skypilot-nightly 1.0.0.dev20250831__py3-none-any.whl → 1.0.0.dev20250901__py3-none-any.whl

sky/server/requests/executor.py

@@ -41,6 +41,7 @@ from sky import skypilot_config
 from sky.server import common as server_common
 from sky.server import config as server_config
 from sky.server import constants as server_constants
+from sky.server import metrics as metrics_lib
 from sky.server.requests import payloads
 from sky.server.requests import preconditions
 from sky.server.requests import process
@@ -373,6 +374,7 @@ def _request_execution_wrapper(request_id: str,
         request_task.status = api_requests.RequestStatus.RUNNING
         func = request_task.entrypoint
         request_body = request_task.request_body
+        request_name = request_task.name
 
     # Append to the log file instead of overwriting it since there might be
     # logs from previous retries.
@@ -390,7 +392,9 @@ def _request_execution_wrapper(request_id: str,
                     config = skypilot_config.to_dict()
                     logger.debug(f'request config: \n'
                                  f'{yaml_utils.dump_yaml_str(dict(config))}')
-                return_value = func(**request_body.to_kwargs())
+                with metrics_lib.time_it(name=request_name,
+                                         group='request_execution'):
+                    return_value = func(**request_body.to_kwargs())
                 f.flush()
         except KeyboardInterrupt:
             logger.info(f'Request {request_id} cancelled by user')
@@ -453,7 +457,7 @@ async def execute_request_coroutine(request: api_requests.Request):
                                                   **request_body.to_kwargs())
 
     async def poll_task(request_id: str) -> bool:
-        request = api_requests.get_request(request_id)
+        request = await api_requests.get_request_async(request_id)
         if request is None:
             raise RuntimeError('Request not found')
 

sky/server/requests/preconditions.py

@@ -98,7 +98,7 @@ class Precondition(abc.ABC):
                 return False
 
             # Check if the request has been cancelled
-            request = api_requests.get_request(self.request_id)
+            request = await api_requests.get_request_async(self.request_id)
             if request is None:
                 logger.error(f'Request {self.request_id} not found')
                 return False
@@ -112,7 +112,8 @@ class Precondition(abc.ABC):
                     return True
                 if status_msg is not None and status_msg != last_status_msg:
                     # Update the status message if it has changed.
-                    with api_requests.update_request(self.request_id) as req:
+                    async with api_requests.update_request_async(
+                            self.request_id) as req:
                         assert req is not None, self.request_id
                         req.status_msg = status_msg
                     last_status_msg = status_msg

sky/server/requests/requests.py

@@ -13,7 +13,8 @@ import sqlite3
 import threading
 import time
 import traceback
-from typing import Any, Callable, Dict, Generator, List, Optional, Tuple
+from typing import (Any, AsyncContextManager, Callable, Dict, Generator, List,
+                    Optional, Tuple)
 
 import colorama
 import filelock
@@ -25,6 +26,7 @@ from sky import skypilot_config
 from sky.server import common as server_common
 from sky.server import constants as server_constants
 from sky.server import daemons
+from sky.server import metrics as metrics_lib
 from sky.server.requests import payloads
 from sky.server.requests.serializers import decoders
 from sky.server.requests.serializers import encoders
@@ -402,26 +404,46 @@ _DB = None
 _init_db_lock = threading.Lock()
 
 
+def _init_db_within_lock():
+    global _DB
+    if _DB is None:
+        db_path = os.path.expanduser(
+            server_constants.API_SERVER_REQUEST_DB_PATH)
+        pathlib.Path(db_path).parents[0].mkdir(parents=True, exist_ok=True)
+        _DB = db_utils.SQLiteConn(db_path, create_table)
+
+
 def init_db(func):
     """Initialize the database."""
 
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
-        global _DB
         if _DB is not None:
             return func(*args, **kwargs)
         with _init_db_lock:
-            if _DB is None:
-                db_path = os.path.expanduser(
-                    server_constants.API_SERVER_REQUEST_DB_PATH)
-                pathlib.Path(db_path).parents[0].mkdir(parents=True,
-                                                       exist_ok=True)
-                _DB = db_utils.SQLiteConn(db_path, create_table)
+            _init_db_within_lock()
         return func(*args, **kwargs)
 
     return wrapper
 
 
+def init_db_async(func):
+    """Async version of init_db."""
+
+    @functools.wraps(func)
+    async def wrapper(*args, **kwargs):
+        if _DB is not None:
+            return await func(*args, **kwargs)
+        # If _DB is not initialized, init_db_async will be blocked if there
+        # is a thread initializing _DB, this is fine since it occurs on process
+        # startup.
+        with _init_db_lock:
+            _init_db_within_lock()
+        return await func(*args, **kwargs)
+
+    return wrapper
+
+
 def reset_db_and_logs():
     """Create the database."""
     server_common.clear_local_api_server_database()
@@ -439,30 +461,66 @@ def request_lock_path(request_id: str) -> str:
 
 @contextlib.contextmanager
 @init_db
+@metrics_lib.time_me
 def update_request(request_id: str) -> Generator[Optional[Request], None, None]:
-    """Get a SkyPilot API request."""
+    """Get and update a SkyPilot API request."""
     request = _get_request_no_lock(request_id)
     yield request
     if request is not None:
         _add_or_update_request_no_lock(request)
 
 
+@init_db
+@metrics_lib.time_me
+def update_request_async(
+        request_id: str) -> AsyncContextManager[Optional[Request]]:
+    """Async version of update_request.
+
+    Returns an async context manager that yields the request record and
+    persists any in-place updates upon exit.
+    """
+
+    @contextlib.asynccontextmanager
+    async def _cm():
+        request = await _get_request_no_lock_async(request_id)
+        try:
+            yield request
+        finally:
+            if request is not None:
+                await _add_or_update_request_no_lock_async(request)
+
+    return _cm()
+
+
+_get_request_sql = (f'SELECT {", ".join(REQUEST_COLUMNS)} FROM {REQUEST_TABLE} '
+                    'WHERE request_id LIKE ?')
+
+
 def _get_request_no_lock(request_id: str) -> Optional[Request]:
     """Get a SkyPilot API request."""
     assert _DB is not None
-    columns_str = ', '.join(REQUEST_COLUMNS)
     with _DB.conn:
         cursor = _DB.conn.cursor()
-        cursor.execute(
-            f'SELECT {columns_str} FROM {REQUEST_TABLE} '
-            'WHERE request_id LIKE ?', (request_id + '%',))
+        cursor.execute(_get_request_sql, (request_id + '%',))
         row = cursor.fetchone()
         if row is None:
             return None
     return Request.from_row(row)
 
 
+async def _get_request_no_lock_async(request_id: str) -> Optional[Request]:
+    """Async version of _get_request_no_lock."""
+    assert _DB is not None
+    async with _DB.execute_fetchall_async(_get_request_sql,
+                                          (request_id + '%',)) as rows:
+        row = rows[0] if rows else None
+        if row is None:
+            return None
+    return Request.from_row(row)
+
+
 @init_db
+@metrics_lib.time_me
 def get_latest_request_id() -> Optional[str]:
     """Get the latest request ID."""
     assert _DB is not None
@@ -475,13 +533,23 @@ def get_latest_request_id() -> Optional[str]:
 
 
 @init_db
+@metrics_lib.time_me
 def get_request(request_id: str) -> Optional[Request]:
     """Get a SkyPilot API request."""
     with filelock.FileLock(request_lock_path(request_id)):
         return _get_request_no_lock(request_id)
 
 
+@init_db_async
+@metrics_lib.time_me_async
+async def get_request_async(request_id: str) -> Optional[Request]:
+    """Async version of get_request."""
+    async with filelock.AsyncFileLock(request_lock_path(request_id)):
+        return await _get_request_no_lock_async(request_id)
+
+
 @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)):
@@ -491,7 +559,19 @@ def create_if_not_exists(request: Request) -> bool:
         return True
 
 
+@init_db_async
+@metrics_lib.time_me_async
+async def create_if_not_exists_async(request: Request) -> bool:
+    """Async version of create_if_not_exists."""
+    async with filelock.AsyncFileLock(request_lock_path(request.request_id)):
+        if await _get_request_no_lock_async(request.request_id) is not None:
+            return False
+        await _add_or_update_request_no_lock_async(request)
+        return True
+
+
 @init_db
+@metrics_lib.time_me
 def get_request_tasks(
     status: Optional[List[RequestStatus]] = None,
     cluster_names: Optional[List[str]] = None,
@@ -565,16 +645,15 @@ def get_request_tasks(
     return requests
 
 
-@init_db
-def get_api_request_ids_start_with(incomplete: str) -> List[str]:
+@init_db_async
+@metrics_lib.time_me_async
+async def get_api_request_ids_start_with(incomplete: str) -> List[str]:
     """Get a list of API request ids for shell completion."""
     assert _DB is not None
-    with _DB.conn:
-        cursor = _DB.conn.cursor()
-        # Prioritize alive requests (PENDING, RUNNING) over finished ones,
-        # then order by creation time (newest first) within each category.
-        cursor.execute(
-            f"""SELECT request_id FROM {REQUEST_TABLE}
+    # Prioritize alive requests (PENDING, RUNNING) over finished ones,
+    # then order by creation time (newest first) within each category.
+    async with _DB.execute_fetchall_async(
+        f"""SELECT request_id FROM {REQUEST_TABLE}
                 WHERE request_id LIKE ?
                 ORDER BY
                     CASE
@@ -582,21 +661,30 @@ def get_api_request_ids_start_with(incomplete: str) -> List[str]:
                         ELSE 1
                     END,
                     created_at DESC
-                LIMIT 1000""", (f'{incomplete}%',))
-        return [row[0] for row in cursor.fetchall()]
+                LIMIT 1000""", (f'{incomplete}%',)) as rows:
+        if not rows:
+            return []
+    return [row[0] for row in rows]
+
+
+_add_or_update_request_sql = (f'INSERT OR REPLACE INTO {REQUEST_TABLE} '
+                              f'({", ".join(REQUEST_COLUMNS)}) VALUES '
+                              f'({", ".join(["?"] * len(REQUEST_COLUMNS))})')
 
 
 def _add_or_update_request_no_lock(request: Request):
     """Add or update a REST request into the database."""
-    row = request.to_row()
-    key_str = ', '.join(REQUEST_COLUMNS)
-    fill_str = ', '.join(['?'] * len(row))
     assert _DB is not None
     with _DB.conn:
         cursor = _DB.conn.cursor()
-        cursor.execute(
-            f'INSERT OR REPLACE INTO {REQUEST_TABLE} ({key_str}) '
-            f'VALUES ({fill_str})', row)
+        cursor.execute(_add_or_update_request_sql, request.to_row())
+
+
+async def _add_or_update_request_no_lock_async(request: Request):
+    """Async version of _add_or_update_request_no_lock."""
+    assert _DB is not None
+    await _DB.execute_and_commit_async(_add_or_update_request_sql,
+                                       request.to_row())
 
 
 def set_request_failed(request_id: str, e: BaseException) -> None:
@@ -630,6 +718,7 @@ def set_request_cancelled(request_id: str) -> None:
 
 
 @init_db
+@metrics_lib.time_me
 def _delete_requests(requests: List[Request]):
     """Clean up requests by their IDs."""
     id_list_str = ','.join(repr(req.request_id) for req in requests)

sky/server/server.py

@@ -68,6 +68,7 @@ from sky.utils import common_utils
 from sky.utils import context
 from sky.utils import context_utils
 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.volumes.server import server as volumes_rest
@@ -421,6 +422,28 @@ async def cleanup_upload_ids():
                 upload_ids_to_cleanup.pop((upload_id, user_hash))
 
 
+async def loop_lag_monitor(loop: asyncio.AbstractEventLoop,
+                           interval: float = 0.1) -> None:
+    target = loop.time() + interval
+
+    pid = str(os.getpid())
+    lag_threshold = perf_utils.get_loop_lag_threshold()
+
+    def tick():
+        nonlocal target
+        now = loop.time()
+        lag = max(0.0, now - target)
+        if lag_threshold is not None and lag > lag_threshold:
+            logger.warning(f'Event loop lag {lag} seconds exceeds threshold '
+                           f'{lag_threshold} seconds.')
+        metrics.SKY_APISERVER_EVENT_LOOP_LAG_SECONDS.labels(
+            pid=pid).observe(lag)
+        target = now + interval
+        loop.call_at(target, tick)
+
+    loop.call_at(target, tick)
+
+
 @contextlib.asynccontextmanager
 async def lifespan(app: fastapi.FastAPI):  # pylint: disable=redefined-outer-name
     """FastAPI lifespan context manager."""
@@ -446,6 +469,10 @@ async def lifespan(app: fastapi.FastAPI):  # pylint: disable=redefined-outer-nam
             # can safely ignore the error if the task is already scheduled.
             logger.debug(f'Request {event.id} already exists.')
     asyncio.create_task(cleanup_upload_ids())
+    if metrics.METRICS_ENABLED:
+        # Start monitoring the event loop lag in each server worker
+        # event loop (process).
+        asyncio.create_task(loop_lag_monitor(asyncio.get_event_loop()))
     yield
     # Shutdown: Add any cleanup code here if needed
 
@@ -1254,20 +1281,25 @@ async def download(download_body: payloads.DownloadBody,
         logs_dir_on_api_server).expanduser().resolve() / zip_filename
 
     try:
-        folders = [
-            str(folder_path.expanduser().resolve())
-            for folder_path in folder_paths
-        ]
-        # Check for optional query parameter to control zip entry structure
-        relative = request.query_params.get('relative', 'home')
-        if relative == 'items':
-            # Dashboard-friendly: entries relative to selected folders
-            storage_utils.zip_files_and_folders(folders,
-                                                zip_path,
-                                                relative_to_items=True)
-        else:
-            # CLI-friendly (default): entries with full paths for mapping
-            storage_utils.zip_files_and_folders(folders, zip_path)
+
+        def _zip_files_and_folders(folder_paths, zip_path):
+            folders = [
+                str(folder_path.expanduser().resolve())
+                for folder_path in folder_paths
+            ]
+            # Check for optional query parameter to control zip entry structure
+            relative = request.query_params.get('relative', 'home')
+            if relative == 'items':
+                # Dashboard-friendly: entries relative to selected folders
+                storage_utils.zip_files_and_folders(folders,
+                                                    zip_path,
+                                                    relative_to_items=True)
+            else:
+                # CLI-friendly (default): entries with full paths for mapping
+                storage_utils.zip_files_and_folders(folders, zip_path)
+
+        await context_utils.to_thread(_zip_files_and_folders, folder_paths,
+                                      zip_path)
 
         # Add home path to the response headers, so that the client can replace
         # the remote path in the zip file to the local path.
@@ -1397,7 +1429,7 @@ 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 = requests_lib.get_request(request_id)
+        request_task = await requests_lib.get_request_async(request_id)
         if request_task is None:
             print(f'No task with request ID {request_id}', flush=True)
             raise fastapi.HTTPException(
@@ -1486,7 +1518,7 @@ async def stream(
 
     # Original plain text streaming logic
     if request_id is not None:
-        request_task = requests_lib.get_request(request_id)
+        request_task = await requests_lib.get_request_async(request_id)
         if request_task is None:
             print(f'No task with request ID {request_id}')
             raise fastapi.HTTPException(
@@ -1581,7 +1613,7 @@ async def api_status(
     else:
         encoded_request_tasks = []
         for request_id in request_ids:
-            request_task = requests_lib.get_request(request_id)
+            request_task = await requests_lib.get_request_async(request_id)
             if request_task is None:
                 continue
             encoded_request_tasks.append(request_task.readable_encode())
@@ -1791,7 +1823,7 @@ async def complete_volume_name(incomplete: str,) -> List[str]:
 
 @app.get('/api/completion/api_request')
 async def complete_api_request(incomplete: str,) -> List[str]:
-    return requests_lib.get_api_request_ids_start_with(incomplete)
+    return await requests_lib.get_api_request_ids_start_with(incomplete)
 
 
 @app.get('/dashboard/{full_path:path}')

sky/server/stream_utils.py

@@ -56,7 +56,7 @@ async def log_streamer(request_id: Optional[str],
     if request_id is not None:
         status_msg = rich_utils.EncodedStatusMessage(
             f'[dim]Checking request: {request_id}[/dim]')
-        request_task = requests_lib.get_request(request_id)
+        request_task = await requests_lib.get_request_async(request_id)
 
         if request_task is None:
             raise fastapi.HTTPException(
@@ -86,10 +86,12 @@ async def log_streamer(request_id: Optional[str],
                 # Use smaller padding (1024 bytes) to force browser rendering
                 yield f'{waiting_msg}' + ' ' * 4096 + '\n'
             # Sleep shortly to avoid storming the DB and CPU and allow other
-            # coroutines to run. This busy waiting loop is performance critical
-            # for short-running requests, so we do not want to yield too long.
+            # coroutines to run.
+            # TODO(aylei): we should use a better mechanism to avoid busy
+            # polling the DB, which can be a bottleneck for high-concurrency
+            # requests.
             await asyncio.sleep(0.1)
-            request_task = requests_lib.get_request(request_id)
+            request_task = await requests_lib.get_request_async(request_id)
             if not follow:
                 break
         if show_request_waiting_spinner:
@@ -151,7 +153,7 @@ 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 = requests_lib.get_request(request_id)
+                request_task = await requests_lib.get_request_async(request_id)
                 if request_task.status > requests_lib.RequestStatus.RUNNING:
                     if (request_task.status ==
                             requests_lib.RequestStatus.CANCELLED):

sky/server/uvicorn.py

@@ -24,6 +24,7 @@ from sky.server.requests import requests as requests_lib
 from sky.skylet import constants
 from sky.utils import context_utils
 from sky.utils import env_options
+from sky.utils import perf_utils
 from sky.utils import subprocess_utils
 
 logger = sky_logging.init_logger(__name__)
@@ -198,6 +199,12 @@ class Server(uvicorn.Server):
         context_utils.hijack_sys_attrs()
         # Use default loop policy of uvicorn (use uvloop if available).
         self.config.setup_event_loop()
+        lag_threshold = perf_utils.get_loop_lag_threshold()
+        if lag_threshold is not None:
+            event_loop = asyncio.get_event_loop()
+            # Same as set PYTHONASYNCIODEBUG=1, but with custom threshold.
+            event_loop.set_debug(True)
+            event_loop.slow_callback_duration = lag_threshold
         with self.capture_signals():
             asyncio.run(self.serve(*args, **kwargs))
 

sky/setup_files/dependencies.py

@@ -39,7 +39,8 @@ install_requires = [
     # Light weight requirement, can be replaced with "typing" once
     # we deprecate Python 3.7 (this will take a while).
     'typing_extensions',
-    'filelock >= 3.6.0',
+    # filelock 3.15.0 or higher is required for async file locking.
+    'filelock >= 3.15.0',
     'packaging',
     'psutil',
     'pulp',
@@ -75,6 +76,7 @@ install_requires = [
     'types-paramiko',
     'alembic',
     'aiohttp',
+    'aiosqlite',
     'anyio',
 ]
 
@@ -100,6 +102,7 @@ server_dependencies = [
     'anyio',
     GRPC,
     PROTOBUF,
+    'aiosqlite',
 ]
 
 local_ray = [

sky/skylet/constants.py

@@ -505,3 +505,6 @@ COST_REPORT_DEFAULT_DAYS = 30
 
 # The directory for file locks.
 SKY_LOCKS_DIR = os.path.expanduser('~/.sky/locks')
+
+ENV_VAR_LOOP_LAG_THRESHOLD_MS = (SKYPILOT_ENV_VAR_PREFIX +
+                                 'DEBUG_LOOP_LAG_THRESHOLD_MS')

sky/utils/db/db_utils.py

@@ -1,11 +1,14 @@
 """Utils for sky databases."""
+import asyncio
 import contextlib
 import enum
 import sqlite3
 import threading
 import typing
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Iterable, Optional
 
+import aiosqlite
+import aiosqlite.context
 import sqlalchemy
 from sqlalchemy import exc as sqlalchemy_exc
 
@@ -283,3 +286,63 @@ class SQLiteConn(threading.local):
         self.conn = sqlite3.connect(db_path, timeout=_DB_TIMEOUT_S)
         self.cursor = self.conn.cursor()
         create_table(self.cursor, self.conn)
+        self._async_conn: Optional[aiosqlite.Connection] = None
+        self._async_conn_lock: Optional[asyncio.Lock] = None
+
+    async def _get_async_conn(self) -> aiosqlite.Connection:
+        """Get the shared aiosqlite connection for current thread.
+
+        Typically, external caller should not get the connection directly,
+        instead, SQLiteConn.{operation}_async methods should be used. This
+        is to avoid txn interleaving on the shared aiosqlite connection.
+        E.g.
+        coroutine 1:
+            A: await write(row1)
+            B: cursor = await conn.execute(read_row1)
+            C: await cursor.fetchall()
+        coroutine 2:
+            D: await write(row2)
+            E: cursor = await conn.execute(read_row2)
+            F: await cursor.fetchall()
+        The A -> B -> D -> E -> C time sequence will cause B and D read at the
+        same snapshot point when B started, thus cause coroutine2 lost the
+        read-after-write consistency. When you are adding new async operations
+        to SQLiteConn, make sure the txn pattern does not cause this issue.
+        """
+        # Python 3.8 binds current event loop to asyncio.Lock(), which requires
+        # a loop available in current thread. Lazy-init the lock to avoid this
+        # dependency. The correctness is guranteed since SQLiteConn is
+        # thread-local so there is no race condition between check and init.
+        if self._async_conn_lock is None:
+            self._async_conn_lock = asyncio.Lock()
+        if self._async_conn is None:
+            async with self._async_conn_lock:
+                if self._async_conn is None:
+                    # Init logic like requests.init_db_within_lock will handle
+                    # initialization like setting the WAL mode, so we do not
+                    # duplicate that logic here.
+                    self._async_conn = await aiosqlite.connect(self.db_path)
+        return self._async_conn
+
+    async def execute_and_commit_async(self,
+                                       sql: str,
+                                       parameters: Optional[
+                                           Iterable[Any]] = None) -> None:
+        """Execute the sql and commit the transaction in a sync block."""
+        conn = await self._get_async_conn()
+
+        def exec_and_commit(sql: str, parameters: Optional[Iterable[Any]]):
+            # pylint: disable=protected-access
+            conn._conn.execute(sql, parameters)
+            conn._conn.commit()
+
+        # pylint: disable=protected-access
+        await conn._execute(exec_and_commit, sql, parameters)
+
+    @aiosqlite.context.contextmanager
+    async def execute_fetchall_async(self,
+                                     sql: str,
+                                     parameters: Optional[Iterable[Any]] = None
+                                    ) -> Iterable[sqlite3.Row]:
+        conn = await self._get_async_conn()
+        return await conn.execute_fetchall(sql, parameters)

sky/utils/perf_utils.py

@@ -0,0 +1,22 @@
+"""Utility functions for performance monitoring."""
+import os
+from typing import Optional
+
+from sky import sky_logging
+from sky.skylet import constants
+
+logger = sky_logging.init_logger(__name__)
+
+
+def get_loop_lag_threshold() -> Optional[float]:
+    """Get the loop lag threshold from the environment variable."""
+    lag_threshold = os.getenv(constants.ENV_VAR_LOOP_LAG_THRESHOLD_MS, None)
+    if lag_threshold is not None:
+        try:
+            return float(lag_threshold) / 1000.0
+        except ValueError:
+            logger.warning(
+                f'Invalid value for {constants.ENV_VAR_LOOP_LAG_THRESHOLD_MS}:'
+                f' {lag_threshold}')
+            return None
+    return None