[jobs] catch jobs that are DONE but non-terminal (#4641)

* [jobs] catch jobs that are DONE but non-terminal

Usually, a job should always transition to a terminal status before it
is set to DONE. However, if a bug makes that fail for some
reason (e.g. the issues fixed in #4637), we should still catch this
state.

* address PR comments

* address more PR comments

* Update sky/jobs/utils.py

Co-authored-by: Zhanghao Wu <[email protected]>

* Update sky/jobs/utils.py

Co-authored-by: Zhanghao Wu <[email protected]>

---------

Co-authored-by: Zhanghao Wu <[email protected]>

Author: cg505

sky/jobs/state.py

@@ -758,8 +758,9 @@ def get_jobs_to_check_status(job_id: Optional[int] = None) -> List[int]:
         job_id: Optional job ID to check. If None, checks all jobs.
 
     Returns a list of job_ids, including the following:
-    - For jobs with schedule state: jobs that have schedule state not DONE
-    - For legacy jobs (no schedule state): jobs that are in non-terminal status
+    - Jobs that have a schedule_state that is not DONE
+    - Jobs have schedule_state DONE but are in a non-terminal status
+    - Legacy jobs (that is, no schedule state) that are in non-terminal status
     """
     job_filter = '' if job_id is None else 'AND spot.spot_job_id=(?)'
     job_value = () if job_id is None else (job_id,)
@@ -772,7 +773,9 @@ def get_jobs_to_check_status(job_id: Optional[int] = None) -> List[int]:
 
     # Get jobs that are either:
     # 1. Have schedule state that is not DONE, or
-    # 2. Have no schedule state (legacy) AND are in non-terminal status
+    # 2. Have schedule state DONE AND are in non-terminal status (unexpected
+    #    inconsistent state), or
+    # 3. Have no schedule state (legacy) AND are in non-terminal status
     with db_utils.safe_cursor(_DB_PATH) as cursor:
         rows = cursor.execute(
             f"""\
@@ -781,14 +784,23 @@ def get_jobs_to_check_status(job_id: Optional[int] = None) -> List[int]:
             LEFT OUTER JOIN job_info
             ON spot.spot_job_id=job_info.spot_job_id
             WHERE (
+                -- non-legacy jobs that are not DONE
                 (job_info.schedule_state IS NOT NULL AND
                  job_info.schedule_state IS NOT ?)
                 OR
-                (job_info.schedule_state IS NULL AND
+                -- legacy or that are in non-terminal status or
+                -- DONE jobs that are in non-terminal status
+                ((-- legacy jobs
+                  job_info.schedule_state IS NULL OR
+                  -- non-legacy DONE jobs
+                  job_info.schedule_state IS ?
+                 ) AND
+                 -- non-terminal
                  status NOT IN ({status_filter_str}))
             )
             {job_filter}
             ORDER BY spot.spot_job_id DESC""", [
+                ManagedJobScheduleState.DONE.value,
                 ManagedJobScheduleState.DONE.value, *terminal_status_values,
                 *job_value
             ]).fetchall()

sky/jobs/utils.py

@@ -245,16 +245,35 @@ def _handle_legacy_job(job_id: int):
         schedule_state = tasks[0]['schedule_state']
 
         # Backwards compatibility: this job was submitted when ray was still
-        # used for managing the parallelism of job controllers.
+        # used for managing the parallelism of job controllers, before #4485.
         # TODO(cooperc): Remove before 0.11.0.
         if (schedule_state is
                 managed_job_state.ManagedJobScheduleState.INVALID):
             _handle_legacy_job(job_id)
             continue
 
-        # For jobs with schedule state:
+        # Handle jobs with schedule state (non-legacy jobs):
         pid = tasks[0]['controller_pid']
-        if pid is None:
+        if schedule_state == managed_job_state.ManagedJobScheduleState.DONE:
+            # There are two cases where we could get a job that is DONE.
+            # 1. At query time (get_jobs_to_check_status), the job was not yet
+            #    DONE, but since then (before get_managed_jobs is called) it has
+            #    hit a terminal status, marked itself done, and exited. This is
+            #    fine.
+            # 2. The job is DONE, but in a non-terminal status. This is
+            #    unexpected. For instance, the task status is RUNNING, but the
+            #    job schedule_state is DONE.
+            if all(task['status'].is_terminal() for task in tasks):
+                # Turns out this job is fine, even though it got pulled by
+                # get_jobs_to_check_status. Probably case #1 above.
+                continue
+
+            logger.error(f'Job {job_id} has DONE schedule state, but some '
+                         f'tasks are not terminal. Task statuses: '
+                         f'{", ".join(task["status"].value for task in tasks)}')
+            failure_reason = ('Inconsistent internal job state. This is a bug.')
+        elif pid is None:
+            # Non-legacy job and controller process has not yet started.
             if schedule_state in (
                     managed_job_state.ManagedJobScheduleState.INACTIVE,
                     managed_job_state.ManagedJobScheduleState.WAITING):

sky/skylet/constants.py

@@ -86,7 +86,7 @@
 # cluster yaml is updated.
 #
 # TODO(zongheng,zhanghao): make the upgrading of skylet automatic?
-SKYLET_VERSION = '11'
+SKYLET_VERSION = '12'
 # The version of the lib files that skylet/jobs use. Whenever there is an API
 # change for the job_lib or log_lib, we need to bump this version, so that the
 # user can be notified to update their SkyPilot version on the remote cluster.