skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/server/requests/requests.py

@@ -0,0 +1,567 @@
+"""Utilities for REST API."""
+import contextlib
+import dataclasses
+import enum
+import functools
+import json
+import os
+import pathlib
+import shutil
+import signal
+import sqlite3
+import time
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import colorama
+import filelock
+
+from sky import exceptions
+from sky import global_user_state
+from sky import sky_logging
+from sky.server import common as server_common
+from sky.server import constants as server_constants
+from sky.server.requests import payloads
+from sky.server.requests.serializers import decoders
+from sky.server.requests.serializers import encoders
+from sky.utils import common
+from sky.utils import common_utils
+from sky.utils import db_utils
+from sky.utils import env_options
+
+logger = sky_logging.init_logger(__name__)
+
+# Tables in task.db.
+REQUEST_TABLE = 'requests'
+COL_CLUSTER_NAME = 'cluster_name'
+COL_USER_ID = 'user_id'
+REQUEST_LOG_PATH_PREFIX = '~/sky_logs/api_server/requests'
+
+# TODO(zhwu): For scalability, there are several TODOs:
+# [x] Have a way to queue requests.
+# [ ] Move logs to persistent place.
+# [ ] Deploy API server in a autoscaling fashion.
+
+
+class RequestStatus(enum.Enum):
+    """The status of a request."""
+
+    PENDING = 'PENDING'
+    RUNNING = 'RUNNING'
+    SUCCEEDED = 'SUCCEEDED'
+    FAILED = 'FAILED'
+    CANCELLED = 'CANCELLED'
+
+    def __gt__(self, other):
+        return (list(RequestStatus).index(self) >
+                list(RequestStatus).index(other))
+
+    def colored_str(self):
+        color = _STATUS_TO_COLOR[self]
+        return f'{color}{self.value}{colorama.Style.RESET_ALL}'
+
+
+_STATUS_TO_COLOR = {
+    RequestStatus.PENDING: colorama.Fore.BLUE,
+    RequestStatus.RUNNING: colorama.Fore.GREEN,
+    RequestStatus.SUCCEEDED: colorama.Fore.GREEN,
+    RequestStatus.FAILED: colorama.Fore.RED,
+    RequestStatus.CANCELLED: colorama.Fore.WHITE,
+}
+
+REQUEST_COLUMNS = [
+    'request_id',
+    'name',
+    'entrypoint',
+    'request_body',
+    'status',
+    'return_value',
+    'error',
+    'pid',
+    'created_at',
+    COL_CLUSTER_NAME,
+    'schedule_type',
+    COL_USER_ID,
+]
+
+
+class ScheduleType(enum.Enum):
+    """The schedule type for the requests."""
+    LONG = 'long'
+    # Queue for requests that should be executed quickly for a quick response.
+    SHORT = 'short'
+
+
+@dataclasses.dataclass
+class RequestPayload:
+    """The payload for the requests."""
+
+    request_id: str
+    name: str
+    entrypoint: str
+    request_body: str
+    status: str
+    created_at: float
+    user_id: str
+    return_value: str
+    error: str
+    pid: Optional[int]
+    schedule_type: str
+    user_name: Optional[str] = None
+    # Resources the request operates on.
+    cluster_name: Optional[str] = None
+
+
+@dataclasses.dataclass
+class Request:
+    """A SkyPilot API request."""
+
+    request_id: str
+    name: str
+    entrypoint: Callable
+    request_body: payloads.RequestBody
+    status: RequestStatus
+    created_at: float
+    user_id: str
+    return_value: Any = None
+    error: Optional[Dict[str, Any]] = None
+    # The pid of the request worker that is(was) running this request.
+    pid: Optional[int] = None
+    schedule_type: ScheduleType = ScheduleType.LONG
+    # Resources the request operates on.
+    cluster_name: Optional[str] = None
+
+    @property
+    def log_path(self) -> pathlib.Path:
+        log_path_prefix = pathlib.Path(
+            REQUEST_LOG_PATH_PREFIX).expanduser().absolute()
+        log_path_prefix.mkdir(parents=True, exist_ok=True)
+        log_path = (log_path_prefix / self.request_id).with_suffix('.log')
+        return log_path
+
+    def set_error(self, error: Exception) -> None:
+        """Set the error."""
+        # TODO(zhwu): pickle.dump does not work well with custom exceptions if
+        # it has more than 1 arguments.
+        serialized = exceptions.serialize_exception(error)
+        self.error = {
+            'object': encoders.pickle_and_encode(serialized),
+            'type': type(error).__name__,
+            'message': str(error),
+        }
+
+    def get_error(self) -> Optional[Dict[str, Any]]:
+        """Get the error."""
+        if self.error is None:
+            return None
+        unpickled = decoders.decode_and_unpickle(self.error['object'])
+        deserialized = exceptions.deserialize_exception(unpickled)
+        return {
+            'object': deserialized,
+            'type': self.error['type'],
+            'message': self.error['message'],
+        }
+
+    def set_return_value(self, return_value: Any) -> None:
+        """Set the return value."""
+        self.return_value = encoders.get_encoder(self.name)(return_value)
+
+    def get_return_value(self) -> Any:
+        """Get the return value."""
+        return decoders.get_decoder(self.name)(self.return_value)
+
+    @classmethod
+    def from_row(cls, row: Tuple[Any, ...]) -> 'Request':
+        content = dict(zip(REQUEST_COLUMNS, row))
+        return cls.decode(RequestPayload(**content))
+
+    def to_row(self) -> Tuple[Any, ...]:
+        payload = self.encode()
+        row = []
+        for k in REQUEST_COLUMNS:
+            row.append(getattr(payload, k))
+        return tuple(row)
+
+    def readable_encode(self) -> RequestPayload:
+        """Serialize the SkyPilot API request for display purposes.
+
+        This function should be called on the server side to serialize the
+        request body into human readable format, e.g., the entrypoint should
+        be a string, and the pid, error, or return value are not needed.
+
+        The returned value will then be displayed on the client side in request
+        table.
+
+        We do not use `encode` for display to avoid a large amount of data being
+        sent to the client side, especially for the request table could include
+        all the requests.
+        """
+        assert isinstance(self.request_body,
+                          payloads.RequestBody), (self.name, self.request_body)
+        user_name = global_user_state.get_user(self.user_id).name
+        return RequestPayload(
+            request_id=self.request_id,
+            name=self.name,
+            entrypoint=self.entrypoint.__name__,
+            request_body=self.request_body.model_dump_json(),
+            status=self.status.value,
+            return_value=json.dumps(None),
+            error=json.dumps(None),
+            pid=None,
+            created_at=self.created_at,
+            schedule_type=self.schedule_type.value,
+            user_id=self.user_id,
+            user_name=user_name,
+            cluster_name=self.cluster_name,
+        )
+
+    def encode(self) -> RequestPayload:
+        """Serialize the SkyPilot API request."""
+        assert isinstance(self.request_body,
+                          payloads.RequestBody), (self.name, self.request_body)
+        try:
+            return RequestPayload(
+                request_id=self.request_id,
+                name=self.name,
+                entrypoint=encoders.pickle_and_encode(self.entrypoint),
+                request_body=encoders.pickle_and_encode(self.request_body),
+                status=self.status.value,
+                return_value=json.dumps(self.return_value),
+                error=json.dumps(self.error),
+                pid=self.pid,
+                created_at=self.created_at,
+                schedule_type=self.schedule_type.value,
+                user_id=self.user_id,
+                cluster_name=self.cluster_name,
+            )
+        except (TypeError, ValueError) as e:
+            # The error is unexpected, so we don't suppress the stack trace.
+            logger.error(
+                f'Error encoding: {e}\n'
+                f'  {self.request_id}\n'
+                f'  {self.name}\n'
+                f'  {self.request_body}\n'
+                f'  {self.return_value}\n'
+                f'  {self.created_at}\n',
+                exc_info=e)
+            raise
+
+    @classmethod
+    def decode(cls, payload: RequestPayload) -> 'Request':
+        """Deserialize the SkyPilot API request."""
+        try:
+            return cls(
+                request_id=payload.request_id,
+                name=payload.name,
+                entrypoint=decoders.decode_and_unpickle(payload.entrypoint),
+                request_body=decoders.decode_and_unpickle(payload.request_body),
+                status=RequestStatus(payload.status),
+                return_value=json.loads(payload.return_value),
+                error=json.loads(payload.error),
+                pid=payload.pid,
+                created_at=payload.created_at,
+                schedule_type=ScheduleType(payload.schedule_type),
+                user_id=payload.user_id,
+                cluster_name=payload.cluster_name,
+            )
+        except (TypeError, ValueError) as e:
+            logger.error(
+                f'Error decoding: {e}\n'
+                f'  {payload.request_id}\n'
+                f'  {payload.name}\n'
+                f'  {payload.entrypoint}\n'
+                f'  {payload.request_body}\n'
+                f'  {payload.created_at}\n',
+                exc_info=e)
+            # The error is unexpected, so we don't suppress the stack trace.
+            raise
+
+
+def kill_cluster_requests(cluster_name: str, exclude_request_name: str):
+    """Kill all pending and running requests for a cluster.
+
+    Args:
+        cluster_name: the name of the cluster.
+        exclude_request_names: exclude requests with these names. This is to
+            prevent killing the caller request.
+    """
+    request_ids = [
+        request_task.request_id for request_task in get_request_tasks(
+            cluster_names=[cluster_name],
+            status=[RequestStatus.PENDING, RequestStatus.RUNNING],
+            exclude_request_names=[exclude_request_name])
+    ]
+    kill_requests(request_ids)
+
+
+def refresh_cluster_status_event():
+    """Periodically refresh the cluster status."""
+    # pylint: disable=import-outside-toplevel
+    from sky import core
+
+    # Disable logging for periodic refresh to avoid the usage message being
+    # sent multiple times.
+    os.environ[env_options.Options.DISABLE_LOGGING.env_key] = '1'
+
+    while True:
+        logger.info('=== Refreshing cluster status ===')
+        # This periodically refresh will hold the lock for the cluster being
+        # refreshed, but it is OK because other operations will just wait for
+        # the lock and get the just refreshed status without refreshing again.
+        core.status(refresh=common.StatusRefreshMode.FORCE, all_users=True)
+        logger.info(
+            'Status refreshed. Sleeping '
+            f'{server_constants.CLUSTER_REFRESH_DAEMON_INTERVAL_SECONDS}'
+            ' seconds for the next refresh...\n')
+        time.sleep(server_constants.CLUSTER_REFRESH_DAEMON_INTERVAL_SECONDS)
+
+
+@dataclasses.dataclass
+class InternalRequestDaemon:
+    id: str
+    name: str
+    event_fn: Callable[[], None]
+
+
+# Register the events to run in the background.
+INTERNAL_REQUEST_DAEMONS = [
+    # This status refresh daemon can cause the autostopp'ed/autodown'ed cluster
+    # set to updated status automatically, without showing users the hint of
+    # cluster being stopped or down when `sky status -r` is called.
+    InternalRequestDaemon(id='skypilot-status-refresh-daemon',
+                          name='status',
+                          event_fn=refresh_cluster_status_event)
+]
+
+
+def kill_requests(request_ids: Optional[List[str]] = None,
+                  user_id: Optional[str] = None) -> List[str]:
+    """Kill a SkyPilot API request and set its status to cancelled.
+
+    Args:
+        request_ids: The request IDs to kill. If None, all requests for the
+            user are killed.
+        user_id: The user ID to kill requests for. If None, all users are
+            killed.
+
+    Returns:
+        A list of request IDs that were cancelled.
+    """
+    if request_ids is None:
+        request_ids = [
+            request_task.request_id for request_task in get_request_tasks(
+                user_id=user_id,
+                status=[RequestStatus.RUNNING, RequestStatus.PENDING],
+                # Avoid cancelling the cancel request itself.
+                exclude_request_names=['sky.api_cancel'])
+        ]
+    cancelled_request_ids = []
+    for request_id in request_ids:
+        with update_request(request_id) as request_record:
+            if request_record is None:
+                logger.debug(f'No request ID {request_id}')
+                continue
+            # Skip internal requests. The internal requests are scheduled with
+            # request_id in range(len(INTERNAL_REQUEST_EVENTS)).
+            if request_record.request_id in set(
+                    event.id for event in INTERNAL_REQUEST_DAEMONS):
+                continue
+            if request_record.status > RequestStatus.RUNNING:
+                logger.debug(f'Request {request_id} already finished')
+                continue
+            if request_record.pid is not None:
+                logger.debug(f'Killing request process {request_record.pid}')
+                # Use SIGTERM instead of SIGKILL:
+                # - The executor can handle SIGTERM gracefully
+                # - After SIGTERM, the executor can reuse the request process
+                #   for other requests, avoiding the overhead of forking a new
+                #   process for each request.
+                os.kill(request_record.pid, signal.SIGTERM)
+            request_record.status = RequestStatus.CANCELLED
+            cancelled_request_ids.append(request_id)
+    return cancelled_request_ids
+
+
+_DB_PATH = os.path.expanduser(server_constants.API_SERVER_REQUEST_DB_PATH)
+pathlib.Path(_DB_PATH).parents[0].mkdir(parents=True, exist_ok=True)
+
+
+def create_table(cursor, conn):
+    del conn
+    # Enable WAL mode to avoid locking issues.
+    # See: issue #1441 and PR #1509
+    # https://github.com/microsoft/WSL/issues/2395
+    # TODO(romilb): We do not enable WAL for WSL because of known issue in WSL.
+    #  This may cause the database locked problem from WSL issue #1441.
+    if not common_utils.is_wsl():
+        try:
+            cursor.execute('PRAGMA journal_mode=WAL')
+        except sqlite3.OperationalError as e:
+            if 'database is locked' not in str(e):
+                raise
+            # If the database is locked, it is OK to continue, as the WAL mode
+            # is not critical and is likely to be enabled by other processes.
+
+    # Table for Requests
+    cursor.execute(f"""\
+        CREATE TABLE IF NOT EXISTS {REQUEST_TABLE} (
+        request_id TEXT PRIMARY KEY,
+        name TEXT,
+        entrypoint TEXT,
+        request_body TEXT,
+        status TEXT,
+        created_at REAL,
+        return_value TEXT,
+        error BLOB,
+        pid INTEGER,
+        {COL_CLUSTER_NAME} TEXT,
+        schedule_type TEXT,
+        {COL_USER_ID} TEXT)""")
+
+
+_DB = None
+
+
+def init_db(func):
+    """Initialize the database."""
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        global _DB
+        if _DB is None:
+            _DB = db_utils.SQLiteConn(_DB_PATH, create_table)
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
+def reset_db_and_logs():
+    """Create the database."""
+    common_utils.remove_file_if_exists(_DB_PATH)
+    shutil.rmtree(pathlib.Path(REQUEST_LOG_PATH_PREFIX).expanduser(),
+                  ignore_errors=True)
+    shutil.rmtree(server_common.API_SERVER_CLIENT_DIR.expanduser(),
+                  ignore_errors=True)
+
+
+def request_lock_path(request_id: str) -> str:
+    lock_path = os.path.expanduser(REQUEST_LOG_PATH_PREFIX)
+    os.makedirs(lock_path, exist_ok=True)
+    return os.path.join(lock_path, f'.{request_id}.lock')
+
+
+@contextlib.contextmanager
+@init_db
+def update_request(request_id: str):
+    """Get 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)
+
+
+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 + '%',))
+        row = 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."""
+    assert _DB is not None
+    with _DB.conn:
+        cursor = _DB.conn.cursor()
+        cursor.execute(f'SELECT request_id FROM {REQUEST_TABLE} '
+                       'ORDER BY created_at DESC LIMIT 1')
+        row = cursor.fetchone()
+        return row[0] if row else None
+
+
+@init_db
+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
+def create_if_not_exists(request: Request) -> bool:
+    """Create a SkyPilot API request if it does not exist."""
+    with filelock.FileLock(request_lock_path(request.request_id)):
+        if _get_request_no_lock(request.request_id) is not None:
+            return False
+        _add_or_update_request_no_lock(request)
+        return True
+
+
+@init_db
+def get_request_tasks(
+    status: Optional[List[RequestStatus]] = None,
+    cluster_names: Optional[List[str]] = None,
+    exclude_request_names: Optional[List[str]] = None,
+    user_id: Optional[str] = None,
+) -> List[Request]:
+    """Get a list of requests that match the given filters.
+
+    Args:
+        status: a list of statuses of the requests to filter on.
+        cluster_names: a list of cluster names to filter requests on.
+        exclude_request_names: a list of request names to exclude from results.
+        user_id: the user ID to filter requests on.
+            If None, all users are included.
+    """
+    filters = []
+    filter_params = []
+    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)
+    assert _DB is not None
+    with _DB.conn:
+        cursor = _DB.conn.cursor()
+        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)
+        rows = cursor.fetchall()
+        if rows is None:
+            return []
+    requests = []
+    for row in rows:
+        request = Request.from_row(row)
+        requests.append(request)
+    return requests
+
+
+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)