skypilot-nightly 1.0.0.dev20250905__py3-none-any.whl → 1.0.0.dev20251210__py3-none-any.whl

sky/server/metrics.py

@@ -1,74 +1,33 @@
 """Instrumentation for the API server."""
 
-import contextlib
-import functools
+import asyncio
+import multiprocessing
 import os
+import threading
 import time
+from typing import List
 
 import fastapi
 from prometheus_client import generate_latest
 from prometheus_client import multiprocess
 import prometheus_client as prom
+import psutil
 import starlette.middleware.base
 import uvicorn
 
+from sky import core
 from sky import sky_logging
-from sky.skylet import constants
-
-# Whether the metrics are enabled, cannot be changed at runtime.
-METRICS_ENABLED = os.environ.get(constants.ENV_VAR_SERVER_METRICS_ENABLED,
-                                 'false').lower() == 'true'
+from sky.metrics import utils as metrics_utils
 
 logger = sky_logging.init_logger(__name__)
 
-# Total number of API server requests, grouped by path, method, and status.
-SKY_APISERVER_REQUESTS_TOTAL = prom.Counter(
-    'sky_apiserver_requests_total',
-    'Total number of API server requests',
-    ['path', 'method', 'status'],
-)
-
-# Time spent processing API server requests, grouped by path, method, and
-# status.
-SKY_APISERVER_REQUEST_DURATION_SECONDS = prom.Histogram(
-    'sky_apiserver_request_duration_seconds',
-    'Time spent processing API server requests',
-    ['path', 'method', 'status'],
-    buckets=(0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 20.0, 30.0,
-             60.0, 120.0, float('inf')),
-)
-
-# Time spent processing requests in executor.
-SKY_APISERVER_REQUEST_EXECUTION_DURATION_SECONDS = prom.Histogram(
-    'sky_apiserver_request_execution_duration_seconds',
-    'Time spent executing requests in executor',
-    ['request', 'worker'],
-    buckets=(0.5, 1, 2.5, 5.0, 10.0, 15.0, 25.0, 40.0, 60.0, 90.0, 120.0, 180.0,
-             float('inf')),
-)
-
-# Time spent processing a piece of code, refer to time_it().
-SKY_APISERVER_CODE_DURATION_SECONDS = prom.Histogram(
-    'sky_apiserver_code_duration_seconds',
-    'Time spent processing code',
-    ['name', 'group'],
-    buckets=(0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 20.0, 30.0,
-             60.0, 120.0, float('inf')),
-)
-
-SKY_APISERVER_EVENT_LOOP_LAG_SECONDS = prom.Histogram(
-    'sky_apiserver_event_loop_lag_seconds',
-    'Scheduling delay of the server event loop',
-    ['pid'],
-    buckets=(0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1, 2, 5, 20.0,
-             60.0, float('inf')),
-)
-
 metrics_app = fastapi.FastAPI()
 
 
+# Serve /metrics in dedicated thread to avoid blocking the event loop
+# of metrics server.
 @metrics_app.get('/metrics')
-async def metrics() -> fastapi.Response:
+def metrics() -> fastapi.Response:
     """Expose aggregated Prometheus metrics from all worker processes."""
     if os.environ.get('PROMETHEUS_MULTIPROC_DIR'):
         # In multiprocess mode, we need to collect metrics from all processes.
@@ -82,6 +41,42 @@ async def metrics() -> fastapi.Response:
                             headers={'Cache-Control': 'no-cache'})
 
 
+@metrics_app.get('/gpu-metrics')
+async def gpu_metrics() -> fastapi.Response:
+    """Gets the GPU metrics from multiple external k8s clusters"""
+    contexts = core.get_all_contexts()
+    all_metrics: List[str] = []
+    successful_contexts = 0
+
+    tasks = [
+        asyncio.create_task(metrics_utils.get_metrics_for_context(context))
+        for context in contexts
+        if context != 'in-cluster'
+    ]
+
+    results = await asyncio.gather(*tasks, return_exceptions=True)
+
+    for i, result in enumerate(results):
+        if isinstance(result, Exception):
+            logger.error(
+                f'Failed to get metrics for context {contexts[i]}: {result}')
+        elif isinstance(result, BaseException):
+            # Avoid changing behavior for non-Exception BaseExceptions
+            # like KeyboardInterrupt/SystemExit: re-raise them.
+            raise result
+        else:
+            metrics_text = result
+            all_metrics.append(metrics_text)
+            successful_contexts += 1
+
+    combined_metrics = '\n\n'.join(all_metrics)
+
+    # Return as plain text for Prometheus compatibility
+    return fastapi.Response(
+        content=combined_metrics,
+        media_type='text/plain; version=0.0.4; charset=utf-8')
+
+
 def build_metrics_server(host: str, port: int) -> uvicorn.Server:
     metrics_config = uvicorn.Config(
         'sky.server.metrics:metrics_app',
@@ -125,56 +120,41 @@ class PrometheusMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
             status_code_group = '5xx'
             raise
         finally:
-            SKY_APISERVER_REQUESTS_TOTAL.labels(path=path,
-                                                method=method,
-                                                status=status_code_group).inc()
+            metrics_utils.SKY_APISERVER_REQUESTS_TOTAL.labels(
+                path=path, method=method, status=status_code_group).inc()
             if not streaming:
                 duration = time.time() - start_time
-                SKY_APISERVER_REQUEST_DURATION_SECONDS.labels(
+                metrics_utils.SKY_APISERVER_REQUEST_DURATION_SECONDS.labels(
                     path=path, method=method,
                     status=status_code_group).observe(duration)
 
         return response
 
 
-@contextlib.contextmanager
-def time_it(name: str, group: str = 'default'):
-    """Context manager to measure and record code execution duration."""
-    if not METRICS_ENABLED:
-        yield
-    else:
-        start_time = time.time()
-        try:
-            yield
-        finally:
-            duration = time.time() - start_time
-            SKY_APISERVER_CODE_DURATION_SECONDS.labels(
-                name=name, group=group).observe(duration)
-
-
-def time_me(func):
-    """Measure the duration of decorated function."""
-
-    @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        if not METRICS_ENABLED:
-            return func(*args, **kwargs)
-        name = f'{func.__module__}/{func.__name__}'
-        with time_it(name, group='function'):
-            return func(*args, **kwargs)
-
-    return wrapper
-
-
-def time_me_async(func):
-    """Measure the duration of decorated async function."""
-
-    @functools.wraps(func)
-    async def async_wrapper(*args, **kwargs):
-        if not METRICS_ENABLED:
-            return await func(*args, **kwargs)
-        name = f'{func.__module__}/{func.__name__}'
-        with time_it(name, group='function'):
-            return await func(*args, **kwargs)
-
-    return async_wrapper
+peak_rss_bytes = 0
+
+
+def process_monitor(process_type: str, stop: threading.Event):
+    pid = multiprocessing.current_process().pid
+    proc = psutil.Process(pid)
+    last_bucket_end = time.time()
+    bucket_peak = 0
+    global peak_rss_bytes
+    while not stop.is_set():
+        if time.time() - last_bucket_end >= 30:
+            # Reset peak RSS for the next time bucket.
+            last_bucket_end = time.time()
+            bucket_peak = 0
+        peak_rss_bytes = max(bucket_peak, proc.memory_info().rss)
+        metrics_utils.SKY_APISERVER_PROCESS_PEAK_RSS.labels(
+            pid=pid, type=process_type).set(peak_rss_bytes)
+        ctimes = proc.cpu_times()
+        metrics_utils.SKY_APISERVER_PROCESS_CPU_TOTAL.labels(pid=pid,
+                                                             type=process_type,
+                                                             mode='user').set(
+                                                                 ctimes.user)
+        metrics_utils.SKY_APISERVER_PROCESS_CPU_TOTAL.labels(pid=pid,
+                                                             type=process_type,
+                                                             mode='system').set(
+                                                                 ctimes.system)
+        time.sleep(1)

sky/server/middleware_utils.py

@@ -0,0 +1,166 @@
+"""Utilities for building middlewares."""
+import enum
+import http
+from typing import Type
+
+import fastapi
+import starlette.middleware.base
+import starlette.types
+
+from sky import sky_logging
+
+logger = sky_logging.init_logger(__name__)
+
+
+class WebSocketDecision(enum.Enum):
+    ACCEPT = 'accept'
+    UNAUTHORIZED = 'unauthorized'
+    FORBIDDEN = 'forbidden'
+    ERROR = 'error'
+
+
+def websocket_aware(
+        middleware_cls: Type[starlette.middleware.base.BaseHTTPMiddleware]):
+    """Decorator to adapt BaseHTTPMiddleware to handle WebSockets.
+
+    It assembles an HTTP-style request like the HTTP upgrade request during
+    websocket handshake and then delegates it to the real HTTP middleware.
+    The websocket connection will be rejected if the HTTP middleware returns
+    a 4xx or 5xx status code.
+
+    Note: for websocket connection, the mutation made by the underlying HTTP
+    middleware on the request and response will be discarded.
+    """
+
+    class WebSocketAwareMiddleware:
+        """WebSocket-aware middleware wrapper."""
+
+        def __init__(self, app: starlette.types.ASGIApp, *args, **kwargs):
+            self.app = app
+            self.middleware = middleware_cls(app, *args, **kwargs)
+
+        async def __call__(self, scope: starlette.types.Scope,
+                           receive: starlette.types.Receive,
+                           send: starlette.types.Send):
+            scope_type = scope.get('type')
+            if scope_type == 'websocket':
+                await self._handle_websocket(scope, receive, send)
+            else:
+                # Delegate other scopes to the underlying HTTP middleware.
+                await self.middleware(scope, receive, send)
+
+        async def dispatch(
+                self, request: fastapi.Request,
+                call_next: starlette.middleware.base.RequestResponseEndpoint):
+            """Implement dispatch method to keep compatibility."""
+            return await self.middleware.dispatch(request, call_next)
+
+        async def _handle_websocket(self, scope: starlette.types.Scope,
+                                    receive: starlette.types.Receive,
+                                    send: starlette.types.Send):
+            """Handle websocket connection by delegating to HTTP middleware."""
+            decision = await self._run_websocket_dispatch(scope)
+            if decision == WebSocketDecision.ACCEPT:
+                await self.app(scope, receive, send)
+            elif decision == WebSocketDecision.UNAUTHORIZED:
+                await send({
+                    'type': 'websocket.close',
+                    'code': 4401,
+                    'reason': 'Unauthorized',
+                })
+            elif decision == WebSocketDecision.FORBIDDEN:
+                await send({
+                    'type': 'websocket.close',
+                    'code': 4403,
+                    'reason': 'Forbidden',
+                })
+            else:
+                await send({
+                    'type': 'websocket.close',
+                    'code': 1011,
+                    'reason': 'Internal Server Error',
+                })
+
+        async def _run_websocket_dispatch(
+                self, scope: starlette.types.Scope) -> WebSocketDecision:
+            http_scope = self._build_http_scope(scope)
+            http_receive = self._http_receive_adapter()
+            request = fastapi.Request(http_scope, receive=http_receive)
+            call_next_called = False
+            stub_response = fastapi.Response(status_code=http.HTTPStatus.OK)
+
+            async def call_next(req):
+                del req
+                # Capture whether call_next() is called in the underlying
+                # HTTP middleware to determine if we can proceed with current
+                # websocket connection.
+                nonlocal call_next_called
+                call_next_called = True
+                return stub_response
+
+            try:
+                response = await self.dispatch(request, call_next)
+            except Exception as e:  # pylint: disable=broad-except
+                logger.error('Exception occurred in middleware dispatch for '
+                             f'WebSocket scope: {e}')
+                return WebSocketDecision.ERROR
+
+            if response is None:
+                response = stub_response
+
+            status_code = response.status_code
+
+            if call_next_called and 200 <= status_code < 400:
+                return WebSocketDecision.ACCEPT
+            if status_code == http.HTTPStatus.UNAUTHORIZED:
+                return WebSocketDecision.UNAUTHORIZED
+            if status_code == http.HTTPStatus.FORBIDDEN:
+                return WebSocketDecision.FORBIDDEN
+            return WebSocketDecision.ERROR
+
+        @staticmethod
+        def _build_http_scope(
+                scope: starlette.types.Scope) -> starlette.types.Scope:
+            state = scope.setdefault('state', {})
+            scheme = scope.get('scheme', 'ws')
+            if scheme == 'ws':
+                http_scheme = 'http'
+            elif scheme == 'wss':
+                http_scheme = 'https'
+            else:
+                http_scheme = scheme
+            http_scope = dict(scope)
+            http_scope['type'] = 'http'
+            http_scope['scheme'] = http_scheme
+            http_scope['method'] = 'GET'
+            http_scope['http_version'] = scope.get('http_version', '1.1')
+            http_scope['state'] = state
+            return http_scope
+
+        @staticmethod
+        def _http_receive_adapter() -> starlette.types.Receive:
+            """Adapter thatmimics the sequence produced by Starlette for an HTTP
+            request: a single http.request event followed by a http.disconnect
+            """
+            sent = False
+
+            async def receive():
+                nonlocal sent
+                if not sent:
+                    sent = True
+                    return {
+                        'type': 'http.request',
+                        'body': b'',
+                        'more_body': False,
+                    }
+                return {
+                    'type': 'http.disconnect',
+                }
+
+            return receive
+
+    WebSocketAwareMiddleware.__name__ = middleware_cls.__name__
+    WebSocketAwareMiddleware.__qualname__ = middleware_cls.__qualname__
+    WebSocketAwareMiddleware.__module__ = middleware_cls.__module__
+    WebSocketAwareMiddleware.__doc__ = middleware_cls.__doc__
+    return WebSocketAwareMiddleware

sky/server/plugins.py

@@ -0,0 +1,222 @@
+"""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 js_extension_path(self) -> Optional[str]:
+        """Optional API route to the JavaScript extension to load."""
+        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']
+        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