Repository: skypilot-org/skypilot
Pull request: Optimize the `sky job queue` when large amount of jobs running
Closes #612. This PR optimizes the `sky job queue` time from 24 seconds to 1.4 seconds when there are 30 jobs running. (To have 30 jobs running together, we need to change the original placement_group setting of CPUs to smaller than 1 to allow more jobs than CPU cores.)

```
# With 30 jobs running
# Old sky queue
sky queue old-jq-test  0.95s user 0.27s system 5% cpu 23.932 total

# New sky queue
sky queue  0.98s user 0.25s system 86% cpu 1.422 total
```

The previous implementation will update the job status for each job when `sky queue` is called, using `ray job status`.  It will introduce significant overheads. Since our tasks will actively update their own status in the `job_table` and the job status update will only be used for dealing with the staled jobs caused by the cluster's stop and restart, it is OK to only update the status during the provision.

Problem statement:
`sky queue` takes a long time when there is many jobs running
Quotes from Daniel:
> with 32 jobs running, sky queue takes 40 seconds to print out the job list (87 total jobs, including cancelled/failed/succeeded)
> just logged in to check and CPU is pretty undersubscribed (I thought it might be oversubscribed anyway, but turns out my code has both poor CPU utilization and poor GPU utilization sad)

The reason could be that we are using `ray job` to query the status of each running job, which may cause a large amount of overhead.

Hints from issue discussion:

Commit message: format

Patch to review:
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index e3172683146..f1dd6c9da7e 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -1130,6 +1130,11 @@ def add_job(self, job_name: str, username: str, run_timestamp: str) -> None:
             'print(job_id, flush=True)',
         ]
 
+    def update_status(self) -> None:
+        self._code += [
+            'job_lib.update_status()',
+        ]
+
     def show_jobs(self, username: Optional[str], all_jobs: bool) -> None:
         self._code.append(f'job_lib.show_jobs({username!r}, {all_jobs})')
 
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index 93a73082214..7d5e5db9d71 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -227,7 +227,7 @@ def add_gang_scheduling_placement_group(
         # Set CPU to avoid ray hanging the resources allocation
         # for remote functions, since the task will request 1 CPU
         # by default.
-        bundles = [{'CPU': 1} for _ in range(num_nodes)]
+        bundles = [{'CPU': 0.001} for _ in range(num_nodes)]
 
         if accelerator_dict is not None:
             acc_name = list(accelerator_dict.keys())[0]
@@ -273,7 +273,8 @@ def add_gang_scheduling_placement_group(
                 def check_ip():
                     return ray.util.get_node_ip_address()
                 ip_list = ray.get([
-                    check_ip.options(placement_group=pg,
+                    check_ip.options(num_cpus=0,
+                                     placement_group=pg,
                                      placement_group_bundle_index=i).remote()
                     for i in range(pg.bundle_count)
                 ])
@@ -323,11 +324,11 @@ def add_ray_task(
         if task_name is None:
             # Make the task name more meaningful in ray log.
             name_str = 'name=\'task\''
+        cpu_str = ', num_cpus=0'
 
-        if ray_resources_dict is None:
-            resources_str = ''
-            num_gpus_str = ''
-        else:
+        resources_str = ''
+        num_gpus_str = ''
+        if ray_resources_dict is not None:
             assert len(ray_resources_dict) == 1, \
                 ('There can only be one type of accelerator per instance.'
                  f' Found: {ray_resources_dict}.')
@@ -345,8 +346,8 @@ def add_ray_task(
 
         resources_str = ', placement_group=pg'
         resources_str += f', placement_group_bundle_index={gang_scheduling_id}'
-        logger.debug(
-            f'Added Task with options: {name_str}{resources_str}{num_gpus_str}')
+        logger.debug('Added Task with options: '
+                     f'{name_str}{cpu_str}{resources_str}{num_gpus_str}')
         self._code += [
             textwrap.dedent(f"""\
         script = {bash_script!r}
@@ -357,7 +358,7 @@ def add_ray_task(
             node_export_sky_env_vars = (export_sky_env_vars +
                                         'export SKY_NODE_RANK={gang_scheduling_id}\\n')
             futures.append(run_bash_command_with_log \\
-                    .options({name_str}{resources_str}{num_gpus_str}) \\
+                    .options({name_str}{cpu_str}{resources_str}{num_gpus_str}) \\
                     .remote(
                         script,
                         {log_path!r},
@@ -1277,12 +1278,23 @@ def provision(self,
                 # TPU.
                 tpu_create_script=config_dict.get('tpu-create-script'),
                 tpu_delete_script=config_dict.get('tpu-delete-script'))
+
+            # Update job queue to avoid stale jobs (when restarted), before
+            # setting the cluster to be ready.
+            codegen = backend_utils.JobLibCodeGen()
+            codegen.update_status()
+            cmd = codegen.build()
+            returncode = self.run_on_head(handle, cmd)
+            backend_utils.handle_returncode(returncode, cmd,
+                                            'Failed to update job status.')
+
             global_user_state.add_or_update_cluster(cluster_name,
                                                     handle,
                                                     ready=True)
             auth_config = backend_utils.read_yaml(handle.cluster_yaml)['auth']
             backend_utils.SSHConfigHelper.add_cluster(cluster_name, ip_list,
                                                       auth_config)
+
             os.remove(lock_path)
             return handle
 
diff --git a/sky/skylet/job_lib.py b/sky/skylet/job_lib.py
index 66af82b1b1c..746de27dcdf 100644
--- a/sky/skylet/job_lib.py
+++ b/sky/skylet/job_lib.py
@@ -210,8 +210,9 @@ def query_job_status(job_ids: List[int]) -> List[JobStatus]:
     return job_status_list
 
 
-def _update_status() -> None:
-    running_jobs = _get_jobs(username=None, status_list=[JobStatus.RUNNING])
+def update_status() -> None:
+    running_jobs = _get_jobs(username=None,
+                             status_list=[JobStatus.PENDING, JobStatus.RUNNING])
     running_job_ids = [job['job_id'] for job in running_jobs]
 
     job_status = query_job_status(running_job_ids)
@@ -256,7 +257,6 @@ def show_jobs(username: Optional[str], all_jobs: bool) -> None:
         username: The username to show jobs for. Show all the users if None.
         all_jobs: Whether to show all jobs, not just the pending/running ones.
     """
-    _update_status()
     status_list = [JobStatus.PENDING, JobStatus.RUNNING]
     if all_jobs:
         status_list = None
@@ -273,7 +273,6 @@ def cancel_jobs(jobs: Optional[List[int]]) -> None:
     """
     # Update the status of the jobs to avoid setting the status of staled
     # jobs to CANCELLED.
-    _update_status()
     if jobs is None:
         job_records = _get_jobs(None, [JobStatus.PENDING, JobStatus.RUNNING])
     else:


Code review comments:
diff_hunk: @@ -1277,12 +1278,23 @@ def provision(self,
                 # TPU.
                 tpu_create_script=config_dict.get('tpu-create-script'),
                 tpu_delete_script=config_dict.get('tpu-delete-script'))
+
+            # Update job queue to avoid stale jobs (when restarted), before
+            # setting the cluster to be ready.
+            codegen = backend_utils.JobLibCodeGen()
+            codegen.update_status()
+            cmd = codegen.build()
+            returncode = self.run_on_head(handle, cmd)
+            backend_utils.handle_returncode(returncode, cmd,
+                                            'Failed to update job status.')
path: sky/backends/cloud_vm_ray_backend.py
text: @user1:
It is possible users want to do `sky launch` on an UP cluster. It seems like we should make that case fast too (i.e., not go through this)?

@author:
Good point! Fixed. Thanks!

diff_hunk: @@ -323,11 +324,11 @@ def add_ray_task(
         if task_name is None:
             # Make the task name more meaningful in ray log.
             name_str = 'name=\'task\''
+        cpu_str = ', num_cpus=0'
path: sky/backends/cloud_vm_ray_backend.py
text: @user1:
Hmm, is it a good idea to allow potentially a large amount of tasks?  Is this just for debugging?

@author:
No, I did this on purpose based on the description of Daniel. I feel like it may be possible for the user to have code that has low utilization of the CPU. Making each task taking 1 CPU core may waste some of the resources. 
> my code has both poor CPU utilization and poor GPU utilization

I set the placement_group for each task to take `0.001` CPUs, but yes, it may be too small. I tend to make it `0.1`, so that makes it possible for 10 tasks per CPU core. WDYT?



@user1:
Yeah, I think setting 0.1 here sounds reasonable.  In the future we can let the user set it.

Just to confirm, the value for the bundles and the value for L327 here should both be 0.1?

@author:
After further testing, the ray program will consume relatively large memory, taking more than 30GB of memory after launching 30 tasks. For a `m4.2xlarge` (our cpunode), it has 8 CPU threads and 32 GB memory.  That is to say, we can only launch less than 4 tasks per CPU thread. Maybe let me change that CPU requirement back to `0.5` first.