Expanded Job Queue

sky/backends/cloud_vm_ray_backend.py

@@ -107,6 +107,13 @@
     pathlib.Path(sky.__file__).resolve().parent / 'backends' /
     'monkey_patches' / 'monkey_patch_ray_up.py')
 
+# Restart skylet when the version does not match to keep the skylet up-to-date.
+_MAYBE_SKYLET_RESTART_CMD = (
+    f'[[ $(cat {constants.SKYLET_VERSION_FILE}) = "{constants.SKYLET_VERSION}"'
+    ' ]] || (pkill -f "python3 -m sky.skylet.skylet";'
+    f' echo {constants.SKYLET_VERSION} > {constants.SKYLET_VERSION_FILE};'
+    'nohup python3 -m sky.skylet.skylet >> ~/.sky/skylet.log 2>&1 &);')
+
 
 def _get_cluster_config_template(cloud):
     cloud_to_template = {
@@ -1649,7 +1656,7 @@ def _ensure_cluster_ray_started(self, handle: 'CloudVmRayResourceHandle',
         # For backward compatability and robustness of skylet, it is restarted
         returncode = backend.run_on_head(
             handle,
-            f'ray status',
+            'ray status',
             # At this state, an erroneous cluster may not have cached
             # handle.head_ip (global_user_state.add_or_update_cluster(...,
             # ready=True)).
@@ -2332,6 +2339,12 @@ def _update_after_cluster_provisioned(
         usage_lib.messages.usage.update_final_cluster_status(
             global_user_state.ClusterStatus.UP)
 
+        self.run_on_head(
+            handle,
+            _MAYBE_SKYLET_RESTART_CMD,
+            use_cached_head_ip=False,
+        )
+
         # Update job queue to avoid stale jobs (when restarted), before
         # setting the cluster to be ready.
         if prev_cluster_status == global_user_state.ClusterStatus.INIT:
@@ -2606,7 +2619,7 @@ def _exec_code_on_head(
                                                       require_outputs=True)
 
         if 'has no attribute' in stdout:
-            # Happens when someone calls `sky exec` but remote has not been updated
+            # Happens when someone calls `sky exec` but remote is outdated
             # necessicating calling `sky launch`
             with ux_utils.print_exception_no_traceback():
                 raise RuntimeError(

sky/execution.py

@@ -66,14 +66,6 @@
   sky.spot_launch(task, ...)
 """.strip()
 
-SKYLET_VERSION = 1  # maybe move it to sky/skylet/constants.py
-SKYLET_VERSION_FILE = '~/.sky/skylet_version'  # maybe move it to sky/skylet/constants.py
-
-# Restart skylet when the version does not match to keep the skylet up-to-date.
-_MAYBE_SKYLET_RESTART_CMD = (
-    f'[[ $(cat {SKYLET_VERSION_FILE}) = "{SKYLET_VERSION}" ]] || (pkill -f "python3 -m sky.skylet.skylet"; echo {SKYLET_VERSION} >  {SKYLET_VERSION_FILE}; nohup python3 -m '
-    'sky.skylet.skylet >> ~/.sky/skylet.log 2>&1 &);')
-
 
 def _convert_to_dag(entrypoint: Any) -> 'sky.Dag':
     """Convert the entrypoint to a sky.Dag.
@@ -293,11 +285,6 @@ def _execute(
             logger.info('Setup commands skipped.')
         elif Stage.SETUP in stages:
             backend.setup(handle, task, detach_setup=detach_setup)
-            backend.run_on_head(
-                handle,
-                _MAYBE_SKYLET_RESTART_CMD,
-                use_cached_head_ip=False,
-            )
 
         if Stage.PRE_EXEC in stages:
             if idle_minutes_to_autostop is not None:

sky/skylet/constants.py

@@ -10,3 +10,6 @@
     'command to initialize it locally: sky launch -c {cluster} \'\'')
 
 JOB_ID_ENV_VAR = 'SKYPILOT_JOB_ID'
+
+SKYLET_VERSION = 1
+SKYLET_VERSION_FILE = '~/.sky/skylet_version'

sky/skylet/job_lib.py

@@ -11,7 +11,7 @@
 import subprocess
 import time
 import typing
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 import colorama
 import filelock
@@ -156,7 +156,7 @@ def queue(self, job_id: int, cmd: str) -> None:
         set_status(job_id, JobStatus.PENDING)
         self.schedule_step()
 
-    def remove_job_no_lock(self, job_id: str) -> None:
+    def remove_job_no_lock(self, job_id: int) -> None:
         _CURSOR.execute(f'DELETE FROM pending_jobs WHERE job_id={job_id!r}')
         _CONN.commit()
 
@@ -191,14 +191,14 @@ def async_schedule_step(self) -> None:
         p = multiprocessing.Process(target=self.schedule_step)
         p.start()
 
-    def _get_jobs(self) -> List[str]:
+    def _get_jobs(self) -> List[Tuple[int, str, int, int]]:
         raise NotImplementedError
 
 
 class FIFOScheduler(JobScheduler):
     """First in first out job scheduler"""
 
-    def _get_jobs(self) -> List[str]:
+    def _get_jobs(self) -> List[Tuple[int, str, int, int]]:
         return list(
             _CURSOR.execute('SELECT * FROM pending_jobs ORDER BY job_id'))
 
@@ -219,7 +219,8 @@ def _get_jobs(self) -> List[str]:
 _RAY_TO_JOB_STATUS_MAP = {
     # These are intentionally set this way, because:
     # 1. when the ray status indicates the job is PENDING the generated
-    # python program has left the job queue and is now PENDING
+    # python program has been `ray job submit` from the job queue
+    # and is now PENDING
     # 2. when the ray status indicates the job is RUNNING the job can be in
     # setup or resources may not be allocated yet, i.e. the job should be
     # PENDING.