Repository: skypilot-org/skypilot
Pull request: [Spot] Keep SKYPILOT_JOB_ID the same for the same spot job
Fixes #1399 

Also, fixes the testing for the `SKYPILOT_JOB_ID`.

~~This PR also adds the `SKYPILOT_JOB_ID` to the `sky spot queue -a` to make it easier for the user to find.~~
```
(sky-dev) ➜  sky-experiment-dev (fix-sky-job-id) sky spot queue -a                                                                                                                                                                               ✱
Fetching managed spot job statuses...
Managed spot jobs:
In progress jobs: 1 STARTING

ID  NAME                            RESOURCES     SUBMITTED                      TOT. DURATION  JOB DURATION  #RECOVERIES  STATUS             STARTED                        CLUSTER                    REGION     SKYPILOT_JOB_ID                           
14  sky-9e85-zhwu                   1x [CPU:0.5]  45s ago                        45s            -             0            STARTING           -                              1x AWS(m6i.2xlarge[Spot])  us-east-2  sky-2022-11-12-05-55-13-159500_spot_id-14 
```
Moved the discussion of how to add the JOB_ID to the queue to the issue #1401 , so as to unblock this PR for fixing the bug.

Tested:
- [x] `tests/run_smoke_tests.sh test_spot_recovery`

Future TODO:
- [ ] Add `SKYPILOT_JOB_ID` into the `sky queue` as well.

Problem statement:
[Spot] The SKYPILOT_JOB_ID changes across restarts
From our user:
> It seems that the SKYPILOT_JOB_ID isn’t being preserved across restarts. Any idea what could be happening?
During the initial run it was 2022-11-12-00-17-08-581710_test-managed-spot-6_id-1
After I killed the cluster and the spot controller restarted, it became sky-2022-11-12-02-51-24-649160_test-managed-spot-6_id-1

Hints from issue discussion:

Commit message: don't set job_id_env_var for spot controller task

Patch to review:
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index 8082dbd7cad..367b24ce5ee 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -311,7 +311,6 @@ def add_gang_scheduling_placement_group(
                       file=sys.stderr,
                       flush=True)
                 job_lib.set_job_started({self.job_id!r})
-                sky_env_vars_dict = dict()
                 """)
         ]
 
@@ -337,10 +336,6 @@ 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
                 """),
         ]
 
@@ -365,7 +360,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,
+                     job_run_id: Optional[str],
                      ray_resources_dict: Optional[Dict[str, float]],
                      log_path: str,
                      env_vars: Dict[str, str] = None,
@@ -409,15 +404,26 @@ def add_ray_task(self,
         resources_str += ', placement_group=pg'
         resources_str += f', placement_group_bundle_index={gang_scheduling_id}'
 
-        sky_env_vars_dict_str = ''
+        sky_env_vars_dict_str = [
+            textwrap.dedent("""\
+            sky_env_vars_dict = dict()
+            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
+            """)
+        ]
+
         if env_vars is not None:
-            sky_env_vars_dict_str = '\n'.join(
+            sky_env_vars_dict_str += [
                 f'sky_env_vars_dict[{k!r}] = {v!r}'
-                for k, v in env_vars.items())
+                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}')
+            sky_env_vars_dict_str += [
+                f'sky_env_vars_dict[{constants.JOB_ID_ENV_VAR!r}]'
+                f' = {job_run_id!r}'
+            ]
+        sky_env_vars_dict_str = '\n'.join(sky_env_vars_dict_str)
 
         logger.debug('Added Task with options: '
                      f'{name_str}{cpu_str}{resources_str}{num_gpus_str}')
@@ -3086,8 +3092,13 @@ 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)
+        if task.spot_task is None:
+            # Don't set JOB_ID_ENV_VAR if it is a spot controller task.
+            job_run_id = task.envs.get(
+                constants.JOB_ID_ENV_VAR,
+                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)
@@ -3153,8 +3164,13 @@ def _execute_task_n_nodes(self, handle: ResourceHandle, task: task_lib.Task,
             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)
+        if task.spot_task is None:
+            # Don't set JOB_ID_ENV_VAR if it is a spot controller task.
+            job_run_id = task.envs.get(
+                constants.JOB_ID_ENV_VAR,
+                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.
diff --git a/sky/cli.py b/sky/cli.py
index 9de8414cdd8..0ff074b1a66 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -1787,6 +1787,14 @@ def autostop(
     required=False,
     help=('Retry provisioning infinitely until the cluster is up, '
           'if we fail to start the cluster due to unavailability errors.'))
+@click.option(
+    '--force',
+    '-f',
+    default=False,
+    is_flag=True,
+    required=False,
+    help=('Force start the cluster even if it is already UP. Useful for '
+          'upgrading SkyPilot runtime.'))
 @usage_lib.entrypoint
 # pylint: disable=redefined-builtin
 def start(
@@ -1795,7 +1803,8 @@ def start(
         yes: bool,
         idle_minutes_to_autostop: Optional[int],
         down: bool,  # pylint: disable=redefined-outer-name
-        retry_until_up: bool):
+        retry_until_up: bool,
+        force: bool):
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Restart cluster(s).
 
@@ -1883,7 +1892,8 @@ def start(
             #      INIT state cluster due to head_ip not being cached).
             #
             #      This can be replicated by adding `exit 1` to Task.setup.
-            if cluster_status == global_user_state.ClusterStatus.UP:
+            if (not force and
+                    cluster_status == global_user_state.ClusterStatus.UP):
                 # An UP cluster; skipping 'sky start' because:
                 #  1. For a really up cluster, this has no effects (ray up -y
                 #    --no-restart) anyway.
@@ -1896,7 +1906,8 @@ def start(
                 #    This is dangerous and unwanted behavior!
                 click.echo(f'Cluster {name} already has status UP.')
                 continue
-            assert cluster_status in (
+
+            assert force or cluster_status in (
                 global_user_state.ClusterStatus.INIT,
                 global_user_state.ClusterStatus.STOPPED), cluster_status
             to_start.append(name)
@@ -1918,7 +1929,8 @@ def start(
             core.start(name,
                        idle_minutes_to_autostop,
                        retry_until_up,
-                       down=down)
+                       down=down,
+                       force=force)
         except exceptions.NotSupportedError as e:
             click.echo(str(e))
         click.secho(f'Cluster {name} started.', fg='green')
diff --git a/sky/core.py b/sky/core.py
index a0b370be5d3..e9f1cb333ac 100644
--- a/sky/core.py
+++ b/sky/core.py
@@ -65,20 +65,21 @@ def status(refresh: bool = False) -> List[Dict[str, Any]]:
 
 
 def _start(
-        cluster_name: str,
-        idle_minutes_to_autostop: Optional[int] = None,
-        retry_until_up: bool = False,
-        down: bool = False,  # pylint: disable=redefined-outer-name
+    cluster_name: str,
+    idle_minutes_to_autostop: Optional[int] = None,
+    retry_until_up: bool = False,
+    down: bool = False,  # pylint: disable=redefined-outer-name
+    force: bool = False,
 ) -> backends.Backend.ResourceHandle:
 
     cluster_status, handle = backend_utils.refresh_cluster_status_handle(
         cluster_name)
     if handle is None:
         raise ValueError(f'Cluster {cluster_name!r} does not exist.')
-    if cluster_status == global_user_state.ClusterStatus.UP:
+    if not force and cluster_status == global_user_state.ClusterStatus.UP:
         print(f'Cluster {cluster_name!r} is already up.')
         return
-    assert cluster_status in (
+    assert force or cluster_status in (
         global_user_state.ClusterStatus.INIT,
         global_user_state.ClusterStatus.STOPPED), cluster_status
 
@@ -108,10 +109,11 @@ def _start(
 
 @usage_lib.entrypoint
 def start(
-        cluster_name: str,
-        idle_minutes_to_autostop: Optional[int] = None,
-        retry_until_up: bool = False,
-        down: bool = False,  # pylint: disable=redefined-outer-name
+    cluster_name: str,
+    idle_minutes_to_autostop: Optional[int] = None,
+    retry_until_up: bool = False,
+    down: bool = False,  # pylint: disable=redefined-outer-name
+    force: bool = False,
 ) -> None:
     # NOTE(dev): Keep the docstring consistent between the Python API and CLI.
     """Restart a cluster.
@@ -142,6 +144,8 @@ def start(
         down: Autodown the cluster: tear down the cluster after specified
             minutes of idle time after all jobs finish (successfully or
             abnormally). Requires ``idle_minutes_to_autostop`` to be set.
+        force: whether to force start the cluster even if it is already up.
+            Useful for upgrading SkyPilot runtime.
 
     Raises:
         ValueError: the specified cluster does not exist; or if ``down`` is set
@@ -153,7 +157,11 @@ def start(
     if down and idle_minutes_to_autostop is None:
         raise ValueError(
             '`idle_minutes_to_autostop` must be set if `down` is True.')
-    _start(cluster_name, idle_minutes_to_autostop, retry_until_up, down)
+    _start(cluster_name,
+           idle_minutes_to_autostop,
+           retry_until_up,
+           down,
+           force=force)
 
 
 @usage_lib.entrypoint
diff --git a/sky/spot/controller.py b/sky/spot/controller.py
index 08053093b9c..c4dcead9f02 100644
--- a/sky/spot/controller.py
+++ b/sky/spot/controller.py
@@ -40,15 +40,18 @@ def __init__(self, job_id: int, task_yaml: str,
         # the user can have the same id for multiple recoveries.
         #   Example value: sky-2022-10-04-22-46-52-467694_id-17
         task_envs = self._task.envs or {}
-        task_envs[constants.JOB_ID_ENV_VAR] = common_utils.get_global_job_id(
+        job_id_env_var = common_utils.get_global_job_id(
             self.backend.run_timestamp, 'spot', self._job_id)
+        task_envs[constants.JOB_ID_ENV_VAR] = job_id_env_var
         self._task.set_envs(task_envs)
 
         spot_state.set_submitted(
             self._job_id,
             self._task_name,
             self.backend.run_timestamp,
-            resources_str=backend_utils.get_task_resources_str(self._task))
+            resources_str=backend_utils.get_task_resources_str(self._task),
+            job_id_env_var=job_id_env_var)
+        logger.info(f'Submitted with $SKY_JOB_ID: {job_id_env_var}')
         self._cluster_name = spot_utils.generate_spot_cluster_name(
             self._task_name, self._job_id)
         self._strategy_executor = recovery_strategy.StrategyExecutor.make(
diff --git a/sky/spot/spot_state.py b/sky/spot/spot_state.py
index cef5e0cc8bc..d73357504e4 100644
--- a/sky/spot/spot_state.py
+++ b/sky/spot/spot_state.py
@@ -9,6 +9,7 @@
 
 from sky import sky_logging
 from sky.backends import backend_utils
+from sky.utils import db_utils
 
 logger = sky_logging.init_logger(__name__)
 
@@ -32,6 +33,9 @@
     last_recovered_at FLOAT DEFAULT -1,
     recovery_count INTEGER DEFAULT 0,
     job_duration FLOAT DEFAULT 0)""")
+
+db_utils.add_column_to_table(_CURSOR, _CONN, 'spot', 'job_id_env_var', 'TEXT')
+
 # job_duration is the time a job actually runs before last_recover,
 # excluding the provision and recovery time.
 # If the job is not finished:
@@ -43,7 +47,7 @@
 columns = [
     'job_id', 'job_name', 'resources', 'submitted_at', 'status',
     'run_timestamp', 'start_at', 'end_at', 'last_recovered_at',
-    'recovery_count', 'job_duration'
+    'recovery_count', 'job_duration', 'job_id_env_var'
 ]
 
 
@@ -89,7 +93,7 @@ def set_pending(job_id: int, name: str, resources_str: str):
 
 
 def set_submitted(job_id: int, name: str, run_timestamp: str,
-                  resources_str: str):
+                  resources_str: str, job_id_env_var: str):
     """Set the job to submitted."""
     # Use the timestamp in the `run_timestamp` ('sky-2022-10...'), to make the
     # log directory and submission time align with each other, so as to make
@@ -104,10 +108,11 @@ def set_submitted(job_id: int, name: str, run_timestamp: str,
         resources=(?),
         submitted_at=(?),
         status=(?),
-        run_timestamp=(?)
+        run_timestamp=(?),
+        job_id_env_var=(?)
         WHERE job_id=(?)""",
         (name, resources_str, submit_time, SpotStatus.SUBMITTED.value,
-         run_timestamp, job_id))
+         run_timestamp, job_id_env_var, job_id))
     _CONN.commit()
 
 
diff --git a/sky/spot/spot_utils.py b/sky/spot/spot_utils.py
index 0700c1da751..1709896f5fe 100644
--- a/sky/spot/spot_utils.py
+++ b/sky/spot/spot_utils.py
@@ -342,7 +342,7 @@ def format_job_table(jobs: List[Dict[str, Any]], show_all: bool) -> str:
         '#RECOVERIES', 'STATUS'
     ]
     if show_all:
-        columns += ['STARTED', 'CLUSTER', 'REGION']
+        columns += ['STARTED', 'CLUSTER', 'REGION', 'SKYPILOT_JOB_ID']
     job_table = log_utils.create_table(columns)
 
     status_counts = collections.defaultdict(int)
@@ -352,15 +352,13 @@ def format_job_table(jobs: List[Dict[str, Any]], show_all: bool) -> str:
         job_duration = log_utils.readable_time_duration(0,
                                                         job['job_duration'],
                                                         absolute=True)
-        ago_suffix = ' ago' if show_all else ''
-        submitted = log_utils.readable_time_duration(job['submitted_at'],
-                                                     absolute=show_all)
+        submitted = log_utils.readable_time_duration(job['submitted_at'])
         values = [
             job['job_id'],
             job['job_name'],
             job['resources'],
             # SUBMITTED
-            submitted + ago_suffix if submitted != '-' else submitted,
+            submitted if submitted != '-' else submitted,
             # TOT. DURATION
             log_utils.readable_time_duration(job['submitted_at'],
                                              job['end_at'],
@@ -373,15 +371,13 @@ def format_job_table(jobs: List[Dict[str, Any]], show_all: bool) -> str:
             status_counts[job['status'].value] += 1
         if show_all:
             # STARTED
-            started = log_utils.readable_time_duration(job['start_at'],
-                                                       absolute=True)
-            if started != '-':
-                started += ago_suffix
+            started = log_utils.readable_time_duration(job['start_at'])
             values.append(started)
             values.extend([
                 job['cluster_resources'],
                 job['region'],
             ])
+            values.append(job['job_id_env_var'])
         job_table.add_row(values)
     status_str = ', '.join([
         f'{count} {status}' for status, count in sorted(status_counts.items())
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
index 9f4a736bfb2..67dcd57883c 100644
--- a/tests/test_smoke.py
+++ b/tests/test_smoke.py
@@ -723,7 +723,7 @@ def test_spot_recovery():
             f'sky spot launch --cloud aws --region {region} -n {name} "echo SKYPILOT_JOB_ID: \$SKYPILOT_JOB_ID; sleep 1000"  -y -d',
             'sleep 300',
             f's=$(sky spot queue); printf "$s"; echo; echo; printf "$s" | grep {name} | head -n1 | grep "RUNNING"',
-            f'RUN_ID=$(sky spot logs -n {name} --no-follow | grep SKYPILOT_JOB_ID | cut -d: -f2); echo $RUN_ID',
+            f'RUN_ID=$(sky spot logs -n {name} --no-follow | grep SKYPILOT_JOB_ID | cut -d: -f2); echo $RUN_ID | tee /tmp/{name}-run-id',
             # Terminate the cluster manually.
             (f'aws ec2 terminate-instances --region {region} --instance-ids $('
              f'aws ec2 describe-instances --region {region} '
@@ -734,7 +734,7 @@ def test_spot_recovery():
             f's=$(sky spot queue); printf "$s"; echo; echo; printf "$s" | grep {name} | head -n1 | grep "RECOVERING"',
             'sleep 200',
             f's=$(sky spot queue); printf "$s"; echo; echo; printf "$s" | grep {name} | head -n1 | grep "RUNNING"',
-            f'sky spot logs -n {name} --no-follow | grep SKYPILOT_JOB_ID | grep "$RUN_ID"',
+            f'RUN_ID=$(cat /tmp/{name}-run-id); echo $RUN_ID; sky spot logs -n {name} --no-follow | grep SKYPILOT_JOB_ID | grep "$RUN_ID"',
         ],
         f'sky spot cancel -y -n {name}',
     )


Code review comments:
diff_hunk: @@ -1787,6 +1787,14 @@ def autostop(
     required=False,
     help=('Retry provisioning infinitely until the cluster is up, '
           'if we fail to start the cluster due to unavailability errors.'))
+@click.option(
+    '--force',
+    '-f',
+    default=False,
+    is_flag=True,
+    required=False,
+    help=('Force start the cluster even if it is already UP. Useful for '
+          'upgrading SkyPilot runtime.'))
path: sky/cli.py
text: @user1:
```suggestion
          'upgrading the SkyPilot runtime on the cluster.'))
```

@author:
Done. Thanks!

diff_hunk: @@ -40,15 +40,18 @@ def __init__(self, job_id: int, task_yaml: str,
         # the user can have the same id for multiple recoveries.
         #   Example value: sky-2022-10-04-22-46-52-467694_id-17
         task_envs = self._task.envs or {}
-        task_envs[constants.JOB_ID_ENV_VAR] = common_utils.get_global_job_id(
+        job_id_env_var = common_utils.get_global_job_id(
             self.backend.run_timestamp, 'spot', self._job_id)
+        task_envs[constants.JOB_ID_ENV_VAR] = job_id_env_var
         self._task.set_envs(task_envs)
 
         spot_state.set_submitted(
             self._job_id,
             self._task_name,
             self.backend.run_timestamp,
-            resources_str=backend_utils.get_task_resources_str(self._task))
+            resources_str=backend_utils.get_task_resources_str(self._task),
+            job_id_env_var=job_id_env_var)
+        logger.info(f'Submitted with $SKY_JOB_ID: {job_id_env_var}')
path: sky/spot/controller.py
text: @user1:
```suggestion
        logger.info(f'Submitted spot job; SKYPILOT_JOB_ID: {job_id_env_var}')
```

@author:
Done. Thanks!

diff_hunk: @@ -3086,8 +3092,13 @@ 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)
+        if task.spot_task is None:
+            # Don't set JOB_ID_ENV_VAR if it is a spot controller task.
path: sky/backends/cloud_vm_ray_backend.py
text: @user1:
```suggestion
            # Don't regenerate JOB_ID_ENV_VAR if it is a spot controller task (already has it set).
```

@author:
Ahh, this comment was for the if statement, where we do not set the `JOB_ID_ENV_VAR` for the spot controller process, which is only for reduce the confusion when we add the `SKY_JOB_ID` in the `sky queue`.
For now, since we do not have the `SKY_JOB_ID` in the `sky queue`, I will remove the if statement.