[UX] Add environment variable SKY_NUM_GPUS_PER_NODE (skypilot-org#1337)

* add SKY_NUM_GPUS_PER_NODE

* increase multi-node progress timeout

* pin torch version

* add comment

* address comment

* fix smoke test

* address comments

Author: Michaelvll

docs/source/examples/distributed-jobs.rst

@@ -13,14 +13,15 @@ For example, here is a simple PyTorch Distributed training example:
   name: resnet-distributed-app
 
   resources:
-    accelerators: V100
+    accelerators: V100:4
 
   num_nodes: 2
 
   setup: |
     pip3 install --upgrade pip
     git clone https://github.com/michaelzhiluo/pytorch-distributed-resnet
-    cd pytorch-distributed-resnet && pip3 install -r requirements.txt
+    # SkyPilot's default image on AWS/GCP has CUDA 11.6 (Azure 11.5).
+    cd pytorch-distributed-resnet && pip3 install -r requirements.txt torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
     mkdir -p data  && mkdir -p saved_models && cd data && \
       wget -c --quiet https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz
     tar -xvzf cifar-10-python.tar.gz
@@ -30,20 +31,22 @@ For example, here is a simple PyTorch Distributed training example:
 
     num_nodes=`echo "$SKY_NODE_IPS" | wc -l`
     master_addr=`echo "$SKY_NODE_IPS" | head -n1`
-    python3 -m torch.distributed.launch --nproc_per_node=1 \
+    python3 -m torch.distributed.launch --nproc_per_node=$SKY_NUM_GPUS_PER_NODE \
       --nnodes=$num_nodes --node_rank=${SKY_NODE_RANK} --master_addr=$master_addr \
       --master_port=8008 resnet_ddp.py --num_epochs 20
 
 In the above, :code:`num_nodes: 2` specifies that this task is to be run on 2
 nodes. The :code:`setup` and :code:`run` commands are executed on both nodes.
 
-SkyPilot exposes two environment variables to distinguish per-node commands:
+SkyPilot exposes these environment variables that can be accessed in a task's ``run`` commands:
 
 - :code:`SKY_NODE_RANK`: rank (an integer ID from 0 to :code:`num_nodes-1`) of
-  the node executing the task
+  the node executing the task.
 - :code:`SKY_NODE_IPS`: a string of IP addresses of the nodes reserved to execute
   the task, where each line contains one IP address.
 
   You can retrieve the number of nodes by :code:`echo "$SKY_NODE_IPS" | wc -l`
   and the IP address of the third node by :code:`echo "$SKY_NODE_IPS" | sed -n
   3p`.
+- :code:`SKY_NUM_GPUS_PER_NODE`: number of GPUs reserved on each node to execute the
+  task; the same as the count in ``accelerators: <name>:<count>`` (rounded up if a fraction).

examples/resnet_distributed_torch.yaml

@@ -9,7 +9,7 @@ num_nodes: 2
 setup: |
     pip3 install --upgrade pip
     git clone https://github.com/michaelzhiluo/pytorch-distributed-resnet
-    cd pytorch-distributed-resnet && pip3 install -r requirements.txt
+    cd pytorch-distributed-resnet && pip3 install -r requirements.txt torch==1.12.1+cu113 --extra-index-url https://download.pytorch.org/whl/cu113
     mkdir -p data  && mkdir -p saved_models && cd data && \
     wget -c --quiet https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz
     tar -xvzf cifar-10-python.tar.gz

sky/backends/cloud_vm_ray_backend.py

@@ -3,6 +3,7 @@
 import enum
 import getpass
 import inspect
+import math
 import json
 import os
 import pathlib
@@ -62,7 +63,7 @@
 
 # Timeout (seconds) for provision progress: if in this duration no new nodes
 # are launched, abort and failover.
-_NODES_LAUNCHING_PROGRESS_TIMEOUT = 60
+_NODES_LAUNCHING_PROGRESS_TIMEOUT = 90
 
 # Time gap between retries after failing to provision in all possible places.
 # Used only if --retry-until-up is set.
@@ -339,17 +340,19 @@ def add_ray_task(self,
         cpu_str = f', num_cpus={backend_utils.DEFAULT_TASK_CPU_DEMAND}'
 
         resources_str = ''
+        num_gpus = 0
         num_gpus_str = ''
         if ray_resources_dict is not None:
             assert len(ray_resources_dict) == 1, \
                 ('There can only be one type of accelerator per instance.'
                  f' Found: {ray_resources_dict}.')
+            num_gpus = list(ray_resources_dict.values())[0]
             resources_str = f', resources={json.dumps(ray_resources_dict)}'
 
             # Passing this ensures that the Ray remote task gets
             # CUDA_VISIBLE_DEVICES set correctly.  If not passed, that flag
             # would be force-set to empty by Ray.
-            num_gpus_str = f', num_gpus={list(ray_resources_dict.values())[0]}'
+            num_gpus_str = f', num_gpus={num_gpus}'
             # `num_gpus` should be empty when the accelerator is not GPU.
             # FIXME: use a set of GPU types.
             resources_key = list(ray_resources_dict.keys())[0]
@@ -376,6 +379,7 @@ def add_ray_task(self,
         log_path = os.path.expanduser({log_path!r})
 
         if script is not None:
+            sky_env_vars_dict['SKY_NUM_GPUS_PER_NODE'] = {int(math.ceil(num_gpus))!r}
             ip = gang_scheduling_id_to_ip[{gang_scheduling_id!r}]
             sky_env_vars_dict['SKY_NODE_RANK'] = ip_rank_map[ip]
             sky_env_vars_dict['SKY_JOB_ID'] = {self.job_id}

tests/test_smoke.py

@@ -294,6 +294,11 @@ def test_job_queue():
             f'sky cancel {name} 2',
             'sleep 5',
             f'sky queue {name} | grep {name}-3 | grep RUNNING',
+            f'sky cancel {name} 3',
+            f'sky exec {name} --gpus K80:0.2 "[[ \$SKY_NUM_GPUS_PER_NODE -eq 1 ]] || exit 1"',
+            f'sky exec {name} --gpus K80:1 "[[ \$SKY_NUM_GPUS_PER_NODE -eq 1 ]] || exit 1"',
+            f'sky logs {name} 4 --status',
+            f'sky logs {name} 5 --status',
         ],
         f'sky down -y {name}',
     )
@@ -315,6 +320,13 @@ def test_n_node_job_queue():
             f'sky cancel {name} 1',
             'sleep 5',
             f'sky queue {name} | grep {name}-3 | grep RUNNING',
+            f'sky cancel {name} 3',
+            f'sky exec {name} --gpus K80:0.2 "[[ \$SKY_NUM_GPUS_PER_NODE -eq 1 ]] || exit 1"',
+            f'sky exec {name} --gpus K80:0.2 --num-nodes 2 "[[ \$SKY_NUM_GPUS_PER_NODE -eq 1 ]] || exit 1"',
+            f'sky exec {name} --gpus K80:1 --num-nodes 2 "[[ \$SKY_NUM_GPUS_PER_NODE -eq 1 ]] || exit 1"',
+            f'sky logs {name} 4 --status',
+            f'sky logs {name} 5 --status',
+            f'sky logs {name} 6 --status',
         ],
         f'sky down -y {name}',
     )