skypilot-nightly 1.0.0.dev20251021__py3-none-any.whl → 1.0.0.dev20251022__py3-none-any.whl

sky/server/requests/payloads.py

@@ -542,6 +542,9 @@ class JobsQueueV2Body(RequestBody):
     page: Optional[int] = None
     limit: Optional[int] = None
     statuses: Optional[List[str]] = None
+    # The fields to return in the response.
+    # Refer to the fields in the `class ManagedJobRecord` in `response.py`
+    fields: Optional[List[str]] = None
 
 
 class JobsCancelBody(RequestBody):

sky/server/requests/requests.py

@@ -16,6 +16,7 @@ import time
 import traceback
 from typing import (Any, Callable, Dict, Generator, List, NamedTuple, Optional,
                     Tuple)
+import uuid
 
 import anyio
 import colorama
@@ -293,6 +294,11 @@ class Request:
             raise
 
 
+def get_new_request_id() -> str:
+    """Get a new request ID."""
+    return str(uuid.uuid4())
+
+
 def encode_requests(requests: List[Request]) -> List[payloads.RequestPayload]:
     """Serialize the SkyPilot API request for display purposes.
 
@@ -657,17 +663,15 @@ async def _get_request_no_lock_async(
     return Request.from_row(row)
 
 
-@init_db
+@init_db_async
 @metrics_lib.time_me
-def get_latest_request_id() -> Optional[str]:
+async def get_latest_request_id_async() -> Optional[str]:
     """Get the latest request ID."""
     assert _DB is not None
-    with _DB.conn:
-        cursor = _DB.conn.cursor()
-        cursor.execute(f'SELECT request_id FROM {REQUEST_TABLE} '
-                       'ORDER BY created_at DESC LIMIT 1')
-        row = cursor.fetchone()
-        return row[0] if row else None
+    async with _DB.execute_fetchall_async(
+        (f'SELECT request_id FROM {REQUEST_TABLE} '
+         'ORDER BY created_at DESC LIMIT 1')) as rows:
+        return rows[0][0] if rows else None
 
 
 @init_db
@@ -725,17 +729,6 @@ async def get_request_status_async(
         return StatusWithMsg(status, status_msg)
 
 
-@init_db
-@metrics_lib.time_me
-def create_if_not_exists(request: Request) -> bool:
-    """Create a SkyPilot API request if it does not exist."""
-    with filelock.FileLock(request_lock_path(request.request_id)):
-        if _get_request_no_lock(request.request_id) is not None:
-            return False
-        _add_or_update_request_no_lock(request)
-        return True
-
-
 @init_db_async
 @metrics_lib.time_me_async
 @asyncio_utils.shield

sky/server/requests/serializers/encoders.py

@@ -121,7 +121,7 @@ def encode_status_kubernetes(
         encoded_cluster = dataclasses.asdict(cluster)
         encoded_cluster['status'] = encoded_cluster['status'].value
         encoded_unmanaged_clusters.append(encoded_cluster)
-    all_jobs = [job.model_dump() for job in all_jobs]
+    all_jobs = [job.model_dump(by_alias=True) for job in all_jobs]
     return encoded_all_clusters, encoded_unmanaged_clusters, all_jobs, context
 
 
@@ -151,9 +151,9 @@ def encode_jobs_queue_v2(
     for job in jobs:
         job['status'] = job['status'].value
     if total is None:
-        return [job.model_dump() for job in jobs]
+        return [job.model_dump(by_alias=True) for job in jobs]
     return {
-        'jobs': [job.model_dump() for job in jobs],
+        'jobs': [job.model_dump(by_alias=True) for job in jobs],
         'total': total,
         'total_no_filter': total_no_filter,
         'status_counts': status_counts

sky/server/server.py

@@ -163,7 +163,7 @@ class RequestIDMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
     """Middleware to add a request ID to each request."""
 
     async def dispatch(self, request: fastapi.Request, call_next):
-        request_id = str(uuid.uuid4())
+        request_id = requests_lib.get_new_request_id()
         request.state.request_id = request_id
         response = await call_next(request)
         # TODO(syang): remove X-Request-ID when v0.10.0 is released.
@@ -455,9 +455,9 @@ async def loop_lag_monitor(loop: asyncio.AbstractEventLoop,
     loop.call_at(target, tick)
 
 
-def schedule_on_boot_check():
+async def schedule_on_boot_check_async():
     try:
-        executor.schedule_request(
+        await executor.schedule_request_async(
             request_id='skypilot-server-on-boot-check',
             request_name='check',
             request_body=payloads.CheckBody(),
@@ -480,7 +480,7 @@ async def lifespan(app: fastapi.FastAPI):  # pylint: disable=redefined-outer-nam
         if event.should_skip():
             continue
         try:
-            executor.schedule_request(
+            await executor.schedule_request_async(
                 request_id=event.id,
                 request_name=event.name,
                 request_body=payloads.RequestBody(),
@@ -495,7 +495,7 @@ async def lifespan(app: fastapi.FastAPI):  # pylint: disable=redefined-outer-nam
             # Lifespan will be executed in each uvicorn worker process, we
             # can safely ignore the error if the task is already scheduled.
             logger.debug(f'Request {event.id} already exists.')
-    schedule_on_boot_check()
+    await schedule_on_boot_check_async()
     asyncio.create_task(cleanup_upload_ids())
     if metrics_utils.METRICS_ENABLED:
         # Start monitoring the event loop lag in each server worker
@@ -729,7 +729,7 @@ async def token(request: fastapi.Request,
 async def check(request: fastapi.Request,
                 check_body: payloads.CheckBody) -> None:
     """Checks enabled clouds."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='check',
         request_body=check_body,
@@ -743,7 +743,7 @@ async def enabled_clouds(request: fastapi.Request,
                          workspace: Optional[str] = None,
                          expand: bool = False) -> None:
     """Gets enabled clouds on the server."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='enabled_clouds',
         request_body=payloads.EnabledCloudsBody(workspace=workspace,
@@ -759,7 +759,7 @@ async def realtime_kubernetes_gpu_availability(
     realtime_gpu_availability_body: payloads.RealtimeGpuAvailabilityRequestBody
 ) -> None:
     """Gets real-time Kubernetes GPU availability."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='realtime_kubernetes_gpu_availability',
         request_body=realtime_gpu_availability_body,
@@ -774,7 +774,7 @@ async def kubernetes_node_info(
         kubernetes_node_info_body: payloads.KubernetesNodeInfoRequestBody
 ) -> None:
     """Gets Kubernetes nodes information and hints."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='kubernetes_node_info',
         request_body=kubernetes_node_info_body,
@@ -786,7 +786,7 @@ async def kubernetes_node_info(
 @app.get('/status_kubernetes')
 async def status_kubernetes(request: fastapi.Request) -> None:
     """Gets Kubernetes status."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='status_kubernetes',
         request_body=payloads.RequestBody(),
@@ -800,7 +800,7 @@ async def list_accelerators(
         request: fastapi.Request,
         list_accelerator_counts_body: payloads.ListAcceleratorsBody) -> None:
     """Gets list of accelerators from cloud catalog."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='list_accelerators',
         request_body=list_accelerator_counts_body,
@@ -815,7 +815,7 @@ async def list_accelerator_counts(
         list_accelerator_counts_body: payloads.ListAcceleratorCountsBody
 ) -> None:
     """Gets list of accelerator counts from cloud catalog."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='list_accelerator_counts',
         request_body=list_accelerator_counts_body,
@@ -872,7 +872,7 @@ async def validate(validate_body: payloads.ValidateBody) -> None:
 async def optimize(optimize_body: payloads.OptimizeBody,
                    request: fastapi.Request) -> None:
     """Optimizes the user's DAG."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='optimize',
         request_body=optimize_body,
@@ -1082,7 +1082,7 @@ async def launch(launch_body: payloads.LaunchBody,
     """Launches a cluster or task."""
     request_id = request.state.request_id
     logger.info(f'Launching request: {request_id}')
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id,
         request_name='launch',
         request_body=launch_body,
@@ -1098,7 +1098,7 @@ async def launch(launch_body: payloads.LaunchBody,
 async def exec(request: fastapi.Request, exec_body: payloads.ExecBody) -> None:
     """Executes a task on an existing cluster."""
     cluster_name = exec_body.cluster_name
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='exec',
         request_body=exec_body,
@@ -1116,7 +1116,7 @@ async def exec(request: fastapi.Request, exec_body: payloads.ExecBody) -> None:
 async def stop(request: fastapi.Request,
                stop_body: payloads.StopOrDownBody) -> None:
     """Stops a cluster."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='stop',
         request_body=stop_body,
@@ -1136,7 +1136,7 @@ async def status(
         raise fastapi.HTTPException(
             status_code=503,
             detail='Server is shutting down, please try again later.')
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='status',
         request_body=status_body,
@@ -1151,7 +1151,7 @@ async def status(
 async def endpoints(request: fastapi.Request,
                     endpoint_body: payloads.EndpointsBody) -> None:
     """Gets the endpoint for a given cluster and port number (endpoint)."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='endpoints',
         request_body=endpoint_body,
@@ -1165,7 +1165,7 @@ async def endpoints(request: fastapi.Request,
 async def down(request: fastapi.Request,
                down_body: payloads.StopOrDownBody) -> None:
     """Tears down a cluster."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='down',
         request_body=down_body,
@@ -1179,7 +1179,7 @@ async def down(request: fastapi.Request,
 async def start(request: fastapi.Request,
                 start_body: payloads.StartBody) -> None:
     """Restarts a cluster."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='start',
         request_body=start_body,
@@ -1193,7 +1193,7 @@ async def start(request: fastapi.Request,
 async def autostop(request: fastapi.Request,
                    autostop_body: payloads.AutostopBody) -> None:
     """Schedules an autostop/autodown for a cluster."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='autostop',
         request_body=autostop_body,
@@ -1207,7 +1207,7 @@ async def autostop(request: fastapi.Request,
 async def queue(request: fastapi.Request,
                 queue_body: payloads.QueueBody) -> None:
     """Gets the job queue of a cluster."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='queue',
         request_body=queue_body,
@@ -1221,7 +1221,7 @@ async def queue(request: fastapi.Request,
 async def job_status(request: fastapi.Request,
                      job_status_body: payloads.JobStatusBody) -> None:
     """Gets the status of a job."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='job_status',
         request_body=job_status_body,
@@ -1235,7 +1235,7 @@ async def job_status(request: fastapi.Request,
 async def cancel(request: fastapi.Request,
                  cancel_body: payloads.CancelBody) -> None:
     """Cancels jobs on a cluster."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='cancel',
         request_body=cancel_body,
@@ -1255,7 +1255,7 @@ async def logs(
     # launch, to finish, so that a user does not need to manually pull the
     # request status.
     executor.check_request_thread_executor_available()
-    request_task = executor.prepare_request(
+    request_task = await executor.prepare_request_async(
         request_id=request.state.request_id,
         request_name='logs',
         request_body=cluster_job_body,
@@ -1286,7 +1286,7 @@ async def download_logs(
     # We should reuse the original request body, so that the env vars, such as
     # user hash, are kept the same.
     cluster_jobs_body.local_dir = str(logs_dir_on_api_server)
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='download_logs',
         request_body=cluster_jobs_body,
@@ -1437,7 +1437,7 @@ def provision_logs(provision_logs_body: payloads.ProvisionLogsBody,
 async def cost_report(request: fastapi.Request,
                       cost_report_body: payloads.CostReportBody) -> None:
     """Gets the cost report of a cluster."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='cost_report',
         request_body=cost_report_body,
@@ -1449,7 +1449,7 @@ async def cost_report(request: fastapi.Request,
 @app.get('/storage/ls')
 async def storage_ls(request: fastapi.Request) -> None:
     """Gets the storages."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='storage_ls',
         request_body=payloads.RequestBody(),
@@ -1462,7 +1462,7 @@ async def storage_ls(request: fastapi.Request) -> None:
 async def storage_delete(request: fastapi.Request,
                          storage_body: payloads.StorageBody) -> None:
     """Deletes a storage."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='storage_delete',
         request_body=storage_body,
@@ -1475,7 +1475,7 @@ async def storage_delete(request: fastapi.Request,
 async def local_up(request: fastapi.Request,
                    local_up_body: payloads.LocalUpBody) -> None:
     """Launches a Kubernetes cluster on API server."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='local_up',
         request_body=local_up_body,
@@ -1488,7 +1488,7 @@ async def local_up(request: fastapi.Request,
 async def local_down(request: fastapi.Request,
                      local_down_body: payloads.LocalDownBody) -> None:
     """Tears down the Kubernetes cluster started by local_up."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='local_down',
         request_body=local_down_body,
@@ -1566,7 +1566,7 @@ async def stream(
             detail='Only one of request_id and log_path can be provided')
 
     if request_id is None and log_path is None:
-        request_id = requests_lib.get_latest_request_id()
+        request_id = await requests_lib.get_latest_request_id_async()
         if request_id is None:
             raise fastapi.HTTPException(status_code=404,
                                         detail='No request found')
@@ -1672,7 +1672,7 @@ async def stream(
 async def api_cancel(request: fastapi.Request,
                      request_cancel_body: payloads.RequestCancelBody) -> None:
     """Cancels requests."""
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='api_cancel',
         request_body=request_cancel_body,
@@ -1908,7 +1908,7 @@ async def kubernetes_pod_ssh_proxy(websocket: fastapi.WebSocket,
 async def all_contexts(request: fastapi.Request) -> None:
     """Gets all Kubernetes and SSH node pool contexts."""
 
-    executor.schedule_request(
+    await executor.schedule_request_async(
         request_id=request.state.request_id,
         request_name='all_contexts',
         request_body=payloads.RequestBody(),

sky/setup_files/alembic.ini

@@ -98,6 +98,10 @@ version_table = alembic_version_spot_jobs_db
 version_locations = %(here)s/../schemas/db/serve_state
 version_table = alembic_version_serve_state_db
 
+[sky_config_db]
+version_locations = %(here)s/../schemas/db/skypilot_config
+version_table = alembic_version_sky_config_db
+
 [post_write_hooks]
 # post_write_hooks defines scripts or Python functions that are run
 # on newly generated revision scripts.  See the documentation for further

sky/skylet/services.py

@@ -408,17 +408,17 @@ class ManagedJobsServiceImpl(managed_jobsv1_pb2_grpc.ManagedJobsServiceServicer
     ) -> managed_jobsv1_pb2.GetJobTableResponse:
         try:
             accessible_workspaces = list(request.accessible_workspaces)
-            job_ids = list(request.job_ids.ids) if request.job_ids else None
+            job_ids = (list(request.job_ids.ids)
+                       if request.HasField('job_ids') else None)
             user_hashes: Optional[List[Optional[str]]] = None
-            if request.user_hashes:
+            if request.HasField('user_hashes'):
                 user_hashes = list(request.user_hashes.hashes)
                 # For backwards compatibility, we show jobs that do not have a
                 # user_hash. TODO: Remove before 0.12.0.
                 if request.show_jobs_without_user_hash:
                     user_hashes.append(None)
-            statuses = list(
-                request.statuses.statuses) if request.statuses else None
-
+            statuses = (list(request.statuses.statuses)
+                        if request.HasField('statuses') else None)
             job_queue = managed_job_utils.get_managed_job_queue(
                 skip_finished=request.skip_finished,
                 accessible_workspaces=accessible_workspaces,

sky/skypilot_config.py

@@ -64,7 +64,6 @@ from sqlalchemy import orm
 from sqlalchemy.dialects import postgresql
 from sqlalchemy.dialects import sqlite
 from sqlalchemy.ext import declarative
-from sqlalchemy.pool import NullPool
 
 from sky import exceptions
 from sky import sky_logging
@@ -77,6 +76,7 @@ from sky.utils import schemas
 from sky.utils import ux_utils
 from sky.utils import yaml_utils
 from sky.utils.db import db_utils
+from sky.utils.db import migration_utils
 from sky.utils.kubernetes import config_map_utils
 
 if typing.TYPE_CHECKING:
@@ -121,7 +121,8 @@ _PROJECT_CONFIG_PATH = '.sky.yaml'
 
 API_SERVER_CONFIG_KEY = 'api_server_config'
 
-_DB_USE_LOCK = threading.Lock()
+_SQLALCHEMY_ENGINE: Optional[sqlalchemy.engine.Engine] = None
+_SQLALCHEMY_ENGINE_LOCK = threading.Lock()
 
 Base = declarative.declarative_base()
 
@@ -481,7 +482,7 @@ def safe_reload_config() -> None:
         reload_config()
 
 
-def reload_config() -> None:
+def reload_config(init_db: bool = False) -> None:
     internal_config_path = os.environ.get(ENV_VAR_SKYPILOT_CONFIG)
     if internal_config_path is not None:
         # {ENV_VAR_SKYPILOT_CONFIG} is used internally.
@@ -493,7 +494,7 @@ def reload_config() -> None:
         return
 
     if os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None:
-        _reload_config_as_server()
+        _reload_config_as_server(init_db=init_db)
     else:
         _reload_config_as_client()
 
@@ -564,7 +565,43 @@ def _reload_config_from_internal_file(internal_config_path: str) -> None:
     _set_loaded_config_path(config_path)
 
 
-def _reload_config_as_server() -> None:
+def _create_table(engine: sqlalchemy.engine.Engine):
+    """Initialize the config database with migrations."""
+    migration_utils.safe_alembic_upgrade(
+        engine, migration_utils.SKYPILOT_CONFIG_DB_NAME,
+        migration_utils.SKYPILOT_CONFIG_VERSION)
+
+
+def _initialize_and_get_db() -> sqlalchemy.engine.Engine:
+    """Initialize and return the config database engine.
+
+    This function should only be called by the API Server during initialization.
+    Client-side code should never call this function.
+    """
+    assert os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None, (
+        'initialize_and_get_db() can only be called by the API Server')
+
+    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
+
+        # We only store config in the DB when using Postgres,
+        # so no need to pass in db_name here.
+        engine = db_utils.get_engine(None)
+
+        # Run migrations if needed
+        _create_table(engine)
+
+        _SQLALCHEMY_ENGINE = engine
+        return _SQLALCHEMY_ENGINE
+
+
+def _reload_config_as_server(init_db: bool = False) -> None:
     # Reset the global variables, to avoid using stale values.
     _set_loaded_config(config_utils.Config())
     _set_loaded_config_path(None)
@@ -580,37 +617,24 @@ def _reload_config_as_server() -> None:
             raise ValueError(
                 'If db config is specified, no other config is allowed')
         logger.debug('retrieving config from database')
-        with _DB_USE_LOCK:
-            dispose_engine = False
-            if db_utils.get_max_connections() == 0:
-                dispose_engine = True
-                sqlalchemy_engine = sqlalchemy.create_engine(db_url,
-                                                             poolclass=NullPool)
-            else:
-                sqlalchemy_engine = db_utils.get_engine('config')
-            db_utils.add_all_tables_to_db_sqlalchemy(Base.metadata,
-                                                     sqlalchemy_engine)
-
-            def _get_config_yaml_from_db(
-                    key: str) -> Optional[config_utils.Config]:
-                assert sqlalchemy_engine is not None
-                with orm.Session(sqlalchemy_engine) as session:
-                    row = session.query(config_yaml_table).filter_by(
-                        key=key).first()
-                if row:
-                    db_config = config_utils.Config(
-                        yaml_utils.safe_load(row.value))
-                    db_config.pop_nested(('db',), None)
-                    return db_config
-                return None
-
-            db_config = _get_config_yaml_from_db(API_SERVER_CONFIG_KEY)
-            if db_config:
-                server_config = overlay_skypilot_config(server_config,
-                                                        db_config)
-            # Close the engine to avoid connection leaks
-            if dispose_engine:
-                sqlalchemy_engine.dispose()
+
+        if init_db:
+            _initialize_and_get_db()
+
+        def _get_config_yaml_from_db(key: str) -> Optional[config_utils.Config]:
+            assert _SQLALCHEMY_ENGINE is not None
+            with orm.Session(_SQLALCHEMY_ENGINE) as session:
+                row = session.query(config_yaml_table).filter_by(
+                    key=key).first()
+            if row:
+                db_config = config_utils.Config(yaml_utils.safe_load(row.value))
+                db_config.pop_nested(('db',), None)
+                return db_config
+            return None
+
+        db_config = _get_config_yaml_from_db(API_SERVER_CONFIG_KEY)
+        if db_config:
+            server_config = overlay_skypilot_config(server_config, db_config)
     if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
         logger.debug(f'server config: \n'
                      f'{yaml_utils.dump_yaml_str(dict(server_config))}')
@@ -666,7 +690,7 @@ def loaded_config_path_serialized() -> Optional[str]:
 
 
 # Load on import, synchronization is guaranteed by python interpreter.
-reload_config()
+reload_config(init_db=True)
 
 
 def loaded() -> bool:
@@ -880,44 +904,32 @@ def update_api_server_config_no_lock(config: config_utils.Config) -> None:
         if new_db_url and new_db_url != existing_db_url:
             raise ValueError('Cannot change db url while server is running')
         if existing_db_url:
-            with _DB_USE_LOCK:
-                dispose_engine = False
-                if db_utils.get_max_connections() == 0:
-                    dispose_engine = True
-                    sqlalchemy_engine = sqlalchemy.create_engine(
-                        existing_db_url, poolclass=NullPool)
-                else:
-                    sqlalchemy_engine = db_utils.get_engine('config')
-                db_utils.add_all_tables_to_db_sqlalchemy(
-                    Base.metadata, sqlalchemy_engine)
-
-                def _set_config_yaml_to_db(key: str,
-                                           config: config_utils.Config):
-                    assert sqlalchemy_engine is not None
-                    config_str = yaml_utils.dump_yaml_str(dict(config))
-                    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_stmnt = insert_func(config_yaml_table).values(
-                            key=key, value=config_str)
-                        do_update_stmt = insert_stmnt.on_conflict_do_update(
-                            index_elements=[config_yaml_table.c.key],
-                            set_={config_yaml_table.c.value: config_str})
-                        session.execute(do_update_stmt)
-                        session.commit()
-
-                logger.debug('saving api_server config to db')
-                _set_config_yaml_to_db(API_SERVER_CONFIG_KEY, config)
-                db_updated = True
-                # Close the engine to avoid connection leaks
-                if dispose_engine:
-                    sqlalchemy_engine.dispose()
+
+            def _set_config_yaml_to_db(key: str, config: config_utils.Config):
+                # reload_config(init_db=True) is called when this module is
+                # imported, so the database engine must already be initialized.
+                assert _SQLALCHEMY_ENGINE is not None
+                config_str = yaml_utils.dump_yaml_str(dict(config))
+                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_stmnt = insert_func(config_yaml_table).values(
+                        key=key, value=config_str)
+                    do_update_stmt = insert_stmnt.on_conflict_do_update(
+                        index_elements=[config_yaml_table.c.key],
+                        set_={config_yaml_table.c.value: config_str})
+                    session.execute(do_update_stmt)
+                    session.commit()
+
+            logger.debug('saving api_server config to db')
+            _set_config_yaml_to_db(API_SERVER_CONFIG_KEY, config)
+            db_updated = True
 
     if not db_updated:
         # save to the local file (PVC in Kubernetes, local file otherwise)