skypilot-nightly 1.0.0.dev20250909__py3-none-any.whl → 1.0.0.dev20250910__py3-none-any.whl

sky/jobs/controller.py

@@ -1,31 +1,26 @@
-"""Controller: handles the life cycle of a managed job.
-
-TODO(cooperc): Document lifecycle, and multiprocess layout.
+"""Controller: handles scheduling and the life cycle of a managed job.
 """
-import argparse
-import multiprocessing
+import asyncio
+import logging
 import os
-import pathlib
+import resource
 import shutil
+import sys
 import time
 import traceback
 import typing
-from typing import Optional, Tuple
+from typing import Dict, Optional, Set, Tuple
 
-import filelock
+import dotenv
 
-# This import ensures backward compatibility. Controller processes may not have
-# imported this module initially, but will attempt to import it during job
-# termination on the fly. If a job was launched with an old SkyPilot runtime
-# and a new job is launched with a newer runtime, the old job's termination
-# will try to import code from a different SkyPilot runtime, causing exceptions.
-# pylint: disable=unused-import
+import sky
 from sky import core
 from sky import exceptions
 from sky import sky_logging
 from sky.backends import backend_utils
 from sky.backends import cloud_vm_ray_backend
 from sky.data import data_utils
+from sky.jobs import constants as jobs_constants
 from sky.jobs import recovery_strategy
 from sky.jobs import scheduler
 from sky.jobs import state as managed_job_state
@@ -35,20 +30,34 @@ from sky.skylet import job_lib
 from sky.usage import usage_lib
 from sky.utils import common
 from sky.utils import common_utils
+from sky.utils import context
+from sky.utils import context_utils
 from sky.utils import controller_utils
 from sky.utils import dag_utils
 from sky.utils import status_lib
-from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 
-if typing.TYPE_CHECKING:
-    import sky
-
-# Use the explicit logger name so that the logger is under the
-# `sky.jobs.controller` namespace when executed directly, so as
-# to inherit the setup from the `sky` logger.
 logger = sky_logging.init_logger('sky.jobs.controller')
 
+_background_tasks: Set[asyncio.Task] = set()
+_background_tasks_lock: asyncio.Lock = asyncio.Lock()
+
+
+async def create_background_task(coro: typing.Coroutine) -> None:
+    """Create a background task and add it to the set of background tasks.
+
+    Main reason we do this is since tasks are only held as a weak reference in
+    the executor, we need to keep a strong reference to the task to avoid it
+    being garbage collected.
+
+    Args:
+        coro: The coroutine to create a task for.
+    """
+    async with _background_tasks_lock:
+        task = asyncio.create_task(coro)
+        _background_tasks.add(task)
+        task.add_done_callback(_background_tasks.discard)
+
 
 def _get_dag_and_name(dag_yaml: str) -> Tuple['sky.Dag', str]:
     dag = dag_utils.load_chain_dag_from_yaml(dag_yaml)
@@ -58,13 +67,79 @@ def _get_dag_and_name(dag_yaml: str) -> Tuple['sky.Dag', str]:
 
 
 class JobsController:
-    """Each jobs controller manages the life cycle of one managed job."""
+    """Controls the lifecycle of a single managed job.
+
+    This controller executes a chain DAG defined in ``dag_yaml`` by:
+    - Loading the DAG and preparing per-task environment variables so each task
+      has a stable global job identifier across recoveries.
+    - Launching the task on the configured backend (``CloudVmRayBackend``),
+      optionally via a cluster pool.
+    - Persisting state transitions to the managed jobs state store
+      (e.g., STARTING → RUNNING → SUCCEEDED/FAILED/CANCELLED).
+    - Monitoring execution, downloading/streaming logs, detecting failures or
+      preemptions, and invoking recovery through
+      ``recovery_strategy.StrategyExecutor``.
+    - Cleaning up clusters and ephemeral resources when tasks finish.
+
+    Concurrency and coordination:
+    - Runs inside an ``asyncio`` event loop.
+    - Shares a ``starting`` set, guarded by ``starting_lock`` and signaled via
+      ``starting_signal``, to throttle concurrent launches across jobs that the
+      top-level ``Controller`` manages.
+
+    Key attributes:
+    - ``_job_id``: Integer identifier of this managed job.
+    - ``_dag_yaml`` / ``_dag`` / ``_dag_name``: The job definition and metadata.
+    - ``_backend``: Backend used to launch and manage clusters.
+    - ``_pool``: Optional pool name if using a cluster pool.
+    - ``_logger``: Job-scoped logger for progress and diagnostics.
+    - ``starting`` / ``starting_lock`` / ``starting_signal``: Shared scheduler
+      coordination primitives. ``starting_lock`` must be used for accessing
+      ``starting_signal`` and ``starting``
+    - ``_strategy_executor``: Recovery/launch strategy executor (created per
+      task).
+    """
+
+    def __init__(
+        self,
+        job_id: int,
+        dag_yaml: str,
+        job_logger: logging.Logger,
+        starting: Set[int],
+        starting_lock: asyncio.Lock,
+        starting_signal: asyncio.Condition,
+        pool: Optional[str] = None,
+    ) -> None:
+        """Initialize a ``JobsController``.
+
+        Args:
+            job_id: Integer ID of the managed job.
+            dag_yaml: Path to the YAML file containing the chain DAG to run.
+            job_logger: Logger instance dedicated to this job.
+            starting: Shared set of job IDs currently in the STARTING phase,
+                used to limit concurrent launches.
+            starting_lock: ``asyncio.Lock`` guarding access to the shared
+                scheduler state (e.g., the ``starting`` set).
+            starting_signal: ``asyncio.Condition`` used to notify when a job
+                exits STARTING so more jobs can be admitted.
+            pool: Optional cluster pool name. When provided, the job is
+                submitted to the pool rather than launching a dedicated
+                cluster.
+        """
+
+        self.starting = starting
+        self.starting_lock = starting_lock
+        self.starting_signal = starting_signal
+
+        self._logger = job_logger
+        self._logger.info(f'Initializing JobsController for job_id={job_id}, '
+                          f'dag_yaml={dag_yaml}')
 
-    def __init__(self, job_id: int, dag_yaml: str, pool: Optional[str]) -> None:
         self._job_id = job_id
+        self._dag_yaml = dag_yaml
         self._dag, self._dag_name = _get_dag_and_name(dag_yaml)
-        logger.info(self._dag)
-        # TODO(zhwu): this assumes the specific backend.
+        self._logger.info(f'Loaded DAG: {self._dag}')
+
         self._backend = cloud_vm_ray_backend.CloudVmRayBackend()
         self._pool = pool
 
@@ -84,6 +159,7 @@ class JobsController:
                 # dag_utils.maybe_infer_and_fill_dag_and_task_names.
                 assert task_name is not None, self._dag
                 task_name = f'{self._dag_name}_{task_name}'
+
             job_id_env_var = common_utils.get_global_job_id(
                 self._backend.run_timestamp,
                 f'{task_name}',
@@ -102,7 +178,7 @@ class JobsController:
     def _download_log_and_stream(
         self,
         task_id: Optional[int],
-        handle: Optional[cloud_vm_ray_backend.CloudVmRayResourceHandle],
+        handle: Optional['cloud_vm_ray_backend.CloudVmRayResourceHandle'],
         job_id_on_pool_cluster: Optional[int],
     ) -> None:
         """Downloads and streams the logs of the current job with given task ID.
@@ -112,9 +188,10 @@ class JobsController:
         preemptions or ssh disconnection during the streaming.
         """
         if handle is None:
-            logger.info(f'Cluster for job {self._job_id} is not found. '
-                        'Skipping downloading and streaming the logs.')
+            self._logger.info(f'Cluster for job {self._job_id} is not found. '
+                              'Skipping downloading and streaming the logs.')
             return
+
         managed_job_logs_dir = os.path.join(constants.SKY_LOGS_DIRECTORY,
                                             'managed_jobs',
                                             f'job-id-{self._job_id}')
@@ -125,19 +202,25 @@ class JobsController:
             job_ids=[str(job_id_on_pool_cluster)]
             if job_id_on_pool_cluster is not None else None)
         if log_file is not None:
-            # Set the path of the log file for the current task, so it can be
-            # accessed even after the job is finished
+            # Set the path of the log file for the current task, so it can
+            # be accessed even after the job is finished
             managed_job_state.set_local_log_file(self._job_id, task_id,
                                                  log_file)
-        logger.info(f'\n== End of logs (ID: {self._job_id}) ==')
+        else:
+            self._logger.warning(
+                f'No log file was downloaded for job {self._job_id}, '
+                f'task {task_id}')
+
+        self._logger.info(f'\n== End of logs (ID: {self._job_id}) ==')
 
-    def _cleanup_cluster(self, cluster_name: Optional[str]) -> None:
+    async def _cleanup_cluster(self, cluster_name: Optional[str]) -> None:
         if cluster_name is None:
             return
         if self._pool is None:
-            managed_job_utils.terminate_cluster(cluster_name)
+            await context_utils.to_thread(managed_job_utils.terminate_cluster,
+                                          cluster_name)
 
-    def _run_one_task(self, task_id: int, task: 'sky.Task') -> bool:
+    async def _run_one_task(self, task_id: int, task: 'sky.Task') -> bool:
         """Busy loop monitoring cluster status and handling recovery.
 
         When the task is successfully completed, this function returns True,
@@ -172,38 +255,52 @@ class JobsController:
                 3. Any unexpected error happens during the `sky.launch`.
         Other exceptions may be raised depending on the backend.
         """
+        task_start_time = time.time()
+        self._logger.info(
+            f'Starting task {task_id} ({task.name}) for job {self._job_id}')
 
         latest_task_id, last_task_prev_status = (
-            managed_job_state.get_latest_task_id_status(self._job_id))
+            await
+            managed_job_state.get_latest_task_id_status_async(self._job_id))
+
         is_resume = False
         if (latest_task_id is not None and last_task_prev_status !=
                 managed_job_state.ManagedJobStatus.PENDING):
             assert latest_task_id >= task_id, (latest_task_id, task_id)
             if latest_task_id > task_id:
-                logger.info(f'Task {task_id} ({task.name}) has already '
-                            'been executed. Skipping...')
+                self._logger.info(f'Task {task_id} ({task.name}) has already '
+                                  'been executed. Skipping...')
                 return True
             if latest_task_id == task_id:
                 # Start recovery.
                 is_resume = True
+                self._logger.info(
+                    f'Resuming task {task_id} from previous execution')
 
         callback_func = managed_job_utils.event_callback_func(
             job_id=self._job_id, task_id=task_id, task=task)
+
         if task.run is None:
-            logger.info(f'Skip running task {task_id} ({task.name}) due to its '
-                        'run commands being empty.')
+            self._logger.info(
+                f'Skip running task {task_id} ({task.name}) due to its '
+                'run commands being empty.')
             # Call set_started first to initialize columns in the state table,
             # including start_at and last_recovery_at to avoid issues for
             # uninitialized columns.
-            managed_job_state.set_started(job_id=self._job_id,
-                                          task_id=task_id,
-                                          start_time=time.time(),
-                                          callback_func=callback_func)
-            managed_job_state.set_succeeded(job_id=self._job_id,
-                                            task_id=task_id,
-                                            end_time=time.time(),
-                                            callback_func=callback_func)
+            await managed_job_state.set_started_async(
+                job_id=self._job_id,
+                task_id=task_id,
+                start_time=time.time(),
+                callback_func=callback_func)
+            await managed_job_state.set_succeeded_async(
+                job_id=self._job_id,
+                task_id=task_id,
+                end_time=time.time(),
+                callback_func=callback_func)
+            self._logger.info(
+                f'Empty task {task_id} marked as succeeded immediately')
             return True
+
         usage_lib.messages.usage.update_task_id(task_id)
         task_id_env_var = task.envs[constants.TASK_ID_ENV_VAR]
         assert task.name is not None, task
@@ -214,61 +311,97 @@ class JobsController:
             task.name, self._job_id) if self._pool is None else None
         self._strategy_executor = recovery_strategy.StrategyExecutor.make(
             cluster_name, self._backend, task, self._job_id, task_id,
-            self._pool)
+            self._logger, self._pool, self.starting, self.starting_lock,
+            self.starting_signal)
         if not is_resume:
             submitted_at = time.time()
             if task_id == 0:
                 submitted_at = backend_utils.get_timestamp_from_run_timestamp(
                     self._backend.run_timestamp)
-            managed_job_state.set_starting(
+
+            resources_str = backend_utils.get_task_resources_str(
+                task, is_managed_job=True)
+
+            await managed_job_state.set_starting_async(
                 self._job_id,
                 task_id,
                 self._backend.run_timestamp,
                 submitted_at,
-                resources_str=backend_utils.get_task_resources_str(
-                    task, is_managed_job=True),
+                resources_str=resources_str,
                 specs={
                     'max_restarts_on_errors':
                         self._strategy_executor.max_restarts_on_errors
                 },
                 callback_func=callback_func)
-            logger.info(f'Submitted managed job {self._job_id} '
-                        f'(task: {task_id}, name: {task.name!r}); '
-                        f'{constants.TASK_ID_ENV_VAR}: {task_id_env_var}')
+            self._logger.info(f'Submitted managed job {self._job_id} '
+                              f'(task: {task_id}, name: {task.name!r}); '
+                              f'{constants.TASK_ID_ENV_VAR}: {task_id_env_var}')
 
-        logger.info('Started monitoring.')
+        self._logger.info('Started monitoring.')
 
         # Only do the initial cluster launch if not resuming from a controller
         # failure. Otherwise, we will transit to recovering immediately.
         remote_job_submitted_at = time.time()
         if not is_resume:
-            remote_job_submitted_at = self._strategy_executor.launch()
+            launch_start = time.time()
+
+            # Run the launch in a separate thread to avoid blocking the event
+            # loop. The scheduler functions used internally already have their
+            # own file locks.
+            remote_job_submitted_at = await self._strategy_executor.launch()
+
+            launch_time = time.time() - launch_start
+            self._logger.info(f'Cluster launch completed in {launch_time:.2f}s')
             assert remote_job_submitted_at is not None, remote_job_submitted_at
         if self._pool is None:
             job_id_on_pool_cluster = None
         else:
             # Update the cluster name when using cluster pool.
             cluster_name, job_id_on_pool_cluster = (
-                managed_job_state.get_pool_submit_info(self._job_id))
+                await
+                managed_job_state.get_pool_submit_info_async(self._job_id))
         assert cluster_name is not None, (cluster_name, job_id_on_pool_cluster)
 
         if not is_resume:
-            managed_job_state.set_started(job_id=self._job_id,
-                                          task_id=task_id,
-                                          start_time=remote_job_submitted_at,
-                                          callback_func=callback_func)
+            await managed_job_state.set_started_async(
+                job_id=self._job_id,
+                task_id=task_id,
+                start_time=remote_job_submitted_at,
+                callback_func=callback_func)
+
+        monitoring_start_time = time.time()
+        status_check_count = 0
+
+        async with self.starting_lock:
+            try:
+                self.starting.remove(self._job_id)
+                # its fine if we notify again, better to wake someone up
+                # and have them go to sleep again, then have some stuck
+                # sleeping.
+                # ps. this shouldn't actually happen because if its been
+                # removed from the set then we would get a key error.
+                self.starting_signal.notify()
+            except KeyError:
+                pass
 
         while True:
+            status_check_count += 1
+
             # NOTE: if we are resuming from a controller failure, we only keep
             # monitoring if the job is in RUNNING state. For all other cases,
             # we will directly transit to recovering since we have no idea what
             # the cluster status is.
             force_transit_to_recovering = False
             if is_resume:
-                prev_status = managed_job_state.get_job_status_with_task_id(
-                    job_id=self._job_id, task_id=task_id)
+                prev_status = await (
+                    managed_job_state.get_job_status_with_task_id_async(
+                        job_id=self._job_id, task_id=task_id))
+
                 if prev_status is not None:
                     if prev_status.is_terminal():
+                        self._logger.info(
+                            f'Task {task_id} already in terminal state: '
+                            f'{prev_status}')
                         return (prev_status ==
                                 managed_job_state.ManagedJobStatus.SUCCEEDED)
                     if (prev_status ==
@@ -276,23 +409,26 @@ class JobsController:
                         # If the controller is down when cancelling the job,
                         # we re-raise the error to run the `_cleanup` function
                         # again to clean up any remaining resources.
-                        raise exceptions.ManagedJobUserCancelledError(
-                            'Recovering cancel signal.')
+                        self._logger.info(
+                            f'Task {task_id} was being cancelled, '
+                            're-raising cancellation')
+                        raise asyncio.CancelledError()
                 if prev_status != managed_job_state.ManagedJobStatus.RUNNING:
                     force_transit_to_recovering = True
                 # This resume logic should only be triggered once.
                 is_resume = False
 
-            time.sleep(managed_job_utils.JOB_STATUS_CHECK_GAP_SECONDS)
+            await asyncio.sleep(managed_job_utils.JOB_STATUS_CHECK_GAP_SECONDS)
 
             # Check the network connection to avoid false alarm for job failure.
             # Network glitch was observed even in the VM.
             try:
-                backend_utils.check_network_connection()
+                await backend_utils.async_check_network_connection()
             except exceptions.NetworkError:
-                logger.info('Network is not available. Retrying again in '
-                            f'{managed_job_utils.JOB_STATUS_CHECK_GAP_SECONDS} '
-                            'seconds.')
+                self._logger.info(
+                    'Network is not available. Retrying again in '
+                    f'{managed_job_utils.JOB_STATUS_CHECK_GAP_SECONDS} '
+                    'seconds.')
                 continue
 
             # NOTE: we do not check cluster status first because race condition
@@ -303,32 +439,47 @@ class JobsController:
             job_status = None
             if not force_transit_to_recovering:
                 try:
-                    job_status = managed_job_utils.get_job_status(
+                    job_status = await managed_job_utils.get_job_status(
                         self._backend,
                         cluster_name,
-                        job_id=job_id_on_pool_cluster)
+                        job_id=job_id_on_pool_cluster,
+                        job_logger=self._logger,
+                    )
                 except exceptions.FetchClusterInfoError as fetch_e:
-                    logger.info(
+                    self._logger.info(
                         'Failed to fetch the job status. Start recovery.\n'
                         f'Exception: {common_utils.format_exception(fetch_e)}\n'
                         f'Traceback: {traceback.format_exc()}')
 
             if job_status == job_lib.JobStatus.SUCCEEDED:
-                success_end_time = managed_job_utils.try_to_get_job_end_time(
-                    self._backend, cluster_name, job_id_on_pool_cluster)
+                self._logger.info(f'Task {task_id} succeeded! '
+                                  'Getting end time and cleaning up')
+                try:
+                    success_end_time = await context_utils.to_thread(
+                        managed_job_utils.try_to_get_job_end_time,
+                        self._backend, cluster_name, job_id_on_pool_cluster)
+                except Exception as e:  # pylint: disable=broad-except
+                    self._logger.warning(
+                        f'Failed to get job end time: '
+                        f'{common_utils.format_exception(e)}',
+                        exc_info=True)
+                    success_end_time = 0
+
                 # The job is done. Set the job to SUCCEEDED first before start
                 # downloading and streaming the logs to make it more responsive.
-                managed_job_state.set_succeeded(self._job_id,
-                                                task_id,
-                                                end_time=success_end_time,
-                                                callback_func=callback_func)
-                logger.info(
+                await managed_job_state.set_succeeded_async(
+                    self._job_id,
+                    task_id,
+                    end_time=success_end_time,
+                    callback_func=callback_func)
+                self._logger.info(
                     f'Managed job {self._job_id} (task: {task_id}) SUCCEEDED. '
                     f'Cleaning up the cluster {cluster_name}.')
                 try:
                     logger.info(f'Downloading logs on cluster {cluster_name} '
                                 f'and job id {job_id_on_pool_cluster}.')
-                    clusters = backend_utils.get_clusters(
+                    clusters = await context_utils.to_thread(
+                        backend_utils.get_clusters,
                         cluster_names=[cluster_name],
                         refresh=common.StatusRefreshMode.NONE,
                         all_users=True,
@@ -337,17 +488,25 @@ class JobsController:
                         assert len(clusters) == 1, (clusters, cluster_name)
                         handle = clusters[0].get('handle')
                         # Best effort to download and stream the logs.
-                        self._download_log_and_stream(task_id, handle,
-                                                      job_id_on_pool_cluster)
+                        await context_utils.to_thread(
+                            self._download_log_and_stream, task_id, handle,
+                            job_id_on_pool_cluster)
                 except Exception as e:  # pylint: disable=broad-except
                     # We don't want to crash here, so just log and continue.
-                    logger.warning(
+                    self._logger.warning(
                         f'Failed to download and stream logs: '
                         f'{common_utils.format_exception(e)}',
                         exc_info=True)
                 # Only clean up the cluster, not the storages, because tasks may
                 # share storages.
-                self._cleanup_cluster(cluster_name)
+                await self._cleanup_cluster(cluster_name)
+
+                task_total_time = time.time() - task_start_time
+                monitoring_time = time.time() - monitoring_start_time
+                self._logger.info(f'Task {task_id} completed successfully in '
+                                  f'{task_total_time:.2f}s '
+                                  f'(monitoring time: {monitoring_time:.2f}s, '
+                                  f'status checks: {status_check_count})')
                 return True
 
             # For single-node jobs, non-terminated job_status indicates a
@@ -363,7 +522,7 @@ class JobsController:
             if job_status in job_lib.JobStatus.user_code_failure_states():
                 # Add a grace period before the check of preemption to avoid
                 # false alarm for job failure.
-                time.sleep(5)
+                await asyncio.sleep(5)
 
             # Pull the actual cluster status from the cloud provider to
             # determine whether the cluster is preempted or failed.
@@ -383,7 +542,7 @@ class JobsController:
                 # code).
                 cluster_status_str = ('' if cluster_status is None else
                                       f' (status: {cluster_status.value})')
-                logger.info(
+                self._logger.info(
                     f'Cluster is preempted or failed{cluster_status_str}. '
                     'Recovering...')
             else:
@@ -394,14 +553,18 @@ class JobsController:
                       in job_lib.JobStatus.user_code_failure_states() or
                       job_status == job_lib.JobStatus.FAILED_DRIVER):
                     # The user code has probably crashed, fail immediately.
-                    end_time = managed_job_utils.try_to_get_job_end_time(
+                    self._logger.info(
+                        f'Task {task_id} failed with status: {job_status}')
+                    end_time = await context_utils.to_thread(
+                        managed_job_utils.try_to_get_job_end_time,
                         self._backend, cluster_name, job_id_on_pool_cluster)
-                    logger.info(
+                    self._logger.info(
                         f'The user job failed ({job_status}). Please check the '
                         'logs below.\n'
                         f'== Logs of the user job (ID: {self._job_id}) ==\n')
 
-                    self._download_log_and_stream(task_id, handle,
+                    await context_utils.to_thread(self._download_log_and_stream,
+                                                  task_id, handle,
                                                   job_id_on_pool_cluster)
 
                     failure_reason = (
@@ -430,7 +593,7 @@ class JobsController:
                     if should_restart_on_failure:
                         max_restarts = (
                             self._strategy_executor.max_restarts_on_errors)
-                        logger.info(
+                        self._logger.info(
                             f'User program crashed '
                             f'({managed_job_status.value}). '
                             f'Retry the job as max_restarts_on_errors is '
@@ -438,7 +601,9 @@ class JobsController:
                             f'[{self._strategy_executor.restart_cnt_on_failure}'
                             f'/{max_restarts}]')
                     else:
-                        managed_job_state.set_failed(
+                        self._logger.info(
+                            f'Task {task_id} failed and will not be retried')
+                        await managed_job_state.set_failed_async(
                             self._job_id,
                             task_id,
                             failure_type=managed_job_status,
@@ -449,11 +614,11 @@ class JobsController:
                 elif job_status is not None:
                     # Either the job is cancelled (should not happen) or in some
                     # unknown new state that we do not handle.
-                    logger.error(f'Unknown job status: {job_status}')
+                    self._logger.error(f'Unknown job status: {job_status}')
                     failure_reason = (
                         f'Unknown job status {job_status}. To see the details, '
                         f'run: sky jobs logs --controller {self._job_id}')
-                    managed_job_state.set_failed(
+                    await managed_job_state.set_failed_async(
                         self._job_id,
                         task_id,
                         failure_type=managed_job_state.ManagedJobStatus.
@@ -466,9 +631,10 @@ class JobsController:
                     # job status. Try to recover the job (will not restart the
                     # cluster, if the cluster is healthy).
                     assert job_status is None, job_status
-                    logger.info('Failed to fetch the job status while the '
-                                'cluster is healthy. Try to recover the job '
-                                '(the cluster will not be restarted).')
+                    self._logger.info(
+                        'Failed to fetch the job status while the '
+                        'cluster is healthy. Try to recover the job '
+                        '(the cluster will not be restarted).')
             # When the handle is None, the cluster should be cleaned up already.
             if handle is not None:
                 resources = handle.launched_resources
@@ -487,86 +653,121 @@ class JobsController:
                     # Some spot resource (e.g., Spot TPU VM) may need to be
                     # cleaned up after preemption, as running launch again on
                     # those clusters again may fail.
-                    logger.info('Cleaning up the preempted or failed cluster'
-                                '...')
-                    self._cleanup_cluster(cluster_name)
+                    self._logger.info(
+                        'Cleaning up the preempted or failed cluster'
+                        '...')
+                    await self._cleanup_cluster(cluster_name)
 
             # Try to recover the managed jobs, when the cluster is preempted or
             # failed or the job status is failed to be fetched.
-            managed_job_state.set_recovering(
+            self._logger.info(f'Starting recovery for task {task_id}, '
+                              f'it is currently {job_status}')
+            await managed_job_state.set_recovering_async(
                 job_id=self._job_id,
                 task_id=task_id,
                 force_transit_to_recovering=force_transit_to_recovering,
                 callback_func=callback_func)
-            recovered_time = self._strategy_executor.recover()
+
+            recovered_time = await self._strategy_executor.recover()
+
             if self._pool is not None:
                 cluster_name, job_id_on_pool_cluster = (
-                    managed_job_state.get_pool_submit_info(self._job_id))
+                    await
+                    managed_job_state.get_pool_submit_info_async(self._job_id))
                 assert cluster_name is not None
-            managed_job_state.set_recovered(self._job_id,
-                                            task_id,
-                                            recovered_time=recovered_time,
-                                            callback_func=callback_func)
+            await managed_job_state.set_recovered_async(
+                self._job_id,
+                task_id,
+                recovered_time=recovered_time,
+                callback_func=callback_func)
 
-    def run(self):
+    async def run(self):
         """Run controller logic and handle exceptions."""
+        self._logger.info(f'Starting JobsController run for job {self._job_id}')
         task_id = 0
+        cancelled = False
+
         try:
             succeeded = True
             # We support chain DAGs only for now.
             for task_id, task in enumerate(self._dag.tasks):
-                succeeded = self._run_one_task(task_id, task)
+                self._logger.info(
+                    f'Processing task {task_id}/{len(self._dag.tasks)-1}: '
+                    f'{task.name}')
+                task_start = time.time()
+                succeeded = await self._run_one_task(task_id, task)
+                task_time = time.time() - task_start
+                self._logger.info(
+                    f'Task {task_id} completed in {task_time:.2f}s '
+                    f'with success={succeeded}')
+
                 if not succeeded:
+                    self._logger.info(
+                        f'Task {task_id} failed, stopping execution')
                     break
+
         except exceptions.ProvisionPrechecksError as e:
             # Please refer to the docstring of self._run for the cases when
             # this exception can occur.
+            self._logger.error(f'Provision prechecks failed for task {task_id}')
             failure_reason = ('; '.join(
                 common_utils.format_exception(reason, use_bracket=True)
                 for reason in e.reasons))
-            logger.error(failure_reason)
-            self._update_failed_task_state(
+            self._logger.error(failure_reason)
+            await self._update_failed_task_state(
                 task_id, managed_job_state.ManagedJobStatus.FAILED_PRECHECKS,
                 failure_reason)
         except exceptions.ManagedJobReachedMaxRetriesError as e:
             # Please refer to the docstring of self._run for the cases when
             # this exception can occur.
+            self._logger.error(
+                f'Managed job reached max retries for task {task_id}')
             failure_reason = common_utils.format_exception(e)
-            logger.error(failure_reason)
+            self._logger.error(failure_reason)
             # The managed job should be marked as FAILED_NO_RESOURCE, as the
             # managed job may be able to launch next time.
-            self._update_failed_task_state(
+            await self._update_failed_task_state(
                 task_id, managed_job_state.ManagedJobStatus.FAILED_NO_RESOURCE,
                 failure_reason)
+        except asyncio.CancelledError:  # pylint: disable=try-except-raise
+            # have this here to avoid getting caught by the general except block
+            # below.
+            cancelled = True
+            raise
         except (Exception, SystemExit) as e:  # pylint: disable=broad-except
+            self._logger.error(
+                f'Unexpected error in JobsController run for task {task_id}')
             with ux_utils.enable_traceback():
-                logger.error(traceback.format_exc())
+                self._logger.error(traceback.format_exc())
             msg = ('Unexpected error occurred: ' +
                    common_utils.format_exception(e, use_bracket=True))
-            logger.error(msg)
-            self._update_failed_task_state(
+            self._logger.error(msg)
+            await self._update_failed_task_state(
                 task_id, managed_job_state.ManagedJobStatus.FAILED_CONTROLLER,
                 msg)
         finally:
-            # This will set all unfinished tasks to CANCELLING, and will not
-            # affect the jobs in terminal states.
-            # We need to call set_cancelling before set_cancelled to make sure
-            # the table entries are correctly set.
             callback_func = managed_job_utils.event_callback_func(
                 job_id=self._job_id,
                 task_id=task_id,
                 task=self._dag.tasks[task_id])
-            managed_job_state.set_cancelling(job_id=self._job_id,
-                                             callback_func=callback_func)
-            managed_job_state.set_cancelled(job_id=self._job_id,
-                                            callback_func=callback_func)
+            await managed_job_state.set_cancelling_async(
+                job_id=self._job_id, callback_func=callback_func)
+            if not cancelled:
+                # the others haven't been run yet so we can set them to
+                # cancelled immediately (no resources to clean up).
+                # if we are running and get cancelled, we need to clean up the
+                # resources first so this will be done later.
+                await managed_job_state.set_cancelled_async(
+                    job_id=self._job_id, callback_func=callback_func)
 
-    def _update_failed_task_state(
+    async def _update_failed_task_state(
             self, task_id: int,
             failure_type: managed_job_state.ManagedJobStatus,
             failure_reason: str):
         """Update the state of the failed task."""
-        managed_job_state.set_failed(
+        self._logger.info(f'Updating failed task state: task_id={task_id}, '
+                          f'failure_type={failure_type}')
+        await managed_job_state.set_failed_async(
             self._job_id,
             task_id=task_id,
             failure_type=failure_type,
@@ -577,199 +778,420 @@ class JobsController:
                 task=self._dag.tasks[task_id]))
 
 
-def _run_controller(job_id: int, dag_yaml: str, pool: Optional[str]):
-    """Runs the controller in a remote process for interruption."""
-    # The controller needs to be instantiated in the remote process, since
-    # the controller is not serializable.
-    jobs_controller = JobsController(job_id, dag_yaml, pool)
-    jobs_controller.run()
-
-
-def _handle_signal(job_id):
-    """Handle the signal if the user sent it."""
-    signal_file = pathlib.Path(
-        managed_job_utils.SIGNAL_FILE_PREFIX.format(job_id))
-    user_signal = None
-    if signal_file.exists():
-        # Filelock is needed to prevent race condition with concurrent
-        # signal writing.
-        with filelock.FileLock(str(signal_file) + '.lock'):
-            with signal_file.open(mode='r', encoding='utf-8') as f:
-                user_signal = f.read().strip()
-                try:
-                    user_signal = managed_job_utils.UserSignal(user_signal)
-                except ValueError:
-                    logger.warning(
-                        f'Unknown signal received: {user_signal}. Ignoring.')
-                    user_signal = None
-            # Remove the signal file, after reading the signal.
-            signal_file.unlink()
-    if user_signal is None:
-        # None or empty string.
-        return
-    assert user_signal == managed_job_utils.UserSignal.CANCEL, (
-        f'Only cancel signal is supported, but {user_signal} got.')
-    raise exceptions.ManagedJobUserCancelledError(
-        f'User sent {user_signal.value} signal.')
-
-
-def _cleanup(job_id: int, dag_yaml: str, pool: Optional[str]):
-    """Clean up the cluster(s) and storages.
-
-    (1) Clean up the succeeded task(s)' ephemeral storage. The storage has
-        to be cleaned up after the whole job is finished, as the tasks
-        may share the same storage.
-    (2) Clean up the cluster(s) that are not cleaned up yet, which can happen
-        when the task failed or cancelled. At most one cluster should be left
-        when reaching here, as we currently only support chain DAGs, and only
-        task is executed at a time.
-    """
-    # Cleanup the HA recovery script first as it is possible that some error
-    # was raised when we construct the task object (e.g.,
-    # sky.exceptions.ResourcesUnavailableError).
-    managed_job_state.remove_ha_recovery_script(job_id)
-    dag, _ = _get_dag_and_name(dag_yaml)
-    for task in dag.tasks:
-        assert task.name is not None, task
-        if pool is None:
-            cluster_name = managed_job_utils.generate_managed_job_cluster_name(
-                task.name, job_id)
-            managed_job_utils.terminate_cluster(cluster_name)
-        else:
-            cluster_name, job_id_on_pool_cluster = (
-                managed_job_state.get_pool_submit_info(job_id))
-            if cluster_name is not None:
-                if job_id_on_pool_cluster is not None:
-                    core.cancel(cluster_name=cluster_name,
-                                job_ids=[job_id_on_pool_cluster],
-                                _try_cancel_if_cluster_is_init=True)
-
-        # Clean up Storages with persistent=False.
-        # TODO(zhwu): this assumes the specific backend.
-        backend = cloud_vm_ray_backend.CloudVmRayBackend()
-        # Need to re-construct storage object in the controller process
-        # because when SkyPilot API server machine sends the yaml config to the
-        # controller machine, only storage metadata is sent, not the storage
-        # object itself.
-        for storage in task.storage_mounts.values():
-            storage.construct()
-        backend.teardown_ephemeral_storage(task)
-
-        # Clean up any files mounted from the local disk, such as two-hop file
-        # mounts.
-        for file_mount in (task.file_mounts or {}).values():
+class Controller:
+    """Controller for managing jobs."""
+
+    def __init__(self):
+        # Global state for active jobs
+        self.job_tasks: Dict[int, asyncio.Task] = {}
+        self.starting: Set[int] = set()
+
+        # Lock for synchronizing access to global state dictionary
+        # Must always hold _job_tasks_lock when accessing the _starting_signal.
+        self._job_tasks_lock = asyncio.Lock()
+        # We signal whenever a job leaves the api server launching state. Feel
+        # free to signal as much as you want to be safe from leaks (if you
+        # do not signal enough there may be some jobs forever waiting to
+        # launch).
+        self._starting_signal = asyncio.Condition(lock=self._job_tasks_lock)
+
+    async def _cleanup(self,
+                       job_id: int,
+                       dag_yaml: str,
+                       job_logger: logging.Logger,
+                       pool: Optional[str] = None):
+        """Clean up the cluster(s) and storages.
+
+        (1) Clean up the succeeded task(s)' ephemeral storage. The storage has
+            to be cleaned up after the whole job is finished, as the tasks
+            may share the same storage.
+        (2) Clean up the cluster(s) that are not cleaned up yet, which can
+            happen when the task failed or cancelled. At most one cluster
+            should be left when reaching here, as we currently only support
+            chain DAGs, and only one task is executed at a time.
+        """
+        # Cleanup the HA recovery script first as it is possible that some error
+        # was raised when we construct the task object (e.g.,
+        # sky.exceptions.ResourcesUnavailableError).
+        await managed_job_state.remove_ha_recovery_script_async(job_id)
+
+        def task_cleanup(task: 'sky.Task', job_id: int):
+            assert task.name is not None, task
+            error = None
+
             try:
-                # For consolidation mode, there is no two-hop file mounts
-                # and the file path here represents the real user data.
-                # We skip the cleanup for consolidation mode.
-                if (not data_utils.is_cloud_store_url(file_mount) and
-                        not managed_job_utils.is_consolidation_mode()):
-                    path = os.path.expanduser(file_mount)
-                    if os.path.isdir(path):
-                        shutil.rmtree(path)
-                    else:
-                        os.remove(path)
+                if pool is None:
+                    cluster_name = (
+                        managed_job_utils.generate_managed_job_cluster_name(
+                            task.name, job_id))
+                    managed_job_utils.terminate_cluster(cluster_name,
+                                                        _logger=job_logger)
+                    status = core.status(cluster_names=[cluster_name],
+                                         all_users=True)
+                    assert (len(status) == 0 or
+                            status[0]['status'] == sky.ClusterStatus.STOPPED), (
+                                f'{cluster_name} is not down: {status}')
+                    job_logger.info(f'{cluster_name} is down')
+                else:
+                    cluster_name, job_id_on_pool_cluster = (
+                        managed_job_state.get_pool_submit_info(job_id))
+                    if cluster_name is not None:
+                        if job_id_on_pool_cluster is not None:
+                            core.cancel(cluster_name=cluster_name,
+                                        job_ids=[job_id_on_pool_cluster],
+                                        _try_cancel_if_cluster_is_init=True)
             except Exception as e:  # pylint: disable=broad-except
-                logger.warning(
-                    f'Failed to clean up file mount {file_mount}: {e}')
+                error = e
+                job_logger.warning(
+                    f'Failed to terminate cluster {cluster_name}: {e}')
+                # we continue to try cleaning up whatever else we can.
+            # Clean up Storages with persistent=False.
+            # TODO(zhwu): this assumes the specific backend.
+            backend = cloud_vm_ray_backend.CloudVmRayBackend()
+            # Need to re-construct storage object in the controller process
+            # because when SkyPilot API server machine sends the yaml config to
+            # the controller machine, only storage metadata is sent, not the
+            # storage object itself.
+            for storage in task.storage_mounts.values():
+                storage.construct()
+            try:
+                backend.teardown_ephemeral_storage(task)
+            except Exception as e:  # pylint: disable=broad-except
+                error = e
+                job_logger.warning(f'Failed to teardown ephemeral storage: {e}')
+                # we continue to try cleaning up whatever else we can.
 
+            # Clean up any files mounted from the local disk, such as two-hop
+            # file mounts.
+            for file_mount in (task.file_mounts or {}).values():
+                try:
+                    # For consolidation mode, there is no two-hop file mounts
+                    # and the file path here represents the real user data.
+                    # We skip the cleanup for consolidation mode.
+                    if (not data_utils.is_cloud_store_url(file_mount) and
+                            not managed_job_utils.is_consolidation_mode()):
+                        path = os.path.expanduser(file_mount)
+                        if os.path.isdir(path):
+                            shutil.rmtree(path)
+                        else:
+                            os.remove(path)
+                except Exception as e:  # pylint: disable=broad-except
+                    job_logger.warning(
+                        f'Failed to clean up file mount {file_mount}: {e}')
+
+            if error is not None:
+                raise error
 
-def start(job_id, dag_yaml, pool):
-    """Start the controller."""
-    controller_process = None
-    cancelling = False
-    task_id = None
-    try:
-        _handle_signal(job_id)
-        # TODO(suquark): In theory, we should make controller process a
-        #  daemon process so it will be killed after this process exits,
-        #  however daemon process cannot launch subprocesses, explained here:
-        #  https://docs.python.org/3/library/multiprocessing.html#multiprocessing.Process.daemon  # pylint: disable=line-too-long
-        #  So we can only enable daemon after we no longer need to
-        #  start daemon processes like Ray.
-        controller_process = multiprocessing.Process(target=_run_controller,
-                                                     args=(job_id, dag_yaml,
-                                                           pool))
-        controller_process.start()
-        while controller_process.is_alive():
-            _handle_signal(job_id)
-            time.sleep(1)
-    except exceptions.ManagedJobUserCancelledError:
         dag, _ = _get_dag_and_name(dag_yaml)
-        task_id, _ = managed_job_state.get_latest_task_id_status(job_id)
-        assert task_id is not None, job_id
-        logger.info(
-            f'Cancelling managed job, job_id: {job_id}, task_id: {task_id}')
-        managed_job_state.set_cancelling(
-            job_id=job_id,
-            callback_func=managed_job_utils.event_callback_func(
-                job_id=job_id, task_id=task_id, task=dag.tasks[task_id]))
-        cancelling = True
-    finally:
-        if controller_process is not None:
-            logger.info(f'Killing controller process {controller_process.pid}.')
-            # NOTE: it is ok to kill or join a killed process.
-            # Kill the controller process first; if its child process is
-            # killed first, then the controller process will raise errors.
-            # Kill any possible remaining children processes recursively.
-            subprocess_utils.kill_children_processes(
-                parent_pids=[controller_process.pid], force=True)
-            controller_process.join()
-            logger.info(f'Controller process {controller_process.pid} killed.')
-
-        logger.info(f'Cleaning up any cluster for job {job_id}.')
-        # NOTE: Originally, we send an interruption signal to the controller
-        # process and the controller process handles cleanup. However, we
-        # figure out the behavior differs from cloud to cloud
-        # (e.g., GCP ignores 'SIGINT'). A possible explanation is
-        # https://unix.stackexchange.com/questions/356408/strange-problem-with-trap-and-sigint
-        # But anyway, a clean solution is killing the controller process
-        # directly, and then cleanup the cluster job_state.
-        _cleanup(job_id, dag_yaml=dag_yaml, pool=pool)
-        logger.info(f'Cluster of managed job {job_id} has been cleaned up.')
-
-        if cancelling:
-            assert task_id is not None, job_id  # Since it's set with cancelling
-            managed_job_state.set_cancelled(
+        error = None
+        for task in dag.tasks:
+            # most things in this function are blocking
+            try:
+                await context_utils.to_thread(task_cleanup, task, job_id)
+            except Exception as e:  # pylint: disable=broad-except
+                error = e
+
+        if error is not None:
+            # we only raise the last error that occurred, but its fine to lose
+            # some data here.
+            raise error
+
+    async def run_job_loop(self,
+                           job_id: int,
+                           dag_yaml: str,
+                           job_logger: logging.Logger,
+                           log_file: str,
+                           env_file_path: Optional[str] = None,
+                           pool: Optional[str] = None):
+        """Background task that runs the job loop."""
+        # Replace os.environ with ContextualEnviron to enable per-job
+        # environment isolation. This allows each job to have its own
+        # environment variables without affecting other jobs or the main
+        # process.
+        context.initialize()
+        ctx = context.get()
+        ctx.redirect_log(log_file)  # type: ignore
+
+        # Load and apply environment variables from the job's environment file
+        if env_file_path and os.path.exists(env_file_path):
+            try:
+                # Load environment variables from the file
+                env_vars = dotenv.dotenv_values(env_file_path)
+                job_logger.info(f'Loading environment from {env_file_path}: '
+                                f'{list(env_vars.keys())}')
+
+                # Apply environment variables to the job's context
+                ctx = context.get()
+                if ctx is not None:
+                    for key, value in env_vars.items():
+                        if value is not None:
+                            ctx.override_envs({key: value})
+                            job_logger.debug(
+                                f'Set environment variable: {key}={value}')
+                else:
+                    job_logger.error(
+                        'Context is None, cannot set environment variables')
+            except Exception as e:  # pylint: disable=broad-except
+                job_logger.error(
+                    f'Failed to load environment file {env_file_path}: {e}')
+        elif env_file_path:
+            job_logger.error(f'Environment file not found: {env_file_path}')
+
+        cancelling = False
+        try:
+            job_logger.info(f'Starting job loop for {job_id}')
+
+            controller = JobsController(job_id, dag_yaml, job_logger,
+                                        self.starting, self._job_tasks_lock,
+                                        self._starting_signal, pool)
+
+            async with self._job_tasks_lock:
+                if job_id in self.job_tasks:
+                    job_logger.error(
+                        f'Job {job_id} already exists in job_tasks')
+                    raise ValueError(f'Job {job_id} already exists')
+
+                # Create the task and store it
+                # This function should return instantly and run the job loop in
+                # the background.
+                task = asyncio.create_task(controller.run())
+                self.job_tasks[job_id] = task
+            await task
+        except asyncio.CancelledError:
+            job_logger.info(f'Job {job_id} was cancelled')
+            dag, _ = _get_dag_and_name(dag_yaml)
+            task_id, _ = await (
+                managed_job_state.get_latest_task_id_status_async(job_id))
+            assert task_id is not None, job_id
+            job_logger.info(f'Cancelling managed job, job_id: {job_id}, '
+                            f'task_id: {task_id}')
+            await managed_job_state.set_cancelling_async(
                 job_id=job_id,
                 callback_func=managed_job_utils.event_callback_func(
                     job_id=job_id, task_id=task_id, task=dag.tasks[task_id]))
+            cancelling = True
+            raise
+        except Exception as e:
+            job_logger.error(f'Unexpected error in job loop for {job_id}: '
+                             f'{common_utils.format_exception(e)}')
+            raise
+        finally:
+            try:
+                await self._cleanup(job_id,
+                                    dag_yaml=dag_yaml,
+                                    job_logger=job_logger,
+                                    pool=pool)
+                job_logger.info(
+                    f'Cluster of managed job {job_id} has been cleaned up.')
+            except Exception as e:  # pylint: disable=broad-except
+                await managed_job_state.set_failed_async(
+                    job_id,
+                    task_id=None,
+                    failure_type=managed_job_state.ManagedJobStatus.
+                    FAILED_CONTROLLER,
+                    failure_reason=e,
+                    override_terminal=True)
+
+            if cancelling:
+                # Since it's set with cancelling
+                assert task_id is not None, job_id
+                await managed_job_state.set_cancelled_async(
+                    job_id=job_id,
+                    callback_func=managed_job_utils.event_callback_func(
+                        job_id=job_id, task_id=task_id,
+                        task=dag.tasks[task_id]))
+
+            # We should check job status after 'set_cancelled', otherwise
+            # the job status is not terminal.
+            job_status = await managed_job_state.get_status_async(job_id)
+            assert job_status is not None
+            # The job can be non-terminal if the controller exited abnormally,
+            # e.g. failed to launch cluster after reaching the MAX_RETRY.
+            if not job_status.is_terminal():
+                job_logger.info(f'Previous job status: {job_status.value}')
+                await managed_job_state.set_failed_async(
+                    job_id,
+                    task_id=None,
+                    failure_type=managed_job_state.ManagedJobStatus.
+                    FAILED_CONTROLLER,
+                    failure_reason=(
+                        'Unexpected error occurred. For details, '
+                        f'run: sky jobs logs --controller {job_id}'))
+
+            await scheduler.job_done_async(job_id)
+
+            async with self._job_tasks_lock:
+                try:
+                    # just in case we were cancelled or some other error
+                    # occurred during launch
+                    self.starting.remove(job_id)
+                    # its fine if we notify again, better to wake someone up
+                    # and have them go to sleep again, then have some stuck
+                    # sleeping.
+                    self._starting_signal.notify()
+                except KeyError:
+                    pass
+
+            # Remove the job from the job_tasks dictionary.
+            async with self._job_tasks_lock:
+                if job_id in self.job_tasks:
+                    del self.job_tasks[job_id]
+
+    async def start_job(
+        self,
+        job_id: int,
+        dag_yaml: str,
+        env_file_path: Optional[str] = None,
+        pool: Optional[str] = None,
+    ):
+        """Start a new job.
+
+        Args:
+            job_id: The ID of the job to start.
+            dag_yaml: Path to the YAML file containing the DAG definition.
+            env_file_path: Optional path to environment file for the job.
+        """
+        # Create a job-specific logger
+        log_dir = os.path.expanduser(jobs_constants.JOBS_CONTROLLER_LOGS_DIR)
+        os.makedirs(log_dir, exist_ok=True)
+        log_file = os.path.join(log_dir, f'{job_id}.log')
+
+        job_logger = logging.getLogger(f'sky.jobs.{job_id}')
+        job_logger.setLevel(logging.DEBUG)
+
+        # Create file handler
+        file_handler = logging.FileHandler(log_file)
+        file_handler.setLevel(logging.DEBUG)
+
+        # Use Sky's standard formatter
+        file_handler.setFormatter(sky_logging.FORMATTER)
 
-        # We should check job status after 'set_cancelled', otherwise
-        # the job status is not terminal.
-        job_status = managed_job_state.get_status(job_id)
-        assert job_status is not None
-        # The job can be non-terminal if the controller exited abnormally,
-        # e.g. failed to launch cluster after reaching the MAX_RETRY.
-        if not job_status.is_terminal():
-            logger.info(f'Previous job status: {job_status.value}')
-            managed_job_state.set_failed(
-                job_id,
-                task_id=None,
-                failure_type=managed_job_state.ManagedJobStatus.
-                FAILED_CONTROLLER,
-                failure_reason=('Unexpected error occurred. For details, '
-                                f'run: sky jobs logs --controller {job_id}'))
-
-        scheduler.job_done(job_id)
+        # Add the handler to the logger
+        job_logger.addHandler(file_handler)
+
+        # Prevent log propagation to avoid duplicate logs
+        job_logger.propagate = False
+
+        job_logger.info(f'Starting job {job_id} with dag_yaml={dag_yaml}, '
+                        f'env_file_path={env_file_path}')
+
+        async with self._job_tasks_lock:
+            self.starting.add(job_id)
+        await create_background_task(
+            self.run_job_loop(job_id, dag_yaml, job_logger, log_file,
+                              env_file_path, pool))
+
+        job_logger.info(f'Job {job_id} started successfully')
+
+    async def cancel_job(self):
+        """Cancel an existing job."""
+        while True:
+            cancels = os.listdir(jobs_constants.CONSOLIDATED_SIGNAL_PATH)
+            for cancel in cancels:
+                async with self._job_tasks_lock:
+                    job_id = int(cancel)
+                    if job_id in self.job_tasks:
+                        logger.info(f'Cancelling job {job_id}')
+
+                        task = self.job_tasks[job_id]
+
+                        # Run the cancellation in the background, so we can
+                        # return immediately.
+                        task.cancel()
+                        logger.info(f'Job {job_id} cancelled successfully')
+
+                        os.remove(f'{jobs_constants.CONSOLIDATED_SIGNAL_PATH}/'
+                                  f'{job_id}')
+            await asyncio.sleep(15)
+
+    async def monitor_loop(self):
+        """Monitor the job loop."""
+        logger.info(f'Starting monitor loop for pid {os.getpid()}...')
+
+        while True:
+            async with self._job_tasks_lock:
+                running_tasks = [
+                    task for task in self.job_tasks.values() if not task.done()
+                ]
+
+            async with self._job_tasks_lock:
+                starting_count = len(self.starting)
+
+            if starting_count >= scheduler.LAUNCHES_PER_WORKER:
+                # launching a job takes around 1 minute, so lets wait half that
+                # time
+                await asyncio.sleep(30)
+                continue
+
+            if len(running_tasks) >= scheduler.JOBS_PER_WORKER:
+                await asyncio.sleep(60)
+                continue
+
+            # Check if there are any jobs that are waiting to launch
+            try:
+                waiting_job = await managed_job_state.get_waiting_job_async(
+                    pid=-os.getpid())
+            except Exception as e:  # pylint: disable=broad-except
+                logger.error(f'Failed to get waiting job: {e}')
+                await asyncio.sleep(5)
+                continue
+
+            if waiting_job is None:
+                await asyncio.sleep(10)
+                continue
+
+            job_id = waiting_job['job_id']
+            dag_yaml_path = waiting_job['dag_yaml_path']
+            env_file_path = waiting_job.get('env_file_path')
+            pool = waiting_job.get('pool', None)
+
+            cancels = os.listdir(jobs_constants.CONSOLIDATED_SIGNAL_PATH)
+            if str(job_id) in cancels:
+                status = await managed_job_state.get_status_async(job_id)
+                if status == managed_job_state.ManagedJobStatus.PENDING:
+                    logger.info(f'Job {job_id} cancelled')
+                    os.remove(f'{jobs_constants.CONSOLIDATED_SIGNAL_PATH}/'
+                              f'{job_id}')
+                    await managed_job_state.set_cancelling_async(
+                        job_id=job_id,
+                        callback_func=managed_job_utils.event_callback_func(
+                            job_id=job_id, task_id=None, task=None))
+                    await managed_job_state.set_cancelled_async(
+                        job_id=job_id,
+                        callback_func=managed_job_utils.event_callback_func(
+                            job_id=job_id, task_id=None, task=None))
+                    continue
+
+            await self.start_job(job_id, dag_yaml_path, env_file_path, pool)
+
+
+async def main():
+    context_utils.hijack_sys_attrs()
+
+    controller = Controller()
+
+    # Will happen multiple times, who cares though
+    os.makedirs(jobs_constants.CONSOLIDATED_SIGNAL_PATH, exist_ok=True)
+
+    # Increase number of files we can open
+    soft = None
+    try:
+        soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE)
+        resource.setrlimit(resource.RLIMIT_NOFILE, (hard, hard))
+    except OSError as e:
+        logger.warning(f'Failed to increase number of files we can open: {e}\n'
+                       f'Current soft limit: {soft}, hard limit: {hard}')
+
+    # Will loop forever, do it in the background
+    cancel_job_task = asyncio.create_task(controller.cancel_job())
+    monitor_loop_task = asyncio.create_task(controller.monitor_loop())
+
+    try:
+        await asyncio.gather(cancel_job_task, monitor_loop_task)
+    except Exception as e:  # pylint: disable=broad-except
+        logger.error(f'Controller server crashed: {e}')
+        sys.exit(1)
 
 
 if __name__ == '__main__':
-    parser = argparse.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.')
-    parser.add_argument('--pool',
-                        required=False,
-                        default=None,
-                        type=str,
-                        help='The pool to use for the controller job.')
-    args = parser.parse_args()
-    # We start process with 'spawn', because 'fork' could result in weird
-    # behaviors; 'spawn' is also cross-platform.
-    multiprocessing.set_start_method('spawn', force=True)
-    start(args.job_id, args.dag_yaml, args.pool)
+    asyncio.run(main())