generate-docker-compose.py

import sys
import argparse

#image = 'image-dfl-bad'

def generate_docker_compose(image_name, num_nodes):
    docker_compose_template = """\
version: '3.8'
services:
"""

    service_template = """
  node{i}:
    image: {img}
    container_name: node{i}
    environment:
      - RANK={i}
      - BASE_PORT=8000
      - IP_ADDRESS=node{i}
    volumes:
      - ./neighbors/neighbors{i}.txt:/app/neighbors.txt
    networks:
      - my_network

"""

    network_section = """
networks:
  my_network:
    driver: bridge
"""


    docker_compose_content = docker_compose_template
    for i in range(0, int(num_nodes)):
        docker_compose_content += service_template.format(i=i,img=image_name)

    docker_compose_content += network_section

    with open('docker-compose.yml', 'w') as file:
        file.write(docker_compose_content)

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Upload PDF files to an S3 bucket")
    parser.add_argument("--image_name", help="The name of the image")
    parser.add_argument("--nb_nodes", help="The number of nodes")
    args = parser.parse_args()
    # image_name= args.image_name
    # if len(sys.argv) != 2:
    #     print("Usage: python generate_docker_compose.py <num_nodes>")
    #     sys.exit(1)

    # num_nodes = int(sys.argv[1])
    generate_docker_compose(args.image_name, args.nb_nodes)