Repository: skypilot-org/skypilot
Pull request: Fix race condition for skylet daemon
Closes #745. Changed the behavior of the status update in skylet daemon: only update the status of the jobs that are submitted more than 2 seconds ago (need to be tuned based on heuristics), to make sure `ray job` has been correctly updated.

The problem is hard to reproduce, so I chose a reasonable gap (2 seconds) for the job status update.

Problem statement:
Possible race condition of job status update in sky daemon 
This is observed by @concretevitamin. When a job is just submitted, the `sky logs` called by `sky launch` may return status FAILED immediately. This can be caused by a race condition, where status update in skylet daemon set the job to FAILED, as the `ray job status` has not been updated yet and returned empty.

Possible solution: for the status update in skylet daemon, we only update the status of the jobs that are submitted longer than 30 seconds ago.

Hints from issue discussion:

Commit message: update comment

Patch to review:
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index 751cfc489ef..92ad527ebae 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -1885,7 +1885,7 @@ def _exec_code_on_head(
 
         try:
             if not detach_run:
-                # Wait for the job being sucessfully submitted to ray job.
+                # Wait for the job being successfully submitted to ray job.
                 time.sleep(1)
                 # Sky logs. Not using subprocess.run since it will make the
                 # ssh keep connected after ctrl-c.
diff --git a/sky/skylet/events.py b/sky/skylet/events.py
index 6bfd3a9bf97..de6c68cfd48 100644
--- a/sky/skylet/events.py
+++ b/sky/skylet/events.py
@@ -52,8 +52,14 @@ class JobUpdateEvent(SkyletEvent):
     """Skylet event for updating job status."""
     EVENT_INTERVAL_SECONDS = 20
 
+    # Only update status of the jobs after this many seconds of job submission,
+    # to avoid race condition with `ray job` to make sure it job has been
+    # correctly updated.
+    # TODO(zhwu): This number should be tuned based on heuristics.
+    _SUBMITTED_GAP_SECONDS = 2
+
     def _run(self):
-        job_lib.update_status()
+        job_lib.update_status(submitted_gap_sec=self._SUBMITTED_GAP_SECONDS)
 
 
 class AutostopEvent(SkyletEvent):
diff --git a/sky/skylet/job_lib.py b/sky/skylet/job_lib.py
index f07d02aa23a..444f0f86ad6 100644
--- a/sky/skylet/job_lib.py
+++ b/sky/skylet/job_lib.py
@@ -149,9 +149,9 @@ def _get_records_from_rows(rows) -> List[Dict[str, Any]]:
     return records
 
 
-def _get_jobs(
-        username: Optional[str],
-        status_list: Optional[List[JobStatus]] = None) -> List[Dict[str, Any]]:
+def _get_jobs(username: Optional[str],
+              status_list: Optional[List[JobStatus]] = None,
+              submitted_gap_sec: int = 0) -> List[Dict[str, Any]]:
     if status_list is None:
         status_list = list(JobStatus)
     status_str_list = [status.value for status in status_list]
@@ -160,17 +160,18 @@ def _get_jobs(
             f"""\
             SELECT * FROM jobs
             WHERE status IN ({','.join(['?'] * len(status_list))})
+            AND submitted_at >= (?)
             ORDER BY job_id DESC""",
-            (*status_str_list,),
+            (*status_str_list, submitted_gap_sec + time.time()),
         )
     else:
         rows = _CURSOR.execute(
             f"""\
             SELECT * FROM jobs
             WHERE status IN ({','.join(['?'] * len(status_list))})
-            AND username=(?)
+            AND username=(?) AND submitted_at >= (?)
             ORDER BY job_id DESC""",
-            (*status_str_list, username),
+            (*status_str_list, username, submitted_gap_sec + time.time()),
         )
 
     records = _get_records_from_rows(rows)
@@ -252,7 +253,7 @@ def fail_all_jobs_in_progress() -> None:
     _CONN.commit()
 
 
-def update_status() -> None:
+def update_status(submitted_gap_sec: int = 0) -> None:
     # This will be called periodically by the skylet to update the status
     # of the jobs in the database, to avoid stale job status.
     # NOTE: there might be a INIT job in the database set to FAILED by this
@@ -261,7 +262,8 @@ def update_status() -> None:
     # app starts.
     running_jobs = _get_jobs(
         username=None,
-        status_list=[JobStatus.INIT, JobStatus.PENDING, JobStatus.RUNNING])
+        status_list=[JobStatus.INIT, JobStatus.PENDING, JobStatus.RUNNING],
+        submitted_gap_sec=submitted_gap_sec)
     running_job_ids = [job['job_id'] for job in running_jobs]
 
     job_status = query_job_status(running_job_ids)


Code review comments:
diff_hunk: @@ -160,17 +160,18 @@ def _get_jobs(
             f"""\
             SELECT * FROM jobs
             WHERE status IN ({','.join(['?'] * len(status_list))})
+            AND submitted_at >= (?)
path: sky/skylet/job_lib.py
text: @user1:
Should this be <= time.time() - submitted_gap_sec?

@author:
Ahh, good catch! Fixed. Thanks!