skypilot-nightly 1.0.0.dev20250509__py3-none-any.whl → 1.0.0.dev20251107__py3-none-any.whl

sky/jobs/state.py

@@ -1,28 +1,59 @@
 """The database for managed jobs status."""
 # TODO(zhwu): maybe use file based status instead of database, so
 # that we can easily switch to a s3-based storage.
+import asyncio
 import enum
+import functools
+import ipaddress
 import json
-import pathlib
 import sqlite3
+import threading
 import time
 import typing
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple, Union
+import urllib.parse
 
 import colorama
+import sqlalchemy
+from sqlalchemy import exc as sqlalchemy_exc
+from sqlalchemy import orm
+from sqlalchemy.dialects import postgresql
+from sqlalchemy.dialects import sqlite
+from sqlalchemy.ext import asyncio as sql_async
+from sqlalchemy.ext import declarative
 
 from sky import exceptions
 from sky import sky_logging
+from sky import skypilot_config
+from sky.adaptors import common as adaptors_common
+from sky.skylet import constants
 from sky.utils import common_utils
-from sky.utils import db_utils
+from sky.utils import context_utils
+from sky.utils.db import db_utils
+from sky.utils.db import migration_utils
 
 if typing.TYPE_CHECKING:
-    import sky
+    from sqlalchemy.engine import row
 
-CallbackType = Callable[[str], None]
+    from sky.schemas.generated import managed_jobsv1_pb2
+else:
+    managed_jobsv1_pb2 = adaptors_common.LazyImport(
+        'sky.schemas.generated.managed_jobsv1_pb2')
+
+# Separate callback types for sync and async contexts
+SyncCallbackType = Callable[[str], None]
+AsyncCallbackType = Callable[[str], Awaitable[Any]]
+CallbackType = Union[SyncCallbackType, AsyncCallbackType]
 
 logger = sky_logging.init_logger(__name__)
 
+_SQLALCHEMY_ENGINE: Optional[sqlalchemy.engine.Engine] = None
+_SQLALCHEMY_ENGINE_ASYNC: Optional[sql_async.AsyncEngine] = None
+_SQLALCHEMY_ENGINE_LOCK = threading.Lock()
+
+_DB_RETRY_TIMES = 30
+
+Base = declarative.declarative_base()
 
 # === Database schema ===
 # `spot` table contains all the finest-grained tasks, including all the
@@ -35,122 +66,248 @@ logger = sky_logging.init_logger(__name__)
 # identifier/primary key for all the tasks. We will use `spot_job_id`
 # to identify the job.
 # TODO(zhwu): schema migration may be needed.
-def create_table(cursor, conn):
+
+spot_table = sqlalchemy.Table(
+    'spot',
+    Base.metadata,
+    sqlalchemy.Column('job_id',
+                      sqlalchemy.Integer,
+                      primary_key=True,
+                      autoincrement=True),
+    sqlalchemy.Column('job_name', sqlalchemy.Text),
+    sqlalchemy.Column('resources', sqlalchemy.Text),
+    sqlalchemy.Column('submitted_at', sqlalchemy.Float),
+    sqlalchemy.Column('status', sqlalchemy.Text),
+    sqlalchemy.Column('run_timestamp', sqlalchemy.Text),
+    sqlalchemy.Column('start_at', sqlalchemy.Float, server_default=None),
+    sqlalchemy.Column('end_at', sqlalchemy.Float, server_default=None),
+    sqlalchemy.Column('last_recovered_at',
+                      sqlalchemy.Float,
+                      server_default='-1'),
+    sqlalchemy.Column('recovery_count', sqlalchemy.Integer, server_default='0'),
+    sqlalchemy.Column('job_duration', sqlalchemy.Float, server_default='0'),
+    sqlalchemy.Column('failure_reason', sqlalchemy.Text),
+    sqlalchemy.Column('spot_job_id', sqlalchemy.Integer, index=True),
+    sqlalchemy.Column('task_id', sqlalchemy.Integer, server_default='0'),
+    sqlalchemy.Column('task_name', sqlalchemy.Text),
+    sqlalchemy.Column('specs', sqlalchemy.Text),
+    sqlalchemy.Column('local_log_file', sqlalchemy.Text, server_default=None),
+    sqlalchemy.Column('metadata', sqlalchemy.Text, server_default='{}'),
+    sqlalchemy.Column('logs_cleaned_at', sqlalchemy.Float, server_default=None),
+)
+
+job_info_table = sqlalchemy.Table(
+    'job_info',
+    Base.metadata,
+    sqlalchemy.Column('spot_job_id',
+                      sqlalchemy.Integer,
+                      primary_key=True,
+                      autoincrement=True),
+    sqlalchemy.Column('name', sqlalchemy.Text),
+    sqlalchemy.Column('schedule_state', sqlalchemy.Text),
+    sqlalchemy.Column('controller_pid', sqlalchemy.Integer,
+                      server_default=None),
+    sqlalchemy.Column('dag_yaml_path', sqlalchemy.Text),
+    sqlalchemy.Column('env_file_path', sqlalchemy.Text),
+    sqlalchemy.Column('dag_yaml_content', sqlalchemy.Text, server_default=None),
+    sqlalchemy.Column('env_file_content', sqlalchemy.Text, server_default=None),
+    sqlalchemy.Column('user_hash', sqlalchemy.Text),
+    sqlalchemy.Column('workspace', sqlalchemy.Text, server_default=None),
+    sqlalchemy.Column('priority',
+                      sqlalchemy.Integer,
+                      server_default=str(constants.DEFAULT_PRIORITY)),
+    sqlalchemy.Column('entrypoint', sqlalchemy.Text, server_default=None),
+    sqlalchemy.Column('original_user_yaml_path',
+                      sqlalchemy.Text,
+                      server_default=None),
+    sqlalchemy.Column('original_user_yaml_content',
+                      sqlalchemy.Text,
+                      server_default=None),
+    sqlalchemy.Column('pool', sqlalchemy.Text, server_default=None),
+    sqlalchemy.Column('current_cluster_name',
+                      sqlalchemy.Text,
+                      server_default=None),
+    sqlalchemy.Column('job_id_on_pool_cluster',
+                      sqlalchemy.Integer,
+                      server_default=None),
+    sqlalchemy.Column('pool_hash', sqlalchemy.Text, server_default=None),
+    sqlalchemy.Column('controller_logs_cleaned_at',
+                      sqlalchemy.Float,
+                      server_default=None),
+)
+
+ha_recovery_script_table = sqlalchemy.Table(
+    'ha_recovery_script',
+    Base.metadata,
+    sqlalchemy.Column('job_id', sqlalchemy.Integer, primary_key=True),
+    sqlalchemy.Column('script', sqlalchemy.Text),
+)
+
+
+def create_table(engine: sqlalchemy.engine.Engine):
     # Enable WAL mode to avoid locking issues.
     # See: issue #3863, #1441 and PR #1509
     # https://github.com/microsoft/WSL/issues/2395
     # TODO(romilb): We do not enable WAL for WSL because of known issue in WSL.
     #  This may cause the database locked problem from WSL issue #1441.
-    if not common_utils.is_wsl():
+    if (engine.dialect.name == db_utils.SQLAlchemyDialect.SQLITE.value and
+            not common_utils.is_wsl()):
         try:
-            cursor.execute('PRAGMA journal_mode=WAL')
-        except sqlite3.OperationalError as e:
+            with orm.Session(engine) as session:
+                session.execute(sqlalchemy.text('PRAGMA journal_mode=WAL'))
+                session.execute(sqlalchemy.text('PRAGMA synchronous=1'))
+                session.commit()
+        except sqlalchemy_exc.OperationalError as e:
             if 'database is locked' not in str(e):
                 raise
             # If the database is locked, it is OK to continue, as the WAL mode
             # is not critical and is likely to be enabled by other processes.
 
-    cursor.execute("""\
-        CREATE TABLE IF NOT EXISTS spot (
-        job_id INTEGER PRIMARY KEY AUTOINCREMENT,
-        job_name TEXT,
-        resources TEXT,
-        submitted_at FLOAT,
-        status TEXT,
-        run_timestamp TEXT CANDIDATE KEY,
-        start_at FLOAT DEFAULT NULL,
-        end_at FLOAT DEFAULT NULL,
-        last_recovered_at FLOAT DEFAULT -1,
-        recovery_count INTEGER DEFAULT 0,
-        job_duration FLOAT DEFAULT 0,
-        failure_reason TEXT,
-        spot_job_id INTEGER,
-        task_id INTEGER DEFAULT 0,
-        task_name TEXT,
-        specs TEXT,
-        local_log_file TEXT DEFAULT NULL)""")
-    conn.commit()
-
-    db_utils.add_column_to_table(cursor, conn, 'spot', 'failure_reason', 'TEXT')
-    # Create a new column `spot_job_id`, which is the same for tasks of the
-    # same managed job.
-    # The original `job_id` no longer has an actual meaning, but only a legacy
-    # identifier for all tasks in database.
-    db_utils.add_column_to_table(cursor,
-                                 conn,
-                                 'spot',
-                                 'spot_job_id',
-                                 'INTEGER',
-                                 copy_from='job_id')
-    db_utils.add_column_to_table(cursor,
-                                 conn,
-                                 'spot',
-                                 'task_id',
-                                 'INTEGER DEFAULT 0',
-                                 value_to_replace_existing_entries=0)
-    db_utils.add_column_to_table(cursor,
-                                 conn,
-                                 'spot',
-                                 'task_name',
-                                 'TEXT',
-                                 copy_from='job_name')
-
-    # Specs is some useful information about the task, e.g., the
-    # max_restarts_on_errors value. It is stored in JSON format.
-    db_utils.add_column_to_table(cursor,
-                                 conn,
-                                 'spot',
-                                 'specs',
-                                 'TEXT',
-                                 value_to_replace_existing_entries=json.dumps({
-                                     'max_restarts_on_errors': 0,
-                                 }))
-    db_utils.add_column_to_table(cursor, conn, 'spot', 'local_log_file',
-                                 'TEXT DEFAULT NULL')
-
-    # `job_info` contains the mapping from job_id to the job_name, as well as
-    # information used by the scheduler.
-    cursor.execute("""\
-        CREATE TABLE IF NOT EXISTS job_info (
-        spot_job_id INTEGER PRIMARY KEY AUTOINCREMENT,
-        name TEXT,
-        schedule_state TEXT,
-        controller_pid INTEGER DEFAULT NULL,
-        dag_yaml_path TEXT,
-        env_file_path TEXT,
-        user_hash TEXT)""")
-
-    db_utils.add_column_to_table(cursor, conn, 'job_info', 'schedule_state',
-                                 'TEXT')
-
-    db_utils.add_column_to_table(cursor, conn, 'job_info', 'controller_pid',
-                                 'INTEGER DEFAULT NULL')
-
-    db_utils.add_column_to_table(cursor, conn, 'job_info', 'dag_yaml_path',
-                                 'TEXT')
-
-    db_utils.add_column_to_table(cursor, conn, 'job_info', 'env_file_path',
-                                 'TEXT')
-
-    db_utils.add_column_to_table(cursor, conn, 'job_info', 'user_hash', 'TEXT')
-
-    conn.commit()
-
-
-# Module-level connection/cursor; thread-safe as the module is only imported
-# once.
-def _get_db_path() -> str:
-    """Workaround to collapse multi-step Path ops for type checker.
-    Ensures _DB_PATH is str, avoiding Union[Path, str] inference.
-    """
-    path = pathlib.Path('~/.sky/spot_jobs.db')
-    path = path.expanduser().absolute()
-    path.parents[0].mkdir(parents=True, exist_ok=True)
-    return str(path)
+    migration_utils.safe_alembic_upgrade(engine,
+                                         migration_utils.SPOT_JOBS_DB_NAME,
+                                         migration_utils.SPOT_JOBS_VERSION)
+
 
+def force_no_postgres() -> bool:
+    """Force no postgres.
+
+    If the db is localhost on the api server, and we are not in consolidation
+    mode, we must force using sqlite and not using the api server on the jobs
+    controller.
+    """
+    conn_string = skypilot_config.get_nested(('db',), None)
+
+    if conn_string:
+        parsed = urllib.parse.urlparse(conn_string)
+        # it freezes if we use the normal get_consolidation_mode function
+        consolidation_mode = skypilot_config.get_nested(
+            ('jobs', 'controller', 'consolidation_mode'), default_value=False)
+        if ((parsed.hostname == 'localhost' or
+             ipaddress.ip_address(parsed.hostname).is_loopback) and
+                not consolidation_mode):
+            return True
+    return False
+
+
+def initialize_and_get_db_async() -> sql_async.AsyncEngine:
+    global _SQLALCHEMY_ENGINE_ASYNC
+    if _SQLALCHEMY_ENGINE_ASYNC is not None:
+        return _SQLALCHEMY_ENGINE_ASYNC
+    with _SQLALCHEMY_ENGINE_LOCK:
+        if _SQLALCHEMY_ENGINE_ASYNC is not None:
+            return _SQLALCHEMY_ENGINE_ASYNC
+
+        _SQLALCHEMY_ENGINE_ASYNC = db_utils.get_engine('spot_jobs',
+                                                       async_engine=True)
+
+    # to create the table in case an async function gets called first
+    initialize_and_get_db()
+    return _SQLALCHEMY_ENGINE_ASYNC
+
+
+# We wrap the sqlalchemy engine initialization in a thread
+# lock to ensure that multiple threads do not initialize the
+# engine which could result in a rare race condition where
+# a session has already been created with _SQLALCHEMY_ENGINE = e1,
+# and then another thread overwrites _SQLALCHEMY_ENGINE = e2
+# which could result in e1 being garbage collected unexpectedly.
+def initialize_and_get_db() -> sqlalchemy.engine.Engine:
+    global _SQLALCHEMY_ENGINE
+    if _SQLALCHEMY_ENGINE is not None:
+        return _SQLALCHEMY_ENGINE
+
+    with _SQLALCHEMY_ENGINE_LOCK:
+        if _SQLALCHEMY_ENGINE is not None:
+            return _SQLALCHEMY_ENGINE
+        # get an engine to the db
+        engine = db_utils.get_engine('spot_jobs')
+
+        # run migrations if needed
+        create_table(engine)
+
+        # return engine
+        _SQLALCHEMY_ENGINE = engine
+        return _SQLALCHEMY_ENGINE
+
+
+def _init_db_async(func):
+    """Initialize the async database. Add backoff to the function call."""
+
+    @functools.wraps(func)
+    async def wrapper(*args, **kwargs):
+        if _SQLALCHEMY_ENGINE_ASYNC is None:
+            # this may happen multiple times since there is no locking
+            # here but thats fine, this is just a short circuit for the
+            # common case.
+            await context_utils.to_thread(initialize_and_get_db_async)
+
+        backoff = common_utils.Backoff(initial_backoff=1, max_backoff_factor=5)
+        last_exc = None
+        for _ in range(_DB_RETRY_TIMES):
+            try:
+                return await func(*args, **kwargs)
+            except (sqlalchemy_exc.OperationalError,
+                    asyncio.exceptions.TimeoutError, OSError,
+                    sqlalchemy_exc.TimeoutError, sqlite3.OperationalError,
+                    sqlalchemy_exc.InterfaceError, sqlite3.InterfaceError) as e:
+                last_exc = e
+            logger.debug(f'DB error: {last_exc}')
+            await asyncio.sleep(backoff.current_backoff())
+        assert last_exc is not None
+        raise last_exc
+
+    return wrapper
+
+
+def _init_db(func):
+    """Initialize the database. Add backoff to the function call."""
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        if _SQLALCHEMY_ENGINE is None:
+            # this may happen multiple times since there is no locking
+            # here but thats fine, this is just a short circuit for the
+            # common case.
+            initialize_and_get_db()
+
+        backoff = common_utils.Backoff(initial_backoff=1, max_backoff_factor=10)
+        last_exc = None
+        for _ in range(_DB_RETRY_TIMES):
+            try:
+                return func(*args, **kwargs)
+            except (sqlalchemy_exc.OperationalError,
+                    asyncio.exceptions.TimeoutError, OSError,
+                    sqlalchemy_exc.TimeoutError, sqlite3.OperationalError,
+                    sqlalchemy_exc.InterfaceError, sqlite3.InterfaceError) as e:
+                last_exc = e
+            logger.debug(f'DB error: {last_exc}')
+            time.sleep(backoff.current_backoff())
+        assert last_exc is not None
+        raise last_exc
+
+    return wrapper
+
+
+async def _describe_task_transition_failure(session: sql_async.AsyncSession,
+                                            job_id: int, task_id: int) -> str:
+    """Return a human-readable description when a task transition fails."""
+    details = 'Couldn\'t fetch the task details.'
+    try:
+        debug_result = await session.execute(
+            sqlalchemy.select(spot_table.c.status, spot_table.c.end_at).where(
+                sqlalchemy.and_(spot_table.c.spot_job_id == job_id,
+                                spot_table.c.task_id == task_id)))
+        rows = debug_result.mappings().all()
+        details = (f'{len(rows)} rows matched job {job_id} and task '
+                   f'{task_id}.')
+        for row in rows:
+            status = row['status']
+            end_at = row['end_at']
+            details += f' Status: {status}, End time: {end_at}.'
+    except Exception as exc:  # pylint: disable=broad-except
+        details += f' Error fetching task details: {exc}'
+    return details
 
-_DB_PATH = _get_db_path()
-db_utils.SQLiteConn(_DB_PATH, create_table)
 
 # job_duration is the time a job actually runs (including the
 # setup duration) before last_recover, excluding the provision
@@ -164,33 +321,52 @@ db_utils.SQLiteConn(_DB_PATH, create_table)
 # e.g., via sky jobs queue. These may not correspond to actual
 # column names in the DB and it corresponds to the combined view
 # by joining the spot and job_info tables.
-columns = [
-    '_job_id',
-    '_task_name',
-    'resources',
-    'submitted_at',
-    'status',
-    'run_timestamp',
-    'start_at',
-    'end_at',
-    'last_recovered_at',
-    'recovery_count',
-    'job_duration',
-    'failure_reason',
-    'job_id',
-    'task_id',
-    'task_name',
-    'specs',
-    'local_log_file',
-    # columns from the job_info table
-    '_job_info_job_id',  # This should be the same as job_id
-    'job_name',
-    'schedule_state',
-    'controller_pid',
-    'dag_yaml_path',
-    'env_file_path',
-    'user_hash',
-]
+def _get_jobs_dict(r: 'row.RowMapping') -> Dict[str, Any]:
+    # WARNING: If you update these you may also need to update GetJobTable in
+    # the skylet ManagedJobsServiceImpl.
+    return {
+        '_job_id': r.get('job_id'),  # from spot table
+        '_task_name': r.get('job_name'),  # deprecated, from spot table
+        'resources': r.get('resources'),
+        'submitted_at': r.get('submitted_at'),
+        'status': r.get('status'),
+        'run_timestamp': r.get('run_timestamp'),
+        'start_at': r.get('start_at'),
+        'end_at': r.get('end_at'),
+        'last_recovered_at': r.get('last_recovered_at'),
+        'recovery_count': r.get('recovery_count'),
+        'job_duration': r.get('job_duration'),
+        'failure_reason': r.get('failure_reason'),
+        'job_id': r.get(spot_table.c.spot_job_id
+                       ),  # ambiguous, use table.column
+        'task_id': r.get('task_id'),
+        'task_name': r.get('task_name'),
+        'specs': r.get('specs'),
+        'local_log_file': r.get('local_log_file'),
+        'metadata': r.get('metadata'),
+        # columns from job_info table (some may be None for legacy jobs)
+        '_job_info_job_id': r.get(job_info_table.c.spot_job_id
+                                 ),  # ambiguous, use table.column
+        'job_name': r.get('name'),  # from job_info table
+        'schedule_state': r.get('schedule_state'),
+        'controller_pid': r.get('controller_pid'),
+        # the _path columns are for backwards compatibility, use the _content
+        # columns instead
+        'dag_yaml_path': r.get('dag_yaml_path'),
+        'env_file_path': r.get('env_file_path'),
+        'dag_yaml_content': r.get('dag_yaml_content'),
+        'env_file_content': r.get('env_file_content'),
+        'user_hash': r.get('user_hash'),
+        'workspace': r.get('workspace'),
+        'priority': r.get('priority'),
+        'entrypoint': r.get('entrypoint'),
+        'original_user_yaml_path': r.get('original_user_yaml_path'),
+        'original_user_yaml_content': r.get('original_user_yaml_content'),
+        'pool': r.get('pool'),
+        'current_cluster_name': r.get('current_cluster_name'),
+        'job_id_on_pool_cluster': r.get('job_id_on_pool_cluster'),
+        'pool_hash': r.get('pool_hash'),
+    }
 
 
 class ManagedJobStatus(enum.Enum):
@@ -206,7 +382,7 @@ class ManagedJobStatus(enum.Enum):
     reset to INIT or SETTING_UP multiple times (depending on the preemptions).
 
     However, a managed job only has one ManagedJobStatus on the jobs controller.
-        ManagedJobStatus = [PENDING, SUBMITTED, STARTING, RUNNING, ...]
+        ManagedJobStatus = [PENDING, STARTING, RUNNING, ...]
     Mapping from JobStatus to ManagedJobStatus:
         INIT            ->  STARTING/RECOVERING
         SETTING_UP      ->  RUNNING
@@ -226,10 +402,14 @@ class ManagedJobStatus(enum.Enum):
     # PENDING: Waiting for the jobs controller to have a slot to run the
     # controller process.
     PENDING = 'PENDING'
+    # SUBMITTED: This state used to be briefly set before immediately changing
+    # to STARTING. Its use was removed in #5682. We keep it for backwards
+    # compatibility, so we can still parse old jobs databases that may have jobs
+    # in this state.
+    # TODO(cooperc): remove this in v0.12.0
+    DEPRECATED_SUBMITTED = 'SUBMITTED'
     # The submitted_at timestamp of the managed job in the 'spot' table will be
     # set to the time when the job controller begins running.
-    # SUBMITTED: The jobs controller starts the controller process.
-    SUBMITTED = 'SUBMITTED'
     # STARTING: The controller process is launching the cluster for the managed
     # job.
     STARTING = 'STARTING'
@@ -302,10 +482,88 @@ class ManagedJobStatus(enum.Enum):
             cls.FAILED_NO_RESOURCE, cls.FAILED_CONTROLLER
         ]
 
+    @classmethod
+    def processing_statuses(cls) -> List['ManagedJobStatus']:
+        # Any status that is not terminal and is not CANCELLING.
+        return [
+            cls.PENDING,
+            cls.STARTING,
+            cls.RUNNING,
+            cls.RECOVERING,
+        ]
+
+    @classmethod
+    def from_protobuf(
+        cls, protobuf_value: 'managed_jobsv1_pb2.ManagedJobStatus'
+    ) -> Optional['ManagedJobStatus']:
+        """Convert protobuf ManagedJobStatus enum to Python enum value."""
+        protobuf_to_enum = {
+            managed_jobsv1_pb2.MANAGED_JOB_STATUS_UNSPECIFIED: None,
+            managed_jobsv1_pb2.MANAGED_JOB_STATUS_PENDING: cls.PENDING,
+            managed_jobsv1_pb2.MANAGED_JOB_STATUS_SUBMITTED:
+                cls.DEPRECATED_SUBMITTED,
+            managed_jobsv1_pb2.MANAGED_JOB_STATUS_STARTING: cls.STARTING,
+            managed_jobsv1_pb2.MANAGED_JOB_STATUS_RUNNING: cls.RUNNING,
+            managed_jobsv1_pb2.MANAGED_JOB_STATUS_SUCCEEDED: cls.SUCCEEDED,
+            managed_jobsv1_pb2.MANAGED_JOB_STATUS_FAILED: cls.FAILED,
+            managed_jobsv1_pb2.MANAGED_JOB_STATUS_FAILED_CONTROLLER:
+                cls.FAILED_CONTROLLER,
+            managed_jobsv1_pb2.MANAGED_JOB_STATUS_FAILED_SETUP:
+                cls.FAILED_SETUP,
+            managed_jobsv1_pb2.MANAGED_JOB_STATUS_CANCELLED: cls.CANCELLED,
+            managed_jobsv1_pb2.MANAGED_JOB_STATUS_RECOVERING: cls.RECOVERING,
+            managed_jobsv1_pb2.MANAGED_JOB_STATUS_CANCELLING: cls.CANCELLING,
+            managed_jobsv1_pb2.MANAGED_JOB_STATUS_FAILED_PRECHECKS:
+                cls.FAILED_PRECHECKS,
+            managed_jobsv1_pb2.MANAGED_JOB_STATUS_FAILED_NO_RESOURCE:
+                cls.FAILED_NO_RESOURCE,
+        }
+
+        if protobuf_value not in protobuf_to_enum:
+            raise ValueError(
+                f'Unknown protobuf ManagedJobStatus value: {protobuf_value}')
+
+        return protobuf_to_enum[protobuf_value]
+
+    def to_protobuf(self) -> 'managed_jobsv1_pb2.ManagedJobStatus':
+        """Convert this Python enum value to protobuf enum value."""
+        enum_to_protobuf = {
+            ManagedJobStatus.PENDING:
+                managed_jobsv1_pb2.MANAGED_JOB_STATUS_PENDING,
+            ManagedJobStatus.DEPRECATED_SUBMITTED:
+                managed_jobsv1_pb2.MANAGED_JOB_STATUS_SUBMITTED,
+            ManagedJobStatus.STARTING:
+                managed_jobsv1_pb2.MANAGED_JOB_STATUS_STARTING,
+            ManagedJobStatus.RUNNING:
+                managed_jobsv1_pb2.MANAGED_JOB_STATUS_RUNNING,
+            ManagedJobStatus.SUCCEEDED:
+                managed_jobsv1_pb2.MANAGED_JOB_STATUS_SUCCEEDED,
+            ManagedJobStatus.FAILED:
+                managed_jobsv1_pb2.MANAGED_JOB_STATUS_FAILED,
+            ManagedJobStatus.FAILED_CONTROLLER:
+                managed_jobsv1_pb2.MANAGED_JOB_STATUS_FAILED_CONTROLLER,
+            ManagedJobStatus.FAILED_SETUP:
+                managed_jobsv1_pb2.MANAGED_JOB_STATUS_FAILED_SETUP,
+            ManagedJobStatus.CANCELLED:
+                managed_jobsv1_pb2.MANAGED_JOB_STATUS_CANCELLED,
+            ManagedJobStatus.RECOVERING:
+                managed_jobsv1_pb2.MANAGED_JOB_STATUS_RECOVERING,
+            ManagedJobStatus.CANCELLING:
+                managed_jobsv1_pb2.MANAGED_JOB_STATUS_CANCELLING,
+            ManagedJobStatus.FAILED_PRECHECKS:
+                managed_jobsv1_pb2.MANAGED_JOB_STATUS_FAILED_PRECHECKS,
+            ManagedJobStatus.FAILED_NO_RESOURCE:
+                managed_jobsv1_pb2.MANAGED_JOB_STATUS_FAILED_NO_RESOURCE,
+        }
+
+        if self not in enum_to_protobuf:
+            raise ValueError(f'Unknown ManagedJobStatus value: {self}')
+
+        return enum_to_protobuf[self]
+
 
 _SPOT_STATUS_TO_COLOR = {
     ManagedJobStatus.PENDING: colorama.Fore.BLUE,
-    ManagedJobStatus.SUBMITTED: colorama.Fore.BLUE,
     ManagedJobStatus.STARTING: colorama.Fore.BLUE,
     ManagedJobStatus.RUNNING: colorama.Fore.GREEN,
     ManagedJobStatus.RECOVERING: colorama.Fore.CYAN,
@@ -317,6 +575,8 @@ _SPOT_STATUS_TO_COLOR = {
     ManagedJobStatus.FAILED_CONTROLLER: colorama.Fore.RED,
     ManagedJobStatus.CANCELLING: colorama.Fore.YELLOW,
     ManagedJobStatus.CANCELLED: colorama.Fore.YELLOW,
+    # TODO(cooperc): backwards compatibility, remove this in v0.12.0
+    ManagedJobStatus.DEPRECATED_SUBMITTED: colorama.Fore.BLUE,
 }
 
 
@@ -333,8 +593,12 @@ class ManagedJobScheduleState(enum.Enum):
     - LAUNCHING -> ALIVE: The launch attempt was completed. It may have
       succeeded or failed. The job controller is not allowed to sky.launch again
       without transitioning to ALIVE_WAITING and then LAUNCHING.
+    - LAUNCHING -> ALIVE_BACKOFF: The launch failed to find resources, and is
+      in backoff waiting for resources.
     - ALIVE -> ALIVE_WAITING: The job controller wants to sky.launch again,
       either for recovery or to launch a subsequent task.
+    - ALIVE_BACKOFF -> ALIVE_WAITING: The backoff period has ended, and the job
+      controller wants to try to launch again.
     - ALIVE_WAITING -> LAUNCHING: The scheduler has determined that the job
       controller may launch again.
     - LAUNCHING, ALIVE, or ALIVE_WAITING -> DONE: The job controller is exiting
@@ -348,6 +612,7 @@ class ManagedJobScheduleState(enum.Enum):
     state or vice versa. (In fact, schedule state is defined on the job and
     status on the task.)
     - INACTIVE or WAITING should only be seen when a job is PENDING.
+    - ALIVE_BACKOFF should only be seen when a job is STARTING.
     - ALIVE_WAITING should only be seen when a job is RECOVERING, has multiple
       tasks, or needs to retry launching.
     - LAUNCHING and ALIVE can be seen in many different statuses.
@@ -359,6 +624,10 @@ class ManagedJobScheduleState(enum.Enum):
     # This job may have been created before scheduler was introduced in #4458.
     # This state is not used by scheduler but just for backward compatibility.
     # TODO(cooperc): remove this in v0.11.0
+    # TODO(luca): the only states we need are INACTIVE, WAITING, ALIVE, and
+    # DONE. ALIVE = old LAUNCHING + ALIVE + ALIVE_BACKOFF + ALIVE_WAITING and
+    # will represent jobs that are claimed by a controller. Delete the rest
+    # in v0.13.0
     INVALID = None
     # The job should be ignored by the scheduler.
     INACTIVE = 'INACTIVE'
@@ -373,194 +642,205 @@ class ManagedJobScheduleState(enum.Enum):
     # The job is running sky.launch, or soon will, using a limited number of
     # allowed launch slots.
     LAUNCHING = 'LAUNCHING'
+    # The job is alive, but is in backoff waiting for resources - a special case
+    # of ALIVE.
+    ALIVE_BACKOFF = 'ALIVE_BACKOFF'
     # The controller for the job is running, but it's not currently launching.
     ALIVE = 'ALIVE'
     # The job is in a terminal state. (Not necessarily SUCCEEDED.)
     DONE = 'DONE'
 
+    @classmethod
+    def from_protobuf(
+        cls, protobuf_value: 'managed_jobsv1_pb2.ManagedJobScheduleState'
+    ) -> Optional['ManagedJobScheduleState']:
+        """Convert protobuf ManagedJobScheduleState enum to Python enum value.
+        """
+        protobuf_to_enum = {
+            managed_jobsv1_pb2.MANAGED_JOB_SCHEDULE_STATE_UNSPECIFIED: None,
+            managed_jobsv1_pb2.MANAGED_JOB_SCHEDULE_STATE_INVALID: cls.INVALID,
+            managed_jobsv1_pb2.MANAGED_JOB_SCHEDULE_STATE_INACTIVE:
+                cls.INACTIVE,
+            managed_jobsv1_pb2.MANAGED_JOB_SCHEDULE_STATE_WAITING: cls.WAITING,
+            managed_jobsv1_pb2.MANAGED_JOB_SCHEDULE_STATE_ALIVE_WAITING:
+                cls.ALIVE_WAITING,
+            managed_jobsv1_pb2.MANAGED_JOB_SCHEDULE_STATE_LAUNCHING:
+                cls.LAUNCHING,
+            managed_jobsv1_pb2.MANAGED_JOB_SCHEDULE_STATE_ALIVE_BACKOFF:
+                cls.ALIVE_BACKOFF,
+            managed_jobsv1_pb2.MANAGED_JOB_SCHEDULE_STATE_ALIVE: cls.ALIVE,
+            managed_jobsv1_pb2.MANAGED_JOB_SCHEDULE_STATE_DONE: cls.DONE,
+        }
+
+        if protobuf_value not in protobuf_to_enum:
+            raise ValueError('Unknown protobuf ManagedJobScheduleState value: '
+                             f'{protobuf_value}')
+
+        return protobuf_to_enum[protobuf_value]
+
+    def to_protobuf(self) -> 'managed_jobsv1_pb2.ManagedJobScheduleState':
+        """Convert this Python enum value to protobuf enum value."""
+        enum_to_protobuf = {
+            ManagedJobScheduleState.INVALID:
+                managed_jobsv1_pb2.MANAGED_JOB_SCHEDULE_STATE_INVALID,
+            ManagedJobScheduleState.INACTIVE:
+                managed_jobsv1_pb2.MANAGED_JOB_SCHEDULE_STATE_INACTIVE,
+            ManagedJobScheduleState.WAITING:
+                managed_jobsv1_pb2.MANAGED_JOB_SCHEDULE_STATE_WAITING,
+            ManagedJobScheduleState.ALIVE_WAITING:
+                managed_jobsv1_pb2.MANAGED_JOB_SCHEDULE_STATE_ALIVE_WAITING,
+            ManagedJobScheduleState.LAUNCHING:
+                managed_jobsv1_pb2.MANAGED_JOB_SCHEDULE_STATE_LAUNCHING,
+            ManagedJobScheduleState.ALIVE_BACKOFF:
+                managed_jobsv1_pb2.MANAGED_JOB_SCHEDULE_STATE_ALIVE_BACKOFF,
+            ManagedJobScheduleState.ALIVE:
+                managed_jobsv1_pb2.MANAGED_JOB_SCHEDULE_STATE_ALIVE,
+            ManagedJobScheduleState.DONE:
+                managed_jobsv1_pb2.MANAGED_JOB_SCHEDULE_STATE_DONE,
+        }
+
+        if self not in enum_to_protobuf:
+            raise ValueError(f'Unknown ManagedJobScheduleState value: {self}')
+
+        return enum_to_protobuf[self]
+
 
 # === Status transition functions ===
-def set_job_info(job_id: int, name: str):
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        cursor.execute(
-            """\
-            INSERT INTO job_info
-            (spot_job_id, name, schedule_state)
-            VALUES (?, ?, ?)""",
-            (job_id, name, ManagedJobScheduleState.INACTIVE.value))
+@_init_db
+def set_job_info_without_job_id(name: str, workspace: str, entrypoint: str,
+                                pool: Optional[str], pool_hash: Optional[str],
+                                user_hash: Optional[str]) -> int:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        if (_SQLALCHEMY_ENGINE.dialect.name ==
+                db_utils.SQLAlchemyDialect.SQLITE.value):
+            insert_func = sqlite.insert
+        elif (_SQLALCHEMY_ENGINE.dialect.name ==
+              db_utils.SQLAlchemyDialect.POSTGRESQL.value):
+            insert_func = postgresql.insert
+        else:
+            raise ValueError('Unsupported database dialect')
+
+        insert_stmt = insert_func(job_info_table).values(
+            name=name,
+            schedule_state=ManagedJobScheduleState.INACTIVE.value,
+            workspace=workspace,
+            entrypoint=entrypoint,
+            pool=pool,
+            pool_hash=pool_hash,
+            user_hash=user_hash,
+        )
 
+        if (_SQLALCHEMY_ENGINE.dialect.name ==
+                db_utils.SQLAlchemyDialect.SQLITE.value):
+            result = session.execute(insert_stmt)
+            ret = result.lastrowid
+            session.commit()
+            return ret
+        elif (_SQLALCHEMY_ENGINE.dialect.name ==
+              db_utils.SQLAlchemyDialect.POSTGRESQL.value):
+            result = session.execute(
+                insert_stmt.returning(job_info_table.c.spot_job_id))
+            ret = result.scalar()
+            session.commit()
+            return ret
+        else:
+            raise ValueError('Unsupported database dialect')
 
-def set_pending(job_id: int, task_id: int, task_name: str, resources_str: str):
-    """Set the task to pending state."""
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        cursor.execute(
-            """\
-            INSERT INTO spot
-            (spot_job_id, task_id, task_name, resources, status)
-            VALUES (?, ?, ?, ?, ?)""",
-            (job_id, task_id, task_name, resources_str,
-             ManagedJobStatus.PENDING.value))
-
-
-def set_submitted(job_id: int, task_id: int, run_timestamp: str,
-                  submit_time: float, resources_str: str,
-                  specs: Dict[str, Union[str,
-                                         int]], callback_func: CallbackType):
-    """Set the task to submitted.
 
-    Args:
-        job_id: The managed job ID.
-        task_id: The task ID.
-        run_timestamp: The run_timestamp of the run. This will be used to
-        determine the log directory of the managed task.
-        submit_time: The time when the managed task is submitted.
-        resources_str: The resources string of the managed task.
-        specs: The specs of the managed task.
-        callback_func: The callback function.
+@_init_db
+def set_pending(
+    job_id: int,
+    task_id: int,
+    task_name: str,
+    resources_str: str,
+    metadata: str,
+):
+    """Set the task to pending state."""
+    assert _SQLALCHEMY_ENGINE is not None
+
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        session.execute(
+            sqlalchemy.insert(spot_table).values(
+                spot_job_id=job_id,
+                task_id=task_id,
+                task_name=task_name,
+                resources=resources_str,
+                metadata=metadata,
+                status=ManagedJobStatus.PENDING.value,
+            ))
+        session.commit()
+
+
+@_init_db_async
+async def set_backoff_pending_async(job_id: int, task_id: int):
+    """Set the task to PENDING state if it is in backoff.
+
+    This should only be used to transition from STARTING or RECOVERING back to
+    PENDING.
     """
-    # Use the timestamp in the `run_timestamp` ('sky-2022-10...'), to make
-    # the log directory and submission time align with each other, so as to
-    # make it easier to find them based on one of the values.
-    # Also, using the earlier timestamp should be closer to the term
-    # `submit_at`, which represents the time the managed task is submitted.
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        cursor.execute(
-            """\
-            UPDATE spot SET
-            resources=(?),
-            submitted_at=(?),
-            status=(?),
-            run_timestamp=(?),
-            specs=(?)
-            WHERE spot_job_id=(?) AND
-            task_id=(?) AND
-            status=(?) AND
-            end_at IS null""",
-            (resources_str, submit_time, ManagedJobStatus.SUBMITTED.value,
-             run_timestamp, json.dumps(specs), job_id, task_id,
-             ManagedJobStatus.PENDING.value))
-        if cursor.rowcount != 1:
-            raise exceptions.ManagedJobStatusError(
-                f'Failed to set the task to submitted. '
-                f'({cursor.rowcount} rows updated)')
-    callback_func('SUBMITTED')
-
-
-def set_starting(job_id: int, task_id: int, callback_func: CallbackType):
-    """Set the task to starting state."""
-    logger.info('Launching the spot cluster...')
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        cursor.execute(
-            """\
-            UPDATE spot SET status=(?)
-            WHERE spot_job_id=(?) AND
-            task_id=(?) AND
-            status=(?) AND
-            end_at IS null""", (ManagedJobStatus.STARTING.value, job_id,
-                                task_id, ManagedJobStatus.SUBMITTED.value))
-        if cursor.rowcount != 1:
-            raise exceptions.ManagedJobStatusError(
-                f'Failed to set the task to starting. '
-                f'({cursor.rowcount} rows updated)')
-    callback_func('STARTING')
-
-
-def set_started(job_id: int, task_id: int, start_time: float,
-                callback_func: CallbackType):
-    """Set the task to started state."""
-    logger.info('Job started.')
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        cursor.execute(
-            """\
-            UPDATE spot SET status=(?), start_at=(?), last_recovered_at=(?)
-            WHERE spot_job_id=(?) AND
-            task_id=(?) AND
-            status IN (?, ?) AND
-            end_at IS null""",
-            (
-                ManagedJobStatus.RUNNING.value,
-                start_time,
-                start_time,
-                job_id,
-                task_id,
-                ManagedJobStatus.STARTING.value,
-                # If the task is empty, we will jump straight from PENDING to
-                # RUNNING
-                ManagedJobStatus.PENDING.value,
-            ),
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        result = await session.execute(
+            sqlalchemy.update(spot_table).where(
+                sqlalchemy.and_(
+                    spot_table.c.spot_job_id == job_id,
+                    spot_table.c.task_id == task_id,
+                    spot_table.c.status.in_([
+                        ManagedJobStatus.STARTING.value,
+                        ManagedJobStatus.RECOVERING.value
+                    ]),
+                    spot_table.c.end_at.is_(None),
+                )).values({spot_table.c.status: ManagedJobStatus.PENDING.value})
         )
-        if cursor.rowcount != 1:
-            raise exceptions.ManagedJobStatusError(
-                f'Failed to set the task to started. '
-                f'({cursor.rowcount} rows updated)')
-    callback_func('STARTED')
-
-
-def set_recovering(job_id: int, task_id: int, callback_func: CallbackType):
-    """Set the task to recovering state, and update the job duration."""
-    logger.info('=== Recovering... ===')
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        cursor.execute(
-            """\
-                UPDATE spot SET
-                status=(?), job_duration=job_duration+(?)-last_recovered_at
-                WHERE spot_job_id=(?) AND
-                task_id=(?) AND
-                status=(?) AND
-                end_at IS null""",
-            (ManagedJobStatus.RECOVERING.value, time.time(), job_id, task_id,
-             ManagedJobStatus.RUNNING.value))
-        if cursor.rowcount != 1:
-            raise exceptions.ManagedJobStatusError(
-                f'Failed to set the task to recovering. '
-                f'({cursor.rowcount} rows updated)')
-    callback_func('RECOVERING')
-
-
-def set_recovered(job_id: int, task_id: int, recovered_time: float,
-                  callback_func: CallbackType):
-    """Set the task to recovered."""
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        cursor.execute(
-            """\
-            UPDATE spot SET
-            status=(?), last_recovered_at=(?), recovery_count=recovery_count+1
-            WHERE spot_job_id=(?) AND
-            task_id=(?) AND
-            status=(?) AND
-            end_at IS null""",
-            (ManagedJobStatus.RUNNING.value, recovered_time, job_id, task_id,
-             ManagedJobStatus.RECOVERING.value))
-        if cursor.rowcount != 1:
-            raise exceptions.ManagedJobStatusError(
-                f'Failed to set the task to recovered. '
-                f'({cursor.rowcount} rows updated)')
-    logger.info('==== Recovered. ====')
-    callback_func('RECOVERED')
-
-
-def set_succeeded(job_id: int, task_id: int, end_time: float,
-                  callback_func: CallbackType):
-    """Set the task to succeeded, if it is in a non-terminal state."""
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        cursor.execute(
-            """\
-            UPDATE spot SET
-            status=(?), end_at=(?)
-            WHERE spot_job_id=(?) AND
-            task_id=(?) AND
-            status=(?) AND
-            end_at IS null""",
-            (ManagedJobStatus.SUCCEEDED.value, end_time, job_id, task_id,
-             ManagedJobStatus.RUNNING.value))
-        if cursor.rowcount != 1:
-            raise exceptions.ManagedJobStatusError(
-                f'Failed to set the task to succeeded. '
-                f'({cursor.rowcount} rows updated)')
-    callback_func('SUCCEEDED')
-    logger.info('Job succeeded.')
-
-
+        count = result.rowcount
+        await session.commit()
+        if count != 1:
+            details = await _describe_task_transition_failure(
+                session, job_id, task_id)
+            message = ('Failed to set the task back to pending. '
+                       f'({count} rows updated. {details})')
+            logger.error(message)
+            raise exceptions.ManagedJobStatusError(message)
+    # Do not call callback_func here, as we don't use the callback for PENDING.
+
+
+@_init_db
+async def set_restarting_async(job_id: int, task_id: int, recovering: bool):
+    """Set the task back to STARTING or RECOVERING from PENDING.
+
+    This should not be used for the initial transition from PENDING to STARTING.
+    In that case, use set_starting instead. This function should only be used
+    after using set_backoff_pending to transition back to PENDING during
+    launch retry backoff.
+    """
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    target_status = ManagedJobStatus.STARTING.value
+    if recovering:
+        target_status = ManagedJobStatus.RECOVERING.value
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        result = await session.execute(
+            sqlalchemy.update(spot_table).where(
+                sqlalchemy.and_(
+                    spot_table.c.spot_job_id == job_id,
+                    spot_table.c.task_id == task_id,
+                    spot_table.c.end_at.is_(None),
+                )).values({spot_table.c.status: target_status}))
+        count = result.rowcount
+        await session.commit()
+        logger.debug(f'back to {target_status}')
+        if count != 1:
+            details = await _describe_task_transition_failure(
+                session, job_id, task_id)
+            message = (f'Failed to set the task back to {target_status}. '
+                       f'({count} rows updated. {details})')
+            logger.error(message)
+            raise exceptions.ManagedJobStatusError(message)
+    # Do not call callback_func here, as it should only be invoked for the
+    # initial (pre-`set_backoff_pending`) transition to STARTING or RECOVERING.
+
+
+@_init_db
 def set_failed(
     job_id: int,
     task_id: Optional[int],
@@ -585,188 +865,158 @@ def set_failed(
         override_terminal: If True, override the current status even if end_at
             is already set.
     """
+    assert _SQLALCHEMY_ENGINE is not None
     assert failure_type.is_failed(), failure_type
     end_time = time.time() if end_time is None else end_time
 
     fields_to_set: Dict[str, Any] = {
-        'status': failure_type.value,
-        'failure_reason': failure_reason,
+        spot_table.c.status: failure_type.value,
+        spot_table.c.failure_reason: failure_reason,
     }
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        previous_status = cursor.execute(
-            'SELECT status FROM spot WHERE spot_job_id=(?)',
-            (job_id,)).fetchone()[0]
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        # Get previous status
+        previous_status = session.execute(
+            sqlalchemy.select(spot_table.c.status).where(
+                spot_table.c.spot_job_id == job_id)).fetchone()[0]
         previous_status = ManagedJobStatus(previous_status)
         if previous_status == ManagedJobStatus.RECOVERING:
             # If the job is recovering, we should set the last_recovered_at to
             # the end_time, so that the end_at - last_recovered_at will not be
             # affect the job duration calculation.
-            fields_to_set['last_recovered_at'] = end_time
-        set_str = ', '.join(f'{k}=(?)' for k in fields_to_set)
-        task_query_str = '' if task_id is None else 'AND task_id=(?)'
-        task_value = [] if task_id is None else [
-            task_id,
-        ]
+            fields_to_set[spot_table.c.last_recovered_at] = end_time
+        where_conditions = [spot_table.c.spot_job_id == job_id]
+        if task_id is not None:
+            where_conditions.append(spot_table.c.task_id == task_id)
 
+        # Handle failure_reason prepending when override_terminal is True
         if override_terminal:
+            # Get existing failure_reason with row lock to prevent race
+            # conditions
+            existing_reason_result = session.execute(
+                sqlalchemy.select(spot_table.c.failure_reason).where(
+                    sqlalchemy.and_(*where_conditions)).with_for_update())
+            existing_reason_row = existing_reason_result.fetchone()
+            if existing_reason_row and existing_reason_row[0]:
+                # Prepend new failure reason to existing one
+                fields_to_set[spot_table.c.failure_reason] = (
+                    failure_reason + '. Previously: ' + existing_reason_row[0])
             # Use COALESCE for end_at to avoid overriding the existing end_at if
             # it's already set.
-            cursor.execute(
-                f"""\
-                UPDATE spot SET
-                end_at = COALESCE(end_at, ?),
-                {set_str}
-                WHERE spot_job_id=(?) {task_query_str}""",
-                (end_time, *list(fields_to_set.values()), job_id, *task_value))
+            fields_to_set[spot_table.c.end_at] = sqlalchemy.func.coalesce(
+                spot_table.c.end_at, end_time)
         else:
-            # Only set if end_at is null, i.e. the previous status is not
-            # terminal.
-            cursor.execute(
-                f"""\
-                UPDATE spot SET
-                end_at = (?),
-                {set_str}
-                WHERE spot_job_id=(?) {task_query_str} AND end_at IS null""",
-                (end_time, *list(fields_to_set.values()), job_id, *task_value))
-
-        updated = cursor.rowcount > 0
+            fields_to_set[spot_table.c.end_at] = end_time
+            where_conditions.append(spot_table.c.end_at.is_(None))
+        count = session.query(spot_table).filter(
+            sqlalchemy.and_(*where_conditions)).update(fields_to_set)
+        session.commit()
+        updated = count > 0
     if callback_func and updated:
         callback_func('FAILED')
     logger.info(failure_reason)
 
 
-def set_cancelling(job_id: int, callback_func: CallbackType):
-    """Set tasks in the job as cancelling, if they are in non-terminal states.
-
-    task_id is not needed, because we expect the job should be cancelled
-    as a whole, and we should not cancel a single task.
-    """
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        rows = cursor.execute(
-            """\
-            UPDATE spot SET
-            status=(?)
-            WHERE spot_job_id=(?) AND end_at IS null""",
-            (ManagedJobStatus.CANCELLING.value, job_id))
-        updated = rows.rowcount > 0
-    if updated:
-        logger.info('Cancelling the job...')
-        callback_func('CANCELLING')
-    else:
-        logger.info('Cancellation skipped, job is already terminal')
-
+@_init_db
+def set_pending_cancelled(job_id: int):
+    """Set the job as pending cancelled, if it is in non-terminal states."""
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        # Subquery to get the spot_job_ids that match the joined condition
+        subquery = session.query(spot_table.c.job_id).join(
+            job_info_table,
+            spot_table.c.spot_job_id == job_info_table.c.spot_job_id
+        ).filter(
+            spot_table.c.spot_job_id == job_id,
+            spot_table.c.status == ManagedJobStatus.PENDING.value,
+            # Note: it's possible that a WAITING job actually needs to be
+            # cleaned up, if we are in the middle of an upgrade/recovery and
+            # the job is waiting to be reclaimed by a new controller. But,
+            # in this case the status will not be PENDING.
+            sqlalchemy.or_(
+                job_info_table.c.schedule_state ==
+                ManagedJobScheduleState.WAITING.value,
+                job_info_table.c.schedule_state ==
+                ManagedJobScheduleState.INACTIVE.value,
+            ),
+        ).subquery()
 
-def set_cancelled(job_id: int, callback_func: CallbackType):
-    """Set tasks in the job as cancelled, if they are in CANCELLING state.
-
-    The set_cancelling should be called before this function.
-    """
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        rows = cursor.execute(
-            """\
-            UPDATE spot SET
-            status=(?), end_at=(?)
-            WHERE spot_job_id=(?) AND status=(?)""",
-            (ManagedJobStatus.CANCELLED.value, time.time(), job_id,
-             ManagedJobStatus.CANCELLING.value))
-        updated = rows.rowcount > 0
-    if updated:
-        logger.info('Job cancelled.')
-        callback_func('CANCELLED')
-    else:
-        logger.info('Cancellation skipped, job is not CANCELLING')
+        count = session.query(spot_table).filter(
+            spot_table.c.job_id.in_(subquery)).update(
+                {spot_table.c.status: ManagedJobStatus.CANCELLED.value},
+                synchronize_session=False)
+        session.commit()
+        return count > 0
 
 
+@_init_db
 def set_local_log_file(job_id: int, task_id: Optional[int],
                        local_log_file: str):
     """Set the local log file for a job."""
-    filter_str = 'spot_job_id=(?)'
-    filter_args = [local_log_file, job_id]
-    if task_id is not None:
-        filter_str += ' AND task_id=(?)'
-        filter_args.append(task_id)
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        cursor.execute(
-            'UPDATE spot SET local_log_file=(?) '
-            f'WHERE {filter_str}', filter_args)
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        where_conditions = [spot_table.c.spot_job_id == job_id]
+        if task_id is not None:
+            where_conditions.append(spot_table.c.task_id == task_id)
+        session.query(spot_table).filter(
+            sqlalchemy.and_(*where_conditions)).update(
+                {spot_table.c.local_log_file: local_log_file})
+        session.commit()
 
 
 # ======== utility functions ========
+@_init_db
 def get_nonterminal_job_ids_by_name(name: Optional[str],
+                                    user_hash: Optional[str] = None,
                                     all_users: bool = False) -> List[int]:
-    """Get non-terminal job ids by name."""
-    statuses = ', '.join(['?'] * len(ManagedJobStatus.terminal_statuses()))
-    field_values = [
-        status.value for status in ManagedJobStatus.terminal_statuses()
-    ]
-
-    job_filter = ''
-    if name is None and not all_users:
-        job_filter += 'AND (job_info.user_hash=(?)) '
-        field_values.append(common_utils.get_user_hash())
-    if name is not None:
-        # We match the job name from `job_info` for the jobs submitted after
-        # #1982, and from `spot` for the jobs submitted before #1982, whose
-        # job_info is not available.
-        job_filter += ('AND (job_info.name=(?) OR '
-                       '(job_info.name IS NULL AND spot.task_name=(?))) ')
-        field_values.extend([name, name])
-
-    # Left outer join is used here instead of join, because the job_info does
-    # not contain the managed jobs submitted before #1982.
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        rows = cursor.execute(
-            f"""\
-            SELECT DISTINCT spot.spot_job_id
-            FROM spot
-            LEFT OUTER JOIN job_info
-            ON spot.spot_job_id=job_info.spot_job_id
-            WHERE status NOT IN
-            ({statuses})
-            {job_filter}
-            ORDER BY spot.spot_job_id DESC""", field_values).fetchall()
-        job_ids = [row[0] for row in rows if row[0] is not None]
-        return job_ids
-
+    """Get non-terminal job ids by name.
 
-def get_schedule_live_jobs(job_id: Optional[int]) -> List[Dict[str, Any]]:
-    """Get jobs from the database that have a live schedule_state.
-
-    This should return job(s) that are not INACTIVE, WAITING, or DONE.  So a
-    returned job should correspond to a live job controller process, with one
-    exception: the job may have just transitioned from WAITING to LAUNCHING, but
-    the controller process has not yet started.
+    If name is None:
+    1. if all_users is False, get for the given user_hash
+    2. otherwise, get for all users
     """
-    job_filter = '' if job_id is None else 'AND spot_job_id=(?)'
-    job_value = (job_id,) if job_id is not None else ()
-
-    # Join spot and job_info tables to get the job name for each task.
-    # We use LEFT OUTER JOIN mainly for backward compatibility, as for an
-    # existing controller before #1982, the job_info table may not exist,
-    # and all the managed jobs created before will not present in the
-    # job_info.
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        rows = cursor.execute(
-            f"""\
-            SELECT spot_job_id, schedule_state, controller_pid
-            FROM job_info
-            WHERE schedule_state not in (?, ?, ?)
-            {job_filter}
-            ORDER BY spot_job_id DESC""",
-            (ManagedJobScheduleState.INACTIVE.value,
-             ManagedJobScheduleState.WAITING.value,
-             ManagedJobScheduleState.DONE.value, *job_value)).fetchall()
-        jobs = []
-        for row in rows:
-            job_dict = {
-                'job_id': row[0],
-                'schedule_state': ManagedJobScheduleState(row[1]),
-                'controller_pid': row[2],
-            }
-            jobs.append(job_dict)
-        return jobs
+    assert _SQLALCHEMY_ENGINE is not None
+
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        # Build the query using SQLAlchemy core
+        query = sqlalchemy.select(
+            spot_table.c.spot_job_id.distinct()).select_from(
+                spot_table.outerjoin(
+                    job_info_table,
+                    spot_table.c.spot_job_id == job_info_table.c.spot_job_id,
+                ))
+        where_conditions = [
+            ~spot_table.c.status.in_([
+                status.value for status in ManagedJobStatus.terminal_statuses()
+            ])
+        ]
+        if name is None and not all_users:
+            if user_hash is None:
+                # For backwards compatibility. With codegen, USER_ID_ENV_VAR
+                # was set to the correct value by the jobs controller, as
+                # part of ManagedJobCodeGen._build(). This is no longer the
+                # case for the Skylet gRPC server, which is why we need to
+                # pass it explicitly through the request body.
+                logger.debug('user_hash is None, using current user hash')
+                user_hash = common_utils.get_user_hash()
+            where_conditions.append(job_info_table.c.user_hash == user_hash)
+        if name is not None:
+            # We match the job name from `job_info` for the jobs submitted after
+            # #1982, and from `spot` for the jobs submitted before #1982, whose
+            # job_info is not available.
+            where_conditions.append(
+                sqlalchemy.or_(
+                    job_info_table.c.name == name,
+                    sqlalchemy.and_(job_info_table.c.name.is_(None),
+                                    spot_table.c.task_name == name),
+                ))
+        query = query.where(sqlalchemy.and_(*where_conditions)).order_by(
+            spot_table.c.spot_job_id.desc())
+        rows = session.execute(query).fetchall()
+        job_ids = [row[0] for row in rows if row[0] is not None]
+        return job_ids
 
 
+@_init_db
 def get_jobs_to_check_status(job_id: Optional[int] = None) -> List[int]:
     """Get jobs that need controller process checking.
 
@@ -778,89 +1028,87 @@ def get_jobs_to_check_status(job_id: Optional[int] = None) -> List[int]:
     - Jobs have schedule_state DONE but are in a non-terminal status
     - Legacy jobs (that is, no schedule state) that are in non-terminal status
     """
-    job_filter = '' if job_id is None else 'AND spot.spot_job_id=(?)'
-    job_value = () if job_id is None else (job_id,)
-
-    status_filter_str = ', '.join(['?'] *
-                                  len(ManagedJobStatus.terminal_statuses()))
-    terminal_status_values = [
-        status.value for status in ManagedJobStatus.terminal_statuses()
-    ]
-
-    # Get jobs that are either:
-    # 1. Have schedule state that is not DONE, or
-    # 2. Have schedule state DONE AND are in non-terminal status (unexpected
-    #    inconsistent state), or
-    # 3. Have no schedule state (legacy) AND are in non-terminal status
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        rows = cursor.execute(
-            f"""\
-            SELECT DISTINCT spot.spot_job_id
-            FROM spot
-            LEFT OUTER JOIN job_info
-            ON spot.spot_job_id=job_info.spot_job_id
-            WHERE (
-                -- non-legacy jobs that are not DONE
-                (job_info.schedule_state IS NOT NULL AND
-                 job_info.schedule_state IS NOT ?)
-                OR
-                -- legacy or that are in non-terminal status or
-                -- DONE jobs that are in non-terminal status
-                ((-- legacy jobs
-                  job_info.schedule_state IS NULL OR
-                  -- non-legacy DONE jobs
-                  job_info.schedule_state IS ?
-                 ) AND
-                 -- non-terminal
-                 status NOT IN ({status_filter_str}))
-            )
-            {job_filter}
-            ORDER BY spot.spot_job_id DESC""", [
-                ManagedJobScheduleState.DONE.value,
-                ManagedJobScheduleState.DONE.value, *terminal_status_values,
-                *job_value
-            ]).fetchall()
-        return [row[0] for row in rows if row[0] is not None]
+    assert _SQLALCHEMY_ENGINE is not None
 
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        terminal_status_values = [
+            status.value for status in ManagedJobStatus.terminal_statuses()
+        ]
 
-def get_all_job_ids_by_name(name: Optional[str]) -> List[int]:
-    """Get all job ids by name."""
-    name_filter = ''
-    field_values = []
-    if name is not None:
-        # We match the job name from `job_info` for the jobs submitted after
-        # #1982, and from `spot` for the jobs submitted before #1982, whose
-        # job_info is not available.
-        name_filter = ('WHERE (job_info.name=(?) OR '
-                       '(job_info.name IS NULL AND spot.task_name=(?)))')
-        field_values = [name, name]
-
-    # Left outer join is used here instead of join, because the job_info does
-    # not contain the managed jobs submitted before #1982.
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        rows = cursor.execute(
-            f"""\
-            SELECT DISTINCT spot.spot_job_id
-            FROM spot
-            LEFT OUTER JOIN job_info
-            ON spot.spot_job_id=job_info.spot_job_id
-            {name_filter}
-            ORDER BY spot.spot_job_id DESC""", field_values).fetchall()
-        job_ids = [row[0] for row in rows if row[0] is not None]
-        return job_ids
+        query = sqlalchemy.select(
+            spot_table.c.spot_job_id.distinct()).select_from(
+                spot_table.outerjoin(
+                    job_info_table,
+                    spot_table.c.spot_job_id == job_info_table.c.spot_job_id))
+
+        # Get jobs that are either:
+        # 1. Have schedule state that is not DONE, or
+        # 2. Have schedule state DONE AND are in non-terminal status (unexpected
+        #    inconsistent state), or
+        # 3. Have no schedule state (legacy) AND are in non-terminal status
+
+        # non-legacy jobs that are not DONE
+        condition1 = sqlalchemy.and_(
+            job_info_table.c.schedule_state.is_not(None),
+            job_info_table.c.schedule_state !=
+            ManagedJobScheduleState.DONE.value)
+        # legacy or that are in non-terminal status or
+        # DONE jobs that are in non-terminal status
+        condition2 = sqlalchemy.and_(
+            sqlalchemy.or_(
+                # legacy jobs
+                job_info_table.c.schedule_state.is_(None),
+                # non-legacy DONE jobs
+                job_info_table.c.schedule_state ==
+                ManagedJobScheduleState.DONE.value),
+            # non-terminal
+            ~spot_table.c.status.in_(terminal_status_values),
+        )
+        where_condition = sqlalchemy.or_(condition1, condition2)
+        if job_id is not None:
+            where_condition = sqlalchemy.and_(
+                where_condition, spot_table.c.spot_job_id == job_id)
+
+        query = query.where(where_condition).order_by(
+            spot_table.c.spot_job_id.desc())
+
+        rows = session.execute(query).fetchall()
+        return [row[0] for row in rows if row[0] is not None]
 
 
+@_init_db
 def _get_all_task_ids_statuses(
         job_id: int) -> List[Tuple[int, ManagedJobStatus]]:
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        id_statuses = cursor.execute(
-            """\
-            SELECT task_id, status FROM spot
-            WHERE spot_job_id=(?)
-            ORDER BY task_id ASC""", (job_id,)).fetchall()
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        id_statuses = session.execute(
+            sqlalchemy.select(
+                spot_table.c.task_id,
+                spot_table.c.status,
+            ).where(spot_table.c.spot_job_id == job_id).order_by(
+                spot_table.c.task_id.asc())).fetchall()
         return [(row[0], ManagedJobStatus(row[1])) for row in id_statuses]
 
 
+@_init_db
+def get_all_task_ids_names_statuses_logs(
+    job_id: int
+) -> List[Tuple[int, str, ManagedJobStatus, str, Optional[float]]]:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        id_names = session.execute(
+            sqlalchemy.select(
+                spot_table.c.task_id,
+                spot_table.c.task_name,
+                spot_table.c.status,
+                spot_table.c.local_log_file,
+                spot_table.c.logs_cleaned_at,
+            ).where(spot_table.c.spot_job_id == job_id).order_by(
+                spot_table.c.task_id.asc())).fetchall()
+        return [(row[0], row[1], ManagedJobStatus(row[2]), row[3], row[4])
+                for row in id_names]
+
+
 def get_num_tasks(job_id: int) -> int:
     return len(_get_all_task_ids_statuses(job_id))
 
@@ -888,31 +1136,43 @@ def get_latest_task_id_status(
     return task_id, status
 
 
+@_init_db
+def get_job_controller_pid(job_id: int) -> Optional[int]:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        pid = session.execute(
+            sqlalchemy.select(job_info_table.c.controller_pid).where(
+                job_info_table.c.spot_job_id == job_id)).fetchone()
+        return pid[0] if pid else None
+
+
 def get_status(job_id: int) -> Optional[ManagedJobStatus]:
     _, status = get_latest_task_id_status(job_id)
     return status
 
 
+@_init_db
 def get_failure_reason(job_id: int) -> Optional[str]:
     """Get the failure reason of a job.
 
     If the job has multiple tasks, we return the first failure reason.
     """
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        reason = cursor.execute(
-            """\
-            SELECT failure_reason FROM spot
-            WHERE spot_job_id=(?)
-            ORDER BY task_id ASC""", (job_id,)).fetchall()
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        reason = session.execute(
+            sqlalchemy.select(spot_table.c.failure_reason).where(
+                spot_table.c.spot_job_id == job_id).order_by(
+                    spot_table.c.task_id.asc())).fetchall()
         reason = [r[0] for r in reason if r[0] is not None]
         if not reason:
             return None
         return reason[0]
 
 
-def get_managed_jobs(job_id: Optional[int] = None) -> List[Dict[str, Any]]:
-    """Get managed jobs from the database."""
-    job_filter = '' if job_id is None else f'WHERE spot.spot_job_id={job_id}'
+@_init_db
+def get_managed_job_tasks(job_id: int) -> List[Dict[str, Any]]:
+    """Get managed job tasks for a specific managed job id from the database."""
+    assert _SQLALCHEMY_ENGINE is not None
 
     # Join spot and job_info tables to get the job name for each task.
     # We use LEFT OUTER JOIN mainly for backward compatibility, as for an
@@ -922,197 +1182,1328 @@ def get_managed_jobs(job_id: Optional[int] = None) -> List[Dict[str, Any]]:
     # Note: we will get the user_hash here, but don't try to call
     # global_user_state.get_user() on it. This runs on the controller, which may
     # not have the user info. Prefer to do it on the API server side.
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        rows = cursor.execute(f"""\
-            SELECT *
-            FROM spot
-            LEFT OUTER JOIN job_info
-            ON spot.spot_job_id=job_info.spot_job_id
-            {job_filter}
-            ORDER BY spot.spot_job_id DESC, spot.task_id ASC""").fetchall()
-        jobs = []
-        for row in rows:
-            job_dict = dict(zip(columns, row))
+    query = sqlalchemy.select(spot_table, job_info_table).select_from(
+        spot_table.outerjoin(
+            job_info_table,
+            spot_table.c.spot_job_id == job_info_table.c.spot_job_id))
+    query = query.where(spot_table.c.spot_job_id == job_id)
+    query = query.order_by(spot_table.c.task_id.asc())
+    rows = None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        rows = session.execute(query).fetchall()
+    jobs = []
+    for row in rows:
+        job_dict = _get_jobs_dict(row._mapping)  # pylint: disable=protected-access
+        job_dict['status'] = ManagedJobStatus(job_dict['status'])
+        job_dict['schedule_state'] = ManagedJobScheduleState(
+            job_dict['schedule_state'])
+        if job_dict['job_name'] is None:
+            job_dict['job_name'] = job_dict['task_name']
+        job_dict['metadata'] = json.loads(job_dict['metadata'])
+
+        # Add user YAML content for managed jobs.
+        job_dict['user_yaml'] = job_dict.get('original_user_yaml_content')
+        if job_dict['user_yaml'] is None:
+            # Backwards compatibility - try to read from file path
+            yaml_path = job_dict.get('original_user_yaml_path')
+            if yaml_path:
+                try:
+                    with open(yaml_path, 'r', encoding='utf-8') as f:
+                        job_dict['user_yaml'] = f.read()
+                except (FileNotFoundError, IOError, OSError) as e:
+                    logger.debug('Failed to read original user YAML for job '
+                                 f'{job_id} from {yaml_path}: {e}')
+
+        jobs.append(job_dict)
+    return jobs
+
+
+def _map_response_field_to_db_column(field: str):
+    """Map the response field name to an actual SQLAlchemy ColumnElement.
+
+    This ensures we never pass plain strings to SQLAlchemy 2.0 APIs like
+    Select.with_only_columns().
+    """
+    # Explicit aliases differing from actual DB column names
+    alias_mapping = {
+        '_job_id': spot_table.c.job_id,  # spot.job_id
+        '_task_name': spot_table.c.job_name,  # deprecated, from spot table
+        'job_id': spot_table.c.spot_job_id,  # public job id -> spot.spot_job_id
+        '_job_info_job_id': job_info_table.c.spot_job_id,
+        'job_name': job_info_table.c.name,  # public job name -> job_info.name
+    }
+    if field in alias_mapping:
+        return alias_mapping[field]
+
+    # Try direct match on the `spot` table columns
+    if field in spot_table.c:
+        return spot_table.c[field]
+
+    # Try direct match on the `job_info` table columns
+    if field in job_info_table.c:
+        return job_info_table.c[field]
+
+    raise ValueError(f'Unknown field: {field}')
+
+
+@_init_db
+def get_managed_jobs_total() -> int:
+    """Get the total number of managed jobs."""
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        result = session.execute(
+            sqlalchemy.select(sqlalchemy.func.count()  # pylint: disable=not-callable
+                             ).select_from(spot_table)).fetchone()
+        return result[0] if result else 0
+
+
+@_init_db
+def get_managed_jobs_highest_priority() -> int:
+    """Get the highest priority of the managed jobs."""
+    assert _SQLALCHEMY_ENGINE is not None
+    query = sqlalchemy.select(sqlalchemy.func.max(
+        job_info_table.c.priority)).where(
+            sqlalchemy.and_(
+                job_info_table.c.schedule_state.in_([
+                    ManagedJobScheduleState.LAUNCHING.value,
+                    ManagedJobScheduleState.ALIVE_BACKOFF.value,
+                    ManagedJobScheduleState.WAITING.value,
+                    ManagedJobScheduleState.ALIVE_WAITING.value,
+                ]),
+                job_info_table.c.priority.is_not(None),
+            ))
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        priority = session.execute(query).fetchone()
+        return priority[0] if priority and priority[
+            0] is not None else constants.MIN_PRIORITY
+
+
+def build_managed_jobs_with_filters_no_status_query(
+    fields: Optional[List[str]] = None,
+    job_ids: Optional[List[int]] = None,
+    accessible_workspaces: Optional[List[str]] = None,
+    workspace_match: Optional[str] = None,
+    name_match: Optional[str] = None,
+    pool_match: Optional[str] = None,
+    user_hashes: Optional[List[Optional[str]]] = None,
+    skip_finished: bool = False,
+    count_only: bool = False,
+    status_count: bool = False,
+) -> sqlalchemy.Select:
+    """Build a query to get managed jobs from the database with filters."""
+    # Join spot and job_info tables to get the job name for each task.
+    # We use LEFT OUTER JOIN mainly for backward compatibility, as for an
+    # existing controller before #1982, the job_info table may not exist,
+    # and all the managed jobs created before will not present in the
+    # job_info.
+    # Note: we will get the user_hash here, but don't try to call
+    # global_user_state.get_user() on it. This runs on the controller, which may
+    # not have the user info. Prefer to do it on the API server side.
+    if count_only:
+        query = sqlalchemy.select(sqlalchemy.func.count().label('count'))  # pylint: disable=not-callable
+    elif status_count:
+        query = sqlalchemy.select(spot_table.c.status,
+                                  sqlalchemy.func.count().label('count'))  # pylint: disable=not-callable
+    else:
+        query = sqlalchemy.select(spot_table, job_info_table)
+    query = query.select_from(
+        spot_table.outerjoin(
+            job_info_table,
+            spot_table.c.spot_job_id == job_info_table.c.spot_job_id))
+    if skip_finished:
+        # Filter out finished jobs at the DB level. If a multi-task job is
+        # partially finished, include all its tasks. We do this by first
+        # selecting job_ids that have at least one non-terminal task, then
+        # restricting the main query to those job_ids.
+        terminal_status_values = [
+            s.value for s in ManagedJobStatus.terminal_statuses()
+        ]
+        non_terminal_job_ids_subquery = (sqlalchemy.select(
+            spot_table.c.spot_job_id).where(
+                sqlalchemy.or_(
+                    spot_table.c.status.is_(None),
+                    sqlalchemy.not_(
+                        spot_table.c.status.in_(terminal_status_values)),
+                )).distinct())
+        query = query.where(
+            spot_table.c.spot_job_id.in_(non_terminal_job_ids_subquery))
+    if not count_only and not status_count and fields:
+        # Resolve requested field names to explicit ColumnElements from
+        # the joined tables.
+        selected_columns = [_map_response_field_to_db_column(f) for f in fields]
+        query = query.with_only_columns(*selected_columns)
+    if job_ids is not None:
+        query = query.where(spot_table.c.spot_job_id.in_(job_ids))
+    if accessible_workspaces is not None:
+        query = query.where(
+            job_info_table.c.workspace.in_(accessible_workspaces))
+    if workspace_match is not None:
+        query = query.where(
+            job_info_table.c.workspace.like(f'%{workspace_match}%'))
+    if name_match is not None:
+        query = query.where(job_info_table.c.name.like(f'%{name_match}%'))
+    if pool_match is not None:
+        query = query.where(job_info_table.c.pool.like(f'%{pool_match}%'))
+    if user_hashes is not None:
+        query = query.where(job_info_table.c.user_hash.in_(user_hashes))
+    return query
+
+
+def build_managed_jobs_with_filters_query(
+    fields: Optional[List[str]] = None,
+    job_ids: Optional[List[int]] = None,
+    accessible_workspaces: Optional[List[str]] = None,
+    workspace_match: Optional[str] = None,
+    name_match: Optional[str] = None,
+    pool_match: Optional[str] = None,
+    user_hashes: Optional[List[Optional[str]]] = None,
+    statuses: Optional[List[str]] = None,
+    skip_finished: bool = False,
+    count_only: bool = False,
+) -> sqlalchemy.Select:
+    """Build a query to get managed jobs from the database with filters."""
+    query = build_managed_jobs_with_filters_no_status_query(
+        fields=fields,
+        job_ids=job_ids,
+        accessible_workspaces=accessible_workspaces,
+        workspace_match=workspace_match,
+        name_match=name_match,
+        pool_match=pool_match,
+        user_hashes=user_hashes,
+        skip_finished=skip_finished,
+        count_only=count_only,
+    )
+    if statuses is not None:
+        query = query.where(spot_table.c.status.in_(statuses))
+    return query
+
+
+@_init_db
+def get_status_count_with_filters(
+    fields: Optional[List[str]] = None,
+    job_ids: Optional[List[int]] = None,
+    accessible_workspaces: Optional[List[str]] = None,
+    workspace_match: Optional[str] = None,
+    name_match: Optional[str] = None,
+    pool_match: Optional[str] = None,
+    user_hashes: Optional[List[Optional[str]]] = None,
+    skip_finished: bool = False,
+) -> Dict[str, int]:
+    """Get the status count of the managed jobs with filters."""
+    query = build_managed_jobs_with_filters_no_status_query(
+        fields=fields,
+        job_ids=job_ids,
+        accessible_workspaces=accessible_workspaces,
+        workspace_match=workspace_match,
+        name_match=name_match,
+        pool_match=pool_match,
+        user_hashes=user_hashes,
+        skip_finished=skip_finished,
+        status_count=True,
+    )
+    query = query.group_by(spot_table.c.status)
+    results: Dict[str, int] = {}
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        rows = session.execute(query).fetchall()
+        for status_value, count in rows:
+            # status_value is already a string (enum value)
+            results[str(status_value)] = int(count)
+    return results
+
+
+@_init_db
+def get_managed_jobs_with_filters(
+    fields: Optional[List[str]] = None,
+    job_ids: Optional[List[int]] = None,
+    accessible_workspaces: Optional[List[str]] = None,
+    workspace_match: Optional[str] = None,
+    name_match: Optional[str] = None,
+    pool_match: Optional[str] = None,
+    user_hashes: Optional[List[Optional[str]]] = None,
+    statuses: Optional[List[str]] = None,
+    skip_finished: bool = False,
+    page: Optional[int] = None,
+    limit: Optional[int] = None,
+) -> Tuple[List[Dict[str, Any]], int]:
+    """Get managed jobs from the database with filters.
+
+    Returns:
+        A tuple containing
+         - the list of managed jobs
+         - the total number of managed jobs
+    """
+    assert _SQLALCHEMY_ENGINE is not None
+
+    count_query = build_managed_jobs_with_filters_query(
+        fields=None,
+        job_ids=job_ids,
+        accessible_workspaces=accessible_workspaces,
+        workspace_match=workspace_match,
+        name_match=name_match,
+        pool_match=pool_match,
+        user_hashes=user_hashes,
+        statuses=statuses,
+        skip_finished=skip_finished,
+        count_only=True,
+    )
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        total = session.execute(count_query).fetchone()[0]
+
+    query = build_managed_jobs_with_filters_query(
+        fields=fields,
+        job_ids=job_ids,
+        accessible_workspaces=accessible_workspaces,
+        workspace_match=workspace_match,
+        name_match=name_match,
+        pool_match=pool_match,
+        user_hashes=user_hashes,
+        statuses=statuses,
+        skip_finished=skip_finished,
+    )
+    query = query.order_by(spot_table.c.spot_job_id.desc(),
+                           spot_table.c.task_id.asc())
+    if page is not None and limit is not None:
+        query = query.offset((page - 1) * limit).limit(limit)
+    rows = None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        rows = session.execute(query).fetchall()
+    jobs = []
+    for row in rows:
+        job_dict = _get_jobs_dict(row._mapping)  # pylint: disable=protected-access
+        if job_dict.get('status') is not None:
             job_dict['status'] = ManagedJobStatus(job_dict['status'])
+        if job_dict.get('schedule_state') is not None:
             job_dict['schedule_state'] = ManagedJobScheduleState(
                 job_dict['schedule_state'])
-            if job_dict['job_name'] is None:
-                job_dict['job_name'] = job_dict['task_name']
-            jobs.append(job_dict)
-        return jobs
-
-
+        if job_dict.get('job_name') is None:
+            job_dict['job_name'] = job_dict.get('task_name')
+        if job_dict.get('metadata') is not None:
+            job_dict['metadata'] = json.loads(job_dict['metadata'])
+
+        # Add user YAML content for managed jobs.
+        job_dict['user_yaml'] = job_dict.get('original_user_yaml_content')
+        if job_dict['user_yaml'] is None:
+            # Backwards compatibility - try to read from file path
+            yaml_path = job_dict.get('original_user_yaml_path')
+            if yaml_path:
+                try:
+                    with open(yaml_path, 'r', encoding='utf-8') as f:
+                        job_dict['user_yaml'] = f.read()
+                except (FileNotFoundError, IOError, OSError) as e:
+                    job_id = job_dict.get('job_id')
+                    if job_id is not None:
+                        logger.debug('Failed to read original user YAML for '
+                                     f'job {job_id} from {yaml_path}: {e}')
+                    else:
+                        logger.debug('Failed to read original user YAML from '
+                                     f'{yaml_path}: {e}')
+
+        jobs.append(job_dict)
+    return jobs, total
+
+
+@_init_db
 def get_task_name(job_id: int, task_id: int) -> str:
     """Get the task name of a job."""
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        task_name = cursor.execute(
-            """\
-            SELECT task_name FROM spot
-            WHERE spot_job_id=(?)
-            AND task_id=(?)""", (job_id, task_id)).fetchone()
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        task_name = session.execute(
+            sqlalchemy.select(spot_table.c.task_name).where(
+                sqlalchemy.and_(
+                    spot_table.c.spot_job_id == job_id,
+                    spot_table.c.task_id == task_id,
+                ))).fetchone()
         return task_name[0]
 
 
+@_init_db
 def get_latest_job_id() -> Optional[int]:
     """Get the latest job id."""
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        rows = cursor.execute("""\
-            SELECT spot_job_id FROM spot
-            WHERE task_id=0
-            ORDER BY submitted_at DESC LIMIT 1""")
-        for (job_id,) in rows:
-            return job_id
-        return None
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        job_id = session.execute(
+            sqlalchemy.select(spot_table.c.spot_job_id).where(
+                spot_table.c.task_id == 0).order_by(
+                    spot_table.c.submitted_at.desc()).limit(1)).fetchone()
+        return job_id[0] if job_id else None
 
 
+@_init_db
 def get_task_specs(job_id: int, task_id: int) -> Dict[str, Any]:
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        task_specs = cursor.execute(
-            """\
-            SELECT specs FROM spot
-            WHERE spot_job_id=(?) AND task_id=(?)""",
-            (job_id, task_id)).fetchone()
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        task_specs = session.execute(
+            sqlalchemy.select(spot_table.c.specs).where(
+                sqlalchemy.and_(
+                    spot_table.c.spot_job_id == job_id,
+                    spot_table.c.task_id == task_id,
+                ))).fetchone()
         return json.loads(task_specs[0])
 
 
-def get_local_log_file(job_id: int, task_id: Optional[int]) -> Optional[str]:
-    """Get the local log directory for a job."""
-    filter_str = 'spot_job_id=(?)'
-    filter_args = [job_id]
-    if task_id is not None:
-        filter_str += ' AND task_id=(?)'
-        filter_args.append(task_id)
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        local_log_file = cursor.execute(
-            f'SELECT local_log_file FROM spot '
-            f'WHERE {filter_str}', filter_args).fetchone()
-        return local_log_file[-1] if local_log_file else None
-
-
-# === Scheduler state functions ===
-# Only the scheduler should call these functions. They may require holding the
-# scheduler lock to work correctly.
-
-
-def scheduler_set_waiting(job_id: int, dag_yaml_path: str, env_file_path: str,
-                          user_hash: str) -> None:
-    """Do not call without holding the scheduler lock."""
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        updated_count = cursor.execute(
-            'UPDATE job_info SET '
-            'schedule_state = (?), dag_yaml_path = (?), env_file_path = (?), '
-            '  user_hash = (?) '
-            'WHERE spot_job_id = (?) AND schedule_state = (?)',
-            (ManagedJobScheduleState.WAITING.value, dag_yaml_path,
-             env_file_path, user_hash, job_id,
-             ManagedJobScheduleState.INACTIVE.value)).rowcount
-        assert updated_count == 1, (job_id, updated_count)
-
-
-def scheduler_set_launching(job_id: int,
-                            current_state: ManagedJobScheduleState) -> None:
-    """Do not call without holding the scheduler lock."""
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        updated_count = cursor.execute(
-            'UPDATE job_info SET '
-            'schedule_state = (?) '
-            'WHERE spot_job_id = (?) AND schedule_state = (?)',
-            (ManagedJobScheduleState.LAUNCHING.value, job_id,
-             current_state.value)).rowcount
-        assert updated_count == 1, (job_id, updated_count)
-
+@_init_db
+def scheduler_set_waiting(job_id: int, dag_yaml_content: str,
+                          original_user_yaml_content: str,
+                          env_file_content: str, priority: int):
+    """Do not call without holding the scheduler lock.
 
-def scheduler_set_alive(job_id: int) -> None:
-    """Do not call without holding the scheduler lock."""
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        updated_count = cursor.execute(
-            'UPDATE job_info SET '
-            'schedule_state = (?) '
-            'WHERE spot_job_id = (?) AND schedule_state = (?)',
-            (ManagedJobScheduleState.ALIVE.value, job_id,
-             ManagedJobScheduleState.LAUNCHING.value)).rowcount
-        assert updated_count == 1, (job_id, updated_count)
+    Returns: Whether this is a recovery run or not.
+        If this is a recovery run, the job may already be in the WAITING
+        state and the update will not change the schedule_state (hence the
+        updated_count will be 0). In this case, we return True.
+        Otherwise, we return False.
+    """
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        updated_count = session.query(job_info_table).filter(
+            sqlalchemy.and_(job_info_table.c.spot_job_id == job_id,)).update({
+                job_info_table.c.schedule_state:
+                    ManagedJobScheduleState.WAITING.value,
+                job_info_table.c.dag_yaml_content: dag_yaml_content,
+                job_info_table.c.original_user_yaml_content:
+                    (original_user_yaml_content),
+                job_info_table.c.env_file_content: env_file_content,
+                job_info_table.c.priority: priority,
+            })
+        session.commit()
+        assert updated_count <= 1, (job_id, updated_count)
+
+
+@_init_db
+def get_job_file_contents(job_id: int) -> Dict[str, Optional[str]]:
+    """Return file information and stored contents for a managed job."""
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        row = session.execute(
+            sqlalchemy.select(
+                job_info_table.c.dag_yaml_path,
+                job_info_table.c.env_file_path,
+                job_info_table.c.dag_yaml_content,
+                job_info_table.c.env_file_content,
+            ).where(job_info_table.c.spot_job_id == job_id)).fetchone()
+
+    if row is None:
+        return {
+            'dag_yaml_path': None,
+            'env_file_path': None,
+            'dag_yaml_content': None,
+            'env_file_content': None,
+        }
+
+    return {
+        'dag_yaml_path': row[0],
+        'env_file_path': row[1],
+        'dag_yaml_content': row[2],
+        'env_file_content': row[3],
+    }
 
 
-def scheduler_set_alive_waiting(job_id: int) -> None:
+@_init_db
+def get_pool_from_job_id(job_id: int) -> Optional[str]:
+    """Get the pool from the job id."""
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        pool = session.execute(
+            sqlalchemy.select(job_info_table.c.pool).where(
+                job_info_table.c.spot_job_id == job_id)).fetchone()
+        return pool[0] if pool else None
+
+
+@_init_db
+def set_current_cluster_name(job_id: int, current_cluster_name: str) -> None:
+    """Set the current cluster name for a job."""
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        session.query(job_info_table).filter(
+            job_info_table.c.spot_job_id == job_id).update(
+                {job_info_table.c.current_cluster_name: current_cluster_name})
+        session.commit()
+
+
+@_init_db_async
+async def set_job_id_on_pool_cluster_async(job_id: int,
+                                           job_id_on_pool_cluster: int) -> None:
+    """Set the job id on the pool cluster for a job."""
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        await session.execute(
+            sqlalchemy.update(job_info_table).
+            where(job_info_table.c.spot_job_id == job_id).values({
+                job_info_table.c.job_id_on_pool_cluster: job_id_on_pool_cluster
+            }))
+        await session.commit()
+
+
+@_init_db
+def get_pool_submit_info(job_id: int) -> Tuple[Optional[str], Optional[int]]:
+    """Get the cluster name and job id on the pool from the managed job id."""
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        info = session.execute(
+            sqlalchemy.select(
+                job_info_table.c.current_cluster_name,
+                job_info_table.c.job_id_on_pool_cluster).where(
+                    job_info_table.c.spot_job_id == job_id)).fetchone()
+        if info is None:
+            return None, None
+        return info[0], info[1]
+
+
+@_init_db_async
+async def get_pool_submit_info_async(
+        job_id: int) -> Tuple[Optional[str], Optional[int]]:
+    """Get the cluster name and job id on the pool from the managed job id."""
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        result = await session.execute(
+            sqlalchemy.select(job_info_table.c.current_cluster_name,
+                              job_info_table.c.job_id_on_pool_cluster).where(
+                                  job_info_table.c.spot_job_id == job_id))
+        info = result.fetchone()
+        if info is None:
+            return None, None
+        return info[0], info[1]
+
+
+@_init_db_async
+async def scheduler_set_launching_async(job_id: int):
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        await session.execute(
+            sqlalchemy.update(job_info_table).where(
+                sqlalchemy.and_(job_info_table.c.spot_job_id == job_id)).values(
+                    {
+                        job_info_table.c.schedule_state:
+                            ManagedJobScheduleState.LAUNCHING.value
+                    }))
+        await session.commit()
+
+
+@_init_db_async
+async def scheduler_set_alive_async(job_id: int) -> None:
     """Do not call without holding the scheduler lock."""
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        updated_count = cursor.execute(
-            'UPDATE job_info SET '
-            'schedule_state = (?) '
-            'WHERE spot_job_id = (?) AND schedule_state = (?)',
-            (ManagedJobScheduleState.ALIVE_WAITING.value, job_id,
-             ManagedJobScheduleState.ALIVE.value)).rowcount
-        assert updated_count == 1, (job_id, updated_count)
-
-
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        result = await session.execute(
+            sqlalchemy.update(job_info_table).where(
+                sqlalchemy.and_(
+                    job_info_table.c.spot_job_id == job_id,
+                    job_info_table.c.schedule_state ==
+                    ManagedJobScheduleState.LAUNCHING.value,
+                )).values({
+                    job_info_table.c.schedule_state:
+                        ManagedJobScheduleState.ALIVE.value
+                }))
+        changes = result.rowcount
+        await session.commit()
+        assert changes == 1, (job_id, changes)
+
+
+@_init_db
 def scheduler_set_done(job_id: int, idempotent: bool = False) -> None:
     """Do not call without holding the scheduler lock."""
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        updated_count = cursor.execute(
-            'UPDATE job_info SET '
-            'schedule_state = (?) '
-            'WHERE spot_job_id = (?) AND schedule_state != (?)',
-            (ManagedJobScheduleState.DONE.value, job_id,
-             ManagedJobScheduleState.DONE.value)).rowcount
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        updated_count = session.query(job_info_table).filter(
+            sqlalchemy.and_(
+                job_info_table.c.spot_job_id == job_id,
+                job_info_table.c.schedule_state !=
+                ManagedJobScheduleState.DONE.value,
+            )).update({
+                job_info_table.c.schedule_state:
+                    ManagedJobScheduleState.DONE.value
+            })
+        session.commit()
         if not idempotent:
             assert updated_count == 1, (job_id, updated_count)
 
 
-def set_job_controller_pid(job_id: int, pid: int):
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        updated_count = cursor.execute(
-            'UPDATE job_info SET '
-            'controller_pid = (?) '
-            'WHERE spot_job_id = (?)', (pid, job_id)).rowcount
-        assert updated_count == 1, (job_id, updated_count)
-
-
+@_init_db
 def get_job_schedule_state(job_id: int) -> ManagedJobScheduleState:
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        state = cursor.execute(
-            'SELECT schedule_state FROM job_info WHERE spot_job_id = (?)',
-            (job_id,)).fetchone()[0]
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        state = session.execute(
+            sqlalchemy.select(job_info_table.c.schedule_state).where(
+                job_info_table.c.spot_job_id == job_id)).fetchone()[0]
         return ManagedJobScheduleState(state)
 
 
+@_init_db
 def get_num_launching_jobs() -> int:
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        return cursor.execute(
-            'SELECT COUNT(*) '
-            'FROM job_info '
-            'WHERE schedule_state = (?)',
-            (ManagedJobScheduleState.LAUNCHING.value,)).fetchone()[0]
-
-
-def get_num_alive_jobs() -> int:
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        return cursor.execute(
-            'SELECT COUNT(*) '
-            'FROM job_info '
-            'WHERE schedule_state IN (?, ?, ?)',
-            (ManagedJobScheduleState.ALIVE_WAITING.value,
-             ManagedJobScheduleState.LAUNCHING.value,
-             ManagedJobScheduleState.ALIVE.value)).fetchone()[0]
-
-
-def get_waiting_job() -> Optional[Dict[str, Any]]:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        return session.execute(
+            sqlalchemy.select(
+                sqlalchemy.func.count()  # pylint: disable=not-callable
+            ).select_from(job_info_table).where(
+                sqlalchemy.and_(
+                    job_info_table.c.schedule_state ==
+                    ManagedJobScheduleState.LAUNCHING.value,
+                    # We only count jobs that are not in the pool, because the
+                    # job in the pool does not actually calling the sky.launch.
+                    job_info_table.c.pool.is_(None)))).fetchone()[0]
+
+
+@_init_db
+def get_num_alive_jobs(pool: Optional[str] = None) -> int:
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        where_conditions = [
+            job_info_table.c.schedule_state.in_([
+                ManagedJobScheduleState.ALIVE_WAITING.value,
+                ManagedJobScheduleState.LAUNCHING.value,
+                ManagedJobScheduleState.ALIVE.value,
+                ManagedJobScheduleState.ALIVE_BACKOFF.value,
+            ])
+        ]
+
+        if pool is not None:
+            where_conditions.append(job_info_table.c.pool == pool)
+
+        return session.execute(
+            sqlalchemy.select(
+                sqlalchemy.func.count()  # pylint: disable=not-callable
+            ).select_from(job_info_table).where(
+                sqlalchemy.and_(*where_conditions))).fetchone()[0]
+
+
+@_init_db
+def get_nonterminal_job_ids_by_pool(pool: str,
+                                    cluster_name: Optional[str] = None
+                                   ) -> List[int]:
+    """Get nonterminal job ids in a pool."""
+    assert _SQLALCHEMY_ENGINE is not None
+
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        query = sqlalchemy.select(
+            spot_table.c.spot_job_id.distinct()).select_from(
+                spot_table.outerjoin(
+                    job_info_table,
+                    spot_table.c.spot_job_id == job_info_table.c.spot_job_id))
+        and_conditions = [
+            ~spot_table.c.status.in_([
+                status.value for status in ManagedJobStatus.terminal_statuses()
+            ]),
+            job_info_table.c.pool == pool,
+        ]
+        if cluster_name is not None:
+            and_conditions.append(
+                job_info_table.c.current_cluster_name == cluster_name)
+        query = query.where(sqlalchemy.and_(*and_conditions)).order_by(
+            spot_table.c.spot_job_id.asc())
+        rows = session.execute(query).fetchall()
+        job_ids = [row[0] for row in rows if row[0] is not None]
+        return job_ids
+
+
+@_init_db_async
+async def get_waiting_job_async(pid: int) -> Optional[Dict[str, Any]]:
     """Get the next job that should transition to LAUNCHING.
 
+    Selects the highest-priority WAITING or ALIVE_WAITING job and atomically
+    transitions it to LAUNCHING state to prevent race conditions.
+
+    Returns the job information if a job was successfully transitioned to
+    LAUNCHING, or None if no suitable job was found.
+
     Backwards compatibility note: jobs submitted before #4485 will have no
     schedule_state and will be ignored by this SQL query.
     """
-    with db_utils.safe_cursor(_DB_PATH) as cursor:
-        row = cursor.execute(
-            'SELECT spot_job_id, schedule_state, dag_yaml_path, env_file_path '
-            'FROM job_info '
-            'WHERE schedule_state in (?, ?) '
-            'ORDER BY spot_job_id LIMIT 1',
-            (ManagedJobScheduleState.WAITING.value,
-             ManagedJobScheduleState.ALIVE_WAITING.value)).fetchone()
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        # Select the highest priority waiting job for update (locks the row)
+        select_query = sqlalchemy.select(
+            job_info_table.c.spot_job_id,
+            job_info_table.c.schedule_state,
+            job_info_table.c.dag_yaml_path,
+            job_info_table.c.env_file_path,
+            job_info_table.c.controller_pid,
+            job_info_table.c.pool,
+        ).where(
+            job_info_table.c.schedule_state.in_([
+                ManagedJobScheduleState.WAITING.value,
+            ])).order_by(
+                job_info_table.c.priority.desc(),
+                job_info_table.c.spot_job_id.asc(),
+            ).limit(1).with_for_update()
+
+        # Execute the select with row locking
+        result = await session.execute(select_query)
+        waiting_job_row = result.fetchone()
+
+        if waiting_job_row is None:
+            return None
+
+        job_id = waiting_job_row[0]
+        current_state = ManagedJobScheduleState(waiting_job_row[1])
+        dag_yaml_path = waiting_job_row[2]
+        env_file_path = waiting_job_row[3]
+        controller_pid = waiting_job_row[4]
+        pool = waiting_job_row[5]
+
+        # Update the job state to LAUNCHING
+        update_result = await session.execute(
+            sqlalchemy.update(job_info_table).where(
+                sqlalchemy.and_(
+                    job_info_table.c.spot_job_id == job_id,
+                    job_info_table.c.schedule_state == current_state.value,
+                )).values({
+                    job_info_table.c.schedule_state:
+                        ManagedJobScheduleState.LAUNCHING.value,
+                    job_info_table.c.controller_pid: pid,
+                }))
+
+        if update_result.rowcount != 1:
+            # Update failed, rollback and return None
+            await session.rollback()
+            return None
+
+        # Commit the transaction
+        await session.commit()
+
         return {
+            'job_id': job_id,
+            'schedule_state': current_state,
+            'dag_yaml_path': dag_yaml_path,
+            'env_file_path': env_file_path,
+            'old_pid': controller_pid,
+            'pool': pool,
+        }
+
+
+@_init_db
+def get_workspace(job_id: int) -> str:
+    """Get the workspace of a job."""
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        workspace = session.execute(
+            sqlalchemy.select(job_info_table.c.workspace).where(
+                job_info_table.c.spot_job_id == job_id)).fetchone()
+        job_workspace = workspace[0] if workspace else None
+        if job_workspace is None:
+            return constants.SKYPILOT_DEFAULT_WORKSPACE
+        return job_workspace
+
+
+# === HA Recovery Script functions ===
+
+
+@_init_db
+def get_ha_recovery_script(job_id: int) -> Optional[str]:
+    """Get the HA recovery script for a job."""
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        row = session.query(ha_recovery_script_table).filter_by(
+            job_id=job_id).first()
+    if row is None:
+        return None
+    return row.script
+
+
+@_init_db
+def set_ha_recovery_script(job_id: int, script: str) -> None:
+    """Set the HA recovery script for a job."""
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        if (_SQLALCHEMY_ENGINE.dialect.name ==
+                db_utils.SQLAlchemyDialect.SQLITE.value):
+            insert_func = sqlite.insert
+        elif (_SQLALCHEMY_ENGINE.dialect.name ==
+              db_utils.SQLAlchemyDialect.POSTGRESQL.value):
+            insert_func = postgresql.insert
+        else:
+            raise ValueError('Unsupported database dialect')
+        insert_stmt = insert_func(ha_recovery_script_table).values(
+            job_id=job_id, script=script)
+        do_update_stmt = insert_stmt.on_conflict_do_update(
+            index_elements=[ha_recovery_script_table.c.job_id],
+            set_={ha_recovery_script_table.c.script: script})
+        session.execute(do_update_stmt)
+        session.commit()
+
+
+@_init_db
+def remove_ha_recovery_script(job_id: int) -> None:
+    """Remove the HA recovery script for a job."""
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        session.query(ha_recovery_script_table).filter_by(
+            job_id=job_id).delete()
+        session.commit()
+
+
+@_init_db_async
+async def get_latest_task_id_status_async(
+        job_id: int) -> Union[Tuple[int, ManagedJobStatus], Tuple[None, None]]:
+    """Returns the (task id, status) of the latest task of a job."""
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        result = await session.execute(
+            sqlalchemy.select(
+                spot_table.c.task_id,
+                spot_table.c.status,
+            ).where(spot_table.c.spot_job_id == job_id).order_by(
+                spot_table.c.task_id.asc()))
+        id_statuses = [
+            (row[0], ManagedJobStatus(row[1])) for row in result.fetchall()
+        ]
+
+    if not id_statuses:
+        return None, None
+    task_id, status = next(
+        ((tid, st) for tid, st in id_statuses if not st.is_terminal()),
+        id_statuses[-1],
+    )
+    return task_id, status
+
+
+@_init_db_async
+async def set_starting_async(job_id: int, task_id: int, run_timestamp: str,
+                             submit_time: float, resources_str: str,
+                             specs: Dict[str, Union[str, int]],
+                             callback_func: AsyncCallbackType):
+    """Set the task to starting state."""
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    logger.info('Launching the spot cluster...')
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        result = await session.execute(
+            sqlalchemy.update(spot_table).where(
+                sqlalchemy.and_(
+                    spot_table.c.spot_job_id == job_id,
+                    spot_table.c.task_id == task_id,
+                    spot_table.c.status == ManagedJobStatus.PENDING.value,
+                    spot_table.c.end_at.is_(None),
+                )).values({
+                    spot_table.c.resources: resources_str,
+                    spot_table.c.submitted_at: submit_time,
+                    spot_table.c.status: ManagedJobStatus.STARTING.value,
+                    spot_table.c.run_timestamp: run_timestamp,
+                    spot_table.c.specs: json.dumps(specs),
+                }))
+        count = result.rowcount
+        await session.commit()
+        if count != 1:
+            details = await _describe_task_transition_failure(
+                session, job_id, task_id)
+            message = ('Failed to set the task to starting. '
+                       f'({count} rows updated. {details})')
+            logger.error(message)
+            raise exceptions.ManagedJobStatusError(message)
+    await callback_func('SUBMITTED')
+    await callback_func('STARTING')
+
+
+@_init_db_async
+async def set_started_async(job_id: int, task_id: int, start_time: float,
+                            callback_func: AsyncCallbackType):
+    """Set the task to started state."""
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    logger.info('Job started.')
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        result = await session.execute(
+            sqlalchemy.update(spot_table).where(
+                sqlalchemy.and_(
+                    spot_table.c.spot_job_id == job_id,
+                    spot_table.c.task_id == task_id,
+                    spot_table.c.status.in_([
+                        ManagedJobStatus.STARTING.value,
+                        ManagedJobStatus.PENDING.value
+                    ]),
+                    spot_table.c.end_at.is_(None),
+                )).values({
+                    spot_table.c.status: ManagedJobStatus.RUNNING.value,
+                    spot_table.c.start_at: start_time,
+                    spot_table.c.last_recovered_at: start_time,
+                }))
+        count = result.rowcount
+        await session.commit()
+        if count != 1:
+            details = await _describe_task_transition_failure(
+                session, job_id, task_id)
+            message = (f'Failed to set the task to started. '
+                       f'({count} rows updated. {details})')
+            logger.error(message)
+            raise exceptions.ManagedJobStatusError(message)
+    await callback_func('STARTED')
+
+
+@_init_db_async
+async def get_job_status_with_task_id_async(
+        job_id: int, task_id: int) -> Optional[ManagedJobStatus]:
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        result = await session.execute(
+            sqlalchemy.select(spot_table.c.status).where(
+                sqlalchemy.and_(spot_table.c.spot_job_id == job_id,
+                                spot_table.c.task_id == task_id)))
+        status = result.fetchone()
+        return ManagedJobStatus(status[0]) if status else None
+
+
+@_init_db_async
+async def set_recovering_async(job_id: int, task_id: int,
+                               force_transit_to_recovering: bool,
+                               callback_func: AsyncCallbackType):
+    """Set the task to recovering state, and update the job duration."""
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    logger.info('=== Recovering... ===')
+    current_time = time.time()
+
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        if force_transit_to_recovering:
+            status_condition = spot_table.c.status.in_(
+                [s.value for s in ManagedJobStatus.processing_statuses()])
+        else:
+            status_condition = (
+                spot_table.c.status == ManagedJobStatus.RUNNING.value)
+
+        result = await session.execute(
+            sqlalchemy.update(spot_table).where(
+                sqlalchemy.and_(
+                    spot_table.c.spot_job_id == job_id,
+                    spot_table.c.task_id == task_id,
+                    status_condition,
+                    spot_table.c.end_at.is_(None),
+                )).values({
+                    spot_table.c.status: ManagedJobStatus.RECOVERING.value,
+                    spot_table.c.job_duration: sqlalchemy.case(
+                        (spot_table.c.last_recovered_at >= 0,
+                         spot_table.c.job_duration + current_time -
+                         spot_table.c.last_recovered_at),
+                        else_=spot_table.c.job_duration),
+                    spot_table.c.last_recovered_at: sqlalchemy.case(
+                        (spot_table.c.last_recovered_at < 0, current_time),
+                        else_=spot_table.c.last_recovered_at),
+                }))
+        count = result.rowcount
+        await session.commit()
+        if count != 1:
+            details = await _describe_task_transition_failure(
+                session, job_id, task_id)
+            message = ('Failed to set the task to recovering with '
+                       'force_transit_to_recovering='
+                       f'{force_transit_to_recovering}. '
+                       f'({count} rows updated. {details})')
+            logger.error(message)
+            raise exceptions.ManagedJobStatusError(message)
+    await callback_func('RECOVERING')
+
+
+@_init_db_async
+async def set_recovered_async(job_id: int, task_id: int, recovered_time: float,
+                              callback_func: AsyncCallbackType):
+    """Set the task to recovered."""
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        result = await session.execute(
+            sqlalchemy.update(spot_table).where(
+                sqlalchemy.and_(
+                    spot_table.c.spot_job_id == job_id,
+                    spot_table.c.task_id == task_id,
+                    spot_table.c.status == ManagedJobStatus.RECOVERING.value,
+                    spot_table.c.end_at.is_(None),
+                )).values({
+                    spot_table.c.status: ManagedJobStatus.RUNNING.value,
+                    spot_table.c.last_recovered_at: recovered_time,
+                    spot_table.c.recovery_count: spot_table.c.recovery_count +
+                                                 1,
+                }))
+        count = result.rowcount
+        await session.commit()
+        if count != 1:
+            details = await _describe_task_transition_failure(
+                session, job_id, task_id)
+            message = (f'Failed to set the task to recovered. '
+                       f'({count} rows updated. {details})')
+            logger.error(message)
+            raise exceptions.ManagedJobStatusError(message)
+    logger.info('==== Recovered. ====')
+    await callback_func('RECOVERED')
+
+
+@_init_db_async
+async def set_succeeded_async(job_id: int, task_id: int, end_time: float,
+                              callback_func: AsyncCallbackType):
+    """Set the task to succeeded, if it is in a non-terminal state."""
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        result = await session.execute(
+            sqlalchemy.update(spot_table).where(
+                sqlalchemy.and_(
+                    spot_table.c.spot_job_id == job_id,
+                    spot_table.c.task_id == task_id,
+                    spot_table.c.status == ManagedJobStatus.RUNNING.value,
+                    spot_table.c.end_at.is_(None),
+                )).values({
+                    spot_table.c.status: ManagedJobStatus.SUCCEEDED.value,
+                    spot_table.c.end_at: end_time,
+                }))
+        count = result.rowcount
+        await session.commit()
+        if count != 1:
+            details = await _describe_task_transition_failure(
+                session, job_id, task_id)
+            message = (f'Failed to set the task to succeeded. '
+                       f'({count} rows updated. {details})')
+            logger.error(message)
+            raise exceptions.ManagedJobStatusError(message)
+    await callback_func('SUCCEEDED')
+    logger.info('Job succeeded.')
+
+
+@_init_db_async
+async def set_failed_async(
+    job_id: int,
+    task_id: Optional[int],
+    failure_type: ManagedJobStatus,
+    failure_reason: str,
+    callback_func: Optional[AsyncCallbackType] = None,
+    end_time: Optional[float] = None,
+    override_terminal: bool = False,
+):
+    """Set an entire job or task to failed."""
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    assert failure_type.is_failed(), failure_type
+    end_time = time.time() if end_time is None else end_time
+
+    fields_to_set: Dict[str, Any] = {
+        spot_table.c.status: failure_type.value,
+        spot_table.c.failure_reason: failure_reason,
+    }
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        # Get previous status
+        result = await session.execute(
+            sqlalchemy.select(
+                spot_table.c.status).where(spot_table.c.spot_job_id == job_id))
+        previous_status_row = result.fetchone()
+        previous_status = ManagedJobStatus(previous_status_row[0])
+        if previous_status == ManagedJobStatus.RECOVERING:
+            fields_to_set[spot_table.c.last_recovered_at] = end_time
+        where_conditions = [spot_table.c.spot_job_id == job_id]
+        if task_id is not None:
+            where_conditions.append(spot_table.c.task_id == task_id)
+
+        # Handle failure_reason prepending when override_terminal is True
+        if override_terminal:
+            # Get existing failure_reason with row lock to prevent race
+            # conditions
+            existing_reason_result = await session.execute(
+                sqlalchemy.select(spot_table.c.failure_reason).where(
+                    sqlalchemy.and_(*where_conditions)).with_for_update())
+            existing_reason_row = existing_reason_result.fetchone()
+            if existing_reason_row and existing_reason_row[0]:
+                # Prepend new failure reason to existing one
+                fields_to_set[spot_table.c.failure_reason] = (
+                    failure_reason + '. Previously: ' + existing_reason_row[0])
+            fields_to_set[spot_table.c.end_at] = sqlalchemy.func.coalesce(
+                spot_table.c.end_at, end_time)
+        else:
+            fields_to_set[spot_table.c.end_at] = end_time
+            where_conditions.append(spot_table.c.end_at.is_(None))
+        result = await session.execute(
+            sqlalchemy.update(spot_table).where(
+                sqlalchemy.and_(*where_conditions)).values(fields_to_set))
+        count = result.rowcount
+        await session.commit()
+        updated = count > 0
+    if callback_func and updated:
+        await callback_func('FAILED')
+    logger.info(failure_reason)
+
+
+@_init_db_async
+async def set_cancelling_async(job_id: int, callback_func: AsyncCallbackType):
+    """Set tasks in the job as cancelling, if they are in non-terminal
+    states."""
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        result = await session.execute(
+            sqlalchemy.update(spot_table).where(
+                sqlalchemy.and_(
+                    spot_table.c.spot_job_id == job_id,
+                    spot_table.c.end_at.is_(None),
+                )).values(
+                    {spot_table.c.status: ManagedJobStatus.CANCELLING.value}))
+        count = result.rowcount
+        await session.commit()
+        updated = count > 0
+    if updated:
+        logger.info('Cancelling the job...')
+        await callback_func('CANCELLING')
+    else:
+        logger.info('Cancellation skipped, job is already terminal')
+
+
+@_init_db_async
+async def set_cancelled_async(job_id: int, callback_func: AsyncCallbackType):
+    """Set tasks in the job as cancelled, if they are in CANCELLING state."""
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        result = await session.execute(
+            sqlalchemy.update(spot_table).where(
+                sqlalchemy.and_(
+                    spot_table.c.spot_job_id == job_id,
+                    spot_table.c.status == ManagedJobStatus.CANCELLING.value,
+                )).values({
+                    spot_table.c.status: ManagedJobStatus.CANCELLED.value,
+                    spot_table.c.end_at: time.time(),
+                }))
+        count = result.rowcount
+        await session.commit()
+        updated = count > 0
+    if updated:
+        logger.info('Job cancelled.')
+        await callback_func('CANCELLED')
+    else:
+        logger.info('Cancellation skipped, job is not CANCELLING')
+
+
+@_init_db_async
+async def remove_ha_recovery_script_async(job_id: int) -> None:
+    """Remove the HA recovery script for a job."""
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        await session.execute(
+            sqlalchemy.delete(ha_recovery_script_table).where(
+                ha_recovery_script_table.c.job_id == job_id))
+        await session.commit()
+
+
+async def get_status_async(job_id: int) -> Optional[ManagedJobStatus]:
+    _, status = await get_latest_task_id_status_async(job_id)
+    return status
+
+
+@_init_db_async
+async def get_job_schedule_state_async(job_id: int) -> ManagedJobScheduleState:
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        result = await session.execute(
+            sqlalchemy.select(job_info_table.c.schedule_state).where(
+                job_info_table.c.spot_job_id == job_id))
+        state = result.fetchone()[0]
+        return ManagedJobScheduleState(state)
+
+
+@_init_db_async
+async def scheduler_set_done_async(job_id: int,
+                                   idempotent: bool = False) -> None:
+    """Do not call without holding the scheduler lock."""
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        result = await session.execute(
+            sqlalchemy.update(job_info_table).where(
+                sqlalchemy.and_(
+                    job_info_table.c.spot_job_id == job_id,
+                    job_info_table.c.schedule_state !=
+                    ManagedJobScheduleState.DONE.value,
+                )).values({
+                    job_info_table.c.schedule_state:
+                        ManagedJobScheduleState.DONE.value
+                }))
+        updated_count = result.rowcount
+        await session.commit()
+        if not idempotent:
+            assert updated_count == 1, (job_id, updated_count)
+
+
+# ==== needed for codegen ====
+# functions have no use outside of codegen, remove at your own peril
+
+
+@_init_db
+def set_job_info(job_id: int,
+                 name: str,
+                 workspace: str,
+                 entrypoint: str,
+                 pool: Optional[str],
+                 pool_hash: Optional[str],
+                 user_hash: Optional[str] = None):
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        if (_SQLALCHEMY_ENGINE.dialect.name ==
+                db_utils.SQLAlchemyDialect.SQLITE.value):
+            insert_func = sqlite.insert
+        elif (_SQLALCHEMY_ENGINE.dialect.name ==
+              db_utils.SQLAlchemyDialect.POSTGRESQL.value):
+            insert_func = postgresql.insert
+        else:
+            raise ValueError('Unsupported database dialect')
+        insert_stmt = insert_func(job_info_table).values(
+            spot_job_id=job_id,
+            name=name,
+            schedule_state=ManagedJobScheduleState.INACTIVE.value,
+            workspace=workspace,
+            entrypoint=entrypoint,
+            pool=pool,
+            pool_hash=pool_hash,
+            user_hash=user_hash,
+        )
+        session.execute(insert_stmt)
+        session.commit()
+
+
+@_init_db
+def reset_jobs_for_recovery() -> None:
+    """Remove controller PIDs for live jobs, allowing them to be recovered."""
+    assert _SQLALCHEMY_ENGINE is not None
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        session.query(job_info_table).filter(
+            # PID should be set.
+            job_info_table.c.controller_pid.isnot(None),
+            # Schedule state should be alive.
+            job_info_table.c.schedule_state.isnot(None),
+            (job_info_table.c.schedule_state !=
+             ManagedJobScheduleState.INVALID.value),
+            (job_info_table.c.schedule_state !=
+             ManagedJobScheduleState.WAITING.value),
+            (job_info_table.c.schedule_state !=
+             ManagedJobScheduleState.DONE.value),
+        ).update({
+            job_info_table.c.controller_pid: None,
+            job_info_table.c.schedule_state:
+                (ManagedJobScheduleState.WAITING.value)
+        })
+        session.commit()
+
+
+@_init_db
+def get_all_job_ids_by_name(name: Optional[str]) -> List[int]:
+    """Get all job ids by name."""
+    assert _SQLALCHEMY_ENGINE is not None
+
+    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+        query = sqlalchemy.select(
+            spot_table.c.spot_job_id.distinct()).select_from(
+                spot_table.outerjoin(
+                    job_info_table,
+                    spot_table.c.spot_job_id == job_info_table.c.spot_job_id))
+        if name is not None:
+            # We match the job name from `job_info` for the jobs submitted after
+            # #1982, and from `spot` for the jobs submitted before #1982, whose
+            # job_info is not available.
+            name_condition = sqlalchemy.or_(
+                job_info_table.c.name == name,
+                sqlalchemy.and_(job_info_table.c.name.is_(None),
+                                spot_table.c.task_name == name))
+            query = query.where(name_condition)
+        query = query.order_by(spot_table.c.spot_job_id.desc())
+        rows = session.execute(query).fetchall()
+        job_ids = [row[0] for row in rows if row[0] is not None]
+        return job_ids
+
+
+@_init_db_async
+async def get_task_logs_to_clean_async(retention_seconds: int,
+                                       batch_size) -> List[Dict[str, Any]]:
+    """Get the logs of job tasks to clean.
+
+    The logs of a task will only cleaned when:
+    - the job schedule state is DONE
+    - AND the end time of the task is older than the retention period
+    """
+
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        now = time.time()
+        result = await session.execute(
+            sqlalchemy.select(
+                spot_table.c.spot_job_id,
+                spot_table.c.task_id,
+                spot_table.c.local_log_file,
+            ).select_from(
+                spot_table.join(
+                    job_info_table,
+                    spot_table.c.spot_job_id == job_info_table.c.spot_job_id,
+                )).
+            where(
+                sqlalchemy.and_(
+                    job_info_table.c.schedule_state.is_(
+                        ManagedJobScheduleState.DONE.value),
+                    spot_table.c.end_at.isnot(None),
+                    spot_table.c.end_at < (now - retention_seconds),
+                    spot_table.c.logs_cleaned_at.is_(None),
+                    # The local log file is set AFTER the task is finished,
+                    # add this condition to ensure the entire log file has
+                    # been written.
+                    spot_table.c.local_log_file.isnot(None),
+                )).limit(batch_size))
+        rows = result.fetchall()
+        return [{
             'job_id': row[0],
-            'schedule_state': ManagedJobScheduleState(row[1]),
-            'dag_yaml_path': row[2],
-            'env_file_path': row[3],
-        } if row is not None else None
+            'task_id': row[1],
+            'local_log_file': row[2]
+        } for row in rows]
+
+
+@_init_db_async
+async def get_controller_logs_to_clean_async(
+        retention_seconds: int, batch_size: int) -> List[Dict[str, Any]]:
+    """Get the controller logs to clean.
+
+    The controller logs will only cleaned when:
+    - the job schedule state is DONE
+    - AND the end time of the latest task is older than the retention period
+    """
+
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        now = time.time()
+
+        result = await session.execute(
+            sqlalchemy.select(job_info_table.c.spot_job_id,).select_from(
+                job_info_table.join(
+                    spot_table,
+                    job_info_table.c.spot_job_id == spot_table.c.spot_job_id,
+                )).where(
+                    sqlalchemy.and_(
+                        job_info_table.c.schedule_state.is_(
+                            ManagedJobScheduleState.DONE.value),
+                        spot_table.c.local_log_file.isnot(None),
+                        job_info_table.c.controller_logs_cleaned_at.is_(None),
+                    )).group_by(
+                        job_info_table.c.spot_job_id,
+                        job_info_table.c.current_cluster_name,
+                    ).having(
+                        sqlalchemy.func.max(
+                            spot_table.c.end_at).isnot(None),).having(
+                                sqlalchemy.func.max(spot_table.c.end_at) < (
+                                    now - retention_seconds)).limit(batch_size))
+        rows = result.fetchall()
+        return [{'job_id': row[0]} for row in rows]
+
+
+@_init_db_async
+async def set_task_logs_cleaned_async(tasks: List[Tuple[int, int]],
+                                      logs_cleaned_at: float):
+    """Set the task logs cleaned at."""
+    if not tasks:
+        return
+    # Deduplicate
+    task_keys = list(dict.fromkeys(tasks))
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        await session.execute(
+            sqlalchemy.update(spot_table).where(
+                sqlalchemy.tuple_(spot_table.c.spot_job_id,
+                                  spot_table.c.task_id).in_(task_keys)).values(
+                                      logs_cleaned_at=logs_cleaned_at))
+        await session.commit()
+
+
+@_init_db_async
+async def set_controller_logs_cleaned_async(job_ids: List[int],
+                                            logs_cleaned_at: float):
+    """Set the controller logs cleaned at."""
+    if not job_ids:
+        return
+    # Deduplicate
+    job_ids = list(dict.fromkeys(job_ids))
+    assert _SQLALCHEMY_ENGINE_ASYNC is not None
+    async with sql_async.AsyncSession(_SQLALCHEMY_ENGINE_ASYNC) as session:
+        await session.execute(
+            sqlalchemy.update(job_info_table).where(
+                job_info_table.c.spot_job_id.in_(job_ids)).values(
+                    controller_logs_cleaned_at=logs_cleaned_at))
+        await session.commit()