# An unique identifier for the head node and workers of this cluster.
cluster_name: horovod-pytorch

# The maximum number of workers nodes to launch in addition to the head
# node. This takes precedence over min_workers. min_workers default to 0.
min_workers: 1
initial_workers: 1
max_workers: 1

target_utilization_fraction: 0.9

# If a node is idle for this many minutes, it will be removed.
idle_timeout_minutes: 50
# docker:
#     image: tensorflow/tensorflow:1.5.0-py3
#     container_name: ray_docker

# Cloud-provider specific configuration.
provider:
    type: aws
    region: us-east-1
    availability_zone: us-east-1c

# How Ray will authenticate with newly launched nodes.
auth:
    ssh_user: ubuntu


head_node:
    InstanceType: p3dn.24xlarge
    ImageId: ami-0698bcaf8bd9ef56d
    InstanceMarketOptions:
        MarketType: spot
    BlockDeviceMappings:
        - DeviceName: /dev/sda1
          Ebs:
              VolumeSize: 250
           # SpotOptions:
           #     MaxPrice: "9.0"


worker_nodes:
    InstanceType: p3dn.24xlarge
    ImageId: ami-0698bcaf8bd9ef56d
    InstanceMarketOptions:
        MarketType: spot
    BlockDeviceMappings:
        - DeviceName: /dev/sda1
          Ebs:
              VolumeSize: 250
        # SpotOptions:
        #     MaxPrice: "9.0"
    #     # Run workers on spot by default. Comment this out to use on-demand.
    #     InstanceMarketOptions:
    #         MarketType: spot

setup_commands:
    - pip install torch torchvision ipdb
    - pip install ray[rllib]  # enable autoscaling
    - git clone https://github.com/horovod/horovod || true
    - git clone https://github.com/NVIDIA/apex && cd apex && pip install -v --no-cache-dir  ./ || true
    - tmux new -d -s my-session "HOROVOD_GPU_ALLREDUCE=NCCL HOROVOD_GPU_BROADCAST=NCCL pip install horovod"


file_mounts: {}

# Custom commands that will be run on the head node after common setup.
head_setup_commands:
    - cat ~/ray_bootstrap_key.pem > ~/.ssh/id_rsa

# Custom commands that will be run on worker nodes after common setup.
worker_setup_commands:
    - pip install horovod

# # Command to start ray on the head node. You don't need to change this.
head_start_ray_commands:
    - ray stop
    - ray start --head --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml --object-store-memory=1000000000

# Command to start ray on worker nodes. You don't need to change this.
worker_start_ray_commands:
    - ray stop
    - ray start --address=$RAY_HEAD_IP:6379 --object-manager-port=8076
    # - nvidia-docker run -it --network=host -d --rm -p 4321:22 horovod:latest bash -c "pip install Pillow==6.1; sleep infinity"

