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

sky/serve/service_spec.py

@@ -2,11 +2,9 @@
 import json
 import os
 import textwrap
-import typing
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from sky import serve
-from sky.adaptors import common as adaptors_common
 from sky.serve import constants
 from sky.serve import load_balancing_policies as lb_policies
 from sky.serve import serve_utils
@@ -14,11 +12,7 @@ from sky.serve import spot_placer as spot_placer_lib
 from sky.utils import common_utils
 from sky.utils import schemas
 from sky.utils import ux_utils
-
-if typing.TYPE_CHECKING:
-    import yaml
-else:
-    yaml = adaptors_common.LazyImport('yaml')
+from sky.utils import yaml_utils
 
 
 class SkyServiceSpec:
@@ -33,7 +27,7 @@ class SkyServiceSpec:
         max_replicas: Optional[int] = None,
         num_overprovision: Optional[int] = None,
         ports: Optional[str] = None,
-        target_qps_per_replica: Optional[float] = None,
+        target_qps_per_replica: Optional[Union[float, Dict[str, float]]] = None,
         post_data: Optional[Dict[str, Any]] = None,
         tls_credential: Optional[serve_utils.TLSCredential] = None,
         readiness_headers: Optional[Dict[str, str]] = None,
@@ -43,7 +37,33 @@ class SkyServiceSpec:
         upscale_delay_seconds: Optional[int] = None,
         downscale_delay_seconds: Optional[int] = None,
         load_balancing_policy: Optional[str] = None,
+        pool: Optional[bool] = None,
     ) -> None:
+        if pool:
+            for unsupported_field in [
+                    'max_replicas',
+                    'num_overprovision',
+                    'target_qps_per_replica',
+                    'upscale_delay_seconds',
+                    'downscale_delay_seconds',
+                    'base_ondemand_fallback_replicas',
+                    'dynamic_ondemand_fallback',
+                    'spot_placer',
+                    'load_balancing_policy',
+                    'ports',
+                    'post_data',
+                    'tls_credential',
+                    'readiness_headers',
+            ]:
+                if locals()[unsupported_field] is not None:
+                    with ux_utils.print_exception_no_traceback():
+                        raise ValueError(
+                            f'{unsupported_field} is not supported for pool.')
+            if max_replicas is not None and max_replicas != min_replicas:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError('Autoscaling is not supported for pool '
+                                     'for now.')
+
         if max_replicas is not None and max_replicas < min_replicas:
             with ux_utils.print_exception_no_traceback():
                 raise ValueError('max_replicas must be greater than or '
@@ -83,7 +103,8 @@ class SkyServiceSpec:
         self._max_replicas: Optional[int] = max_replicas
         self._num_overprovision: Optional[int] = num_overprovision
         self._ports: Optional[str] = ports
-        self._target_qps_per_replica: Optional[float] = target_qps_per_replica
+        self._target_qps_per_replica: Optional[Union[float, Dict[
+            str, float]]] = target_qps_per_replica
         self._post_data: Optional[Dict[str, Any]] = post_data
         self._tls_credential: Optional[serve_utils.TLSCredential] = (
             tls_credential)
@@ -96,6 +117,7 @@ class SkyServiceSpec:
         self._upscale_delay_seconds: Optional[int] = upscale_delay_seconds
         self._downscale_delay_seconds: Optional[int] = downscale_delay_seconds
         self._load_balancing_policy: Optional[str] = load_balancing_policy
+        self._pool: Optional[bool] = pool
 
         self._use_ondemand_fallback: bool = (
             self.dynamic_ondemand_fallback is not None and
@@ -115,7 +137,7 @@ class SkyServiceSpec:
 
         service_config: Dict[str, Any] = {}
 
-        readiness_section = config['readiness_probe']
+        readiness_section = config.get('readiness_probe', '/')
         if isinstance(readiness_section, str):
             service_config['readiness_path'] = readiness_section
             initial_delay_seconds = None
@@ -157,8 +179,29 @@ class SkyServiceSpec:
                     raise ValueError('Port must be between 1 and 65535.')
         service_config['ports'] = str(ports) if ports is not None else None
 
+        pool_config = config.get('pool', None)
+        if pool_config is not None:
+            service_config['pool'] = pool_config
+
         policy_section = config.get('replica_policy', None)
+        if policy_section is not None and pool_config:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError('Cannot specify `replica_policy` for cluster '
+                                 'pool. Only `workers: <num>` is supported '
+                                 'for pool now.')
+
         simplified_policy_section = config.get('replicas', None)
+        workers_config = config.get('workers', None)
+        if simplified_policy_section is not None and workers_config is not None:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError('Cannot specify both `replicas` and `workers`.'
+                                 ' Please use one of them.')
+        if simplified_policy_section is not None and pool_config:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError('Cannot specify `replicas` for pool. '
+                                 'Please use `workers` instead.')
+        if simplified_policy_section is None:
+            simplified_policy_section = workers_config
         if policy_section is None or simplified_policy_section is not None:
             if simplified_policy_section is not None:
                 min_replicas = simplified_policy_section
@@ -193,6 +236,26 @@ class SkyServiceSpec:
         service_config['load_balancing_policy'] = config.get(
             'load_balancing_policy', None)
 
+        # Validate instance-aware settings
+        target_qps_per_replica = service_config['target_qps_per_replica']
+        load_balancing_policy = service_config['load_balancing_policy']
+
+        if isinstance(target_qps_per_replica, dict):
+            if load_balancing_policy != 'instance_aware_least_load':
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(
+                        'When using dict type target_qps_per_replica, '
+                        'load_balancing_policy must be '
+                        '"instance_aware_least_load".')
+
+        if load_balancing_policy == 'instance_aware_least_load':
+            if not isinstance(target_qps_per_replica, dict):
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError(
+                        'When using "instance_aware_least_load" policy, '
+                        'target_qps_per_replica must be a '
+                        'dict mapping GPU types to QPS values.')
+
         tls_section = config.get('tls', None)
         if tls_section is not None:
             service_config['tls_credential'] = serve_utils.TLSCredential(
@@ -203,14 +266,13 @@ class SkyServiceSpec:
         return SkyServiceSpec(**service_config)
 
     @staticmethod
-    def from_yaml(yaml_path: str) -> 'SkyServiceSpec':
-        with open(os.path.expanduser(yaml_path), 'r', encoding='utf-8') as f:
-            config = yaml.safe_load(f)
+    def from_yaml_str(yaml_str: str) -> 'SkyServiceSpec':
+        config = yaml_utils.safe_load(yaml_str)
 
         if isinstance(config, str):
             with ux_utils.print_exception_no_traceback():
                 raise ValueError('YAML loaded as str, not as dict. '
-                                 f'Is it correct? Path: {yaml_path}')
+                                 f'Is it correct? content:\n{yaml_str}')
 
         if config is None:
             config = {}
@@ -218,10 +280,16 @@ class SkyServiceSpec:
         if 'service' not in config:
             with ux_utils.print_exception_no_traceback():
                 raise ValueError('Service YAML must have a "service" section. '
-                                 f'Is it correct? Path: {yaml_path}')
+                                 f'Is it correct? content:\n{yaml_str}')
 
         return SkyServiceSpec.from_yaml_config(config['service'])
 
+    @staticmethod
+    def from_yaml(yaml_path: str) -> 'SkyServiceSpec':
+        with open(os.path.expanduser(yaml_path), 'r', encoding='utf-8') as f:
+            yaml_content = f.read()
+        return SkyServiceSpec.from_yaml_str(yaml_content)
+
     def to_yaml_config(self) -> Dict[str, Any]:
         config: Dict[str, Any] = {}
 
@@ -239,6 +307,13 @@ class SkyServiceSpec:
                         config[section] = dict()
                     config[section][key] = value
 
+        add_if_not_none('pool', None, self._pool)
+
+        if self.pool:
+            # For pool, currently only `workers: <num>` is supported.
+            add_if_not_none('workers', None, self.min_replicas)
+            return config
+
         add_if_not_none('readiness_probe', 'path', self.readiness_path)
         add_if_not_none('readiness_probe', 'initial_delay_seconds',
                         self.initial_delay_seconds)
@@ -306,10 +381,14 @@ class SkyServiceSpec:
         return ' '.join(policy_strs)
 
     def autoscaling_policy_str(self):
+        if self.pool:
+            # We only support fixed-size pool for now.
+            return f'Fixed-size ({self.min_replicas} workers)'
         # TODO(MaoZiming): Update policy_str
+        noun = 'worker' if self.pool else 'replica'
         min_plural = '' if self.min_replicas == 1 else 's'
         if self.max_replicas == self.min_replicas or self.max_replicas is None:
-            return f'Fixed {self.min_replicas} replica{min_plural}'
+            return f'Fixed {self.min_replicas} {noun}{min_plural}'
         # Already checked in __init__.
         assert self.target_qps_per_replica is not None
         # TODO(tian): Refactor to contain more information
@@ -319,8 +398,8 @@ class SkyServiceSpec:
             overprovision_str = (
                 f' with {self.num_overprovision} overprovisioned replicas')
         return (f'Autoscaling from {self.min_replicas} to {self.max_replicas} '
-                f'replica{max_plural}{overprovision_str} (target QPS per '
-                f'replica: {self.target_qps_per_replica})')
+                f'{noun}{max_plural}{overprovision_str} (target QPS per '
+                f'{noun}: {self.target_qps_per_replica})')
 
     def set_ports(self, ports: str) -> None:
         self._ports = ports
@@ -332,6 +411,10 @@ class SkyServiceSpec:
                 f'Certfile: {self.tls_credential.certfile}')
 
     def __repr__(self) -> str:
+        if self.pool:
+            return textwrap.dedent(f"""\
+                Worker policy:  {self.autoscaling_policy_str()}
+            """)
         return textwrap.dedent(f"""\
             Readiness probe method:           {self.probe_str()}
             Readiness initial delay seconds:  {self.initial_delay_seconds}
@@ -372,7 +455,8 @@ class SkyServiceSpec:
         return self._ports
 
     @property
-    def target_qps_per_replica(self) -> Optional[float]:
+    def target_qps_per_replica(
+            self) -> Optional[Union[float, Dict[str, float]]]:
         return self._target_qps_per_replica
 
     @property
@@ -420,3 +504,43 @@ class SkyServiceSpec:
     def load_balancing_policy(self) -> str:
         return lb_policies.LoadBalancingPolicy.make_policy_name(
             self._load_balancing_policy)
+
+    @property
+    def pool(self) -> bool:
+        # This can happen for backward compatibility.
+        if not hasattr(self, '_pool'):
+            return False
+        return bool(self._pool)
+
+    def copy(self, **override) -> 'SkyServiceSpec':
+        return SkyServiceSpec(
+            readiness_path=override.pop('readiness_path', self._readiness_path),
+            initial_delay_seconds=override.pop('initial_delay_seconds',
+                                               self._initial_delay_seconds),
+            readiness_timeout_seconds=override.pop(
+                'readiness_timeout_seconds', self._readiness_timeout_seconds),
+            min_replicas=override.pop('min_replicas', self._min_replicas),
+            max_replicas=override.pop('max_replicas', self._max_replicas),
+            num_overprovision=override.pop('num_overprovision',
+                                           self._num_overprovision),
+            ports=override.pop('ports', self._ports),
+            target_qps_per_replica=override.pop('target_qps_per_replica',
+                                                self._target_qps_per_replica),
+            post_data=override.pop('post_data', self._post_data),
+            tls_credential=override.pop('tls_credential', self._tls_credential),
+            readiness_headers=override.pop('readiness_headers',
+                                           self._readiness_headers),
+            dynamic_ondemand_fallback=override.pop(
+                'dynamic_ondemand_fallback', self._dynamic_ondemand_fallback),
+            base_ondemand_fallback_replicas=override.pop(
+                'base_ondemand_fallback_replicas',
+                self._base_ondemand_fallback_replicas),
+            spot_placer=override.pop('spot_placer', self._spot_placer),
+            upscale_delay_seconds=override.pop('upscale_delay_seconds',
+                                               self._upscale_delay_seconds),
+            downscale_delay_seconds=override.pop('downscale_delay_seconds',
+                                                 self._downscale_delay_seconds),
+            load_balancing_policy=override.pop('load_balancing_policy',
+                                               self._load_balancing_policy),
+            pool=override.pop('pool', self._pool),
+        )

sky/serve/spot_placer.py

@@ -46,6 +46,8 @@ class Location:
 
     @classmethod
     def from_resources(cls, resources: 'resources_lib.Resources') -> 'Location':
+        assert resources.cloud is not None, 'Cloud must be specified'
+        assert resources.region is not None, 'Region must be specified'
         return cls(resources.cloud, resources.region, resources.zone)
 
     def to_dict(self) -> Dict[str, Any]:
@@ -147,6 +149,7 @@ def _get_possible_location_from_task(task: 'task_lib.Task') -> List[Location]:
                 cloud_str = str(launchable.cloud)
                 region = launchable.region
                 zone = launchable.zone
+                assert region is not None, 'Region must be specified'
                 if (cloud_str not in location_requirements and
                         location_requirements):
                     continue

sky/server/auth/authn.py

@@ -0,0 +1,50 @@
+"""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]):
+    # Skip for upload requests to avoid consuming the body prematurely, which
+    # will break the streaming upload.
+    if request.url.path.startswith('/upload'):
+        return
+    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/loopback.py

@@ -0,0 +1,38 @@
+"""Shared loopback detection utilities for auth middlewares."""
+
+import ipaddress
+
+import fastapi
+
+from sky import sky_logging
+
+logger = sky_logging.init_logger(__name__)
+
+COMMON_PROXY_HEADERS = [
+    'X-Forwarded-For', 'Forwarded', 'X-Real-IP', 'X-Client-IP',
+    'X-Forwarded-Host', 'X-Forwarded-Proto'
+]
+
+
+def _is_loopback_ip(ip_str: str) -> bool:
+    """Check if an IP address is a loopback address."""
+    try:
+        ip = ipaddress.ip_address(ip_str)
+        return ip.is_loopback
+    except ValueError:
+        return False
+
+
+def is_loopback_request(request: fastapi.Request) -> bool:
+    """Determine if a request is coming from localhost."""
+    if request.client is None:
+        return False
+
+    client_host = request.client.host
+    if client_host == 'localhost' or _is_loopback_ip(client_host):
+        # Additional checks: ensure no forwarding headers are present.
+        # If there are any, assume this traffic went through a proxy.
+        return not any(
+            request.headers.get(header) for header in COMMON_PROXY_HEADERS)
+
+    return False

sky/server/auth/oauth2_proxy.py

@@ -0,0 +1,202 @@
+"""Authentication based on oauth2-proxy."""
+
+import asyncio
+import hashlib
+import http
+import os
+import traceback
+from typing import Optional
+import urllib
+
+import aiohttp
+import fastapi
+import starlette.middleware.base
+
+from sky import global_user_state
+from sky import models
+from sky import sky_logging
+from sky.jobs import utils as managed_job_utils
+from sky.server import middleware_utils
+from sky.server.auth import authn
+from sky.server.auth import loopback
+from sky.users import permission
+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'
+
+
+@middleware_utils.websocket_aware
+class OAuth2ProxyMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
+    """Middleware to handle authentication by delegating to OAuth2 Proxy."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        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)
+
+        if managed_job_utils.is_consolidation_mode(
+        ) and loopback.is_loopback_request(request):
+            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}'
+                             f'{traceback.format_exc()}')
+                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 = {}
+        auth_url = f'{self.proxy_base}/oauth2/auth'
+        forwarded_headers['X-Forwarded-Uri'] = str(request.url).rstrip('/')
+        forwarded_headers['Host'] = request.url.hostname
+        logger.debug(f'authenticate request: {auth_url}, '
+                     f'headers: {forwarded_headers}')
+
+        async with session.request(
+                method='GET',
+                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.'
+                        })
+                newly_added = global_user_state.add_or_update_user(auth_user)
+                if newly_added:
+                    permission.permission_service.add_user_if_not_exists(
+                        auth_user.id)
+                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