skypilot-nightly 1.0.0.dev20250902__py3-none-any.whl → 1.0.0.dev20250904__py3-none-any.whl

sky/server/requests/requests.py

@@ -1,5 +1,6 @@
 """Utilities for REST API."""
 import asyncio
+import atexit
 import contextlib
 import dataclasses
 import enum
@@ -14,7 +15,7 @@ import threading
 import time
 import traceback
 from typing import (Any, AsyncContextManager, Callable, Dict, Generator, List,
-                    Optional, Tuple)
+                    NamedTuple, Optional, Tuple)
 
 import colorama
 import filelock
@@ -300,10 +301,11 @@ def kill_cluster_requests(cluster_name: str, exclude_request_name: str):
             prevent killing the caller request.
     """
     request_ids = [
-        request_task.request_id for request_task in get_request_tasks(
+        request_task.request_id
+        for request_task in get_request_tasks(req_filter=RequestTaskFilter(
             cluster_names=[cluster_name],
             status=[RequestStatus.PENDING, RequestStatus.RUNNING],
-            exclude_request_names=[exclude_request_name])
+            exclude_request_names=[exclude_request_name]))
     ]
     kill_requests(request_ids)
 
@@ -323,11 +325,12 @@ def kill_requests(request_ids: Optional[List[str]] = None,
     """
     if request_ids is None:
         request_ids = [
-            request_task.request_id for request_task in get_request_tasks(
+            request_task.request_id
+            for request_task in get_request_tasks(req_filter=RequestTaskFilter(
                 user_id=user_id,
                 status=[RequestStatus.RUNNING, RequestStatus.PENDING],
                 # Avoid cancelling the cancel request itself.
-                exclude_request_names=['sky.api_cancel'])
+                exclude_request_names=['sky.api_cancel']))
         ]
     cancelled_request_ids = []
     for request_id in request_ids:
@@ -548,6 +551,40 @@ async def get_request_async(request_id: str) -> Optional[Request]:
         return await _get_request_no_lock_async(request_id)
 
 
+class StatusWithMsg(NamedTuple):
+    status: RequestStatus
+    status_msg: Optional[str] = None
+
+
+@init_db_async
+@metrics_lib.time_me_async
+async def get_request_status_async(
+    request_id: str,
+    include_msg: bool = False,
+) -> Optional[StatusWithMsg]:
+    """Get the status of a request.
+
+    Args:
+        request_id: The ID of the request.
+        include_msg: Whether to include the status message.
+
+    Returns:
+        The status of the request. If the request is not found, returns
+        None.
+    """
+    assert _DB is not None
+    columns = 'status'
+    if include_msg:
+        columns += ', status_msg'
+    sql = f'SELECT {columns} FROM {REQUEST_TABLE} WHERE request_id LIKE ?'
+    async with _DB.execute_fetchall_async(sql, (request_id + '%',)) as rows:
+        if rows is None or len(rows) == 0:
+            return None
+        status = RequestStatus(rows[0][0])
+        status_msg = rows[0][1] if include_msg else None
+        return StatusWithMsg(status, status_msg)
+
+
 @init_db
 @metrics_lib.time_me
 def create_if_not_exists(request: Request) -> bool:
@@ -570,17 +607,9 @@ async def create_if_not_exists_async(request: Request) -> bool:
         return True
 
 
-@init_db
-@metrics_lib.time_me
-def get_request_tasks(
-    status: Optional[List[RequestStatus]] = None,
-    cluster_names: Optional[List[str]] = None,
-    user_id: Optional[str] = None,
-    exclude_request_names: Optional[List[str]] = None,
-    include_request_names: Optional[List[str]] = None,
-    finished_before: Optional[float] = None,
-) -> List[Request]:
-    """Get a list of requests that match the given filters.
+@dataclasses.dataclass
+class RequestTaskFilter:
+    """Filter for requests.
 
     Args:
         status: a list of statuses of the requests to filter on.
@@ -598,51 +627,87 @@ def get_request_tasks(
         ValueError: If both exclude_request_names and include_request_names are
             provided.
     """
-    if exclude_request_names is not None and include_request_names is not None:
-        raise ValueError(
-            'Only one of exclude_request_names or include_request_names can be '
-            'provided, not both.')
-
-    filters = []
-    filter_params: List[Any] = []
-    if status is not None:
-        status_list_str = ','.join(repr(status.value) for status in status)
-        filters.append(f'status IN ({status_list_str})')
-    if exclude_request_names is not None:
-        exclude_request_names_str = ','.join(
-            repr(name) for name in exclude_request_names)
-        filters.append(f'name NOT IN ({exclude_request_names_str})')
-    if cluster_names is not None:
-        cluster_names_str = ','.join(repr(name) for name in cluster_names)
-        filters.append(f'{COL_CLUSTER_NAME} IN ({cluster_names_str})')
-    if user_id is not None:
-        filters.append(f'{COL_USER_ID} = ?')
-        filter_params.append(user_id)
-    if include_request_names is not None:
-        request_names_str = ','.join(
-            repr(name) for name in include_request_names)
-        filters.append(f'name IN ({request_names_str})')
-    if finished_before is not None:
-        filters.append('finished_at < ?')
-        filter_params.append(finished_before)
-    assert _DB is not None
-    with _DB.conn:
-        cursor = _DB.conn.cursor()
+    status: Optional[List[RequestStatus]] = None
+    cluster_names: Optional[List[str]] = None
+    user_id: Optional[str] = None
+    exclude_request_names: Optional[List[str]] = None
+    include_request_names: Optional[List[str]] = None
+    finished_before: Optional[float] = None
+
+    def __post_init__(self):
+        if (self.exclude_request_names is not None and
+                self.include_request_names is not None):
+            raise ValueError(
+                'Only one of exclude_request_names or include_request_names '
+                'can be provided, not both.')
+
+    def build_query(self) -> Tuple[str, List[Any]]:
+        """Build the SQL query and filter parameters.
+
+        Returns:
+            A tuple of (SQL, SQL parameters).
+        """
+        filters = []
+        filter_params: List[Any] = []
+        if self.status is not None:
+            status_list_str = ','.join(
+                repr(status.value) for status in self.status)
+            filters.append(f'status IN ({status_list_str})')
+        if self.exclude_request_names is not None:
+            exclude_request_names_str = ','.join(
+                repr(name) for name in self.exclude_request_names)
+            filters.append(f'name NOT IN ({exclude_request_names_str})')
+        if self.cluster_names is not None:
+            cluster_names_str = ','.join(
+                repr(name) for name in self.cluster_names)
+            filters.append(f'{COL_CLUSTER_NAME} IN ({cluster_names_str})')
+        if self.user_id is not None:
+            filters.append(f'{COL_USER_ID} = ?')
+            filter_params.append(self.user_id)
+        if self.include_request_names is not None:
+            request_names_str = ','.join(
+                repr(name) for name in self.include_request_names)
+            filters.append(f'name IN ({request_names_str})')
+        if self.finished_before is not None:
+            filters.append('finished_at < ?')
+            filter_params.append(self.finished_before)
         filter_str = ' AND '.join(filters)
         if filter_str:
             filter_str = f' WHERE {filter_str}'
         columns_str = ', '.join(REQUEST_COLUMNS)
-        cursor.execute(
-            f'SELECT {columns_str} FROM {REQUEST_TABLE}{filter_str} '
-            'ORDER BY created_at DESC', filter_params)
+        return (f'SELECT {columns_str} FROM {REQUEST_TABLE}{filter_str} '
+                'ORDER BY created_at DESC'), filter_params
+
+
+@init_db
+@metrics_lib.time_me
+def get_request_tasks(req_filter: RequestTaskFilter) -> List[Request]:
+    """Get a list of requests that match the given filters.
+
+    Args:
+        req_filter: the filter to apply to the requests. Refer to
+            RequestTaskFilter for the details.
+    """
+    assert _DB is not None
+    with _DB.conn:
+        cursor = _DB.conn.cursor()
+        cursor.execute(*req_filter.build_query())
         rows = cursor.fetchall()
         if rows is None:
             return []
-    requests = []
-    for row in rows:
-        request = Request.from_row(row)
-        requests.append(request)
-    return requests
+    return [Request.from_row(row) for row in rows]
+
+
+@init_db_async
+@metrics_lib.time_me_async
+async def get_request_tasks_async(
+        req_filter: RequestTaskFilter) -> List[Request]:
+    """Async version of get_request_tasks."""
+    assert _DB is not None
+    async with _DB.execute_fetchall_async(*req_filter.build_query()) as rows:
+        if not rows:
+            return []
+    return [Request.from_row(row) for row in rows]
 
 
 @init_db_async
@@ -739,8 +804,10 @@ def clean_finished_requests_with_retention(retention_seconds: int):
         retention_seconds: Requests older than this many seconds will be
             deleted.
     """
-    reqs = get_request_tasks(status=RequestStatus.finished_status(),
-                             finished_before=time.time() - retention_seconds)
+    reqs = get_request_tasks(
+        req_filter=RequestTaskFilter(status=RequestStatus.finished_status(),
+                                     finished_before=time.time() -
+                                     retention_seconds))
 
     subprocess_utils.run_in_parallel(
         func=lambda req: req.log_path.unlink(missing_ok=True),
@@ -767,7 +834,7 @@ async def requests_gc_daemon():
         try:
             # Negative value disables the requests GC
             if retention_seconds >= 0:
-                clean_finished_requests_with_retention(retention_seconds)
+                await clean_finished_requests_with_retention(retention_seconds)
         except asyncio.CancelledError:
             logger.info('Requests GC daemon cancelled')
             break
@@ -776,3 +843,11 @@ async def requests_gc_daemon():
         # Run the daemon at most once every hour to avoid too frequent
         # cleanup.
         await asyncio.sleep(max(retention_seconds, 3600))
+
+
+def _cleanup():
+    if _DB is not None:
+        asyncio.run(_DB.close())
+
+
+atexit.register(_cleanup)

sky/server/server.py

@@ -71,6 +71,7 @@ from sky.utils import dag_utils
 from sky.utils import perf_utils
 from sky.utils import status_lib
 from sky.utils import subprocess_utils
+from sky.utils.db import db_utils
 from sky.volumes.server import server as volumes_rest
 from sky.workspaces import server as workspaces_rest
 
@@ -1321,10 +1322,11 @@ async def download(download_body: payloads.DownloadBody,
                                     detail=f'Error creating zip file: {str(e)}')
 
 
+# TODO(aylei): run it asynchronously after global_user_state support async op
 @app.post('/provision_logs')
-async def provision_logs(cluster_body: payloads.ClusterNameBody,
-                         follow: bool = True,
-                         tail: int = 0) -> fastapi.responses.StreamingResponse:
+def provision_logs(cluster_body: payloads.ClusterNameBody,
+                   follow: bool = True,
+                   tail: int = 0) -> fastapi.responses.StreamingResponse:
     """Streams the provision.log for the latest launch request of a cluster."""
     # Prefer clusters table first, then cluster_history as fallback.
     log_path_str = global_user_state.get_cluster_provision_log_path(
@@ -1429,27 +1431,29 @@ async def local_down(request: fastapi.Request) -> None:
 async def api_get(request_id: str) -> payloads.RequestPayload:
     """Gets a request with a given request ID prefix."""
     while True:
-        request_task = await requests_lib.get_request_async(request_id)
-        if request_task is None:
+        req_status = await requests_lib.get_request_status_async(request_id)
+        if req_status is None:
             print(f'No task with request ID {request_id}', flush=True)
             raise fastapi.HTTPException(
                 status_code=404, detail=f'Request {request_id!r} not found')
-        if request_task.status > requests_lib.RequestStatus.RUNNING:
-            if request_task.should_retry:
-                raise fastapi.HTTPException(
-                    status_code=503,
-                    detail=f'Request {request_id!r} should be retried')
-            request_error = request_task.get_error()
-            if request_error is not None:
-                raise fastapi.HTTPException(
-                    status_code=500, detail=request_task.encode().model_dump())
-            return request_task.encode()
-        elif (request_task.status == requests_lib.RequestStatus.RUNNING and
-              daemons.is_daemon_request_id(request_id)):
-            return request_task.encode()
+        if (req_status.status == requests_lib.RequestStatus.RUNNING and
+                daemons.is_daemon_request_id(request_id)):
+            # Daemon requests run forever, break without waiting for complete.
+            break
+        if req_status.status > requests_lib.RequestStatus.RUNNING:
+            break
         # yield control to allow other coroutines to run, sleep shortly
         # to avoid storming the DB and CPU in the meantime
         await asyncio.sleep(0.1)
+    request_task = await requests_lib.get_request_async(request_id)
+    if request_task.should_retry:
+        raise fastapi.HTTPException(
+            status_code=503, detail=f'Request {request_id!r} should be retried')
+    request_error = request_task.get_error()
+    if request_error is not None:
+        raise fastapi.HTTPException(status_code=500,
+                                    detail=request_task.encode().model_dump())
+    return request_task.encode()
 
 
 @app.get('/api/stream')
@@ -1606,10 +1610,9 @@ async def api_status(
                 requests_lib.RequestStatus.PENDING,
                 requests_lib.RequestStatus.RUNNING,
             ]
-        return [
-            request_task.readable_encode()
-            for request_task in requests_lib.get_request_tasks(status=statuses)
-        ]
+        request_tasks = await requests_lib.get_request_tasks_async(
+            req_filter=requests_lib.RequestTaskFilter(status=statuses))
+        return [r.readable_encode() for r in request_tasks]
     else:
         encoded_request_tasks = []
         for request_id in request_ids:
@@ -1808,17 +1811,20 @@ async def gpu_metrics() -> fastapi.Response:
 # === Internal APIs ===
 @app.get('/api/completion/cluster_name')
 async def complete_cluster_name(incomplete: str,) -> List[str]:
-    return global_user_state.get_cluster_names_start_with(incomplete)
+    return await context_utils.to_thread(
+        global_user_state.get_cluster_names_start_with, incomplete)
 
 
 @app.get('/api/completion/storage_name')
 async def complete_storage_name(incomplete: str,) -> List[str]:
-    return global_user_state.get_storage_names_start_with(incomplete)
+    return await context_utils.to_thread(
+        global_user_state.get_storage_names_start_with, incomplete)
 
 
 @app.get('/api/completion/volume_name')
 async def complete_volume_name(incomplete: str,) -> List[str]:
-    return global_user_state.get_volume_names_start_with(incomplete)
+    return await context_utils.to_thread(
+        global_user_state.get_volume_names_start_with, incomplete)
 
 
 @app.get('/api/completion/api_request')
@@ -1901,13 +1907,6 @@ if __name__ == '__main__':
 
     skyuvicorn.add_timestamp_prefix_for_server_logs()
 
-    # Initialize global user state db
-    global_user_state.initialize_and_get_db()
-    # Initialize request db
-    requests_lib.reset_db_and_logs()
-    # Restore the server user hash
-    _init_or_restore_server_user_hash()
-
     parser = argparse.ArgumentParser()
     parser.add_argument('--host', default='127.0.0.1')
     parser.add_argument('--port', default=46580, type=int)
@@ -1923,7 +1922,17 @@ if __name__ == '__main__':
     # that it is shown only when the API server is started.
     usage_lib.maybe_show_privacy_policy()
 
-    config = server_config.compute_server_config(cmd_args.deploy)
+    # Initialize global user state db
+    db_utils.set_max_connections(1)
+    global_user_state.initialize_and_get_db()
+    # Initialize request db
+    requests_lib.reset_db_and_logs()
+    # Restore the server user hash
+    _init_or_restore_server_user_hash()
+    max_db_connections = global_user_state.get_max_db_connections()
+    config = server_config.compute_server_config(cmd_args.deploy,
+                                                 max_db_connections)
+
     num_workers = config.num_server_workers
 
     queue_server: Optional[multiprocessing.Process] = None
@@ -1948,11 +1957,12 @@ if __name__ == '__main__':
         logger.info(f'Starting SkyPilot API server, workers={num_workers}')
         # We don't support reload for now, since it may cause leakage of request
         # workers or interrupt running requests.
-        config = uvicorn.Config('sky.server.server:app',
-                                host=cmd_args.host,
-                                port=cmd_args.port,
-                                workers=num_workers)
-        skyuvicorn.run(config)
+        uvicorn_config = uvicorn.Config('sky.server.server:app',
+                                        host=cmd_args.host,
+                                        port=cmd_args.port,
+                                        workers=num_workers)
+        skyuvicorn.run(uvicorn_config,
+                       max_db_connections=config.num_db_connections_per_worker)
     except Exception as exc:  # pylint: disable=broad-except
         logger.error(f'Failed to start SkyPilot API server: '
                      f'{common_utils.format_exception(exc, use_bracket=True)}')

sky/server/stream_utils.py

@@ -75,8 +75,10 @@ async def log_streamer(request_id: Optional[str],
         last_waiting_msg = ''
         waiting_msg = (f'Waiting for {request_task.name!r} request to be '
                        f'scheduled: {request_id}')
-        while request_task.status < requests_lib.RequestStatus.RUNNING:
-            if request_task.status_msg is not None:
+        req_status = request_task.status
+        req_msg = request_task.status_msg
+        while req_status < requests_lib.RequestStatus.RUNNING:
+            if req_msg is not None:
                 waiting_msg = request_task.status_msg
             if show_request_waiting_spinner:
                 yield status_msg.update(f'[dim]{waiting_msg}[/dim]')
@@ -91,7 +93,10 @@ async def log_streamer(request_id: Optional[str],
             # polling the DB, which can be a bottleneck for high-concurrency
             # requests.
             await asyncio.sleep(0.1)
-            request_task = await requests_lib.get_request_async(request_id)
+            status_with_msg = await requests_lib.get_request_status_async(
+                request_id, include_msg=True)
+            req_status = status_with_msg.status
+            req_msg = status_with_msg.status_msg
             if not follow:
                 break
         if show_request_waiting_spinner:
@@ -153,10 +158,13 @@ async def _tail_log_file(f: aiofiles.threadpool.binary.AsyncBufferedReader,
         line: Optional[bytes] = await f.readline()
         if not line:
             if request_id is not None:
-                request_task = await requests_lib.get_request_async(request_id)
-                if request_task.status > requests_lib.RequestStatus.RUNNING:
-                    if (request_task.status ==
+                req_status = await requests_lib.get_request_status_async(
+                    request_id)
+                if req_status.status > requests_lib.RequestStatus.RUNNING:
+                    if (req_status.status ==
                             requests_lib.RequestStatus.CANCELLED):
+                        request_task = await requests_lib.get_request_async(
+                            request_id)
                         if request_task.should_retry:
                             buffer.append(
                                 message_utils.encode_payload(

sky/server/uvicorn.py

@@ -26,6 +26,7 @@ from sky.utils import context_utils
 from sky.utils import env_options
 from sky.utils import perf_utils
 from sky.utils import subprocess_utils
+from sky.utils.db import db_utils
 
 logger = sky_logging.init_logger(__name__)
 
@@ -88,9 +89,12 @@ class Server(uvicorn.Server):
     - Run the server process with contextually aware.
     """
 
-    def __init__(self, config: uvicorn.Config):
+    def __init__(self,
+                 config: uvicorn.Config,
+                 max_db_connections: Optional[int] = None):
         super().__init__(config=config)
         self.exiting: bool = False
+        self.max_db_connections = max_db_connections
 
     def handle_exit(self, sig: int, frame: Union[FrameType, None]) -> None:
         """Handle exit signal.
@@ -146,7 +150,8 @@ class Server(uvicorn.Server):
                 requests_lib.RequestStatus.PENDING,
                 requests_lib.RequestStatus.RUNNING,
             ]
-            reqs = requests_lib.get_request_tasks(status=statuses)
+            reqs = requests_lib.get_request_tasks(
+                req_filter=requests_lib.RequestTaskFilter(status=statuses))
             if not reqs:
                 break
             logger.info(f'{len(reqs)} on-going requests '
@@ -195,6 +200,8 @@ class Server(uvicorn.Server):
 
     def run(self, *args, **kwargs):
         """Run the server process."""
+        if self.max_db_connections is not None:
+            db_utils.set_max_connections(self.max_db_connections)
         add_timestamp_prefix_for_server_logs()
         context_utils.hijack_sys_attrs()
         # Use default loop policy of uvicorn (use uvloop if available).
@@ -209,14 +216,14 @@ class Server(uvicorn.Server):
             asyncio.run(self.serve(*args, **kwargs))
 
 
-def run(config: uvicorn.Config):
+def run(config: uvicorn.Config, max_db_connections: Optional[int] = None):
     """Run unvicorn server."""
     if config.reload:
         # Reload and multi-workers are mutually exclusive
         # in uvicorn. Since we do not use reload now, simply
         # guard by an exception.
         raise ValueError('Reload is not supported yet.')
-    server = Server(config=config)
+    server = Server(config=config, max_db_connections=max_db_connections)
     try:
         if config.workers is not None and config.workers > 1:
             sock = config.bind_socket()

sky/skylet/constants.py

@@ -362,7 +362,7 @@ SKY_SSH_USER_PLACEHOLDER = 'skypilot:ssh_user'
 
 RCLONE_CONFIG_DIR = '~/.config/rclone'
 RCLONE_CONFIG_PATH = f'{RCLONE_CONFIG_DIR}/rclone.conf'
-RCLONE_LOG_DIR = '~/.sky/rclone_log'
+RCLONE_MOUNT_CACHED_LOG_DIR = '~/.sky/rclone_log'
 RCLONE_CACHE_DIR = '~/.cache/rclone'
 RCLONE_CACHE_REFRESH_INTERVAL = 10
 

sky/skypilot_config.py

@@ -227,7 +227,7 @@ def _get_config_from_path(path: Optional[str]) -> config_utils.Config:
     return parse_and_validate_config_file(path)
 
 
-def _resolve_user_config_path() -> Optional[str]:
+def resolve_user_config_path() -> Optional[str]:
     # find the user config file path, None if not resolved.
     user_config_path = _get_config_file_path(ENV_VAR_GLOBAL_CONFIG)
     if user_config_path:
@@ -252,7 +252,7 @@ def _resolve_user_config_path() -> Optional[str]:
 
 def get_user_config() -> config_utils.Config:
     """Returns the user config."""
-    return _get_config_from_path(_resolve_user_config_path())
+    return _get_config_from_path(resolve_user_config_path())
 
 
 def _resolve_project_config_path() -> Optional[str]:
@@ -574,8 +574,13 @@ def _reload_config_as_server() -> None:
                 'If db config is specified, no other config is allowed')
         logger.debug('retrieving config from database')
         with _DB_USE_LOCK:
-            sqlalchemy_engine = sqlalchemy.create_engine(db_url,
-                                                         poolclass=NullPool)
+            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)
 
@@ -597,7 +602,8 @@ def _reload_config_as_server() -> None:
                 server_config = overlay_skypilot_config(server_config,
                                                         db_config)
             # Close the engine to avoid connection leaks
-            sqlalchemy_engine.dispose()
+            if dispose_engine:
+                sqlalchemy_engine.dispose()
     if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
         logger.debug(f'server config: \n'
                      f'{yaml_utils.dump_yaml_str(dict(server_config))}')
@@ -611,7 +617,7 @@ def _reload_config_as_client() -> None:
     _set_loaded_config_path(None)
 
     overrides: List[config_utils.Config] = []
-    user_config_path = _resolve_user_config_path()
+    user_config_path = resolve_user_config_path()
     user_config = _get_config_from_path(user_config_path)
     if user_config:
         overrides.append(user_config)
@@ -867,8 +873,13 @@ def update_api_server_config_no_lock(config: config_utils.Config) -> None:
             raise ValueError('Cannot change db url while server is running')
         if existing_db_url:
             with _DB_USE_LOCK:
-                sqlalchemy_engine = sqlalchemy.create_engine(existing_db_url,
-                                                             poolclass=NullPool)
+                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)
 
@@ -897,7 +908,8 @@ def update_api_server_config_no_lock(config: config_utils.Config) -> None:
                 _set_config_yaml_to_db(API_SERVER_CONFIG_KEY, config)
                 db_updated = True
                 # Close the engine to avoid connection leaks
-                sqlalchemy_engine.dispose()
+                if dispose_engine:
+                    sqlalchemy_engine.dispose()
 
     if not db_updated:
         # save to the local file (PVC in Kubernetes, local file otherwise)

sky/ssh_node_pools/server.py

@@ -15,7 +15,7 @@ router = fastapi.APIRouter()
 
 
 @router.get('')
-async def get_ssh_node_pools() -> Dict[str, Any]:
+def get_ssh_node_pools() -> Dict[str, Any]:
     """Get all SSH Node Pool configurations."""
     try:
         return ssh_node_pools_core.get_all_pools()
@@ -27,7 +27,7 @@ async def get_ssh_node_pools() -> Dict[str, Any]:
 
 
 @router.post('')
-async def update_ssh_node_pools(pools_config: Dict[str, Any]) -> Dict[str, str]:
+def update_ssh_node_pools(pools_config: Dict[str, Any]) -> Dict[str, str]:
     """Update SSH Node Pool configurations."""
     try:
         ssh_node_pools_core.update_pools(pools_config)
@@ -39,7 +39,7 @@ async def update_ssh_node_pools(pools_config: Dict[str, Any]) -> Dict[str, str]:
 
 
 @router.delete('/{pool_name}')
-async def delete_ssh_node_pool(pool_name: str) -> Dict[str, str]:
+def delete_ssh_node_pool(pool_name: str) -> Dict[str, str]:
     """Delete a SSH Node Pool configuration."""
     try:
         if ssh_node_pools_core.delete_pool(pool_name):
@@ -83,7 +83,7 @@ async def upload_ssh_key(request: fastapi.Request) -> Dict[str, str]:
 
 
 @router.get('/keys')
-async def list_ssh_keys() -> List[str]:
+def list_ssh_keys() -> List[str]:
     """List available SSH keys."""
     try:
         return ssh_node_pools_core.list_ssh_keys()
@@ -200,7 +200,7 @@ async def down_ssh_node_pool_general(
 
 
 @router.get('/{pool_name}/status')
-async def get_ssh_node_pool_status(pool_name: str) -> Dict[str, str]:
+def get_ssh_node_pool_status(pool_name: str) -> Dict[str, str]:
     """Get the status of a specific SSH Node Pool."""
     try:
         # Call ssh_status to check the context