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

# 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: 10
# 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
    # ssh_private_key: ...

head_node:
    InstanceType: p3.2xlarge
    ImageId: ami-0698bcaf8bd9ef56d
    # KeyName: ...
    InstanceMarketOptions:
        MarketType: spot
        SpotOptions:
            BlockDurationMinutes: 360
    BlockDeviceMappings:
        - DeviceName: /dev/sda1
          Ebs:
              VolumeSize: 300
           # SpotOptions:
           #     MaxPrice: "9.0"


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

setup_commands:
    # This replaces the standard anaconda Ray installation
    - ray || pip install -U https://s3-us-west-2.amazonaws.com/ray-wheels/latest/ray-1.2.0.dev0-cp36-cp36m-manylinux2014_x86_64.whl
    # Uncomment this and the filemount to update the Ray installation with your local Ray code
    # - rm -rf ./anaconda3/lib/python3.6/site-packages/ray/util/sgd/
    # - cp -rf ~/sgd ./anaconda3/lib/python3.6/site-packages/ray/util/

    # Installing this without -U to make sure we don't replace the existing Ray installation
    - pip install ray[rllib]
    - pip install -U ipdb torch torchvision tqdm
    # Install Apex
    - rm -rf apex || true
    - git clone https://github.com/NVIDIA/apex && cd apex && pip install -v --no-cache-dir  ./ || true

    # Install timm and get data
    - pip install timm
    - ls data || (curl https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz -O && tar zxvf imagenette2.tgz && mv imagenette2 data)


file_mounts: {
    # This should point to ray/python/ray/util/sgd.
    # ~/anaconda3/lib/python3.6/site-packages/ray/util/sgd/: ../../../../sgd,
}

# Custom commands that will be run on the head node after common setup.
head_setup_commands: []

# Custom commands that will be run on worker nodes after common setup.
worker_setup_commands: []

# # 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 --object-store-memory=1000000000

