skypilot-nightly 1.0.0.dev20250827__py3-none-any.whl → 1.0.0.dev20250829__py3-none-any.whl

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
@@ -402,26 +403,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()
@@ -440,28 +461,61 @@ def request_lock_path(request_id: str) -> str:
 @contextlib.contextmanager
 @init_db
 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
+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
+    conn = await _DB.async_conn()
+    async with conn.execute(_get_request_sql, (request_id + '%',)) as cursor:
+        row = await cursor.fetchone()
+        if row is None:
+            return None
+    return Request.from_row(row)
+
+
 @init_db
 def get_latest_request_id() -> Optional[str]:
     """Get the latest request ID."""
@@ -481,6 +535,13 @@ def get_request(request_id: str) -> Optional[Request]:
         return _get_request_no_lock(request_id)
 
 
+@init_db_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
 def create_if_not_exists(request: Request) -> bool:
     """Create a SkyPilot API request if it does not exist."""
@@ -491,6 +552,16 @@ def create_if_not_exists(request: Request) -> bool:
         return True
 
 
+@init_db_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
 def get_request_tasks(
     status: Optional[List[RequestStatus]] = None,
@@ -565,17 +636,48 @@ def get_request_tasks(
     return requests
 
 
+@init_db_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
+    conn = await _DB.async_conn()
+    # Prioritize alive requests (PENDING, RUNNING) over finished ones,
+    # then order by creation time (newest first) within each category.
+    async with conn.execute(
+        f"""SELECT request_id FROM {REQUEST_TABLE}
+                WHERE request_id LIKE ?
+                ORDER BY
+                    CASE
+                        WHEN status IN ('PENDING', 'RUNNING') THEN 0
+                        ELSE 1
+                    END,
+                    created_at DESC
+                LIMIT 1000""", (f'{incomplete}%',)) as cursor:
+        rows = await cursor.fetchall()
+        if rows is None:
+            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
+    conn = await _DB.async_conn()
+    await conn.execute(_add_or_update_request_sql, request.to_row())
+    await conn.commit()
 
 
 def set_request_failed(request_id: str, e: BaseException) -> None:

sky/server/server.py

@@ -21,6 +21,7 @@ import uuid
 import zipfile
 
 import aiofiles
+import anyio
 import fastapi
 from fastapi.middleware import cors
 import starlette.middleware.base
@@ -847,7 +848,7 @@ async def upload_zip_file(request: fastapi.Request, user_hash: str,
     client_file_mounts_dir = (
         common.API_SERVER_CLIENT_DIR.expanduser().resolve() / user_id /
         'file_mounts')
-    client_file_mounts_dir.mkdir(parents=True, exist_ok=True)
+    await anyio.Path(client_file_mounts_dir).mkdir(parents=True, exist_ok=True)
 
     # Check upload_id to be a valid SkyPilot run_timestamp appended with 8 hex
     # characters, e.g. 'sky-2025-01-17-09-10-13-933602-35d31c22'.
@@ -870,7 +871,7 @@ async def upload_zip_file(request: fastapi.Request, user_hash: str,
         zip_file_path = client_file_mounts_dir / f'{upload_id}.zip'
     else:
         chunk_dir = client_file_mounts_dir / upload_id
-        chunk_dir.mkdir(parents=True, exist_ok=True)
+        await anyio.Path(chunk_dir).mkdir(parents=True, exist_ok=True)
         zip_file_path = chunk_dir / f'part{chunk_index}.incomplete'
 
     try:
@@ -916,9 +917,9 @@ async def upload_zip_file(request: fastapi.Request, user_hash: str,
                         await zip_file.write(data)
 
     logger.info(f'Uploaded zip file: {zip_file_path}')
-    unzip_file(zip_file_path, client_file_mounts_dir)
+    await unzip_file(zip_file_path, client_file_mounts_dir)
     if total_chunks > 1:
-        shutil.rmtree(chunk_dir)
+        await context_utils.to_thread(shutil.rmtree, chunk_dir)
     return payloads.UploadZipFileResponse(
         status=responses.UploadStatus.COMPLETED.value)
 
@@ -933,61 +934,69 @@ def _is_relative_to(path: pathlib.Path, parent: pathlib.Path) -> bool:
         return False
 
 
-def unzip_file(zip_file_path: pathlib.Path,
-               client_file_mounts_dir: pathlib.Path) -> None:
-    """Unzips a zip file."""
-    try:
-        with zipfile.ZipFile(zip_file_path, 'r') as zipf:
-            for member in zipf.infolist():
-                # Determine the new path
-                original_path = os.path.normpath(member.filename)
-                new_path = client_file_mounts_dir / original_path.lstrip('/')
-
-                if (member.external_attr >> 28) == 0xA:
-                    # Symlink. Read the target path and create a symlink.
+async def unzip_file(zip_file_path: pathlib.Path,
+                     client_file_mounts_dir: pathlib.Path) -> None:
+    """Unzips a zip file without blocking the event loop."""
+
+    def _do_unzip() -> None:
+        try:
+            with zipfile.ZipFile(zip_file_path, 'r') as zipf:
+                for member in zipf.infolist():
+                    # Determine the new path
+                    original_path = os.path.normpath(member.filename)
+                    new_path = client_file_mounts_dir / original_path.lstrip(
+                        '/')
+
+                    if (member.external_attr >> 28) == 0xA:
+                        # Symlink. Read the target path and create a symlink.
+                        new_path.parent.mkdir(parents=True, exist_ok=True)
+                        target = zipf.read(member).decode()
+                        assert not os.path.isabs(target), target
+                        # Since target is a relative path, we need to check that
+                        # it is under `client_file_mounts_dir` for security.
+                        full_target_path = (new_path.parent / target).resolve()
+                        if not _is_relative_to(full_target_path,
+                                               client_file_mounts_dir):
+                            raise ValueError(
+                                f'Symlink target {target} leads to a '
+                                'file not in userspace. Aborted.')
+
+                        if new_path.exists() or new_path.is_symlink():
+                            new_path.unlink(missing_ok=True)
+                        new_path.symlink_to(
+                            target,
+                            target_is_directory=member.filename.endswith('/'))
+                        continue
+
+                    # Handle directories
+                    if member.filename.endswith('/'):
+                        new_path.mkdir(parents=True, exist_ok=True)
+                        continue
+
+                    # Handle files
                     new_path.parent.mkdir(parents=True, exist_ok=True)
-                    target = zipf.read(member).decode()
-                    assert not os.path.isabs(target), target
-                    # Since target is a relative path, we need to check that it
-                    # is under `client_file_mounts_dir` for security.
-                    full_target_path = (new_path.parent / target).resolve()
-                    if not _is_relative_to(full_target_path,
-                                           client_file_mounts_dir):
-                        raise ValueError(f'Symlink target {target} leads to a '
-                                         'file not in userspace. Aborted.')
-
-                    if new_path.exists() or new_path.is_symlink():
-                        new_path.unlink(missing_ok=True)
-                    new_path.symlink_to(
-                        target,
-                        target_is_directory=member.filename.endswith('/'))
-                    continue
-
-                # Handle directories
-                if member.filename.endswith('/'):
-                    new_path.mkdir(parents=True, exist_ok=True)
-                    continue
-
-                # Handle files
-                new_path.parent.mkdir(parents=True, exist_ok=True)
-                with zipf.open(member) as member_file, new_path.open('wb') as f:
-                    # Use shutil.copyfileobj to copy files in chunks, so it does
-                    # not load the entire file into memory.
-                    shutil.copyfileobj(member_file, f)
-    except zipfile.BadZipFile as e:
-        logger.error(f'Bad zip file: {zip_file_path}')
-        raise fastapi.HTTPException(
-            status_code=400,
-            detail=f'Invalid zip file: {common_utils.format_exception(e)}')
-    except Exception as e:
-        logger.error(f'Error unzipping file: {zip_file_path}')
-        raise fastapi.HTTPException(
-            status_code=500,
-            detail=(f'Error unzipping file: '
-                    f'{common_utils.format_exception(e)}'))
+                    with zipf.open(member) as member_file, new_path.open(
+                            'wb') as f:
+                        # Use shutil.copyfileobj to copy files in chunks,
+                        # so it does not load the entire file into memory.
+                        shutil.copyfileobj(member_file, f)
+        except zipfile.BadZipFile as e:
+            logger.error(f'Bad zip file: {zip_file_path}')
+            raise fastapi.HTTPException(
+                status_code=400,
+                detail=f'Invalid zip file: {common_utils.format_exception(e)}')
+        except Exception as e:
+            logger.error(f'Error unzipping file: {zip_file_path}')
+            raise fastapi.HTTPException(
+                status_code=500,
+                detail=(f'Error unzipping file: '
+                        f'{common_utils.format_exception(e)}'))
+        finally:
+            # Cleanup the temporary file regardless of
+            # success/failure handling above
+            zip_file_path.unlink(missing_ok=True)
 
-    # Cleanup the temporary file
-    zip_file_path.unlink()
+    await context_utils.to_thread(_do_unzip)
 
 
 @app.post('/launch')
@@ -1388,7 +1397,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(
@@ -1403,6 +1412,9 @@ async def api_get(request_id: str) -> payloads.RequestPayload:
                 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()
         # 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)
@@ -1474,7 +1486,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(
@@ -1491,6 +1503,14 @@ async def stream(
         if log_path == constants.API_SERVER_LOGS:
             resolved_log_path = pathlib.Path(
                 constants.API_SERVER_LOGS).expanduser()
+            if not resolved_log_path.exists():
+                raise fastapi.HTTPException(
+                    status_code=404,
+                    detail='Server log file does not exist. The API server may '
+                    'have been started with `--foreground` - check the '
+                    'stdout of API server process, such as: '
+                    '`kubectl logs -n api-server-namespace '
+                    'api-server-pod-name`')
         else:
             # This should be a log path under ~/sky_logs.
             resolved_logs_directory = pathlib.Path(
@@ -1561,7 +1581,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())
@@ -1769,6 +1789,11 @@ async def complete_volume_name(incomplete: str,) -> List[str]:
     return global_user_state.get_volume_names_start_with(incomplete)
 
 
+@app.get('/api/completion/api_request')
+async def complete_api_request(incomplete: str,) -> List[str]:
+    return await requests_lib.get_api_request_ids_start_with(incomplete)
+
+
 @app.get('/dashboard/{full_path:path}')
 async def serve_dashboard(full_path: str):
     """Serves the Next.js dashboard application.

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/setup_files/dependencies.py

@@ -35,7 +35,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',
@@ -71,6 +72,8 @@ install_requires = [
     'types-paramiko',
     'alembic',
     'aiohttp',
+    'aiosqlite',
+    'anyio',
 ]
 
 # See requirements-dev.txt for the version of grpc and protobuf
@@ -92,8 +95,10 @@ server_dependencies = [
     'passlib',
     'pyjwt',
     'aiohttp',
+    'anyio',
     GRPC,
     PROTOBUF,
+    'aiosqlite',
 ]
 
 local_ray = [

sky/sky_logging.py

@@ -19,6 +19,9 @@ _FORMAT = '%(levelname).1s %(asctime)s %(filename)s:%(lineno)d] %(message)s'
 _DATE_FORMAT = '%m-%d %H:%M:%S'
 _SENSITIVE_LOGGER = ['sky.provisioner', 'sky.optimizer']
 
+_DEBUG_LOG_DIR = os.path.expanduser(
+    os.path.join(constants.SKY_LOGS_DIRECTORY, 'request_debug'))
+
 DEBUG = logging.DEBUG
 INFO = logging.INFO
 WARNING = logging.WARNING
@@ -254,3 +257,28 @@ def generate_tmp_logging_file_path(file_name: str) -> str:
     log_path = os.path.expanduser(os.path.join(log_dir, file_name))
 
     return log_path
+
+
+@contextlib.contextmanager
+def add_debug_log_handler(request_id: str):
+    if os.getenv(constants.ENV_VAR_ENABLE_REQUEST_DEBUG_LOGGING) != 'true':
+        yield
+        return
+
+    os.makedirs(_DEBUG_LOG_DIR, exist_ok=True)
+    log_path = os.path.join(_DEBUG_LOG_DIR, f'{request_id}.log')
+    try:
+        debug_log_handler = logging.FileHandler(log_path)
+        debug_log_handler.setFormatter(FORMATTER)
+        debug_log_handler.setLevel(logging.DEBUG)
+        _root_logger.addHandler(debug_log_handler)
+        # sky.provision sets up its own logger/handler with propogate=False,
+        # so add it there too.
+        provision_logger = logging.getLogger('sky.provision')
+        provision_logger.addHandler(debug_log_handler)
+        provision_logger.setLevel(logging.DEBUG)
+        yield
+    finally:
+        _root_logger.removeHandler(debug_log_handler)
+        provision_logger.removeHandler(debug_log_handler)
+        debug_log_handler.close()

sky/skylet/constants.py

@@ -70,6 +70,7 @@ DEACTIVATE_SKY_REMOTE_PYTHON_ENV = (
 
 # Prefix for SkyPilot environment variables
 SKYPILOT_ENV_VAR_PREFIX = 'SKYPILOT_'
+SKYPILOT_SERVER_ENV_VAR_PREFIX = 'SKYPILOT_SERVER_'
 
 # The name for the environment variable that stores the unique ID of the
 # current task. This will stay the same across multiple recoveries of the
@@ -417,6 +418,7 @@ LOCAL_SKYPILOT_CONFIG_PATH_PLACEHOLDER = 'skypilot:local_skypilot_config_path'
 # Path to the generated cluster config yamls and ssh configs.
 SKY_USER_FILE_PATH = '~/.sky/generated'
 
+# TODO(cooperc): Update all env vars to begin with SKYPILOT_ or SKYPILOT_SERVER_
 # Environment variable that is set to 'true' if this is a skypilot server.
 ENV_VAR_IS_SKYPILOT_SERVER = 'IS_SKYPILOT_SERVER'
 
@@ -436,6 +438,10 @@ ENV_VAR_ENABLE_BASIC_AUTH = 'ENABLE_BASIC_AUTH'
 SKYPILOT_INITIAL_BASIC_AUTH = 'SKYPILOT_INITIAL_BASIC_AUTH'
 ENV_VAR_ENABLE_SERVICE_ACCOUNTS = 'ENABLE_SERVICE_ACCOUNTS'
 
+# Enable debug logging for requests.
+ENV_VAR_ENABLE_REQUEST_DEBUG_LOGGING = (
+    f'{SKYPILOT_SERVER_ENV_VAR_PREFIX}ENABLE_REQUEST_DEBUG_LOGGING')
+
 SKYPILOT_DEFAULT_WORKSPACE = 'default'
 
 # BEGIN constants used for service catalog.

sky/skylet/events.py

@@ -20,7 +20,6 @@ from sky.skylet import constants
 from sky.skylet import job_lib
 from sky.usage import usage_lib
 from sky.utils import cluster_utils
-from sky.utils import common_utils
 from sky.utils import registry
 from sky.utils import ux_utils
 from sky.utils import yaml_utils
@@ -181,7 +180,7 @@ class AutostopEvent(SkyletEvent):
 
             config_path = os.path.abspath(
                 os.path.expanduser(cluster_utils.SKY_CLUSTER_YAML_REMOTE_PATH))
-            config = common_utils.read_yaml(config_path)
+            config = yaml_utils.read_yaml(config_path)
             provider_name = cluster_utils.get_provider_name(config)
             cloud = registry.CLOUD_REGISTRY.from_str(provider_name)
             assert cloud is not None, f'Unknown cloud: {provider_name}'
@@ -326,5 +325,5 @@ class AutostopEvent(SkyletEvent):
         config['auth'].pop('ssh_proxy_command', None)
         # Empty the file_mounts.
         config['file_mounts'] = {}
-        common_utils.dump_yaml(yaml_path, config)
+        yaml_utils.dump_yaml(yaml_path, config)
         logger.debug('Replaced upscaling speed to 0.')