[Env] SKYPILOT_JOB_ID for all tasks (skypilot-org#1377)

Michaelvll · web-flow · commit e8c0772ee0e5 · 2022-11-07T10:49:01.000-08:00
* Add run id for normal job

* add example for the run id

* fix env_check

* fix env_check

* fix

* address comments

* Rename to SKYPILOT_JOB_ID

* rename the controller's job id to avoid confusion

* rename env variables

* fix
diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
@@ -53,4 +53,4 @@ jobs:
           pip install pytest pytest-xdist pytest-env>=0.6
 
       - name: Run tests with pytest
-        run: SKY_DISABLE_USAGE_COLLECTION=1 pytest ${{ matrix.test-path }}
+        run: SKYPILOT_DISABLE_USAGE_COLLECTION=1 pytest ${{ matrix.test-path }}
diff --git a/docs/source/examples/spot-jobs.rst b/docs/source/examples/spot-jobs.rst
@@ -150,7 +150,7 @@ Below we show an `example <https://github.com/skypilot-org/skypilot/blob/master/
     --max_seq_length 384 \
     --doc_stride 128 \
     --report_to wandb \
-    --run_name $SKYPILOT_RUN_ID \
+    --run_name $SKYPILOT_JOB_ID \
     --output_dir /checkpoint/bert_qa/ \
     --save_total_limit 10 \
     --save_steps 1000
@@ -162,11 +162,11 @@ the output directory and frequency of checkpointing (see more
 on `Huggingface API <https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments.save_steps>`_).
 You may also refer to another example `here <https://github.com/skypilot-org/skypilot/tree/master/examples/spot/resnet_ddp>`_ for periodically checkpointing with PyTorch.
 
-We also set :code:`--run_name` to :code:`$SKYPILOT_RUN_ID` so that the loggings will be saved
+We also set :code:`--run_name` to :code:`$SKYPILOT_JOB_ID` so that the loggings will be saved
 to the same run in Weights & Biases.
 
 .. note::
-  The environment variable :code:`$SKYPILOT_RUN_ID` can be used to identify the same job, i.e., it is kept identical across all
+  The environment variable :code:`$SKYPILOT_JOB_ID` (example: "sky-2022-10-06-05-17-09-750781_spot_id-22") can be used to identify the same job, i.e., it is kept identical across all
   recoveries of the job.
   It can be accessed in the task's :code:`run` commands or directly in the program itself (e.g., access
   via :code:`os.environ` and pass to Weights & Biases for tracking purposes in your training script). It is made available to
diff --git a/docs/source/running-jobs/distributed-jobs.rst b/docs/source/running-jobs/distributed-jobs.rst
@@ -30,27 +30,27 @@ For example, here is a simple PyTorch Distributed training example:
   run: |
     cd pytorch-distributed-resnet
 
-    num_nodes=`echo "$SKY_NODE_IPS" | wc -l`
-    master_addr=`echo "$SKY_NODE_IPS" | head -n1`
-    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 \
+    num_nodes=`echo "$SKYPILOT_NODE_IPS" | wc -l`
+    master_addr=`echo "$SKYPILOT_NODE_IPS" | head -n1`
+    python3 -m torch.distributed.launch --nproc_per_node=$SKYPILOT_NUM_GPUS_PER_NODE \
+      --nnodes=$num_nodes --node_rank=${SKYPILOT_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 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
+- :code:`SKYPILOT_NODE_RANK`: rank (an integer ID from 0 to :code:`num_nodes-1`) of
   the node executing the task.
-- :code:`SKY_NODE_IPS`: a string of IP addresses of the nodes reserved to execute
+- :code:`SKYPILOT_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
+  - You can retrieve the number of nodes by :code:`echo "$SKYPILOT_NODE_IPS" | wc -l`
+    and the IP address of the third node by :code:`echo "$SKYPILOT_NODE_IPS" | sed -n
     3p`.
 
   - To manipulate these IP addresses, you can also store them to a file in the
-    :code:`run` command with :code:`echo $SKY_NODE_IPS >> ~/sky_node_ips`.
-- :code:`SKY_NUM_GPUS_PER_NODE`: number of GPUs reserved on each node to execute the
+    :code:`run` command with :code:`echo $SKYPILOT_NODE_IPS >> ~/sky_node_ips`.
+- :code:`SKYPILOT_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).
diff --git a/examples/env_check.yaml b/examples/env_check.yaml
@@ -24,7 +24,27 @@ run: |
     exit 1
   fi
 
-  echo NODE ID: $SKY_NODE_RANK
-  echo NODE IPS: "$SKY_NODE_IPS"
-  worker_addr=`echo "$SKY_NODE_IPS" | sed -n 2p`
+  if [[ -z "$SKYPILOT_NODE_RANK" ]]; then
+    echo "SKYPILOT_NODE_RANK is not set"
+    exit 1
+  else
+    echo "SKYPILOT_NODE_RANK is set to ${SKYPILOT_NODE_RANK}"
+  fi
+
+  if [[ -z "$SKYPILOT_NODE_IPS" ]]; then
+    echo "SKYPILOT_NODE_IPS is not set"
+    exit 1
+  else
+    echo "SKYPILOT_NODE_IPS is set to ${SKYPILOT_NODE_IPS}"
+    echo "${SKYPILOT_NODE_IPS}"
+    echo "${SKYPILOT_NODE_IPS}" | wc -l | grep 2 || exit 1
+  fi
+  worker_addr=`echo "$SKYPILOT_NODE_IPS" | sed -n 2p`
   echo Worker IP: $worker_addr
+
+  if [[ -z "$SKYPILOT_JOB_ID" ]]; then
+    echo "SKYPILOT_JOB_ID is not set"
+    exit 1
+  else
+    echo "SKYPILOT_JOB_ID is set to ${SKYPILOT_JOB_ID}"
+  fi
diff --git a/examples/ray_tune_app.yaml b/examples/ray_tune_app.yaml
@@ -11,6 +11,6 @@ setup: |
     pip3 install ray[tune] pytorch-lightning==1.4.9 lightning-bolts torchvision
 
 run: |
-    if [ "${SKY_NODE_RANK}" == "0" ]; then
+    if [ "${SKYPILOT_NODE_RANK}" == "0" ]; then
         python3 tune_ptl_example.py
     fi
diff --git a/examples/resnet_distributed_torch.yaml b/examples/resnet_distributed_torch.yaml
@@ -19,8 +19,8 @@ setup: |
 run: |
     cd pytorch-distributed-resnet
 
-    num_nodes=`echo "$SKY_NODE_IPS" | wc -l`
-    master_addr=`echo "$SKY_NODE_IPS" | head -n1`
+    num_nodes=`echo "$SKYPILOT_NODE_IPS" | wc -l`
+    master_addr=`echo "$SKYPILOT_NODE_IPS" | head -n1`
     python3 -m torch.distributed.launch --nproc_per_node=1 \
-    --nnodes=$num_nodes --node_rank=${SKY_NODE_RANK} --master_addr=$master_addr \
+    --nnodes=$num_nodes --node_rank=${SKYPILOT_NODE_RANK} --master_addr=$master_addr \
     --master_port=8008 resnet_ddp.py --num_epochs 20
diff --git a/examples/resnet_distributed_torch_scripts/run.sh b/examples/resnet_distributed_torch_scripts/run.sh
@@ -4,9 +4,9 @@ conda activate resnet
 conda env list
 
 cd pytorch-distributed-resnet
-num_nodes=`echo "$SKY_NODE_IPS" | wc -l`
-master_addr=`echo "$SKY_NODE_IPS" | head -n1`
+num_nodes=`echo "$SKYPILOT_NODE_IPS" | wc -l`
+master_addr=`echo "$SKYPILOT_NODE_IPS" | head -n1`
 echo MASTER_ADDR $master_addr
 python -m torch.distributed.launch --nproc_per_node=1 \
---nnodes=$num_nodes --node_rank=${SKY_NODE_RANK} --master_addr=$master_addr \
+--nnodes=$num_nodes --node_rank=${SKYPILOT_NODE_RANK} --master_addr=$master_addr \
 --master_port=8008 resnet_ddp.py --num_epochs 20
diff --git a/examples/spot/bert_qa.yaml b/examples/spot/bert_qa.yaml
@@ -36,7 +36,7 @@ run: |
     --max_seq_length 384 \
     --doc_stride 128 \
     --report_to wandb \
-    --run_name $SKYPILOT_RUN_ID \
+    --run_name $SKYPILOT_JOB_ID \
     --output_dir /checkpoint/bert_qa/ \
     --save_total_limit 10 \
     --save_steps 1000
diff --git a/examples/spot/resnet.yaml b/examples/spot/resnet.yaml
@@ -46,9 +46,9 @@ run: |
     # modify your run id for each different run!
     run_id="resnet-run-1"
 
-    num_nodes=`echo "$SKY_NODE_IPS" | wc -l`
-    master_addr=`echo "$SKY_NODE_IPS" | head -n1`
+    num_nodes=`echo "$SKYPILOT_NODE_IPS" | wc -l`
+    master_addr=`echo "$SKYPILOT_NODE_IPS" | head -n1`
     python3 -m torch.distributed.launch --nproc_per_node=1 \
-    --nnodes=$num_nodes --node_rank=${SKY_NODE_RANK} --master_addr=$master_addr \
+    --nnodes=$num_nodes --node_rank=${SKYPILOT_NODE_RANK} --master_addr=$master_addr \
     --master_port=8008 resnet_ddp.py --num_epochs 100000 --model_dir /checkpoint/torch_ddp_resnet/ \
     --resume --model_filename resnet_distributed-with-epochs.pth --run_id $run_id --wandb_dir /checkpoint/
diff --git a/examples/storage/checkpointed_training.yaml b/examples/storage/checkpointed_training.yaml
@@ -43,9 +43,9 @@ run: |
     cd pytorch-distributed-resnet
     git pull
 
-    num_nodes=`echo "$SKY_NODE_IPS" | wc -l`
-    master_addr=`echo "$SKY_NODE_IPS" | head -n1`
+    num_nodes=`echo "$SKYPILOT_NODE_IPS" | wc -l`
+    master_addr=`echo "$SKYPILOT_NODE_IPS" | head -n1`
     python3 -m torch.distributed.launch --nproc_per_node=1 \
-    --nnodes=$num_nodes --node_rank=${SKY_NODE_RANK} --master_addr=$master_addr \
+    --nnodes=$num_nodes --node_rank=${SKYPILOT_NODE_RANK} --master_addr=$master_addr \
     --master_port=8008 resnet_ddp.py --num_epochs 100 --model_dir /checkpoints/torch_ddp_resnet/ \
     --resume --model_filename resnet_distributed-with-epochs.pth
diff --git a/examples/storage/pingpong.yaml b/examples/storage/pingpong.yaml
@@ -27,5 +27,5 @@ setup: |
   git clone https://github.com/romilbhardwaj/sharedfs-pingpong.git
 
 run: |
-  num_nodes=`echo "$SKY_NODE_IPS" | wc -l`
-  python -u sharedfs-pingpong/main.py --process-id ${SKY_NODE_RANK} --shared-path /sharedfs/ --num-processes $num_nodes
+  num_nodes=`echo "$SKYPILOT_NODE_IPS" | wc -l`
+  python -u sharedfs-pingpong/main.py --process-id ${SKYPILOT_NODE_RANK} --shared-path /sharedfs/ --num-processes $num_nodes
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
@@ -292,7 +292,9 @@ def check_ip():
                 else:
                     ip_rank_map = {{ip: i for i, ip in enumerate(gang_scheduling_id_to_ip)}}
                     ip_list_str = '\\n'.join(gang_scheduling_id_to_ip)
-
+                
+                sky_env_vars_dict['SKYPILOT_NODE_IPS'] = ip_list_str
+                # Environment starting with `SKY_` is deprecated.
                 sky_env_vars_dict['SKY_NODE_IPS'] = ip_list_str
                 """),
         ]
@@ -318,6 +320,7 @@ def register_run_fn(self, run_fn: str, run_fn_name: str) -> None:
     def add_ray_task(self,
                      bash_script: str,
                      task_name: Optional[str],
+                     job_run_id: str,
                      ray_resources_dict: Optional[Dict[str, float]],
                      log_path: str,
                      env_vars: Dict[str, str] = None,
@@ -366,6 +369,10 @@ def add_ray_task(self,
             sky_env_vars_dict_str = '\n'.join(
                 f'sky_env_vars_dict[{k!r}] = {v!r}'
                 for k, v in env_vars.items())
+        if job_run_id is not None:
+            sky_env_vars_dict_str += (
+                f'\nsky_env_vars_dict[{constants.JOB_ID_ENV_VAR!r}]'
+                f' = {job_run_id!r}')
 
         logger.debug('Added Task with options: '
                      f'{name_str}{cpu_str}{resources_str}{num_gpus_str}')
@@ -379,10 +386,18 @@ def add_ray_task(self,
         log_path = os.path.expanduser({log_path!r})
 
         if script is not None:
+            sky_env_vars_dict['SKYPILOT_NUM_GPUS_PER_NODE'] = {int(math.ceil(num_gpus))!r}
+            # Environment starting with `SKY_` is deprecated.
             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['SKYPILOT_NODE_RANK'] = ip_rank_map[ip]
+            # Environment starting with `SKY_` is deprecated.
             sky_env_vars_dict['SKY_NODE_RANK'] = ip_rank_map[ip]
-            sky_env_vars_dict['SKY_JOB_ID'] = {self.job_id}
+            
+            sky_env_vars_dict['SKYPILOT_INTERNAL_JOB_ID'] = {self.job_id}
+            # Environment starting with `SKY_` is deprecated.
+            sky_env_vars_dict['SKY_INTERNAL_JOB_ID'] = {self.job_id}
 
             futures.append(run_bash_command_with_log \\
                     .options({name_str}{cpu_str}{resources_str}{num_gpus_str}) \\
@@ -3002,12 +3017,16 @@ def _execute_task_one_node(self, handle: ResourceHandle,
             run_fn_name = task.run.__name__
             codegen.register_run_fn(run_fn_code, run_fn_name)
 
+        job_run_id = common_utils.get_global_job_id(
+            self.run_timestamp, cluster_name=handle.cluster_name, job_id=job_id)
+
         command_for_node = task.run if isinstance(task.run, str) else None
         use_sudo = isinstance(handle.launched_resources.cloud, clouds.Local)
         codegen.add_ray_task(
             bash_script=command_for_node,
             env_vars=task.envs,
             task_name=task.name,
+            job_run_id=job_run_id,
             ray_resources_dict=backend_utils.get_task_demands_dict(task),
             log_path=log_path,
             use_sudo=use_sudo)
@@ -3061,6 +3080,10 @@ def _execute_task_n_nodes(self, handle: ResourceHandle, task: task_lib.Task,
             run_fn_code = textwrap.dedent(inspect.getsource(task.run))
             run_fn_name = task.run.__name__
             codegen.register_run_fn(run_fn_code, run_fn_name)
+
+        job_run_id = common_utils.get_global_job_id(
+            self.run_timestamp, cluster_name=handle.cluster_name, job_id=job_id)
+
         # TODO(zhwu): The resources limitation for multi-node ray.tune and
         # horovod should be considered.
         for i in range(num_actual_nodes):
@@ -3075,6 +3098,7 @@ def _execute_task_n_nodes(self, handle: ResourceHandle, task: task_lib.Task,
                 bash_script=command_for_node,
                 env_vars=task.envs,
                 task_name=name,
+                job_run_id=job_run_id,
                 ray_resources_dict=accelerator_dict,
                 log_path=log_path,
                 gang_scheduling_id=i,
diff --git a/sky/callbacks/sky_callback/utils.py b/sky/callbacks/sky_callback/utils.py
@@ -1,7 +1,7 @@
 import os
 from typing import Optional
 
-DISABLE_CALLBACK = os.environ.get('SKY_DISABLE_CALLBACK',
+DISABLE_CALLBACK = os.environ.get('SKYPILOT_DISABLE_CALLBACK',
                                   'False').lower() in ('true', '1')
 
 
diff --git a/sky/execution.py b/sky/execution.py
@@ -548,7 +548,7 @@ def _maybe_translate_local_file_mounts_and_sync_up(
     # Translate the workdir and local file mounts to cloud file mounts.
     # ================================================================
     task = copy.deepcopy(task)
-    run_id = common_utils.get_run_id()[:8]
+    run_id = common_utils.get_usage_run_id()[:8]
     original_file_mounts = task.file_mounts if task.file_mounts else {}
     original_storage_mounts = task.storage_mounts if task.storage_mounts else {}
 
diff --git a/sky/skylet/constants.py b/sky/skylet/constants.py
@@ -8,3 +8,5 @@
 UNINITIALIZED_ONPREM_CLUSTER_MESSAGE = (
     'Found uninitialized local cluster {cluster}. Run this '
     'command to initialize it locally: sky launch -c {cluster} \'\'')
+
+JOB_ID_ENV_VAR = 'SKYPILOT_JOB_ID'
diff --git a/sky/spot/controller.py b/sky/spot/controller.py
@@ -13,10 +13,12 @@
 from sky import sky_logging
 from sky.backends import backend_utils
 from sky.backends import cloud_vm_ray_backend
+from sky.skylet import constants
 from sky.skylet import job_lib
 from sky.spot import recovery_strategy
 from sky.spot import spot_state
 from sky.spot import spot_utils
+from sky.utils import common_utils
 
 logger = sky_logging.init_logger(__name__)
 
@@ -37,11 +39,9 @@ def __init__(self, job_id: int, task_yaml: str,
         # Add a unique identifier to the task environment variables, so that
         # the user can have the same id for multiple recoveries.
         #   Example value: sky-2022-10-04-22-46-52-467694_id-17
-        # TODO(zhwu): support SKYPILOT_RUN_ID for normal jobs as well, so
-        # the use can use env_var for normal jobs.
         task_envs = self._task.envs or {}
-        task_envs['SKYPILOT_RUN_ID'] = (f'{self.backend.run_timestamp}'
-                                        f'_id-{self._job_id}')
+        task_envs[constants.JOB_ID_ENV_VAR] = common_utils.get_global_job_id(
+            self.backend.run_timestamp, 'spot', self._job_id)
         self._task.set_envs(task_envs)
 
         spot_state.set_submitted(
diff --git a/sky/templates/gcp-ray.yml.j2 b/sky/templates/gcp-ray.yml.j2
@@ -171,15 +171,15 @@ head_start_ray_commands:
   # Line "which prlimit ..": increase the limit of the number of open files for the raylet process, as the `ulimit` may not take effect at this point, because it requires
   # all the sessions to be reloaded. This is a workaround.
   - ((ps aux | grep -v nohup | grep -v grep | grep -q -- "python3 -m sky.skylet.skylet") || nohup python3 -m sky.skylet.skylet >> ~/.sky/skylet.log 2>&1 &);
-    export SKY_NUM_GPUS=0 && which nvidia-smi > /dev/null && SKY_NUM_GPUS=$(nvidia-smi --query-gpu=index,name --format=csv,noheader | wc -l);
-    ray stop; ray start --disable-usage-stats --head --port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} --num-gpus=$SKY_NUM_GPUS || exit 1;
+    export SKYPILOT_NUM_GPUS=0 && which nvidia-smi > /dev/null && SKYPILOT_NUM_GPUS=$(nvidia-smi --query-gpu=index,name --format=csv,noheader | wc -l);
+    ray stop; ray start --disable-usage-stats --head --port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} --num-gpus=$SKYPILOT_NUM_GPUS || exit 1;
     which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
 
 # Worker commands are needed for TPU VM Pods
 {%- if num_nodes > 1 or tpu_vm %}
 worker_start_ray_commands:
-  - SKY_NUM_GPUS=0 && which nvidia-smi > /dev/null && SKY_NUM_GPUS=$(nvidia-smi --query-gpu=index,name --format=csv,noheader | wc -l);
-    ray stop; ray start --disable-usage-stats --address=$RAY_HEAD_IP:6379 --object-manager-port=8076 {{"--resources='%s'" % custom_resources if custom_resources}} --num-gpus=$SKY_NUM_GPUS || exit 1;
+  - SKYPILOT_NUM_GPUS=0 && which nvidia-smi > /dev/null && SKYPILOT_NUM_GPUS=$(nvidia-smi --query-gpu=index,name --format=csv,noheader | wc -l);
+    ray stop; ray start --disable-usage-stats --address=$RAY_HEAD_IP:6379 --object-manager-port=8076 {{"--resources='%s'" % custom_resources if custom_resources}} --num-gpus=$SKYPILOT_NUM_GPUS || exit 1;
     which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
 {%- else %}
 worker_start_ray_commands: []
diff --git a/sky/templates/spot-controller.yaml.j2 b/sky/templates/spot-controller.yaml.j2
@@ -35,7 +35,7 @@ setup: |
 run: |
   python -u -m sky.spot.controller \
     {{remote_user_yaml_prefix}}/{{cluster_name}}.yaml \
-    --job-id $SKY_JOB_ID {% if retry_until_up %}--retry-until-up{% endif %}
+    --job-id $SKYPILOT_INTERNAL_JOB_ID {% if retry_until_up %}--retry-until-up{% endif %}
 
 envs:
   SKYPILOT_USAGE_USER_ID: {{logging_user_hash}}
diff --git a/sky/usage/usage_lib.py b/sky/usage/usage_lib.py
@@ -77,7 +77,7 @@ def __init__(self) -> None:
         super().__init__(constants.USAGE_MESSAGE_SCHEMA_VERSION)
         # Message identifier.
         self.user: str = _get_user_hash()
-        self.run_id: str = common_utils.get_run_id()
+        self.run_id: str = common_utils.get_usage_run_id()
         self.sky_version: str = sky.__version__
 
         # Entry
diff --git a/sky/utils/common_utils.py b/sky/utils/common_utils.py
diff --git a/tests/test_onprem.py b/tests/test_onprem.py
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
