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

sky/server/uvicorn.py

@@ -3,24 +3,238 @@
 This module is a wrapper around uvicorn to customize the behavior of the
 server.
 """
+import asyncio
+import logging
 import os
+import signal
+import sys
 import threading
-from typing import Optional
+import time
+from types import FrameType
+from typing import Optional, Union
 
+import filelock
 import uvicorn
 from uvicorn.supervisors import multiprocess
 
+from sky import sky_logging
+from sky.server import daemons
+from sky.server import metrics as metrics_lib
+from sky.server import state
+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
+from sky.utils.db import db_utils
 
+logger = sky_logging.init_logger(__name__)
 
-def run(config: uvicorn.Config):
+# File lock path for coordinating graceful shutdown across processes
+_GRACEFUL_SHUTDOWN_LOCK_PATH = '/tmp/skypilot_graceful_shutdown.lock'
+
+# Interval to check for on-going requests.
+_WAIT_REQUESTS_INTERVAL_SECONDS = 5
+
+# Timeout for waiting for on-going requests to finish.
+try:
+    _WAIT_REQUESTS_TIMEOUT_SECONDS = int(
+        os.environ.get(constants.GRACE_PERIOD_SECONDS_ENV_VAR, '60'))
+except ValueError:
+    _WAIT_REQUESTS_TIMEOUT_SECONDS = 60
+
+# TODO(aylei): use decorator to register requests that need to be proactively
+# cancelled instead of hardcoding here.
+_RETRIABLE_REQUEST_NAMES = {
+    'sky.logs',
+    'sky.jobs.logs',
+    'sky.serve.logs',
+}
+
+
+def add_timestamp_prefix_for_server_logs() -> None:
+    """Configure logging for API server.
+
+    Note: we only do this in the main API server process and uvicorn processes,
+    to avoid affecting executor logs (including in modules like
+    sky.server.requests) that may get sent to the client.
+    """
+    server_logger = sky_logging.init_logger('sky.server')
+    # Clear existing handlers first to prevent duplicates
+    server_logger.handlers.clear()
+    # Disable propagation to avoid the root logger of SkyPilot being affected
+    server_logger.propagate = False
+    # Add date prefix to the log message printed by loggers under
+    # server.
+    stream_handler = logging.StreamHandler(sys.stdout)
+    if env_options.Options.SHOW_DEBUG_INFO.get():
+        stream_handler.setLevel(logging.DEBUG)
+    else:
+        stream_handler.setLevel(logging.INFO)
+    stream_handler.flush = sys.stdout.flush  # type: ignore
+    stream_handler.setFormatter(sky_logging.FORMATTER)
+    server_logger.addHandler(stream_handler)
+    # Add date prefix to the log message printed by uvicorn.
+    for name in ['uvicorn', 'uvicorn.access']:
+        uvicorn_logger = logging.getLogger(name)
+        uvicorn_logger.handlers.clear()
+        uvicorn_logger.addHandler(stream_handler)
+
+
+class Server(uvicorn.Server):
+    """Server wrapper for uvicorn.
+
+    Extended functionalities:
+    - Handle exit signal and perform custom graceful shutdown.
+    - Run the server process with contextually aware.
+    """
+
+    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.
+
+        When a server process receives a SIGTERM or SIGINT signal, a graceful
+        shutdown will be initiated. If a SIGINT signal is received again, the
+        server will be forcefully shutdown.
+        """
+        if self.exiting and sig == signal.SIGINT:
+            # The server has been siganled to exit and recieved a SIGINT again,
+            # do force shutdown.
+            logger.info('Force shutdown.')
+            self.should_exit = True
+            super().handle_exit(sig, frame)
+            return
+        if not self.exiting:
+            self.exiting = True
+            # Perform graceful shutdown in a separate thread to avoid blocking
+            # the main thread.
+            threading.Thread(target=self._graceful_shutdown,
+                             args=(sig, frame),
+                             daemon=True).start()
+
+    def _graceful_shutdown(self, sig: int, frame: Union[FrameType,
+                                                        None]) -> None:
+        """Perform graceful shutdown."""
+        # Block new requests so that we can wait until all on-going requests
+        # are finished. Note that /api/$verb operations are still allowed in
+        # this stage to ensure the client can still operate the on-going
+        # requests, e.g. /api/logs, /api/cancel, etc.
+        logger.info('Block new requests being submitted in worker '
+                    f'{os.getpid()}.')
+        state.set_block_requests(True)
+        # Ensure the shutting_down are set on all workers before next step.
+        # TODO(aylei): hacky, need a reliable solution.
+        time.sleep(1)
+
+        lock = filelock.FileLock(_GRACEFUL_SHUTDOWN_LOCK_PATH)
+        # Elect a coordinator process to handle on-going requests check
+        with lock.acquire():
+            logger.info(f'Worker {os.getpid()} elected as shutdown coordinator')
+            self._wait_requests()
+
+        logger.info('Shutting down server...')
+        self.should_exit = True
+        super().handle_exit(sig, frame)
+
+    def _wait_requests(self) -> None:
+        """Wait until all on-going requests are finished or cancelled."""
+        start_time = time.time()
+        while True:
+            statuses = [
+                requests_lib.RequestStatus.PENDING,
+                requests_lib.RequestStatus.RUNNING,
+            ]
+            requests = [(request_task.request_id, request_task.name)
+                        for request_task in requests_lib.get_request_tasks(
+                            req_filter=requests_lib.RequestTaskFilter(
+                                status=statuses, fields=['request_id', 'name']))
+                       ]
+            if not requests:
+                break
+            logger.info(f'{len(requests)} on-going requests '
+                        'found, waiting for them to finish...')
+            # Proactively cancel internal requests and logs requests since
+            # they can run for infinite time.
+            internal_request_ids = {
+                d.id for d in daemons.INTERNAL_REQUEST_DAEMONS
+            }
+            if time.time() - start_time > _WAIT_REQUESTS_TIMEOUT_SECONDS:
+                logger.warning('Timeout waiting for on-going requests to '
+                               'finish, cancelling all on-going requests.')
+                for request_id, _ in requests:
+                    self.interrupt_request_for_retry(request_id)
+                break
+            interrupted = 0
+            for request_id, name in requests:
+                if (name in _RETRIABLE_REQUEST_NAMES or
+                        request_id in internal_request_ids):
+                    self.interrupt_request_for_retry(request_id)
+                    interrupted += 1
+                # TODO(aylei): interrupt pending requests to accelerate the
+                # shutdown.
+            # If some requests are not interrupted, wait for them to finish,
+            # otherwise we just check again immediately to accelerate the
+            # shutdown process.
+            if interrupted < len(requests):
+                time.sleep(_WAIT_REQUESTS_INTERVAL_SECONDS)
+
+    def interrupt_request_for_retry(self, request_id: str) -> None:
+        """Interrupt a request for retry."""
+        with requests_lib.update_request(request_id) as req:
+            if req is None:
+                return
+            if req.pid is not None:
+                try:
+                    os.kill(req.pid, signal.SIGTERM)
+                except ProcessLookupError:
+                    logger.debug(f'Process {req.pid} already finished.')
+            req.status = requests_lib.RequestStatus.CANCELLED
+            req.should_retry = True
+        logger.info(
+            f'Request {request_id} interrupted and will be retried by client.')
+
+    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).
+        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
+        stop_monitor = threading.Event()
+        monitor = threading.Thread(target=metrics_lib.process_monitor,
+                                   args=('server', stop_monitor),
+                                   daemon=True)
+        monitor.start()
+        try:
+            with self.capture_signals():
+                asyncio.run(self.serve(*args, **kwargs))
+        finally:
+            stop_monitor.set()
+            monitor.join()
+
+
+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 = uvicorn.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/server/versions.py

@@ -0,0 +1,270 @@
+"""API versioning module."""
+
+import contextvars
+import functools
+import re
+from typing import Callable, Literal, Mapping, NamedTuple, Optional, Tuple
+
+import colorama
+from packaging import version as version_lib
+
+import sky
+from sky import exceptions
+from sky import sky_logging
+from sky.server import constants
+from sky.utils import ux_utils
+
+logger = sky_logging.init_logger(__name__)
+
+CLIENT_TOO_OLD_ERROR = (
+    f'{colorama.Fore.YELLOW}Your SkyPilot client version is too old:'
+    '{remote_version}\n'
+    f'{colorama.Style.RESET_ALL}'
+    'The server is running on {local_version} and the minimum compatible '
+    'version is {min_version}.\n'
+    f'Upgrade your client with:\n{colorama.Fore.YELLOW}'
+    '{command}'
+    f'{colorama.Style.RESET_ALL}')
+SERVER_TOO_OLD_ERROR = (
+    f'{colorama.Fore.YELLOW}Your SkyPilot API server version is too old: '
+    '{remote_version}\n'
+    f'{colorama.Style.RESET_ALL}'
+    'The client is running on {local_version} and the minimum compatible '
+    'version is {min_version}.\n'
+    'Contact your administrator to upgrade the remote API server or downgrade '
+    f'your client with:\n{colorama.Fore.YELLOW}'
+    '{command}'
+    f'{colorama.Style.RESET_ALL}')
+
+# SkyPilot dev version.
+DEV_VERSION = '1.0.0-dev0'
+
+_REMOTE_TO_ERROR = {
+    'client': CLIENT_TOO_OLD_ERROR,
+    'server': SERVER_TOO_OLD_ERROR,
+}
+
+# Context-local (thread or cooroutine) remote API version, captured during
+# communication with the remote peer.
+_remote_api_version: contextvars.ContextVar[Optional[int]] = \
+    contextvars.ContextVar('remote_api_version', default=None)
+_remote_version: contextvars.ContextVar[str] = \
+    contextvars.ContextVar('remote_version', default='unknown')
+_reminded_for_minor_version_upgrade = False
+
+
+def get_remote_api_version() -> Optional[int]:
+    return _remote_api_version.get()
+
+
+def set_remote_api_version(api_version: int) -> None:
+    _remote_api_version.set(api_version)
+
+
+def get_remote_version() -> str:
+    return _remote_version.get()
+
+
+def set_remote_version(version: str) -> None:
+    _remote_version.set(version)
+
+
+class VersionInfo(NamedTuple):
+    api_version: int
+    version: str
+    error: Optional[str] = None
+
+
+def check_compatibility_at_server(
+        client_headers: Mapping[str, str]) -> Optional[VersionInfo]:
+    """Check API compatibility between client and server."""
+    return _check_version_compatibility(client_headers, 'client')
+
+
+def check_compatibility_at_client(
+        server_headers: Mapping[str, str]) -> Optional[VersionInfo]:
+    """Check API compatibility between client and server."""
+    return _check_version_compatibility(server_headers, 'server')
+
+
+def _check_version_compatibility(
+        remote_headers: Mapping[str, str],
+        remote_type: Literal['client', 'server']) -> Optional[VersionInfo]:
+    """Check API compatibility between client and server.
+
+    This function can be called at both client and server side, where the
+    headers should contain the version info of the remote.
+
+    Args:
+        remote_headers: The headers of the request/response sent from the
+            remote.
+        remote_type: The type of the remote, used to determine the error
+            message. Valid options are 'client' and 'server'.
+
+    Returns:
+        The version info of the remote, None if the version info is not found
+        in the headers for backward compatibility.
+    """
+    api_version_str = remote_headers.get(constants.API_VERSION_HEADER)
+    version = remote_headers.get(constants.VERSION_HEADER)
+    if version is None or api_version_str is None:
+        return None
+    try:
+        api_version = int(api_version_str)
+    except ValueError:
+        # The future change is expected to not break the compatibility of this
+        # header, so we are encountering a bug or a malicious request here,
+        # just raise an error.
+        raise ValueError(
+            f'Header {constants.API_VERSION_HEADER}: '
+            f'{api_version_str} is not a valid API version.') from None
+
+    if api_version < constants.MIN_COMPATIBLE_API_VERSION:
+        if remote_type == 'server':
+            # Hint the user to downgrade to client to the remote server server.
+            server_version, server_commit = parse_readable_version(version)
+            command = install_version_command(server_version, server_commit)
+        else:
+            # Hint the client to upgrade to upgrade the server version
+            command = install_version_command(sky.__version__, sky.__commit__)
+        return VersionInfo(api_version=api_version,
+                           version=version,
+                           error=_REMOTE_TO_ERROR[remote_type].format(
+                               remote_version=version,
+                               local_version=get_local_readable_version(),
+                               min_version=constants.MIN_COMPATIBLE_VERSION,
+                               command=command,
+                           ))
+
+    if remote_type == 'server':
+        # Only print the reminder at client-side.
+        _remind_minor_version_upgrade(version)
+
+    return VersionInfo(api_version=api_version, version=version)
+
+
+def get_local_readable_version() -> str:
+    """Get the readable version of the SkyPilot code loaded in current process.
+
+    For dev version, the version is formatted as: 1.0.0-dev0 (commit: 1234567)
+    to make it meaningful for users.
+    """
+    if sky.__version__ == DEV_VERSION:
+        return f'{sky.__version__} (commit: {sky.__commit__})'
+    else:
+        return sky.__version__
+
+
+def parse_readable_version(version: str) -> Tuple[str, Optional[str]]:
+    """Parse a readable produced by get_local_readable_version.
+
+    Args:
+        version: The version string to parse.
+
+    Returns:
+        A tuple of (version, optional_commit) where:
+        - version: The base version string (e.g., "1.0.0-dev0")
+        - optional_commit: The commit hash if present, None otherwise
+    """
+    # Check if this is a dev version with commit info
+    # Format: "1.0.0-dev0 (commit: 1234567)"
+    commit_pattern = r'^(.+) \(commit: ([a-f0-9]+)\)$'
+    match = re.match(commit_pattern, version)
+
+    if match:
+        base_version = match.group(1)
+        commit = match.group(2)
+        return base_version, commit
+    else:
+        # Regular version without commit info
+        return version, None
+
+
+def install_version_command(version: str, commit: Optional[str] = None) -> str:
+    if version == DEV_VERSION:
+        if commit is not None:
+            return ('pip install git+https://github.com/skypilot-org/skypilot@'
+                    f'{commit}')
+    elif 'dev' in version:
+        return f'pip install -U "skypilot-nightly=={version}"'
+    return f'pip install -U "skypilot=={version}"'
+
+
+def _remind_minor_version_upgrade(remote_version: str) -> None:
+    """Remind the user to upgrade the CLI/SDK."""
+    # Only print the reminder once per process.
+    global _reminded_for_minor_version_upgrade
+    if _reminded_for_minor_version_upgrade:
+        return
+    # Skip for dev versions.
+    if 'dev' in sky.__version__ or 'dev' in remote_version:
+        return
+
+    # Remove the commit info if any.
+    remote_base_version, _ = parse_readable_version(remote_version)
+
+    # Parse semver for both local and remote versions
+    try:
+        local = version_lib.parse(sky.__version__)
+        remote = version_lib.parse(remote_base_version)
+
+        # Check if local version is behind remote version, ignore patch version.
+        if (local.major, local.minor) < (remote.major, remote.minor):
+            logger.warning(
+                f'{colorama.Fore.YELLOW}The SkyPilot API server is running in '
+                f'version {remote_version}, which is newer than your client '
+                f'version {sky.__version__}. The compatibility for your '
+                f'current version might be dropped in the next server upgrade.'
+                f'\nConsider upgrading your client with:\n'
+                f'{install_version_command(remote_version)}'
+                f'{colorama.Style.RESET_ALL}')
+            _reminded_for_minor_version_upgrade = True
+    except version_lib.InvalidVersion:
+        # Skip for non-valid semver (probabely a dev version)
+        pass
+
+
+# TODO(aylei): maybe we can use similiar approach to mark a new argument can
+# only be used in the new server version.
+def minimal_api_version(min_version: int) -> Callable:
+    """Decorator to enforce a minimum remote API version for an SDK function.
+
+    New SDK method must be decorated with this decorator to make sure it raises
+    an readable error when the remote server is not upgraded.
+
+    Args:
+        min_version: The minimum remote API version required to call the
+        function.
+
+    Returns:
+        A decorator function that checks API version before execution.
+
+    Raises:
+        APINotSupportedError: If the remote API version is below the minimum
+            required.
+    """
+
+    def decorator(func: Callable) -> Callable:
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            remote_api_version = get_remote_api_version()
+            if remote_api_version is None:
+                return func(*args, **kwargs)
+            if remote_api_version < min_version:
+                with ux_utils.print_exception_no_traceback():
+                    hint = 'Please upgrade the remote server.'
+                    # The client runs in a released version, do better hint.
+                    if 'dev' not in sky.__version__:
+                        hint = (
+                            f'Upgrade the remote server to {sky.__version__} '
+                            'and re-run the command.')
+                    raise exceptions.APINotSupportedError(
+                        f'Function {func.__name__} is introduced after the '
+                        f'remote server version {get_remote_version()!r} is '
+                        f'released. {hint}')
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

sky/setup_files/MANIFEST.in

@@ -9,6 +9,7 @@ include sky/skylet/providers/ibm/*
 include sky/skylet/providers/scp/*
 include sky/skylet/providers/*.py
 include sky/skylet/ray_patches/*.patch
+include sky/skylet/ray_patches/*.diff
 include sky/jobs/dashboard/*
 include sky/jobs/dashboard/templates/*
 include sky/jobs/dashboard/static/*
@@ -16,3 +17,7 @@ include sky/templates/*
 include sky/utils/kubernetes/*
 include sky/server/html/*
 recursive-include sky/dashboard/out *
+include sky/users/*.conf
+include sky/utils/*.sh
+include sky/setup_files/alembic.ini
+recursive-include sky/schemas/db *

sky/setup_files/alembic.ini

@@ -0,0 +1,156 @@
+# alembic configuration for global user state, jobs state, and sky config db migrations.
+
+[DEFAULT]
+# path to migration scripts.
+# this is typically a path given in POSIX (e.g. forward slashes)
+# format, relative to the token %(here)s which refers to the location of this
+# ini file
+script_location = %(here)s/../schemas/db
+
+# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
+# Uncomment the line below if you want the files to be prepended with date and time
+# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file
+# for all available tokens
+# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
+
+# sys.path path, will be prepended to sys.path if present.
+# defaults to the current working directory.  for multiple paths, the path separator
+# is defined by "path_separator" below.
+prepend_sys_path = .
+
+
+# timezone to use when rendering the date within the migration file
+# as well as the filename.
+# If specified, requires the python>=3.9 or backports.zoneinfo library and tzdata library.
+# Any required deps can installed by adding `alembic[tz]` to the pip requirements
+# string value is passed to ZoneInfo()
+# leave blank for localtime
+# timezone =
+
+# max length of characters to apply to the "slug" field
+# truncate_slug_length = 40
+
+# set to 'true' to run the environment during
+# the 'revision' command, regardless of autogenerate
+# revision_environment = false
+
+# set to 'true' to allow .pyc and .pyo files without
+# a source .py file to be detected as revisions in the
+# versions/ directory
+# sourceless = false
+
+# version location specification; This defaults
+# to <script_location>/versions.  When using multiple version
+# directories, initial revisions must be specified with --version-path.
+# The path separator used here should be the separator specified by "path_separator"
+# below.
+# version_locations = %(here)s/bar:%(here)s/bat:%(here)s/alembic/versions
+
+# path_separator; This indicates what character is used to split lists of file
+# paths, including version_locations and prepend_sys_path within configparser
+# files such as alembic.ini.
+# The default rendered in new alembic.ini files is "os", which uses os.pathsep
+# to provide os-dependent path splitting.
+#
+# Note that in order to support legacy alembic.ini files, this default does NOT
+# take place if path_separator is not present in alembic.ini.  If this
+# option is omitted entirely, fallback logic is as follows:
+#
+# 1. Parsing of the version_locations option falls back to using the legacy
+#    "version_path_separator" key, which if absent then falls back to the legacy
+#    behavior of splitting on spaces and/or commas.
+# 2. Parsing of the prepend_sys_path option falls back to the legacy
+#    behavior of splitting on spaces, commas, or colons.
+#
+# Valid values for path_separator are:
+#
+# path_separator = :
+# path_separator = ;
+# path_separator = space
+# path_separator = newline
+#
+# Use os.pathsep. Default configuration used for new projects.
+path_separator = os
+
+# set to 'true' to search source files recursively
+# in each "version_locations" directory
+# new in Alembic version 1.10
+# recursive_version_locations = false
+
+# the output encoding used when revision files
+# are written from script.py.mako
+# output_encoding = utf-8
+
+# database URL.  This is consumed by the user-maintained env.py script only.
+# other means of configuring database URLs may be customized within the env.py
+# file.
+# sqlalchemy.url = driver://user:pass@localhost/dbname
+
+[state_db]
+version_locations = %(here)s/../schemas/db/global_user_state
+version_table = alembic_version_state_db
+
+[spot_jobs_db]
+version_locations = %(here)s/../schemas/db/spot_jobs
+version_table = alembic_version_spot_jobs_db
+
+[serve_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
+# detail and examples
+
+# format using "black" - use the console_scripts runner, against the "black" entrypoint
+# hooks = black
+# black.type = console_scripts
+# black.entrypoint = black
+# black.options = -l 79 REVISION_SCRIPT_FILENAME
+
+# lint with attempts to fix using "ruff" - use the exec runner, execute a binary
+# hooks = ruff
+# ruff.type = exec
+# ruff.executable = %(here)s/.venv/bin/ruff
+# ruff.options = check --fix REVISION_SCRIPT_FILENAME
+
+# Logging configuration.  This is also consumed by the user-maintained
+# env.py script only.
+[loggers]
+keys = root,sqlalchemy,alembic
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARNING
+handlers = console
+qualname =
+
+[logger_sqlalchemy]
+level = WARNING
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = WARNING
+handlers =
+qualname = alembic
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s
+datefmt = %H:%M:%S