skypilot-nightly 1.0.0.dev20251203__py3-none-any.whl → 1.0.0.dev20260112__py3-none-any.whl

sky/server/plugins.py

@@ -0,0 +1,238 @@
+"""Load plugins for the SkyPilot API server."""
+import abc
+import dataclasses
+import importlib
+import os
+from typing import Dict, List, Optional, Tuple
+
+from fastapi import FastAPI
+
+from sky import sky_logging
+from sky.skylet import constants as skylet_constants
+from sky.utils import common_utils
+from sky.utils import config_utils
+from sky.utils import yaml_utils
+
+logger = sky_logging.init_logger(__name__)
+
+_DEFAULT_PLUGINS_CONFIG_PATH = '~/.sky/plugins.yaml'
+_PLUGINS_CONFIG_ENV_VAR = (
+    f'{skylet_constants.SKYPILOT_SERVER_ENV_VAR_PREFIX}PLUGINS_CONFIG')
+
+
+class ExtensionContext:
+    """Context provided to plugins during installation.
+
+    Attributes:
+        app: The FastAPI application instance.
+        rbac_rules: List of RBAC rules registered by the plugin.
+        Example:
+        [
+            ('user', RBACRule(path='/plugins/api/xx/*', method='POST')),
+            ('user', RBACRule(path='/plugins/api/xx/*', method='DELETE'))
+        ]
+    """
+
+    def __init__(self, app: Optional[FastAPI] = None):
+        self.app = app
+        self.rbac_rules: List[Tuple[str, RBACRule]] = []
+
+    def register_rbac_rule(self,
+                           path: str,
+                           method: str,
+                           description: Optional[str] = None,
+                           role: str = 'user') -> None:
+        """Register an RBAC rule for this plugin.
+
+        This method allows plugins to declare which endpoints should be
+        restricted to admin users during the install phase.
+
+        Args:
+            path: The path pattern to restrict (supports wildcards with
+                  keyMatch2).
+                  Example: '/plugins/api/credentials/*'
+            method: The HTTP method to restrict. Example: 'POST', 'DELETE'
+            description: Optional description of what this rule protects.
+            role: The role to add this rule to (default: 'user').
+                  Rules added to 'user' role block regular users but allow
+                  admins.
+
+        Example:
+            def install(self, ctx: ExtensionContext):
+                # Only admin can upload credentials
+                ctx.register_rbac_rule(
+                    path='/plugins/api/credentials/*',
+                    method='POST',
+                    description='Only admin can upload credentials'
+                )
+        """
+        rule = RBACRule(path=path, method=method, description=description)
+        self.rbac_rules.append((role, rule))
+        logger.debug(f'Registered RBAC rule for {role}: {method} {path}'
+                     f'{f" - {description}" if description else ""}')
+
+
+@dataclasses.dataclass
+class RBACRule:
+    """RBAC rule for a plugin endpoint.
+
+    Attributes:
+        path: The path pattern to match (supports wildcards with keyMatch2).
+              Example: '/plugins/api/credentials/*'
+        method: The HTTP method to restrict. Example: 'POST', 'DELETE'
+        description: Optional description of what this rule protects.
+    """
+    path: str
+    method: str
+    description: Optional[str] = None
+
+
+class BasePlugin(abc.ABC):
+    """Base class for all SkyPilot server plugins."""
+
+    @property
+    def name(self) -> Optional[str]:
+        """Plugin name for display purposes."""
+        return None
+
+    @property
+    def js_extension_path(self) -> Optional[str]:
+        """Optional API route to the JavaScript extension to load."""
+        return None
+
+    @property
+    def version(self) -> Optional[str]:
+        """Plugin version."""
+        return None
+
+    @property
+    def commit(self) -> Optional[str]:
+        """Plugin git commit hash."""
+        return None
+
+    @abc.abstractmethod
+    def install(self, extension_context: ExtensionContext):
+        """Hook called by API server to let the plugin install itself."""
+        raise NotImplementedError
+
+    def shutdown(self):
+        """Hook called by API server to let the plugin shutdown."""
+        pass
+
+
+def _config_schema():
+    plugin_schema = {
+        'type': 'object',
+        'required': ['class'],
+        'additionalProperties': False,
+        'properties': {
+            'class': {
+                'type': 'string',
+            },
+            'parameters': {
+                'type': 'object',
+                'required': [],
+                'additionalProperties': True,
+            },
+        },
+    }
+    return {
+        'type': 'object',
+        'required': [],
+        'additionalProperties': False,
+        'properties': {
+            'plugins': {
+                'type': 'array',
+                'items': plugin_schema,
+                'default': [],
+            },
+        },
+    }
+
+
+def _load_plugin_config() -> Optional[config_utils.Config]:
+    """Load plugin config."""
+    config_path = os.getenv(_PLUGINS_CONFIG_ENV_VAR,
+                            _DEFAULT_PLUGINS_CONFIG_PATH)
+    config_path = os.path.expanduser(config_path)
+    if not os.path.exists(config_path):
+        return None
+    config = yaml_utils.read_yaml(config_path) or {}
+    common_utils.validate_schema(config,
+                                 _config_schema(),
+                                 err_msg_prefix='Invalid plugins config: ')
+    return config_utils.Config.from_dict(config)
+
+
+_PLUGINS: Dict[str, BasePlugin] = {}
+_EXTENSION_CONTEXT: Optional[ExtensionContext] = None
+
+
+def load_plugins(extension_context: ExtensionContext):
+    """Load and initialize plugins from the config."""
+    global _EXTENSION_CONTEXT
+    _EXTENSION_CONTEXT = extension_context
+
+    config = _load_plugin_config()
+    if not config:
+        return
+
+    for plugin_config in config.get('plugins', []):
+        class_path = plugin_config['class']
+        logger.debug(f'Loading plugins: {class_path}')
+        module_path, class_name = class_path.rsplit('.', 1)
+        try:
+            module = importlib.import_module(module_path)
+        except ImportError as e:
+            raise ImportError(
+                f'Failed to import plugin module: {module_path}. '
+                'Please check if the module is installed in your Python '
+                'environment.') from e
+        try:
+            plugin_cls = getattr(module, class_name)
+        except AttributeError as e:
+            raise AttributeError(
+                f'Could not find plugin {class_name} class in module '
+                f'{module_path}. ') from e
+        if not issubclass(plugin_cls, BasePlugin):
+            raise TypeError(
+                f'Plugin {class_path} must inherit from BasePlugin.')
+        parameters = plugin_config.get('parameters') or {}
+        plugin = plugin_cls(**parameters)
+        plugin.install(extension_context)
+        _PLUGINS[class_path] = plugin
+
+
+def get_plugins() -> List[BasePlugin]:
+    """Return shallow copies of the registered plugins."""
+    return list(_PLUGINS.values())
+
+
+def get_plugin_rbac_rules() -> Dict[str, List[Dict[str, str]]]:
+    """Collect RBAC rules from all loaded plugins.
+
+    Collects rules from the ExtensionContext.
+
+    Returns:
+        Dictionary mapping role names to lists of blocklist rules.
+        Example:
+        {
+            'user': [
+                {'path': '/plugins/api/credentials/*', 'method': 'POST'},
+                {'path': '/plugins/api/credentials/*', 'method': 'DELETE'}
+            ]
+        }
+    """
+    rules_by_role: Dict[str, List[Dict[str, str]]] = {}
+
+    # Collect rules registered via ExtensionContext
+    if _EXTENSION_CONTEXT:
+        for role, rule in _EXTENSION_CONTEXT.rbac_rules:
+            if role not in rules_by_role:
+                rules_by_role[role] = []
+            rules_by_role[role].append({
+                'path': rule.path,
+                'method': rule.method,
+            })
+
+    return rules_by_role

sky/server/requests/executor.py

@@ -44,6 +44,7 @@ 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 import plugins
 from sky.server.requests import payloads
 from sky.server.requests import preconditions
 from sky.server.requests import process
@@ -159,6 +160,8 @@ queue_backend = server_config.QueueBackend.MULTIPROCESSING
 def executor_initializer(proc_group: str):
     setproctitle.setproctitle(f'SkyPilot:executor:{proc_group}:'
                               f'{multiprocessing.current_process().pid}')
+    # Load plugins for executor process.
+    plugins.load_plugins(plugins.ExtensionContext())
     # Executor never stops, unless the whole process is killed.
     threading.Thread(target=metrics_lib.process_monitor,
                      args=(f'worker:{proc_group}', threading.Event()),
@@ -533,8 +536,8 @@ def _request_execution_wrapper(request_id: str,
         # so that the "Request xxxx failed due to ..." log message will be
         # written to the original stdout and stderr file descriptors.
         _restore_output()
-        logger.info(f'Request {request_id} failed due to '
-                    f'{common_utils.format_exception(e)}')
+        logger.error(f'Request {request_id} failed due to '
+                     f'{common_utils.format_exception(e)}')
         return
     else:
         api_requests.set_request_succeeded(

sky/server/requests/payloads.py

@@ -82,7 +82,7 @@ def request_body_env_vars() -> dict:
         if common.is_api_server_local() and env_var in EXTERNAL_LOCAL_ENV_VARS:
             env_vars[env_var] = os.environ[env_var]
     env_vars[constants.USER_ID_ENV_VAR] = common_utils.get_user_hash()
-    env_vars[constants.USER_ENV_VAR] = common_utils.get_current_user_name()
+    env_vars[constants.USER_ENV_VAR] = common_utils.get_local_user_name()
     env_vars[
         usage_constants.USAGE_RUN_ID_ENV_VAR] = usage_lib.messages.usage.run_id
     if not common.is_api_server_local():
@@ -482,6 +482,7 @@ class VolumeApplyBody(RequestBody):
 class VolumeDeleteBody(RequestBody):
     """The request body for the volume delete endpoint."""
     names: List[str]
+    purge: bool = False
 
 
 class VolumeListBody(RequestBody):
@@ -670,6 +671,11 @@ class KubernetesNodeInfoRequestBody(RequestBody):
     context: Optional[str] = None
 
 
+class SlurmNodeInfoRequestBody(RequestBody):
+    """The request body for the slurm node info endpoint."""
+    slurm_cluster_name: Optional[str] = None
+
+
 class ListAcceleratorsBody(RequestBody):
     """The request body for the list accelerators endpoint."""
     gpus_only: bool = True
@@ -854,3 +860,26 @@ class RequestPayload(BasePayload):
     status_msg: Optional[str] = None
     should_retry: bool = False
     finished_at: Optional[float] = None
+
+
+class SlurmGpuAvailabilityRequestBody(RequestBody):
+    """Request body for getting Slurm real-time GPU availability."""
+    name_filter: Optional[str] = None
+    quantity_filter: Optional[int] = None
+
+
+class ClusterEventsBody(RequestBody):
+    """The request body for the cluster events endpoint."""
+    cluster_name: Optional[str] = None
+    cluster_hash: Optional[str] = None
+    event_type: str  # 'STATUS_CHANGE' or 'DEBUG'
+    include_timestamps: bool = False
+    limit: Optional[
+        int] = None  # If specified, returns at most this many events
+
+
+class GetJobEventsBody(RequestBody):
+    """The request body for the get job task events endpoint."""
+    job_id: int
+    task_id: Optional[int] = None
+    limit: Optional[int] = 10  # Default to 10 most recent task events

sky/server/requests/request_names.py

@@ -10,6 +10,8 @@ class RequestName(str, enum.Enum):
     REALTIME_KUBERNETES_GPU_AVAILABILITY = (
         'realtime_kubernetes_gpu_availability')
     KUBERNETES_NODE_INFO = 'kubernetes_node_info'
+    REALTIME_SLURM_GPU_AVAILABILITY = 'realtime_slurm_gpu_availability'
+    SLURM_NODE_INFO = 'slurm_node_info'
     STATUS_KUBERNETES = 'status_kubernetes'
     LIST_ACCELERATORS = 'list_accelerators'
     LIST_ACCELERATOR_COUNTS = 'list_accelerator_counts'
@@ -29,6 +31,7 @@ class RequestName(str, enum.Enum):
     CLUSTER_JOB_LOGS = 'logs'
     CLUSTER_JOB_DOWNLOAD_LOGS = 'download_logs'
     CLUSTER_COST_REPORT = 'cost_report'
+    CLUSTER_EVENTS = 'cluster_events'
     # Storage requests
     STORAGE_LS = 'storage_ls'
     STORAGE_DELETE = 'storage_delete'
@@ -50,6 +53,7 @@ class RequestName(str, enum.Enum):
     JOBS_POOL_STATUS = 'jobs.pool_status'
     JOBS_POOL_LOGS = 'jobs.pool_logs'
     JOBS_POOL_SYNC_DOWN_LOGS = 'jobs.pool_sync_down_logs'
+    JOBS_EVENTS = 'jobs.events'
     # Serve requests
     SERVE_UP = 'serve.up'
     SERVE_UPDATE = 'serve.update'

sky/server/requests/requests.py

@@ -33,6 +33,7 @@ from sky.server import daemons
 from sky.server.requests import payloads
 from sky.server.requests.serializers import decoders
 from sky.server.requests.serializers import encoders
+from sky.server.requests.serializers import return_value_serializers
 from sky.utils import asyncio_utils
 from sky.utils import common_utils
 from sky.utils import ux_utils
@@ -231,13 +232,16 @@ class Request:
         assert isinstance(self.request_body,
                           payloads.RequestBody), (self.name, self.request_body)
         try:
+            # Use version-aware serializer to handle backward compatibility
+            # for old clients that don't recognize new fields.
+            serializer = return_value_serializers.get_serializer(self.name)
             return payloads.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=orjson.dumps(self.return_value).decode('utf-8'),
+                return_value=serializer(self.return_value),
                 error=orjson.dumps(self.error).decode('utf-8'),
                 pid=self.pid,
                 created_at=self.created_at,
@@ -869,11 +873,17 @@ async def create_if_not_exists_async(request: Request) -> bool:
         f'({request_columns}) VALUES '
         f'({values_str}) ON CONFLICT(request_id) DO NOTHING RETURNING ROWID')
     request_row = request.to_row()
-    # Execute the SQL statement without getting the request lock.
-    # The request lock is used to prevent racing with cancellation codepath,
-    # but a request cannot be cancelled before it is created.
-    row = await _DB.execute_get_returning_value_async(sql_statement,
-                                                      request_row)
+    if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
+        logger.debug(f'Start creating request {request.request_id}')
+    try:
+        # Execute the SQL statement without getting the request lock.
+        # The request lock is used to prevent racing with cancellation codepath,
+        # but a request cannot be cancelled before it is created.
+        row = await _DB.execute_get_returning_value_async(
+            sql_statement, request_row)
+    finally:
+        if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
+            logger.debug(f'End creating request {request.request_id}')
     return True if row else False
 
 
@@ -1030,9 +1040,15 @@ _add_or_update_request_sql = (f'INSERT OR REPLACE INTO {REQUEST_TABLE} '
 def _add_or_update_request_no_lock(request: Request):
     """Add or update a REST request into the database."""
     assert _DB is not None
-    with _DB.conn:
-        cursor = _DB.conn.cursor()
-        cursor.execute(_add_or_update_request_sql, request.to_row())
+    if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
+        logger.debug(f'Start adding or updating request {request.request_id}')
+    try:
+        with _DB.conn:
+            cursor = _DB.conn.cursor()
+            cursor.execute(_add_or_update_request_sql, request.to_row())
+    finally:
+        if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
+            logger.debug(f'End adding or updating request {request.request_id}')
 
 
 async def _add_or_update_request_no_lock_async(request: Request):
@@ -1121,8 +1137,14 @@ async def _delete_requests(request_ids: List[str]):
     """Clean up requests by their IDs."""
     id_list_str = ','.join(repr(request_id) for request_id in request_ids)
     assert _DB is not None
-    await _DB.execute_and_commit_async(
-        f'DELETE FROM {REQUEST_TABLE} WHERE request_id IN ({id_list_str})')
+    if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
+        logger.debug(f'Start deleting requests {request_ids}')
+    try:
+        await _DB.execute_and_commit_async(
+            f'DELETE FROM {REQUEST_TABLE} WHERE request_id IN ({id_list_str})')
+    finally:
+        if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
+            logger.debug(f'End deleting requests {request_ids}')
 
 
 async def clean_finished_requests_with_retention(retention_seconds: int,

sky/server/requests/serializers/encoders.py

@@ -69,6 +69,11 @@ def encode_status(
             response_cluster['last_use'] = ''
         if 'status_updated_at' not in response_cluster:
             response_cluster['status_updated_at'] = 0
+        # Ensure labels is always included, defaulting to empty dict if None
+        # This is needed because exclude_none=True would exclude None labels
+        if 'labels' not in response_cluster or response_cluster.get(
+                'labels') is None:
+            response_cluster['labels'] = {}
         response_cluster['status'] = cluster['status'].value
         handle = serialize_utils.prepare_handle_for_backwards_compatibility(
             cluster['handle'])
@@ -266,6 +271,23 @@ def encode_realtime_gpu_availability(
     return encoded
 
 
+@register_encoder('realtime_slurm_gpu_availability')
+def encode_realtime_slurm_gpu_availability(
+    return_value: List[Tuple[str,
+                             List[Any]]]) -> List[Tuple[str, List[List[Any]]]]:
+    # Convert RealtimeGpuAvailability namedtuples to lists
+    # for JSON serialization.
+    encoded = []
+    for context, gpu_list in return_value:
+        converted_gpu_list = []
+        for gpu in gpu_list:
+            assert isinstance(gpu, models.RealtimeGpuAvailability), (
+                f'Expected RealtimeGpuAvailability, got {type(gpu)}')
+            converted_gpu_list.append(list(gpu))
+        encoded.append((context, converted_gpu_list))
+    return encoded
+
+
 @register_encoder('list_accelerators')
 def encode_list_accelerators(
         return_value: Dict[str, List[Any]]) -> Dict[str, Any]:

sky/server/requests/serializers/return_value_serializers.py

@@ -0,0 +1,70 @@
+"""Version-aware serializers for request return values.
+
+These serializers run at encode() time when remote_api_version is available,
+to handle backward compatibility for old clients.
+
+The existing encoders.py handles object -> dict conversion at set_return_value()
+time. This module handles dict -> JSON string serialization at encode() time,
+with version-aware field filtering for backward compatibility.
+"""
+from typing import Any, Callable, Dict
+
+import orjson
+
+from sky.server import constants as server_constants
+from sky.server import versions
+
+handlers: Dict[str, Callable[[Any], str]] = {}
+
+
+def register_serializer(*names: str):
+    """Decorator to register a version-aware serializer."""
+
+    def decorator(func):
+        for name in names:
+            if name != server_constants.DEFAULT_HANDLER_NAME:
+                name = server_constants.REQUEST_NAME_PREFIX + name
+            if name in handlers:
+                raise ValueError(f'Serializer {name} already registered: '
+                                 f'{handlers[name]}')
+            handlers[name] = func
+        return func
+
+    return decorator
+
+
+def get_serializer(name: str) -> Callable[[Any], str]:
+    """Get the serializer for a request name."""
+    return handlers.get(name, handlers[server_constants.DEFAULT_HANDLER_NAME])
+
+
+@register_serializer(server_constants.DEFAULT_HANDLER_NAME)
+def default_serializer(return_value: Any) -> str:
+    """The default serializer."""
+    return orjson.dumps(return_value).decode('utf-8')
+
+
+@register_serializer('kubernetes_node_info')
+def serialize_kubernetes_node_info(return_value: Dict[str, Any]) -> str:
+    """Serialize kubernetes node info with version compatibility.
+
+    The is_ready field was added in API version 25. Remove it for old clients
+    that don't recognize it.
+    The cpu_count, memory_gb, cpu_free, and memory_free_gb fields were added
+    in API version 26. Remove them for old clients that don't recognize them.
+    """
+    remote_api_version = versions.get_remote_api_version()
+    if (return_value and remote_api_version is not None):
+        for node_info in return_value.get('node_info_dict', {}).values():
+            if remote_api_version < 25:
+                # Remove is_ready field for old clients that don't recognize it
+                node_info.pop('is_ready', None)
+            if remote_api_version < 26:
+                # Remove cpu_count, memory_gb, cpu_free, and
+                # memory_free_gb fields for old clients that don't
+                # recognize them
+                node_info.pop('cpu_count', None)
+                node_info.pop('memory_gb', None)
+                node_info.pop('cpu_free', None)
+                node_info.pop('memory_free_gb', None)
+    return orjson.dumps(return_value).decode('utf-8')