skypilot-nightly 1.0.0.dev20250116__py3-none-any.whl → 1.0.0.dev20250118__py3-none-any.whl

sky/jobs/recovery_strategy.py

@@ -17,6 +17,7 @@ from sky import global_user_state
 from sky import sky_logging
 from sky import status_lib
 from sky.backends import backend_utils
+from sky.jobs import scheduler
 from sky.jobs import utils as managed_job_utils
 from sky.skylet import job_lib
 from sky.usage import usage_lib
@@ -42,45 +43,20 @@ MAX_JOB_CHECKING_RETRY = 10
 _AUTODOWN_MINUTES = 5
 
 
-def terminate_cluster(cluster_name: str, max_retry: int = 3) -> None:
-    """Terminate the cluster."""
-    retry_cnt = 0
-    while True:
-        try:
-            usage_lib.messages.usage.set_internal()
-            sky.down(cluster_name)
-            return
-        except exceptions.ClusterDoesNotExist:
-            # The cluster is already down.
-            logger.debug(f'The cluster {cluster_name} is already down.')
-            return
-        except Exception as e:  # pylint: disable=broad-except
-            retry_cnt += 1
-            if retry_cnt >= max_retry:
-                raise RuntimeError(
-                    f'Failed to terminate the cluster {cluster_name}.') from e
-            logger.error(
-                f'Failed to terminate the cluster {cluster_name}. Retrying.'
-                f'Details: {common_utils.format_exception(e)}')
-            with ux_utils.enable_traceback():
-                logger.error(f'  Traceback: {traceback.format_exc()}')
-
-
 class StrategyExecutor:
     """Handle the launching, recovery and termination of managed job clusters"""
 
     RETRY_INIT_GAP_SECONDS = 60
 
     def __init__(self, cluster_name: str, backend: 'backends.Backend',
-                 task: 'task_lib.Task', retry_until_up: bool,
-                 max_restarts_on_errors: int) -> None:
+                 task: 'task_lib.Task', max_restarts_on_errors: int,
+                 job_id: int) -> None:
         """Initialize the strategy executor.
 
         Args:
             cluster_name: The name of the cluster.
             backend: The backend to use. Only CloudVMRayBackend is supported.
             task: The task to execute.
-            retry_until_up: Whether to retry until the cluster is up.
         """
         assert isinstance(backend, backends.CloudVmRayBackend), (
             'Only CloudVMRayBackend is supported.')
@@ -88,8 +64,8 @@ class StrategyExecutor:
         self.dag.add(task)
         self.cluster_name = cluster_name
         self.backend = backend
-        self.retry_until_up = retry_until_up
         self.max_restarts_on_errors = max_restarts_on_errors
+        self.job_id = job_id
         self.restart_cnt_on_failure = 0
 
     def __init_subclass__(cls, name: str, default: bool = False):
@@ -102,7 +78,7 @@ class StrategyExecutor:
 
     @classmethod
     def make(cls, cluster_name: str, backend: 'backends.Backend',
-             task: 'task_lib.Task', retry_until_up: bool) -> 'StrategyExecutor':
+             task: 'task_lib.Task', job_id: int) -> 'StrategyExecutor':
         """Create a strategy from a task."""
 
         resource_list = list(task.resources)
@@ -127,8 +103,9 @@ class StrategyExecutor:
             job_recovery_name = job_recovery
             max_restarts_on_errors = 0
         return RECOVERY_STRATEGIES[job_recovery_name](cluster_name, backend,
-                                                      task, retry_until_up,
-                                                      max_restarts_on_errors)
+                                                      task,
+                                                      max_restarts_on_errors,
+                                                      job_id)
 
     def launch(self) -> float:
         """Launch the cluster for the first time.
@@ -142,10 +119,7 @@ class StrategyExecutor:
         Raises: Please refer to the docstring of self._launch().
         """
 
-        if self.retry_until_up:
-            job_submit_at = self._launch(max_retry=None)
-        else:
-            job_submit_at = self._launch()
+        job_submit_at = self._launch(max_retry=None)
         assert job_submit_at is not None
         return job_submit_at
 
@@ -195,7 +169,7 @@ class StrategyExecutor:
                         f'{common_utils.format_exception(e)}\n'
                         'Terminating the cluster explicitly to ensure no '
                         'remaining job process interferes with recovery.')
-            terminate_cluster(self.cluster_name)
+            managed_job_utils.terminate_cluster(self.cluster_name)
 
     def _wait_until_job_starts_on_cluster(self) -> Optional[float]:
         """Wait for MAX_JOB_CHECKING_RETRY times until job starts on the cluster
@@ -304,89 +278,96 @@ class StrategyExecutor:
         backoff = common_utils.Backoff(self.RETRY_INIT_GAP_SECONDS)
         while True:
             retry_cnt += 1
-            try:
-                usage_lib.messages.usage.set_internal()
-                # Detach setup, so that the setup failure can be detected
-                # by the controller process (job_status -> FAILED_SETUP).
-                sky.launch(
-                    self.dag,
-                    cluster_name=self.cluster_name,
-                    # We expect to tear down the cluster as soon as the job is
-                    # finished. However, in case the controller dies, set
-                    # autodown to try and avoid a resource leak.
-                    idle_minutes_to_autostop=_AUTODOWN_MINUTES,
-                    down=True,
-                    detach_setup=True,
-                    detach_run=True,
-                    _is_launched_by_jobs_controller=True)
-                logger.info('Managed job cluster launched.')
-            except (exceptions.InvalidClusterNameError,
-                    exceptions.NoCloudAccessError,
-                    exceptions.ResourcesMismatchError) as e:
-                logger.error('Failure happened before provisioning. '
-                             f'{common_utils.format_exception(e)}')
-                if raise_on_failure:
-                    raise exceptions.ProvisionPrechecksError(reasons=[e])
-                return None
-            except exceptions.ResourcesUnavailableError as e:
-                # This is raised when the launch fails due to prechecks or
-                # after failing over through all the candidates.
-                # Please refer to the docstring of `sky.launch` for more
-                # details of how the exception will be structured.
-                if not any(
-                        isinstance(err, exceptions.ResourcesUnavailableError)
-                        for err in e.failover_history):
-                    # _launch() (this function) should fail/exit directly, if
-                    # none of the failover reasons were because of resource
-                    # unavailability or no failover was attempted (the optimizer
-                    # cannot find feasible resources for requested resources),
-                    # i.e., e.failover_history is empty.
-                    # Failing directly avoids the infinite loop of retrying
-                    # the launch when, e.g., an invalid cluster name is used
-                    # and --retry-until-up is specified.
-                    reasons = (e.failover_history
-                               if e.failover_history else [e])
-                    reasons_str = '; '.join(
-                        common_utils.format_exception(err) for err in reasons)
-                    logger.error(
-                        'Failure happened before provisioning. Failover '
-                        f'reasons: {reasons_str}')
+            with scheduler.scheduled_launch(self.job_id):
+                try:
+                    usage_lib.messages.usage.set_internal()
+                    # Detach setup, so that the setup failure can be detected
+                    # by the controller process (job_status -> FAILED_SETUP).
+                    sky.launch(
+                        self.dag,
+                        cluster_name=self.cluster_name,
+                        # We expect to tear down the cluster as soon as the job
+                        # is finished. However, in case the controller dies, set
+                        # autodown to try and avoid a resource leak.
+                        idle_minutes_to_autostop=_AUTODOWN_MINUTES,
+                        down=True,
+                        detach_setup=True,
+                        detach_run=True,
+                        _is_launched_by_jobs_controller=True)
+                    logger.info('Managed job cluster launched.')
+                except (exceptions.InvalidClusterNameError,
+                        exceptions.NoCloudAccessError,
+                        exceptions.ResourcesMismatchError) as e:
+                    logger.error('Failure happened before provisioning. '
+                                 f'{common_utils.format_exception(e)}')
                     if raise_on_failure:
-                        raise exceptions.ProvisionPrechecksError(reasons)
-                    return None
-                logger.info('Failed to launch a cluster with error: '
-                            f'{common_utils.format_exception(e)})')
-            except Exception as e:  # pylint: disable=broad-except
-                # If the launch fails, it will be recovered by the following
-                # code.
-                logger.info('Failed to launch a cluster with error: '
-                            f'{common_utils.format_exception(e)})')
-                with ux_utils.enable_traceback():
-                    logger.info(f'  Traceback: {traceback.format_exc()}')
-            else:  # No exception, the launch succeeds.
-                # At this point, a sky.launch() has succeeded. Cluster may be
-                # UP (no preemption since) or DOWN (newly preempted).
-                job_submitted_at = self._wait_until_job_starts_on_cluster()
-                if job_submitted_at is not None:
-                    return job_submitted_at
-                # The job fails to start on the cluster, retry the launch.
-                # TODO(zhwu): log the unexpected error to usage collection
-                # for future debugging.
-                logger.info(
-                    'Failed to successfully submit the job to the '
-                    'launched cluster, due to unexpected submission errors or '
-                    'the cluster being preempted during job submission.')
-
-            terminate_cluster(self.cluster_name)
-            if max_retry is not None and retry_cnt >= max_retry:
-                # Retry forever if max_retry is None.
-                if raise_on_failure:
-                    with ux_utils.print_exception_no_traceback():
-                        raise exceptions.ManagedJobReachedMaxRetriesError(
-                            'Resources unavailable: failed to launch clusters '
-                            f'after {max_retry} retries.')
-                else:
+                        raise exceptions.ProvisionPrechecksError(reasons=[e])
                     return None
+                except exceptions.ResourcesUnavailableError as e:
+                    # This is raised when the launch fails due to prechecks or
+                    # after failing over through all the candidates.
+                    # Please refer to the docstring of `sky.launch` for more
+                    # details of how the exception will be structured.
+                    if not any(
+                            isinstance(err,
+                                       exceptions.ResourcesUnavailableError)
+                            for err in e.failover_history):
+                        # _launch() (this function) should fail/exit directly,
+                        # if none of the failover reasons were because of
+                        # resource unavailability or no failover was attempted
+                        # (the optimizer cannot find feasible resources for
+                        # requested resources), i.e., e.failover_history is
+                        # empty. Failing directly avoids the infinite loop of
+                        # retrying the launch when, e.g., an invalid cluster
+                        # name is used and --retry-until-up is specified.
+                        reasons = (e.failover_history
+                                   if e.failover_history else [e])
+                        reasons_str = '; '.join(
+                            common_utils.format_exception(err)
+                            for err in reasons)
+                        logger.error(
+                            'Failure happened before provisioning. Failover '
+                            f'reasons: {reasons_str}')
+                        if raise_on_failure:
+                            raise exceptions.ProvisionPrechecksError(reasons)
+                        return None
+                    logger.info('Failed to launch a cluster with error: '
+                                f'{common_utils.format_exception(e)})')
+                except Exception as e:  # pylint: disable=broad-except
+                    # If the launch fails, it will be recovered by the following
+                    # code.
+                    logger.info('Failed to launch a cluster with error: '
+                                f'{common_utils.format_exception(e)})')
+                    with ux_utils.enable_traceback():
+                        logger.info(f'  Traceback: {traceback.format_exc()}')
+                else:  # No exception, the launch succeeds.
+                    # At this point, a sky.launch() has succeeded. Cluster may
+                    # be UP (no preemption since) or DOWN (newly preempted).
+                    job_submitted_at = self._wait_until_job_starts_on_cluster()
+                    if job_submitted_at is not None:
+                        return job_submitted_at
+                    # The job fails to start on the cluster, retry the launch.
+                    # TODO(zhwu): log the unexpected error to usage collection
+                    # for future debugging.
+                    logger.info(
+                        'Failed to successfully submit the job to the '
+                        'launched cluster, due to unexpected submission errors '
+                        'or the cluster being preempted during job submission.')
+
+                # If we get here, the launch did not succeed. Tear down the
+                # cluster and retry.
+                managed_job_utils.terminate_cluster(self.cluster_name)
+                if max_retry is not None and retry_cnt >= max_retry:
+                    # Retry forever if max_retry is None.
+                    if raise_on_failure:
+                        with ux_utils.print_exception_no_traceback():
+                            raise exceptions.ManagedJobReachedMaxRetriesError(
+                                'Resources unavailable: failed to launch '
+                                f'clusters after {max_retry} retries.')
+                    else:
+                        return None
+            # Exit the scheduled_launch context so that the scheulde state is
+            # ALIVE during the backoff. This allows other jobs to launch.
             gap_seconds = backoff.current_backoff()
             logger.info('Retrying to launch the cluster in '
                         f'{gap_seconds:.1f} seconds.')
@@ -411,10 +392,10 @@ class FailoverStrategyExecutor(StrategyExecutor, name='FAILOVER',
     _MAX_RETRY_CNT = 240  # Retry for 4 hours.
 
     def __init__(self, cluster_name: str, backend: 'backends.Backend',
-                 task: 'task_lib.Task', retry_until_up: bool,
-                 max_restarts_on_errors: int) -> None:
-        super().__init__(cluster_name, backend, task, retry_until_up,
-                         max_restarts_on_errors)
+                 task: 'task_lib.Task', max_restarts_on_errors: int,
+                 job_id: int) -> None:
+        super().__init__(cluster_name, backend, task, max_restarts_on_errors,
+                         job_id)
         # Note down the cloud/region of the launched cluster, so that we can
         # first retry in the same cloud/region. (Inside recover() we may not
         # rely on cluster handle, as it can be None if the cluster is
@@ -468,7 +449,7 @@ class FailoverStrategyExecutor(StrategyExecutor, name='FAILOVER',
             # Step 2
             logger.debug('Terminating unhealthy cluster and reset cloud '
                          'region.')
-            terminate_cluster(self.cluster_name)
+            managed_job_utils.terminate_cluster(self.cluster_name)
 
             # Step 3
             logger.debug('Relaunch the cluster  without constraining to prior '
@@ -478,16 +459,11 @@ class FailoverStrategyExecutor(StrategyExecutor, name='FAILOVER',
                                             raise_on_failure=False)
             if job_submitted_at is None:
                 # Failed to launch the cluster.
-                if self.retry_until_up:
-                    gap_seconds = self.RETRY_INIT_GAP_SECONDS
-                    logger.info('Retrying to recover the cluster in '
-                                f'{gap_seconds:.1f} seconds.')
-                    time.sleep(gap_seconds)
-                    continue
-                with ux_utils.print_exception_no_traceback():
-                    raise exceptions.ResourcesUnavailableError(
-                        f'Failed to recover the cluster after retrying '
-                        f'{self._MAX_RETRY_CNT} times.')
+                gap_seconds = self.RETRY_INIT_GAP_SECONDS
+                logger.info('Retrying to recover the cluster in '
+                            f'{gap_seconds:.1f} seconds.')
+                time.sleep(gap_seconds)
+                continue
 
             return job_submitted_at
 
@@ -531,7 +507,7 @@ class EagerFailoverStrategyExecutor(FailoverStrategyExecutor,
 
         # Step 1
         logger.debug('Terminating unhealthy cluster and reset cloud region.')
-        terminate_cluster(self.cluster_name)
+        managed_job_utils.terminate_cluster(self.cluster_name)
 
         # Step 2
         logger.debug('Relaunch the cluster skipping the previously launched '
@@ -566,15 +542,10 @@ class EagerFailoverStrategyExecutor(FailoverStrategyExecutor,
                                             raise_on_failure=False)
             if job_submitted_at is None:
                 # Failed to launch the cluster.
-                if self.retry_until_up:
-                    gap_seconds = self.RETRY_INIT_GAP_SECONDS
-                    logger.info('Retrying to recover the cluster in '
-                                f'{gap_seconds:.1f} seconds.')
-                    time.sleep(gap_seconds)
-                    continue
-                with ux_utils.print_exception_no_traceback():
-                    raise exceptions.ResourcesUnavailableError(
-                        f'Failed to recover the cluster after retrying '
-                        f'{self._MAX_RETRY_CNT} times.')
+                gap_seconds = self.RETRY_INIT_GAP_SECONDS
+                logger.info('Retrying to recover the cluster in '
+                            f'{gap_seconds:.1f} seconds.')
+                time.sleep(gap_seconds)
+                continue
 
             return job_submitted_at

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)