skypilot-nightly 1.0.0.dev20250804__py3-none-any.whl → 1.0.0.dev20250806__py3-none-any.whl

sky/serve/service.py

@@ -112,6 +112,10 @@ def cleanup_storage(task_yaml: str) -> bool:
 
 def _cleanup(service_name: str) -> bool:
     """Clean up all service related resources, i.e. replicas and storage."""
+    # Cleanup the HA recovery script first as it is possible that some error
+    # was raised when we construct the task object (e.g.,
+    # sky.exceptions.ResourcesUnavailableError).
+    serve_state.remove_ha_recovery_script(service_name)
     failed = False
     replica_infos = serve_state.get_replica_infos(service_name)
     info2proc: Dict[replica_managers.ReplicaInfo,
@@ -223,7 +227,8 @@ def _start(service_name: str, tmp_task_yaml: str, job_id: int):
             load_balancing_policy=service_spec.load_balancing_policy,
             status=serve_state.ServiceStatus.CONTROLLER_INIT,
             tls_encrypted=service_spec.tls_credential is not None,
-            pool=service_spec.pool)
+            pool=service_spec.pool,
+            controller_pid=os.getpid())
         # Directly throw an error here. See sky/serve/api.py::up
         # for more details.
         if not success:
@@ -241,6 +246,8 @@ def _start(service_name: str, tmp_task_yaml: str, job_id: int):
         # sync to a tmp file first and then copy it to the final name
         # if there is no name conflict.
         shutil.copy(tmp_task_yaml, service_task_yaml)
+    else:
+        serve_state.update_service_controller_pid(service_name, os.getpid())
 
     controller_process = None
     load_balancer_process = None

sky/server/auth/authn.py

@@ -0,0 +1,46 @@
+"""Authentication module."""
+import json
+from typing import Optional
+
+import fastapi
+
+from sky import models
+from sky import sky_logging
+from sky.skylet import constants
+
+logger = sky_logging.init_logger(__name__)
+
+
+# TODO(hailong): Remove this function and use request.state.auth_user instead.
+async def override_user_info_in_request_body(request: fastapi.Request,
+                                             auth_user: Optional[models.User]):
+    if auth_user is None:
+        return
+
+    body = await request.body()
+    if body:
+        try:
+            original_json = await request.json()
+        except (json.JSONDecodeError, UnicodeDecodeError) as e:
+            logger.error(f'Error parsing request JSON: {e}')
+        else:
+            logger.debug(f'Overriding user for {request.state.request_id}: '
+                         f'{auth_user.name}, {auth_user.id}')
+            if 'env_vars' in original_json:
+                if isinstance(original_json.get('env_vars'), dict):
+                    original_json['env_vars'][
+                        constants.USER_ID_ENV_VAR] = auth_user.id
+                    original_json['env_vars'][
+                        constants.USER_ENV_VAR] = auth_user.name
+                else:
+                    logger.warning(
+                        f'"env_vars" in request body is not a dictionary '
+                        f'for request {request.state.request_id}. '
+                        'Skipping user info injection into body.')
+            else:
+                original_json['env_vars'] = {}
+                original_json['env_vars'][
+                    constants.USER_ID_ENV_VAR] = auth_user.id
+                original_json['env_vars'][
+                    constants.USER_ENV_VAR] = auth_user.name
+            request._body = json.dumps(original_json).encode('utf-8')  # pylint: disable=protected-access

sky/server/auth/oauth2_proxy.py

@@ -0,0 +1,185 @@
+"""Authentication based on oauth2-proxy."""
+
+import asyncio
+import hashlib
+import http
+import os
+from typing import Optional
+import urllib
+
+import aiohttp
+import fastapi
+import starlette.middleware.base
+
+from sky import models
+from sky import sky_logging
+from sky.server.auth import authn
+from sky.utils import common_utils
+
+logger = sky_logging.init_logger(__name__)
+
+# We do not support setting these in config.yaml because:
+# 1. config.yaml can be updated dynamically, but auth middleware does not
+#    support hot reload yet.
+# 2. If we introduce hot reload for auth middleware, bad config might
+#    invalidate all authenticated sessions and thus cannot be rolled back
+#    by API users.
+# TODO(aylei): we should introduce server.yaml for static server admin config,
+# which is more structured than multiple environment variables and can be less
+# confusing to users.
+OAUTH2_PROXY_BASE_URL_ENV_VAR = 'SKYPILOT_AUTH_OAUTH2_PROXY_BASE_URL'
+OAUTH2_PROXY_ENABLED_ENV_VAR = 'SKYPILOT_AUTH_OAUTH2_PROXY_ENABLED'
+
+
+class OAuth2ProxyMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
+    """Middleware to handle authentication by delegating to OAuth2 Proxy."""
+
+    def __init__(self, application: fastapi.FastAPI):
+        super().__init__(application)
+        self.enabled: bool = (os.getenv(OAUTH2_PROXY_ENABLED_ENV_VAR,
+                                        'false') == 'true')
+        self.proxy_base: str = ''
+        if self.enabled:
+            proxy_base = os.getenv(OAUTH2_PROXY_BASE_URL_ENV_VAR)
+            if not proxy_base:
+                raise ValueError('OAuth2 Proxy is enabled but base_url is not '
+                                 'set')
+            self.proxy_base = proxy_base.rstrip('/')
+
+    async def dispatch(self, request: fastapi.Request, call_next):
+        if not self.enabled:
+            return await call_next(request)
+
+        # Forward /oauth2/* to oauth2-proxy, including /oauth2/start and
+        # /oauth2/callback.
+        if request.url.path.startswith('/oauth2'):
+            return await self.forward_to_oauth2_proxy(request)
+
+        return await self.authenticate(request, call_next)
+
+    async def forward_to_oauth2_proxy(self, request: fastapi.Request):
+        """Forward requests to oauth2-proxy service."""
+        logger.debug(f'forwarding to oauth2-proxy: {request.url.path}')
+        path = request.url.path.lstrip('/')
+        target_url = f'{self.proxy_base}/{path}'
+        body = await request.body()
+        async with aiohttp.ClientSession() as session:
+            try:
+                forwarded_headers = dict(request.headers)
+                async with session.request(
+                        method=request.method,
+                        url=target_url,
+                        headers=forwarded_headers,
+                        data=body,
+                        cookies=request.cookies,
+                        params=request.query_params,
+                        allow_redirects=False,
+                ) as response:
+                    response_body = await response.read()
+                    fastapi_response = fastapi.responses.Response(
+                        content=response_body,
+                        status_code=response.status,
+                        headers=dict(response.headers),
+                    )
+                    # Forward cookies from OAuth2 proxy response to client
+                    for cookie_name, cookie in response.cookies.items():
+                        fastapi_response.set_cookie(
+                            key=cookie_name,
+                            value=cookie.value,
+                            max_age=cookie.get('max-age'),
+                            expires=cookie.get('expires'),
+                            path=cookie.get('path', '/'),
+                            domain=cookie.get('domain'),
+                            secure=cookie.get('secure', False),
+                            httponly=cookie.get('httponly', False),
+                        )
+                    return fastapi_response
+            except (aiohttp.ClientError, asyncio.TimeoutError) as e:
+                logger.error(f'Error forwarding to OAuth2 proxy: {e}')
+                return fastapi.responses.JSONResponse(
+                    status_code=http.HTTPStatus.BAD_GATEWAY,
+                    content={'detail': 'oauth2-proxy service unavailable'})
+
+    async def authenticate(self, request: fastapi.Request, call_next):
+        if request.state.auth_user is not None:
+            # Already authenticated
+            return await call_next(request)
+
+        async with aiohttp.ClientSession() as session:
+            try:
+                return await self._authenticate(request, call_next, session)
+            except (aiohttp.ClientError, asyncio.TimeoutError) as e:
+                logger.error(f'Error communicating with OAuth2 proxy: {e}')
+                # Fail open or closed based on your security requirements
+                return fastapi.responses.JSONResponse(
+                    status_code=http.HTTPStatus.BAD_GATEWAY,
+                    content={'detail': 'oauth2-proxy service unavailable'})
+
+    async def _authenticate(self, request: fastapi.Request, call_next,
+                            session: aiohttp.ClientSession):
+        forwarded_headers = dict(request.headers)
+        auth_url = f'{self.proxy_base}/oauth2/auth'
+        forwarded_headers['X-Forwarded-Uri'] = str(request.url).rstrip('/')
+        logger.debug(f'authenticate request: {request.url.path}')
+
+        async with session.request(
+                method=request.method,
+                url=auth_url,
+                headers=forwarded_headers,
+                cookies=request.cookies,
+                timeout=aiohttp.ClientTimeout(total=10),
+                allow_redirects=False,
+        ) as auth_response:
+
+            if auth_response.status == http.HTTPStatus.ACCEPTED:
+                # User is authenticated, extract user info from headers
+                auth_user = self.get_auth_user(auth_response)
+                if not auth_user:
+                    return fastapi.responses.JSONResponse(
+                        status_code=http.HTTPStatus.INTERNAL_SERVER_ERROR,
+                        content={
+                            'detail':
+                                'oauth2-proxy is enabled but did not'
+                                'return user info, check your oauth2-proxy'
+                                'setup.'
+                        })
+                request.state.auth_user = auth_user
+                await authn.override_user_info_in_request_body(
+                    request, auth_user)
+                return await call_next(request)
+            elif auth_response.status == http.HTTPStatus.UNAUTHORIZED:
+                # For /api/health, we should allow unauthenticated requests to
+                # not break healthz check.
+                # TODO(aylei): remove this to an aggregated login middleware
+                # in favor of the unified authentication.
+                if request.url.path.startswith('/api/health'):
+                    request.state.anonymous_user = True
+                    return await call_next(request)
+
+                # TODO(aylei): in unified authentication, the redirection
+                # or rejection should be done after all the authentication
+                # methods are performed.
+                # Not authenticated, redirect to sign-in
+                redirect_path = request.url.path
+                if request.url.query:
+                    redirect_path += f'?{request.url.query}'
+                rd = urllib.parse.quote(redirect_path)
+                signin_url = (f'{request.base_url}oauth2/start?'
+                              f'rd={rd}')
+                return fastapi.responses.RedirectResponse(url=signin_url)
+            else:
+                logger.error('oauth2-proxy returned unexpected status '
+                             f'{auth_response.status}: {auth_response.text}')
+                return fastapi.responses.JSONResponse(
+                    status_code=auth_response.status,
+                    content={'detail': 'oauth2-proxy error'})
+
+    def get_auth_user(
+            self, response: aiohttp.ClientResponse) -> Optional[models.User]:
+        """Extract user info from OAuth2 proxy response headers."""
+        email_header = response.headers.get('X-Auth-Request-Email')
+        if email_header:
+            user_hash = hashlib.md5(email_header.encode()).hexdigest(
+            )[:common_utils.USER_HASH_LENGTH]
+            return models.User(id=user_hash, name=email_header)
+        return None

sky/server/common.py

@@ -41,12 +41,14 @@ from sky.utils import rich_utils
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
+    import aiohttp
     import pydantic
     import requests
 
     from sky import dag as dag_lib
     from sky import models
 else:
+    aiohttp = adaptors_common.LazyImport('aiohttp')
     pydantic = adaptors_common.LazyImport('pydantic')
     requests = adaptors_common.LazyImport('requests')
 
@@ -175,24 +177,14 @@ def get_cookies_from_response(
     return cookies
 
 
-def make_authenticated_request(method: str,
-                               path: str,
-                               server_url: Optional[str] = None,
-                               retry: bool = True,
-                               **kwargs) -> 'requests.Response':
-    """Make an authenticated HTTP request to the API server.
-
-    Automatically handles service account token authentication or cookie-based
-    authentication based on what's available.
-
-    Args:
-        method: HTTP method (GET, POST, etc.)
-        path: API path (e.g., '/api/v1/status')
-        server_url: Server URL, defaults to configured server
-        **kwargs: Additional arguments to pass to requests
+def _prepare_authenticated_request_params(
+        path: str,
+        server_url: Optional[str] = None,
+        **kwargs) -> Tuple[str, Dict[str, Any]]:
+    """Prepare common parameters for authenticated requests (sync or async).
 
     Returns:
-        requests.Response object
+        Tuple of (url, updated_kwargs)
     """
     if server_url is None:
         server_url = get_server_url()
@@ -214,6 +206,41 @@ def make_authenticated_request(method: str,
     if not headers.get('Authorization') and 'cookies' not in kwargs:
         kwargs['cookies'] = get_api_cookie_jar()
 
+    return url, kwargs
+
+
+def _convert_requests_cookies_to_aiohttp(
+        cookie_jar: requests.cookies.RequestsCookieJar) -> Dict[str, str]:
+    """Convert requests cookie jar to aiohttp-compatible dict format."""
+    cookies = {}
+    for cookie in cookie_jar:
+        cookies[cookie.name] = cookie.value
+    return cookies  # type: ignore
+
+
+def make_authenticated_request(method: str,
+                               path: str,
+                               server_url: Optional[str] = None,
+                               retry: bool = True,
+                               **kwargs) -> 'requests.Response':
+    """Make an authenticated HTTP request to the API server.
+
+    Automatically handles service account token authentication or cookie-based
+    authentication based on what's available.
+
+    Args:
+        method: HTTP method (GET, POST, etc.)
+        path: API path (e.g., '/api/v1/status')
+        server_url: Server URL, defaults to configured server
+        retry: Whether to retry on transient errors
+        **kwargs: Additional arguments to pass to requests
+
+    Returns:
+        requests.Response object
+    """
+    url, kwargs = _prepare_authenticated_request_params(path, server_url,
+                                                        **kwargs)
+
     # Make the request
     if retry:
         return rest.request(method, url, **kwargs)
@@ -222,6 +249,69 @@ def make_authenticated_request(method: str,
         return rest.request_without_retry(method, url, **kwargs)
 
 
+async def make_authenticated_request_async(
+        session: 'aiohttp.ClientSession',
+        method: str,
+        path: str,
+        server_url: Optional[str] = None,
+        retry: bool = True,
+        **kwargs) -> 'aiohttp.ClientResponse':
+    """Make an authenticated async HTTP request to the API server using aiohttp.
+
+    Automatically handles service account token authentication or cookie-based
+    authentication based on what's available.
+
+    Example usage:
+        async with aiohttp.ClientSession() as session:
+            response = await make_authenticated_request_async(
+                session, 'GET', '/api/v1/status')
+            data = await response.json()
+
+    Args:
+        session: aiohttp ClientSession to use for the request
+        method: HTTP method (GET, POST, etc.)
+        path: API path (e.g., '/api/v1/status')
+        server_url: Server URL, defaults to configured server
+        retry: Whether to retry on transient errors
+        **kwargs: Additional arguments to pass to aiohttp
+
+    Returns:
+        aiohttp.ClientResponse object
+
+    Raises:
+        aiohttp.ClientError: For HTTP-related errors
+        exceptions.ServerTemporarilyUnavailableError: When server returns 503
+        exceptions.RequestInterruptedError: When request is interrupted
+    """
+    url, kwargs = _prepare_authenticated_request_params(path, server_url,
+                                                        **kwargs)
+
+    # Convert cookies to aiohttp format if needed
+    if 'cookies' in kwargs and isinstance(kwargs['cookies'],
+                                          requests.cookies.RequestsCookieJar):
+        kwargs['cookies'] = _convert_requests_cookies_to_aiohttp(
+            kwargs['cookies'])
+
+    # Convert params to strings for aiohttp compatibility
+    if 'params' in kwargs and kwargs['params'] is not None:
+        normalized_params = {}
+        for key, value in kwargs['params'].items():
+            if isinstance(value, bool):
+                normalized_params[key] = str(value).lower()
+            elif value is not None:
+                normalized_params[key] = str(value)
+            # Skip None values
+        kwargs['params'] = normalized_params
+
+    # Make the request
+    if retry:
+        return await rest.request_async(session, method, url, **kwargs)
+    else:
+        assert method == 'GET', 'Only GET requests can be done without retry'
+        return await rest.request_without_retry_async(session, method, url,
+                                                      **kwargs)
+
+
 @annotations.lru_cache(scope='global')
 def get_server_url(host: Optional[str] = None) -> str:
     endpoint = DEFAULT_SERVER_URL
@@ -322,13 +412,14 @@ def get_api_server_status(endpoint: Optional[str] = None) -> ApiServerInfo:
         # The response is 200, so we can parse the response.
         try:
             result = response.json()
+            server_status = result.get('status')
             api_version = result.get('api_version')
             version = result.get('version')
             version_on_disk = result.get('version_on_disk')
             commit = result.get('commit')
             user = result.get('user')
             basic_auth_enabled = result.get('basic_auth_enabled')
-            server_info = ApiServerInfo(status=ApiServerStatus.HEALTHY,
+            server_info = ApiServerInfo(status=ApiServerStatus(server_status),
                                         api_version=api_version,
                                         version=version,
                                         version_on_disk=version_on_disk,

sky/server/constants.py

@@ -10,7 +10,7 @@ from sky.skylet import constants
 # based on version info is needed.
 # For more details and code guidelines, refer to:
 # https://docs.skypilot.co/en/latest/developers/CONTRIBUTING.html#backward-compatibility-guidelines
-API_VERSION = 13
+API_VERSION = 14
 
 # The minimum peer API version that the code should still work with.
 # Notes (dev):

sky/server/daemons.py

@@ -14,6 +14,10 @@ from sky.utils import ux_utils
 logger = sky_logging.init_logger(__name__)
 
 
+def _default_should_skip():
+    return False
+
+
 @dataclasses.dataclass
 class InternalRequestDaemon:
     """Internal daemon that runs an event in the background."""
@@ -22,6 +26,7 @@ class InternalRequestDaemon:
     name: str
     event_fn: Callable[[], None]
     default_log_level: str = 'INFO'
+    should_skip: Callable[[], bool] = _default_should_skip
 
     def refresh_log_level(self) -> int:
         # pylint: disable=import-outside-toplevel
@@ -110,14 +115,14 @@ def managed_job_status_refresh_event():
     """Refresh the managed job status for controller consolidation mode."""
     # pylint: disable=import-outside-toplevel
     from sky.jobs import utils as managed_job_utils
-    if not managed_job_utils.is_consolidation_mode():
-        return
+    from sky.utils import controller_utils
+
     # We run the recovery logic before starting the event loop as those two are
     # conflicting. Check PERSISTENT_RUN_RESTARTING_SIGNAL_FILE for details.
-    from sky.utils import controller_utils
     if controller_utils.high_availability_specified(
             controller_utils.Controllers.JOBS_CONTROLLER.value.cluster_name):
         managed_job_utils.ha_recovery_for_consolidation_mode()
+
     # After recovery, we start the event loop.
     from sky.skylet import events
     refresh_event = events.ManagedJobEvent()
@@ -128,20 +133,58 @@ def managed_job_status_refresh_event():
     time.sleep(events.EVENT_CHECKING_INTERVAL_SECONDS)
 
 
-def sky_serve_status_refresh_event():
+def should_skip_managed_job_status_refresh():
+    """Check if the managed job status refresh event should be skipped."""
+    # pylint: disable=import-outside-toplevel
+    from sky.jobs import utils as managed_job_utils
+    return not managed_job_utils.is_consolidation_mode()
+
+
+def _serve_status_refresh_event(pool: bool):
     """Refresh the sky serve status for controller consolidation mode."""
     # pylint: disable=import-outside-toplevel
     from sky.serve import serve_utils
-    if not serve_utils.is_consolidation_mode():
-        return
-    # TODO(tian): Add HA recovery logic.
+    from sky.utils import controller_utils
+
+    # We run the recovery logic before starting the event loop as those two are
+    # conflicting. Check PERSISTENT_RUN_RESTARTING_SIGNAL_FILE for details.
+    controller = controller_utils.get_controller_for_pool(pool)
+    if controller_utils.high_availability_specified(
+            controller.value.cluster_name):
+        serve_utils.ha_recovery_for_consolidation_mode(pool=pool)
+
+    # After recovery, we start the event loop.
     from sky.skylet import events
-    event = events.ServiceUpdateEvent()
-    logger.info('=== Running serve status refresh event ===')
+    event = events.ServiceUpdateEvent(pool=pool)
+    noun = 'pool' if pool else 'serve'
+    logger.info(f'=== Running {noun} status refresh event ===')
     event.run()
     time.sleep(events.EVENT_CHECKING_INTERVAL_SECONDS)
 
 
+def _should_skip_serve_status_refresh_event(pool: bool):
+    """Check if the serve status refresh event should be skipped."""
+    # pylint: disable=import-outside-toplevel
+    from sky.serve import serve_utils
+    return not serve_utils.is_consolidation_mode(pool=pool)
+
+
+def sky_serve_status_refresh_event():
+    _serve_status_refresh_event(pool=False)
+
+
+def should_skip_sky_serve_status_refresh():
+    return _should_skip_serve_status_refresh_event(pool=False)
+
+
+def pool_status_refresh_event():
+    _serve_status_refresh_event(pool=True)
+
+
+def should_skip_pool_status_refresh():
+    return _should_skip_serve_status_refresh_event(pool=True)
+
+
 # Register the events to run in the background.
 INTERNAL_REQUEST_DAEMONS = [
     # This status refresh daemon can cause the autostopp'ed/autodown'ed cluster
@@ -157,8 +200,14 @@ INTERNAL_REQUEST_DAEMONS = [
                           event_fn=refresh_volume_status_event),
     InternalRequestDaemon(id='managed-job-status-refresh-daemon',
                           name='managed-job-status',
-                          event_fn=managed_job_status_refresh_event),
+                          event_fn=managed_job_status_refresh_event,
+                          should_skip=should_skip_managed_job_status_refresh),
     InternalRequestDaemon(id='sky-serve-status-refresh-daemon',
                           name='sky-serve-status',
-                          event_fn=sky_serve_status_refresh_event),
+                          event_fn=sky_serve_status_refresh_event,
+                          should_skip=should_skip_sky_serve_status_refresh),
+    InternalRequestDaemon(id='pool-status-refresh-daemon',
+                          name='pool-status',
+                          event_fn=pool_status_refresh_event,
+                          should_skip=should_skip_pool_status_refresh),
 ]