Repository: skypilot-org/skypilot
Pull request: [Spot] Add spot job name in the `SKYPILOT_TASK_ID` env var
<!-- Describe the changes in this PR -->

Fixes #3391

This PR adds the task name in the `SKYPILOT_TASK_ID`.

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [ ] Code formatting: `bash format.sh`
- [x] Any manual or new tests for this PR (please specify below)
  - [x] `sky spot launch -n test "echo \$SKYPILOT_TASK_ID" --cpus 2`
    ```
    sky-managed-2024-04-05-17-47-48-461038_test_14-0
    ```
  - [x] `sky launch -c my-clus -n test "echo \$SKYPILOT_TASK_ID" --cpus 2`
    ```
    sky-2024-04-05-17-12-35-855707_my-clus_1
    ```
- [ ] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `bash tests/backward_comaptibility_tests.sh`


Problem statement:
[Spot/UX] Make spot job name part of `SKYPILOT_TASK_ID`
```
name: my-task-name

resources:
  cpus: 2+
  use_spot: true

run: |
  echo SKYPILOT_TASK_ID: $SKYPILOT_TASK_ID
```

`sky launch -c my-cluster t-name.yaml --down`
- `SKYPILOT_TASK_ID: sky-2024-03-29-08-33-56-952900_my-cluster_id-1` -- note the cluster (rather than task name) is part of the ID

`sky spot launch t-name.yaml`
- `SKYPILOT_TASK_ID: sky-2024-03-29-15-36-30-985352_spot_id-2-0` -- note `name: my-task-name` is not part of the ID

user:
> For managed spot jobs, would it be possible to make the spot job name influence the task ID? Right now we use the task ID to identify wandb experiments for uniqueness, and for managed spot jobs, the ID is much less informative than if the job is launched on an on-demand or unmanaged spot cluster.

Hints from issue discussion:

Commit message: format

Patch to review:
diff --git a/docs/source/examples/spot-jobs.rst b/docs/source/examples/spot-jobs.rst
index 6aa5606968e..22ddcee58d5 100644
--- a/docs/source/examples/spot-jobs.rst
+++ b/docs/source/examples/spot-jobs.rst
@@ -176,7 +176,7 @@ We also set :code:`--run_name` to :code:`$SKYPILOT_TASK_ID` so that the logs for
 to the same run in Weights & Biases.
 
 .. note::
-  The environment variable :code:`$SKYPILOT_TASK_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
+  The environment variable :code:`$SKYPILOT_TASK_ID` (example: "sky-2022-10-06-05-17-09-750781_bert-qa_spot_id-22-1") 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
@@ -369,6 +369,12 @@ The above YAML file defines a pipeline with two tasks. The first :code:`name: pi
 
 To submit the pipeline, the same command :code:`sky spot launch` is used. The pipeline will be automatically launched and monitored by SkyPilot. You can check the status of the pipeline with :code:`sky spot queue` or :code:`sky spot dashboard`.
 
+.. note::
+
+  The :code:`$SKYPILOT_TASK_ID` environment variable is also available in the :code:`run` section of each task. It is unique for each task in the pipeline.
+  For the :code:`eval` task, the :code:`$SKYPILOT_TASK_ID` will be:
+  :code:`sky-2022-10-06-05-17-09-750781_pipeline_eval_spot_id-8-1`.
+
 .. code-block:: console
 
   $ sky spot launch -n pipeline pipeline.yaml
diff --git a/sky/spot/controller.py b/sky/spot/controller.py
index 3bb8b63d867..159faf832c5 100644
--- a/sky/spot/controller.py
+++ b/sky/spot/controller.py
@@ -58,14 +58,17 @@ def __init__(self, job_id: int, dag_yaml: str,
         # pylint: disable=line-too-long
         # 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_spot_id-17-1
+        #   Example value: sky-2022-10-04-22-46-52-467694_my-spot-name_spot_id-17-1
         job_id_env_vars = []
         for i in range(len(self._dag.tasks)):
+            task_name = self._dag_name
+            if len(self._dag.tasks) > 1:
+                task_name = f'{self._dag_name}_{task_name}'
             job_id_env_var = common_utils.get_global_job_id(
                 self._backend.run_timestamp,
-                'spot',
+                f'{task_name}_spot',
                 str(self._job_id),
-                task_id=i)
+                task_id=i + 1)
             job_id_env_vars.append(job_id_env_var)
 
         for i, task in enumerate(self._dag.tasks):


Code review comments:
diff_hunk: @@ -369,6 +369,12 @@ The above YAML file defines a pipeline with two tasks. The first :code:`name: pi
 
 To submit the pipeline, the same command :code:`sky spot launch` is used. The pipeline will be automatically launched and monitored by SkyPilot. You can check the status of the pipeline with :code:`sky spot queue` or :code:`sky spot dashboard`.
 
+.. note::
+
+  The :code:`$SKYPILOT_TASK_ID` environment variable is also available in the :code:`run` section of each task. It is unique for each task in the pipeline.
+  For the :code:`eval` task, the :code:`$SKYPILOT_TASK_ID` will be:
path: docs/source/examples/spot-jobs.rst
text: @user1:
One UX discussion: I found myself reading these as "spot_id". Took a while to realize that it should be read as "_spot" and appended with "id-xxxx". Shall we skip the `_id` part if backward compatible?

@author:
Oh, it should be read as `spot_id` and appended with `_<job_id>-<task_id>`. I guess there will be no significant backward compatibility error in the case. I am not strongly for keeping the `id` or not. Wdyt?

@user1:
Hmm, reason I found "spot_id" odd is that typically we can just append the IDs without saying IDs first. (E.g., we prepend the timestamp without saying `timestamp-`.)

In this case, how about we remove `-id`? I guess moving `spot-id` as a whole may be less desirable. 

@author:
We do have an `_id` for normal cluster job. To keep it align, we may have to remove it for those jobs as well. 
Removed is for both now.