Expanded Job Queue

sky/backends/cloud_vm_ray_backend.py

@@ -107,6 +107,10 @@
     pathlib.Path(sky.__file__).resolve().parent / 'backends' /
     'monkey_patches' / 'monkey_patch_ray_up.py')
 
+_SKYLET_RESTART_CMD = (
+    '(pkill -f "python3 -m sky.skylet.skylet"; nohup python3 -m '
+    'sky.skylet.skylet >> ~/.sky/skylet.log 2>&1 &);')
+
 
 def _get_cluster_config_template(cloud):
     cloud_to_template = {
@@ -317,6 +321,7 @@ def add_gang_scheduling_placement_group(
                 print('INFO: All task resources reserved.',
                       file=sys.stderr,
                       flush=True)
+                job_lib.scheduler.schedule_step()
                 """)
         ]
 
@@ -331,7 +336,6 @@ def add_gang_scheduling_placement_group(
                 # We unset it so that user setup command may properly use this env var.
                 setup_cmd = 'unset CUDA_VISIBLE_DEVICES; ' + setup_cmd
                 job_lib.set_status({job_id!r}, job_lib.JobStatus.SETTING_UP)
-                print({_CTRL_C_TIP_MESSAGE!r}, file=sys.stderr, flush=True)
                 total_num_nodes = len(ray.nodes())
                 setup_bundles = [{{"CPU": _SETUP_CPUS}} for _ in range(total_num_nodes)]
                 setup_pg = ray.util.placement_group(setup_bundles, strategy='STRICT_SPREAD')
@@ -365,7 +369,6 @@ def add_gang_scheduling_placement_group(
         self._code += [
             textwrap.dedent(f"""\
                 job_lib.set_job_started({self.job_id!r})
-                job_lib.scheduler.schedule_step()
                 """),
         ]
 
@@ -1647,9 +1650,10 @@ def _ensure_cluster_ray_started(self, handle: 'CloudVmRayResourceHandle',
             return
         backend = CloudVmRayBackend()
 
+        # For backward compatability and robustness of skylet, it is restarted
         returncode = backend.run_on_head(
             handle,
-            'ray status',
+            f'{_SKYLET_RESTART_CMD}ray status',
             # At this state, an erroneous cluster may not have cached
             # handle.head_ip (global_user_state.add_or_update_cluster(...,
             # ready=True)).
@@ -2595,16 +2599,22 @@ def _exec_code_on_head(
                 '--no-wait '
                 f'"{executable} -u {script_path} > {remote_log_path} 2>&1"')
 
-            mkdir_code = (
-                f'{cd} && mkdir -p {remote_log_dir} && touch '
-                f'{remote_log_path} && echo START > {remote_log_path} 2>&1')
+            mkdir_code = (f'{cd} && mkdir -p {remote_log_dir} &&'
+                          f'echo START > {remote_log_path} 2>&1')
             code = job_lib.JobLibCodeGen.queue_job(job_id, job_submit_cmd)
             job_submit_cmd = mkdir_code + ' && ' + code
 
         returncode, stdout, stderr = self.run_on_head(handle,
                                                       job_submit_cmd,
                                                       stream_logs=False,
                                                       require_outputs=True)
+
+        if 'has no attribute' in stdout:
+            logger.info(
+                f'{colorama.Fore.RED}SkyPilot must be updated on remote,'
+                f'use `sky launch` instead{colorama.Style.RESET_ALL}')
+            return
+
         subprocess_utils.handle_returncode(returncode,
                                            job_submit_cmd,
                                            f'Failed to submit job {job_id}.',

sky/skylet/job_lib.py

@@ -10,7 +10,7 @@
 import subprocess
 import time
 import typing
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional
 
 import colorama
 import filelock
@@ -92,15 +92,15 @@ class JobStatus(enum.Enum):
     # In the 'jobs' table, the `submitted_at` column will be set to the current
     # time, when the job is firstly created (in the INIT state).
     INIT = 'INIT'
+    # The job is waiting for the required resources. (`ray job status`
+    # shows RUNNING as the generated ray program has started, but blocked
+    # by the placement constraints.)
+    PENDING = 'PENDING'
     # Running the user's setup script (only in effect if --detach-setup is
     # set). Our update_job_status() can temporarily (for a short period) set
     # the status to SETTING_UP, if the generated ray program has not set
     # the status to PENDING or RUNNING yet.
     SETTING_UP = 'SETTING_UP'
-    # The job is waiting for the required resources. (`ray job status`
-    # shows RUNNING as the generated ray program has started, but blocked
-    # by the placement constraints.)
-    PENDING = 'PENDING'
     # The job is running.
     # In the 'jobs' table, the `start_at` column will be set to the current
     # time, when the job is firstly transitioned to RUNNING.
@@ -140,8 +140,9 @@ def colored_str(self):
 # correctly updated.
 # TODO(zhwu): This number should be tuned based on heuristics.
 _SUBMITTED_GAP_SECONDS = 60
+_PENDING_SUBMIT_TIMEOUT = 5
 
-_PRE_RESOURCE_STATUSES = [JobStatus.PENDING, JobStatus.SETTING_UP]
+_PRE_RESOURCE_STATUSES = [JobStatus.PENDING]
 
 
 class JobScheduler:
@@ -154,13 +155,6 @@ def queue(self, job_id: int, cmd: str) -> None:
         set_status(job_id, JobStatus.PENDING)
         self.schedule_step()
 
-    def set_scheduled(self, job_id: str):
-        # TODO(mraheja): remove pylint disabling when filelock
-        # version updated
-        # pylint: disable=abstract-class-instantiated
-        with filelock.FileLock(_get_lock_path(job_id)):
-            self.remove_job_no_lock(job_id)
-
     def remove_job_no_lock(self, job_id: str) -> None:
         _CURSOR.execute(f'DELETE FROM pending_jobs WHERE job_id={job_id!r}')
         _CONN.commit()
@@ -173,7 +167,7 @@ def _run_job(self, job_id: int, run_cmd: str):
 
     def schedule_step(self) -> None:
         job_owner = getpass.getuser()
-        jobs = list(self._get_jobs())
+        jobs = self._get_jobs()
         if len(jobs) > 0:
             update_status(job_owner, _SUBMITTED_GAP_SECONDS)
         for job_id, run_cmd, submit, _ in jobs:
@@ -192,12 +186,16 @@ def schedule_step(self) -> None:
                 self._run_job(job_id, run_cmd)
                 return
 
+    def _get_jobs(self) -> List:
+        raise NotImplementedError
+
 
 class FIFOScheduler(JobScheduler):
     """First in first out job scheduler"""
 
-    def _get_jobs(self) -> Tuple:
-        return _CURSOR.execute('SELECT * FROM pending_jobs ORDER BY job_id')
+    def _get_jobs(self) -> List:
+        return list(
+            _CURSOR.execute('SELECT * FROM pending_jobs ORDER BY job_id'))
 
 
 scheduler = FIFOScheduler()
@@ -229,8 +227,8 @@ def _get_jobs(self) -> Tuple:
     # (ray's status becomes RUNNING), i.e. it will be very rare that the job
     # will be set to SETTING_UP by the update_job_status, as our generated
     # ray program will set the status to PENDING immediately.
-    'PENDING': JobStatus.INIT,
-    'RUNNING': JobStatus.SETTING_UP,
+    'PENDING': JobStatus.PENDING,
+    'RUNNING': JobStatus.PENDING,
     'SUCCEEDED': JobStatus.SUCCEEDED,
     'FAILED': JobStatus.FAILED,
     'STOPPED': JobStatus.CANCELLED,
@@ -496,8 +494,7 @@ def update_job_status(job_owner: str,
     job_detail_lists: List['ray_pydantic.JobDetails'] = job_client.list_jobs()
 
     pending_jobs = _get_pending_jobs()
-
-    job_details = dict()
+    job_details = {}
     ray_job_ids_set = set(ray_job_ids)
     for job_detail in job_detail_lists:
         if job_detail.submission_id in ray_job_ids_set:
@@ -509,10 +506,14 @@ def update_job_status(job_owner: str,
             ray_status = job_details[ray_job_id].status
             job_statuses[i] = _RAY_TO_JOB_STATUS_MAP[ray_status]
         if job_id in pending_jobs:
+            # Gives a 5 second timeout between job being submit from the
+            # pending queue until appearing in ray jobs
             if pending_jobs[job_id]['submit'] > 0 and pending_jobs[job_id][
-                    'submit'] < time.time() - 5:
+                    'submit'] < time.time() - _PENDING_SUBMIT_TIMEOUT:
                 continue
             if pending_jobs[job_id]['created_time'] < psutil.boot_time():
+                # The job is stale as it is created before the instance
+                # is created.
                 job_statuses[i] = JobStatus.FAILED
             else:
                 job_statuses[i] = JobStatus.PENDING
@@ -686,22 +687,24 @@ def cancel_jobs(job_owner: str, jobs: Optional[List[int]]) -> None:
     # ray cluster (tracked in #1262).
     for job in job_records:
         job_id = make_ray_job_id(job['job_id'], job_owner)
-        try:
-            # TODO(mraheja): remove pylint disabling when filelock
-            # version updated
-            # pylint: disable=abstract-class-instantiated
-            with filelock.FileLock(_get_lock_path(job_id)):
+        # Job is locked to ensure that pending queue does not start it while
+        # it is being cancelled
+        with filelock.FileLock(_get_lock_path(job['job_id'],)):
+            try:
+                # TODO(mraheja): remove pylint disabling when filelock
+                # version updated
+                # pylint: disable=abstract-class-instantiated
                 job_client.stop_job(job_id)
-        except RuntimeError as e:
-            # If the job does not exist or if the request to the
-            # job server fails.
-            logger.warning(str(e))
-            continue
+            except RuntimeError as e:
+                # If the job does not exist or if the request to the
+                # job server fails.
+                logger.warning(str(e))
+                continue
 
-        if job['status'] in [
-                JobStatus.SETTING_UP, JobStatus.PENDING, JobStatus.RUNNING
-        ]:
-            set_status(job['job_id'], JobStatus.CANCELLED)
+            if job['status'] in [
+                    JobStatus.SETTING_UP, JobStatus.PENDING, JobStatus.RUNNING
+            ]:
+                set_status(job['job_id'], JobStatus.CANCELLED)
 
 
 def get_run_timestamp(job_id: Optional[int]) -> Optional[str]:
@@ -759,9 +762,9 @@ def add_job(cls, job_name: str, username: str, run_timestamp: str,
         return cls._build(code)
 
     @classmethod
-    def queue_job(cls, job_name: str, cmd: str) -> None:
+    def queue_job(cls, job_id: int, cmd: str) -> str:
         code = ['job_lib.scheduler.queue('
-                f'{job_name!r},'
+                f'{job_id!r},'
                 f'{cmd!r})']
         return cls._build(code)
 

tests/test_smoke.py

@@ -836,7 +836,7 @@ def test_large_job_queue(generic_cloud: str):
             f'sky launch -y -c {name} --cloud {generic_cloud}',
             f'for i in `seq 1 75`; do sky exec {name} -n {name}-$i -d "echo $i; sleep 100000000"; done',
             f'sky cancel -y {name} 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16',
-            'sleep 120',
+            'sleep 90',
             # Each job takes 0.5 CPU and the default VM has 8 CPUs, so there should be 8 / 0.5 = 16 jobs running.
             # The first 16 jobs are canceled, so there should be 75 - 32 = 43 jobs PENDING.
             f'sky queue {name} | grep -v grep | grep PENDING | wc -l | grep 43',