skypilot-nightly 1.0.0.dev20250905__py3-none-any.whl → 1.0.0.dev20251203__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