skypilot-nightly 1.0.0.dev20250114__py3-none-any.whl → 1.0.0.dev20250124__py3-none-any.whl

sky/jobs/scheduler.py

@@ -0,0 +1,283 @@
+"""Scheduler for managed jobs.
+
+Once managed jobs are submitted via submit_job, the scheduler is responsible for
+the business logic of deciding when they are allowed to start, and choosing the
+right one to start. The scheduler will also schedule jobs that are already live
+but waiting to launch a new task or recover.
+
+The scheduler is not its own process - instead, maybe_schedule_next_jobs() can
+be called from any code running on the managed jobs controller instance to
+trigger scheduling of new jobs if possible. This function should be called
+immediately after any state change that could result in jobs newly being able to
+be scheduled.
+
+The scheduling logic limits the number of running jobs according to two limits:
+1. The number of jobs that can be launching (that is, STARTING or RECOVERING) at
+   once, based on the number of CPUs. (See _get_launch_parallelism.) This the
+   most compute-intensive part of the job lifecycle, which is why we have an
+   additional limit.
+2. The number of jobs that can be running at any given time, based on the amount
+   of memory. (See _get_job_parallelism.) Since the job controller is doing very
+   little once a job starts (just checking its status periodically), the most
+   significant resource it consumes is memory.
+
+The state of the scheduler is entirely determined by the schedule_state column
+of all the jobs in the job_info table. This column should only be modified via
+the functions defined in this file. We will always hold the lock while modifying
+this state. See state.ManagedJobScheduleState.
+
+Nomenclature:
+- job: same as managed job (may include multiple tasks)
+- launch/launching: launching a cluster (sky.launch) as part of a job
+- start/run: create the job controller process for a job
+- schedule: transition a job to the LAUNCHING state, whether a new job or a job
+  that is already alive
+- alive: a job controller exists (includes multiple schedule_states: ALIVE,
+  ALIVE_WAITING, LAUNCHING)
+"""
+
+from argparse import ArgumentParser
+import contextlib
+from functools import lru_cache
+import os
+import time
+
+import filelock
+import psutil
+
+from sky import sky_logging
+from sky.jobs import constants as managed_job_constants
+from sky.jobs import state
+from sky.skylet import constants
+from sky.utils import subprocess_utils
+
+logger = sky_logging.init_logger('sky.jobs.controller')
+
+# The _MANAGED_JOB_SCHEDULER_LOCK should be held whenever we are checking the
+# parallelism control or updating the schedule_state of any job.
+# Any code that takes this lock must conclude by calling
+# maybe_schedule_next_jobs.
+_MANAGED_JOB_SCHEDULER_LOCK = '~/.sky/locks/managed_job_scheduler.lock'
+_ALIVE_JOB_LAUNCH_WAIT_INTERVAL = 0.5
+
+
+@lru_cache(maxsize=1)
+def _get_lock_path() -> str:
+    path = os.path.expanduser(_MANAGED_JOB_SCHEDULER_LOCK)
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+    return path
+
+
+def maybe_schedule_next_jobs() -> None:
+    """Determine if any managed jobs can be scheduled, and if so, schedule them.
+
+    Here, "schedule" means to select job that is waiting, and allow it to
+    proceed. It does NOT mean to submit a job to the scheduler.
+
+    For newly submitted jobs, scheduling means updating the state of the jobs,
+    and starting the job controller process. For jobs that are already alive but
+    are waiting to launch a new task or recover, just update the state of the
+    job to indicate that the launch can proceed.
+
+    This function transitions jobs into LAUNCHING on a best-effort basis. That
+    is, if we can start any jobs, we will, but if not, we will exit (almost)
+    immediately. It's expected that if some WAITING or ALIVE_WAITING jobs cannot
+    be started now (either because the lock is held, or because there are not
+    enough resources), another call to this function will be made whenever that
+    situation is resolved. (If the lock is held, the lock holder should start
+    the jobs. If there aren't enough resources, the next controller to exit and
+    free up resources should start the jobs.)
+
+    If this function obtains the lock, it will launch as many jobs as possible
+    before releasing the lock. This is what allows other calls to exit
+    immediately if the lock is held, while ensuring that all jobs are started as
+    soon as possible.
+
+    This uses subprocess_utils.launch_new_process_tree() to start the controller
+    processes, which should be safe to call from pretty much any code running on
+    the jobs controller instance. New job controller processes will be detached
+    from the current process and there will not be a parent/child relationship.
+    See launch_new_process_tree for more.
+    """
+    try:
+        # We must use a global lock rather than a per-job lock to ensure correct
+        # parallelism control. If we cannot obtain the lock, exit immediately.
+        # The current lock holder is expected to launch any jobs it can before
+        # releasing the lock.
+        with filelock.FileLock(_get_lock_path(), blocking=False):
+            while True:
+                maybe_next_job = state.get_waiting_job()
+                if maybe_next_job is None:
+                    # Nothing left to start, break from scheduling loop
+                    break
+
+                current_state = maybe_next_job['schedule_state']
+
+                assert current_state in (
+                    state.ManagedJobScheduleState.ALIVE_WAITING,
+                    state.ManagedJobScheduleState.WAITING), maybe_next_job
+
+                # Note: we expect to get ALIVE_WAITING jobs before WAITING jobs,
+                # since they will have been submitted and therefore started
+                # first. The requirements to launch in an alive job are more
+                # lenient, so there is no way that we wouldn't be able to launch
+                # an ALIVE_WAITING job, but we would be able to launch a WAITING
+                # job.
+                if current_state == state.ManagedJobScheduleState.ALIVE_WAITING:
+                    if not _can_lauch_in_alive_job():
+                        # Can't schedule anything, break from scheduling loop.
+                        break
+                elif current_state == state.ManagedJobScheduleState.WAITING:
+                    if not _can_start_new_job():
+                        # Can't schedule anything, break from scheduling loop.
+                        break
+
+                logger.debug(f'Scheduling job {maybe_next_job["job_id"]}')
+                state.scheduler_set_launching(maybe_next_job['job_id'],
+                                              current_state)
+
+                if current_state == state.ManagedJobScheduleState.WAITING:
+                    # The job controller has not been started yet. We must start
+                    # it.
+
+                    job_id = maybe_next_job['job_id']
+                    dag_yaml_path = maybe_next_job['dag_yaml_path']
+
+                    # If the command line here is changed, please also update
+                    # utils._controller_process_alive. `--job-id X` should be at
+                    # the end.
+                    run_cmd = (f'{constants.ACTIVATE_SKY_REMOTE_PYTHON_ENV};'
+                               'python -u -m sky.jobs.controller '
+                               f'{dag_yaml_path} --job-id {job_id}')
+
+                    logs_dir = os.path.expanduser(
+                        managed_job_constants.JOBS_CONTROLLER_LOGS_DIR)
+                    os.makedirs(logs_dir, exist_ok=True)
+                    log_path = os.path.join(logs_dir, f'{job_id}.log')
+
+                    pid = subprocess_utils.launch_new_process_tree(
+                        run_cmd, log_output=log_path)
+                    state.set_job_controller_pid(job_id, pid)
+
+                    logger.debug(f'Job {job_id} started with pid {pid}')
+
+    except filelock.Timeout:
+        # If we can't get the lock, just exit. The process holding the lock
+        # should launch any pending jobs.
+        pass
+
+
+def submit_job(job_id: int, dag_yaml_path: str) -> None:
+    """Submit an existing job to the scheduler.
+
+    This should be called after a job is created in the `spot` table as
+    PENDING. It will tell the scheduler to try and start the job controller, if
+    there are resources available. It may block to acquire the lock, so it
+    should not be on the critical path for `sky jobs launch -d`.
+    """
+    with filelock.FileLock(_get_lock_path()):
+        state.scheduler_set_waiting(job_id, dag_yaml_path)
+    maybe_schedule_next_jobs()
+
+
+@contextlib.contextmanager
+def scheduled_launch(job_id: int):
+    """Launch as part of an ongoing job.
+
+    A newly started job will already be LAUNCHING, and this will immediately
+    enter the context.
+
+    If a job is ongoing (ALIVE schedule_state), there are two scenarios where we
+    may need to call sky.launch again during the course of a job controller:
+    - for tasks after the first task
+    - for recovery
+
+    This function will mark the job as ALIVE_WAITING, which indicates to the
+    scheduler that it wants to transition back to LAUNCHING. Then, it will wait
+    until the scheduler transitions the job state, before entering the context.
+
+    On exiting the context, the job will transition to ALIVE.
+
+    This should only be used within the job controller for the given job_id. If
+    multiple uses of this context are nested, behavior is undefined. Don't do
+    that.
+    """
+
+    # If we're already in LAUNCHING schedule_state, we don't need to wait.
+    # This may be the case for the first launch of a job.
+    if (state.get_job_schedule_state(job_id) !=
+            state.ManagedJobScheduleState.LAUNCHING):
+        # Since we aren't LAUNCHING, we need to wait to be scheduled.
+        _set_alive_waiting(job_id)
+
+        while (state.get_job_schedule_state(job_id) !=
+               state.ManagedJobScheduleState.LAUNCHING):
+            time.sleep(_ALIVE_JOB_LAUNCH_WAIT_INTERVAL)
+
+    yield
+
+    with filelock.FileLock(_get_lock_path()):
+        state.scheduler_set_alive(job_id)
+    maybe_schedule_next_jobs()
+
+
+def job_done(job_id: int, idempotent: bool = False) -> None:
+    """Transition a job to DONE.
+
+    If idempotent is True, this will not raise an error if the job is already
+    DONE.
+
+    The job could be in any terminal ManagedJobStatus. However, once DONE, it
+    should never transition back to another state.
+    """
+    if idempotent and (state.get_job_schedule_state(job_id)
+                       == state.ManagedJobScheduleState.DONE):
+        return
+
+    with filelock.FileLock(_get_lock_path()):
+        state.scheduler_set_done(job_id, idempotent)
+    maybe_schedule_next_jobs()
+
+
+def _set_alive_waiting(job_id: int) -> None:
+    """Should use wait_until_launch_okay() to transition to this state."""
+    with filelock.FileLock(_get_lock_path()):
+        state.scheduler_set_alive_waiting(job_id)
+    maybe_schedule_next_jobs()
+
+
+def _get_job_parallelism() -> int:
+    # Assume a running job uses 350MB memory.
+    # We observe 230-300 in practice.
+    job_memory = 350 * 1024 * 1024
+    return max(psutil.virtual_memory().total // job_memory, 1)
+
+
+def _get_launch_parallelism() -> int:
+    cpus = os.cpu_count()
+    return cpus * 4 if cpus is not None else 1
+
+
+def _can_start_new_job() -> bool:
+    launching_jobs = state.get_num_launching_jobs()
+    alive_jobs = state.get_num_alive_jobs()
+    return launching_jobs < _get_launch_parallelism(
+    ) and alive_jobs < _get_job_parallelism()
+
+
+def _can_lauch_in_alive_job() -> bool:
+    launching_jobs = state.get_num_launching_jobs()
+    return launching_jobs < _get_launch_parallelism()
+
+
+if __name__ == '__main__':
+    parser = ArgumentParser()
+    parser.add_argument('--job-id',
+                        required=True,
+                        type=int,
+                        help='Job id for the controller job.')
+    parser.add_argument('dag_yaml',
+                        type=str,
+                        help='The path to the user job yaml file.')
+    args = parser.parse_args()
+    submit_job(args.job_id, args.dag_yaml)

sky/jobs/state.py

@@ -107,12 +107,25 @@ def create_table(cursor, conn):
     db_utils.add_column_to_table(cursor, conn, 'spot', 'local_log_file',
                                  'TEXT DEFAULT NULL')
 
-    # `job_info` contains the mapping from job_id to the job_name.
-    # In the future, it may contain more information about each job.
+    # `job_info` contains the mapping from job_id to the job_name, as well as
+    # information used by the scheduler.
     cursor.execute("""\
         CREATE TABLE IF NOT EXISTS job_info (
         spot_job_id INTEGER PRIMARY KEY AUTOINCREMENT,
-        name TEXT)""")
+        name TEXT,
+        schedule_state TEXT,
+        controller_pid INTEGER DEFAULT NULL,
+        dag_yaml_path TEXT)""")
+
+    db_utils.add_column_to_table(cursor, conn, 'job_info', 'schedule_state',
+                                 'TEXT')
+
+    db_utils.add_column_to_table(cursor, conn, 'job_info', 'controller_pid',
+                                 'INTEGER DEFAULT NULL')
+
+    db_utils.add_column_to_table(cursor, conn, 'job_info', 'dag_yaml_path',
+                                 'TEXT')
+
     conn.commit()
 
 
@@ -164,6 +177,9 @@ columns = [
     # columns from the job_info table
     '_job_info_job_id',  # This should be the same as job_id
     'job_name',
+    'schedule_state',
+    'controller_pid',
+    'dag_yaml_path',
 ]
 
 
@@ -189,16 +205,18 @@ class ManagedJobStatus(enum.Enum):
         SUCCEEDED       ->  SUCCEEDED
         FAILED          ->  FAILED
         FAILED_SETUP    ->  FAILED_SETUP
+    Not all statuses are in this list, since some ManagedJobStatuses are only
+    possible while the cluster is INIT/STOPPED/not yet UP.
     Note that the JobStatus will not be stuck in PENDING, because each cluster
     is dedicated to a managed job, i.e. there should always be enough resource
     to run the job and the job will be immediately transitioned to RUNNING.
+
     """
     # PENDING: Waiting for the jobs controller to have a slot to run the
     # controller process.
-    # The submitted_at timestamp of the managed job in the 'spot' table will be
-    # set to the time when the job is firstly submitted by the user (set to
-    # PENDING).
     PENDING = 'PENDING'
+    # The submitted_at timestamp of the managed job in the 'spot' table will be
+    # set to the time when the job controller begins running.
     # SUBMITTED: The jobs controller starts the controller process.
     SUBMITTED = 'SUBMITTED'
     # STARTING: The controller process is launching the cluster for the managed
@@ -292,14 +310,72 @@ _SPOT_STATUS_TO_COLOR = {
 }
 
 
+class ManagedJobScheduleState(enum.Enum):
+    """Captures the state of the job from the scheduler's perspective.
+
+    A job that predates the introduction of the scheduler will be INVALID.
+
+    A newly created job will be INACTIVE.  The following transitions are valid:
+    - INACTIVE -> WAITING: The job is "submitted" to the scheduler, and its job
+      controller can be started.
+    - WAITING -> LAUNCHING: The job controller is starting by the scheduler and
+      may proceed to sky.launch.
+    - LAUNCHING -> ALIVE: The launch attempt was completed. It may have
+      succeeded or failed. The job controller is not allowed to sky.launch again
+      without transitioning to ALIVE_WAITING and then LAUNCHING.
+    - ALIVE -> ALIVE_WAITING: The job controller wants to sky.launch again,
+      either for recovery or to launch a subsequent task.
+    - ALIVE_WAITING -> LAUNCHING: The scheduler has determined that the job
+      controller may launch again.
+    - LAUNCHING, ALIVE, or ALIVE_WAITING -> DONE: The job controller is exiting
+      and the job is in some terminal status. In the future it may be possible
+      to transition directly from WAITING or even INACTIVE to DONE if the job is
+      cancelled.
+
+    There is no well-defined mapping from the managed job status to schedule
+    state or vice versa. (In fact, schedule state is defined on the job and
+    status on the task.)
+    - INACTIVE or WAITING should only be seen when a job is PENDING.
+    - ALIVE_WAITING should only be seen when a job is RECOVERING, has multiple
+      tasks, or needs to retry launching.
+    - LAUNCHING and ALIVE can be seen in many different statuses.
+    - DONE should only be seen when a job is in a terminal status.
+    Since state and status transitions are not atomic, it may be possible to
+    briefly observe inconsistent states, like a job that just finished but
+    hasn't yet transitioned to DONE.
+    """
+    # This job may have been created before scheduler was introduced in #4458.
+    # This state is not used by scheduler but just for backward compatibility.
+    # TODO(cooperc): remove this in v0.11.0
+    INVALID = None
+    # The job should be ignored by the scheduler.
+    INACTIVE = 'INACTIVE'
+    # The job is waiting to transition to LAUNCHING for the first time. The
+    # scheduler should try to transition it, and when it does, it should start
+    # the job controller.
+    WAITING = 'WAITING'
+    # The job is already alive, but wants to transition back to LAUNCHING,
+    # e.g. for recovery, or launching later tasks in the DAG. The scheduler
+    # should try to transition it to LAUNCHING.
+    ALIVE_WAITING = 'ALIVE_WAITING'
+    # The job is running sky.launch, or soon will, using a limited number of
+    # allowed launch slots.
+    LAUNCHING = 'LAUNCHING'
+    # The controller for the job is running, but it's not currently launching.
+    ALIVE = 'ALIVE'
+    # The job is in a terminal state. (Not necessarily SUCCEEDED.)
+    DONE = 'DONE'
+
+
 # === Status transition functions ===
-def set_job_name(job_id: int, name: str):
+def set_job_info(job_id: int, name: str):
     with db_utils.safe_cursor(_DB_PATH) as cursor:
         cursor.execute(
             """\
             INSERT INTO job_info
-            (spot_job_id, name)
-            VALUES (?, ?)""", (job_id, name))
+            (spot_job_id, name, schedule_state)
+            VALUES (?, ?, ?)""",
+            (job_id, name, ManagedJobScheduleState.INACTIVE.value))
 
 
 def set_pending(job_id: int, task_id: int, task_name: str, resources_str: str):
@@ -324,7 +400,7 @@ def set_submitted(job_id: int, task_id: int, run_timestamp: str,
         job_id: The managed job ID.
         task_id: The task ID.
         run_timestamp: The run_timestamp of the run. This will be used to
-            determine the log directory of the managed task.
+        determine the log directory of the managed task.
         submit_time: The time when the managed task is submitted.
         resources_str: The resources string of the managed task.
         specs: The specs of the managed task.
@@ -458,13 +534,12 @@ def set_failed(
     with db_utils.safe_cursor(_DB_PATH) as cursor:
         previous_status = cursor.execute(
             'SELECT status FROM spot WHERE spot_job_id=(?)',
-            (job_id,)).fetchone()
-        previous_status = ManagedJobStatus(previous_status[0])
-        if previous_status in [ManagedJobStatus.RECOVERING]:
-            # If the job is recovering, we should set the
-            # last_recovered_at to the end_time, so that the
-            # end_at - last_recovered_at will not be affect the job duration
-            # calculation.
+            (job_id,)).fetchone()[0]
+        previous_status = ManagedJobStatus(previous_status)
+        if previous_status == ManagedJobStatus.RECOVERING:
+            # If the job is recovering, we should set the last_recovered_at to
+            # the end_time, so that the end_at - last_recovered_at will not be
+            # affect the job duration calculation.
             fields_to_set['last_recovered_at'] = end_time
         set_str = ', '.join(f'{k}=(?)' for k in fields_to_set)
         task_str = '' if task_id is None else f' AND task_id={task_id}'
@@ -564,6 +639,44 @@ def get_nonterminal_job_ids_by_name(name: Optional[str]) -> List[int]:
         return job_ids
 
 
+def get_schedule_live_jobs(job_id: Optional[int]) -> List[Dict[str, Any]]:
+    """Get jobs from the database that have a live schedule_state.
+
+    This should return job(s) that are not INACTIVE, WAITING, or DONE.  So a
+    returned job should correspond to a live job controller process, with one
+    exception: the job may have just transitioned from WAITING to LAUNCHING, but
+    the controller process has not yet started.
+    """
+    job_filter = '' if job_id is None else 'AND spot_job_id=(?)'
+    job_value = (job_id,) if job_id is not None else ()
+
+    # Join spot and job_info tables to get the job name for each task.
+    # We use LEFT OUTER JOIN mainly for backward compatibility, as for an
+    # existing controller before #1982, the job_info table may not exist,
+    # and all the managed jobs created before will not present in the
+    # job_info.
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        rows = cursor.execute(
+            f"""\
+            SELECT spot_job_id, schedule_state, controller_pid
+            FROM job_info
+            WHERE schedule_state not in (?, ?, ?)
+            {job_filter}
+            ORDER BY spot_job_id DESC""",
+            (ManagedJobScheduleState.INACTIVE.value,
+             ManagedJobScheduleState.WAITING.value,
+             ManagedJobScheduleState.DONE.value, *job_value)).fetchall()
+        jobs = []
+        for row in rows:
+            job_dict = {
+                'job_id': row[0],
+                'schedule_state': ManagedJobScheduleState(row[1]),
+                'controller_pid': row[2],
+            }
+            jobs.append(job_dict)
+        return jobs
+
+
 def get_all_job_ids_by_name(name: Optional[str]) -> List[int]:
     """Get all job ids by name."""
     name_filter = ''
@@ -620,10 +733,12 @@ def get_latest_task_id_status(
     id_statuses = _get_all_task_ids_statuses(job_id)
     if not id_statuses:
         return None, None
-    task_id, status = id_statuses[-1]
-    for task_id, status in id_statuses:
-        if not status.is_terminal():
-            break
+    task_id, status = next(
+        ((tid, st) for tid, st in id_statuses if not st.is_terminal()),
+        id_statuses[-1],
+    )
+    # Unpack the tuple first, or it triggers a Pylint's bug on recognizing
+    # the return type.
     return task_id, status
 
 
@@ -670,6 +785,8 @@ def get_managed_jobs(job_id: Optional[int] = None) -> List[Dict[str, Any]]:
         for row in rows:
             job_dict = dict(zip(columns, row))
             job_dict['status'] = ManagedJobStatus(job_dict['status'])
+            job_dict['schedule_state'] = ManagedJobScheduleState(
+                job_dict['schedule_state'])
             if job_dict['job_name'] is None:
                 job_dict['job_name'] = job_dict['task_name']
             jobs.append(job_dict)
@@ -721,3 +838,128 @@ def get_local_log_file(job_id: int, task_id: Optional[int]) -> Optional[str]:
             f'SELECT local_log_file FROM spot '
             f'WHERE {filter_str}', filter_args).fetchone()
         return local_log_file[-1] if local_log_file else None
+
+
+# === Scheduler state functions ===
+# Only the scheduler should call these functions. They may require holding the
+# scheduler lock to work correctly.
+
+
+def scheduler_set_waiting(job_id: int, dag_yaml_path: str) -> None:
+    """Do not call without holding the scheduler lock."""
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        updated_count = cursor.execute(
+            'UPDATE job_info SET '
+            'schedule_state = (?), dag_yaml_path = (?) '
+            'WHERE spot_job_id = (?) AND schedule_state = (?)',
+            (ManagedJobScheduleState.WAITING.value, dag_yaml_path, job_id,
+             ManagedJobScheduleState.INACTIVE.value)).rowcount
+        assert updated_count == 1, (job_id, updated_count)
+
+
+def scheduler_set_launching(job_id: int,
+                            current_state: ManagedJobScheduleState) -> None:
+    """Do not call without holding the scheduler lock."""
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        updated_count = cursor.execute(
+            'UPDATE job_info SET '
+            'schedule_state = (?) '
+            'WHERE spot_job_id = (?) AND schedule_state = (?)',
+            (ManagedJobScheduleState.LAUNCHING.value, job_id,
+             current_state.value)).rowcount
+        assert updated_count == 1, (job_id, updated_count)
+
+
+def scheduler_set_alive(job_id: int) -> None:
+    """Do not call without holding the scheduler lock."""
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        updated_count = cursor.execute(
+            'UPDATE job_info SET '
+            'schedule_state = (?) '
+            'WHERE spot_job_id = (?) AND schedule_state = (?)',
+            (ManagedJobScheduleState.ALIVE.value, job_id,
+             ManagedJobScheduleState.LAUNCHING.value)).rowcount
+        assert updated_count == 1, (job_id, updated_count)
+
+
+def scheduler_set_alive_waiting(job_id: int) -> None:
+    """Do not call without holding the scheduler lock."""
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        updated_count = cursor.execute(
+            'UPDATE job_info SET '
+            'schedule_state = (?) '
+            'WHERE spot_job_id = (?) AND schedule_state = (?)',
+            (ManagedJobScheduleState.ALIVE_WAITING.value, job_id,
+             ManagedJobScheduleState.ALIVE.value)).rowcount
+        assert updated_count == 1, (job_id, updated_count)
+
+
+def scheduler_set_done(job_id: int, idempotent: bool = False) -> None:
+    """Do not call without holding the scheduler lock."""
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        updated_count = cursor.execute(
+            'UPDATE job_info SET '
+            'schedule_state = (?) '
+            'WHERE spot_job_id = (?) AND schedule_state != (?)',
+            (ManagedJobScheduleState.DONE.value, job_id,
+             ManagedJobScheduleState.DONE.value)).rowcount
+        if not idempotent:
+            assert updated_count == 1, (job_id, updated_count)
+
+
+def set_job_controller_pid(job_id: int, pid: int):
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        updated_count = cursor.execute(
+            'UPDATE job_info SET '
+            'controller_pid = (?) '
+            'WHERE spot_job_id = (?)', (pid, job_id)).rowcount
+        assert updated_count == 1, (job_id, updated_count)
+
+
+def get_job_schedule_state(job_id: int) -> ManagedJobScheduleState:
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        state = cursor.execute(
+            'SELECT schedule_state FROM job_info WHERE spot_job_id = (?)',
+            (job_id,)).fetchone()[0]
+        return ManagedJobScheduleState(state)
+
+
+def get_num_launching_jobs() -> int:
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        return cursor.execute(
+            'SELECT COUNT(*) '
+            'FROM job_info '
+            'WHERE schedule_state = (?)',
+            (ManagedJobScheduleState.LAUNCHING.value,)).fetchone()[0]
+
+
+def get_num_alive_jobs() -> int:
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        return cursor.execute(
+            'SELECT COUNT(*) '
+            'FROM job_info '
+            'WHERE schedule_state IN (?, ?, ?)',
+            (ManagedJobScheduleState.ALIVE_WAITING.value,
+             ManagedJobScheduleState.LAUNCHING.value,
+             ManagedJobScheduleState.ALIVE.value)).fetchone()[0]
+
+
+def get_waiting_job() -> Optional[Dict[str, Any]]:
+    """Get the next job that should transition to LAUNCHING.
+
+    Backwards compatibility note: jobs submitted before #4485 will have no
+    schedule_state and will be ignored by this SQL query.
+    """
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        row = cursor.execute(
+            'SELECT spot_job_id, schedule_state, dag_yaml_path '
+            'FROM job_info '
+            'WHERE schedule_state in (?, ?) '
+            'ORDER BY spot_job_id LIMIT 1',
+            (ManagedJobScheduleState.WAITING.value,
+             ManagedJobScheduleState.ALIVE_WAITING.value)).fetchone()
+        return {
+            'job_id': row[0],
+            'schedule_state': ManagedJobScheduleState(row[1]),
+            'dag_yaml_path': row[2],
+        } if row is not None else None