skypilot-nightly 1.0.0.dev20250630__py3-none-any.whl → 1.0.0.dev20250702__py3-none-any.whl

sky/server/server.py

@@ -39,6 +39,7 @@ from sky import models
 from sky import sky_logging
 from sky.data import storage_utils
 from sky.jobs.server import server as jobs_rest
+from sky.metrics import utils as metrics_utils
 from sky.provision.kubernetes import utils as kubernetes_utils
 from sky.serve.server import server as serve_rest
 from sky.server import common
@@ -218,14 +219,26 @@ class RequestIDMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
 
 
 def _get_auth_user_header(request: fastapi.Request) -> Optional[models.User]:
-    if 'X-Auth-Request-Email' not in request.headers:
+    header_name = os.environ.get(constants.ENV_VAR_SERVER_AUTH_USER_HEADER,
+                                 'X-Auth-Request-Email')
+    if header_name not in request.headers:
         return None
-    user_name = request.headers['X-Auth-Request-Email']
+    user_name = request.headers[header_name]
     user_hash = hashlib.md5(
         user_name.encode()).hexdigest()[:common_utils.USER_HASH_LENGTH]
     return models.User(id=user_hash, name=user_name)
 
 
+class InitializeRequestAuthUserMiddleware(
+        starlette.middleware.base.BaseHTTPMiddleware):
+
+    async def dispatch(self, request: fastapi.Request, call_next):
+        # Make sure that request.state.auth_user is set. Otherwise, we may get a
+        # KeyError while trying to read it.
+        request.state.auth_user = None
+        return await call_next(request)
+
+
 class BasicAuthMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
     """Middleware to handle HTTP Basic Auth."""
 
@@ -277,29 +290,51 @@ class BearerTokenMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
     """Middleware to handle Bearer Token Auth (Service Accounts)."""
 
     async def dispatch(self, request: fastapi.Request, call_next):
-        # Only process requests with Bearer token authorization header
+        """Make sure correct bearer token auth is present.
+
+        1. If the request has the X-Skypilot-Auth-Mode: token header, it must
+           have a valid bearer token.
+        2. For backwards compatibility, if the request has a Bearer token
+           beginning with "sky_" (even if X-Skypilot-Auth-Mode is not present),
+           it must be a valid token.
+        3. If X-Skypilot-Auth-Mode is not set to "token", and there is no Bearer
+           token beginning with "sky_", allow the request to continue.
+
+        In conjunction with an auth proxy, the idea is to make the auth proxy
+        bypass requests with bearer tokens, instead setting the
+        X-Skypilot-Auth-Mode header. The auth proxy should either validate the
+        auth or set the header X-Skypilot-Auth-Mode: token.
+        """
+        has_skypilot_auth_header = (
+            request.headers.get('X-Skypilot-Auth-Mode') == 'token')
         auth_header = request.headers.get('authorization')
-        if not auth_header or not auth_header.lower().startswith('bearer '):
+        has_bearer_token_starting_with_sky = (
+            auth_header and auth_header.lower().startswith('bearer ') and
+            auth_header.split(' ', 1)[1].startswith('sky_'))
+
+        if (not has_skypilot_auth_header and
+                not has_bearer_token_starting_with_sky):
+            # This is case #3 above. We do not need to validate the request.
             # No Bearer token, continue with normal processing (OAuth2 cookies,
             # etc.)
             return await call_next(request)
+        # After this point, all requests must be validated.
+
+        if auth_header is None:
+            return fastapi.responses.JSONResponse(
+                status_code=401, content={'detail': 'Authentication required'})
 
         # Extract token
-        sa_token = auth_header.split(' ', 1)[1]
+        split_header = auth_header.split(' ', 1)
+        if split_header[0].lower() != 'bearer':
+            return fastapi.responses.JSONResponse(
+                status_code=401,
+                content={'detail': 'Invalid authentication method'})
+        sa_token = split_header[1]
 
         # Handle SkyPilot service account tokens
-        if sa_token.startswith('sky_'):
-            return await self._handle_service_account_token(
-                request, sa_token, call_next)
-
-        # Handle other Bearer tokens (OAuth2 access tokens, etc.)
-        # These requests bypassed OAuth2 proxy, so let the application decide
-        # how to handle them
-        # For now, we'll let them continue through normal processing
-        logger.debug(
-            'Non-SkyPilot Bearer token detected, continuing with normal '
-            'processing')
-        return await call_next(request)
+        return await self._handle_service_account_token(request, sa_token,
+                                                        call_next)
 
     async def _handle_service_account_token(self, request: fastapi.Request,
                                             sa_token: str, call_next):
@@ -384,6 +419,18 @@ class AuthProxyMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
     async def dispatch(self, request: fastapi.Request, call_next):
         auth_user = _get_auth_user_header(request)
 
+        if request.state.auth_user is not None:
+            # Previous middleware is trusted more than this middleware.  For
+            # instance, a client could set the Authorization and the
+            # X-Auth-Request-Email header. In that case, the auth proxy will be
+            # skipped and we should rely on the Bearer token to authenticate the
+            # user - but that means the user could set X-Auth-Request-Email to
+            # whatever the user wants. We should thus ignore it.
+            if auth_user is not None:
+                logger.debug('Warning: ignoring auth proxy header since the '
+                             'auth user was already set.')
+            return await call_next(request)
+
         # Add user to database if auth_user is present
         if auth_user is not None:
             newly_added = global_user_state.add_or_update_user(auth_user)
@@ -394,8 +441,6 @@ class AuthProxyMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
         # Store user info in request.state for access by GET endpoints
         if auth_user is not None:
             request.state.auth_user = auth_user
-        else:
-            request.state.auth_user = None
 
         await _override_user_info_in_request_body(request, auth_user)
         return await call_next(request)
@@ -514,10 +559,17 @@ class GracefulShutdownMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
 
 
 app = fastapi.FastAPI(prefix='/api/v1', debug=True, lifespan=lifespan)
+# Middleware wraps in the order defined here. E.g., given
+#   app.add_middleware(Middleware1)
+#   app.add_middleware(Middleware2)
+#   app.add_middleware(Middleware3)
+# The effect will be like:
+#   Middleware3(Middleware2(Middleware1(request)))
+# If MiddlewareN does something like print(n); call_next(); print(n), you'll get
+#   3; 2; 1; <request>; 1; 2; 3
 # Use environment variable to make the metrics middleware optional.
 if os.environ.get(constants.ENV_VAR_SERVER_METRICS_ENABLED):
     app.add_middleware(metrics.PrometheusMiddleware)
-app.add_middleware(RBACMiddleware)
 app.add_middleware(InternalDashboardPrefixMiddleware)
 app.add_middleware(GracefulShutdownMiddleware)
 app.add_middleware(PathCleanMiddleware)
@@ -530,15 +582,26 @@ app.add_middleware(
     allow_credentials=True,
     allow_methods=['*'],
     allow_headers=['*'],
-    # TODO(syang): remove X-Request-ID when v0.10.0 is released.
+    # TODO(syang): remove X-Request-ID \when v0.10.0 is released.
     expose_headers=['X-Request-ID', 'X-Skypilot-Request-ID'])
+# The order of all the authentication-related middleware is important.
+# RBACMiddleware must precede all the auth middleware, so it can access
+# request.state.auth_user.
+app.add_middleware(RBACMiddleware)
+# AuthProxyMiddleware should precede BasicAuthMiddleware and
+# BearerTokenMiddleware, since it should be skipped if either of those set the
+# auth user.
+app.add_middleware(AuthProxyMiddleware)
 enable_basic_auth = os.environ.get(constants.ENV_VAR_ENABLE_BASIC_AUTH, 'false')
 if str(enable_basic_auth).lower() == 'true':
     app.add_middleware(BasicAuthMiddleware)
 # Bearer token middleware should always be present to handle service account
 # authentication
 app.add_middleware(BearerTokenMiddleware)
-app.add_middleware(AuthProxyMiddleware)
+# InitializeRequestAuthUserMiddleware must be the last added middleware so that
+# request.state.auth_user is always set, but can be overridden by the auth
+# middleware above.
+app.add_middleware(InitializeRequestAuthUserMiddleware)
 app.add_middleware(RequestIDMiddleware)
 app.include_router(jobs_rest.router, prefix='/jobs', tags=['jobs'])
 app.include_router(serve_rest.router, prefix='/serve', tags=['serve'])
@@ -1554,6 +1617,38 @@ async def all_contexts(request: fastapi.Request) -> None:
     )
 
 
+@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 = []
+    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}')
+        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')
+
+
 # === Internal APIs ===
 @app.get('/api/completion/cluster_name')
 async def complete_cluster_name(incomplete: str,) -> List[str]:

sky/setup_files/MANIFEST.in

@@ -17,3 +17,4 @@ include sky/utils/kubernetes/*
 include sky/server/html/*
 recursive-include sky/dashboard/out *
 include sky/users/*.conf
+include sky/metrics/*

sky/setup_files/dependencies.py

@@ -65,12 +65,14 @@ install_requires = [
     # Required for API server metrics
     'prometheus_client>=0.8.0',
     'passlib',
+    'pyjwt',
 ]
 
 server_dependencies = [
     'casbin',
     'sqlalchemy_adapter',
     'passlib',
+    'pyjwt',
 ]
 
 local_ray = [

sky/skylet/constants.py

@@ -421,6 +421,9 @@ ENV_VAR_IS_SKYPILOT_SERVER = 'IS_SKYPILOT_SERVER'
 # Environment variable that is set to 'true' if metrics are enabled.
 ENV_VAR_SERVER_METRICS_ENABLED = 'SKY_API_SERVER_METRICS_ENABLED'
 
+# If set, overrides the header that we can use to get the user name.
+ENV_VAR_SERVER_AUTH_USER_HEADER = f'{SKYPILOT_ENV_VAR_PREFIX}AUTH_USER_HEADER'
+
 # Environment variable that is used as the DB connection string for the
 # skypilot server.
 ENV_VAR_DB_CONNECTION_URI = (f'{SKYPILOT_ENV_VAR_PREFIX}DB_CONNECTION_URI')

sky/skypilot_config.py

@@ -63,6 +63,7 @@ from sqlalchemy import orm
 from sqlalchemy.dialects import postgresql
 from sqlalchemy.dialects import sqlite
 from sqlalchemy.ext import declarative
+from sqlalchemy.pool import NullPool
 
 from sky import exceptions
 from sky import sky_logging
@@ -116,9 +117,10 @@ ENV_VAR_PROJECT_CONFIG = f'{constants.SKYPILOT_ENV_VAR_PREFIX}PROJECT_CONFIG'
 _GLOBAL_CONFIG_PATH = '~/.sky/config.yaml'
 _PROJECT_CONFIG_PATH = '.sky.yaml'
 
-_SQLALCHEMY_ENGINE: Optional[sqlalchemy.engine.Engine] = None
 API_SERVER_CONFIG_KEY = 'api_server_config'
 
+_DB_USE_LOCK = threading.Lock()
+
 Base = declarative.declarative_base()
 
 config_yaml_table = sqlalchemy.Table(
@@ -129,44 +131,6 @@ config_yaml_table = sqlalchemy.Table(
 )
 
 
-def create_table():
-    # Create tables if they don't exist
-    Base.metadata.create_all(bind=_SQLALCHEMY_ENGINE)
-
-
-def _get_config_yaml_from_db(key: str) -> Optional[config_utils.Config]:
-    assert _SQLALCHEMY_ENGINE is not None
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        row = session.query(config_yaml_table).filter_by(key=key).first()
-    if row:
-        db_config = config_utils.Config(yaml.safe_load(row.value))
-        db_config.pop_nested(('db',), None)
-        return db_config
-    return None
-
-
-def _set_config_yaml_to_db(key: str, config: config_utils.Config):
-    assert _SQLALCHEMY_ENGINE is not None
-    config.pop_nested(('db',), None)
-    config_str = common_utils.dump_yaml_str(dict(config))
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
-        if (_SQLALCHEMY_ENGINE.dialect.name ==
-                db_utils.SQLAlchemyDialect.SQLITE.value):
-            insert_func = sqlite.insert
-        elif (_SQLALCHEMY_ENGINE.dialect.name ==
-              db_utils.SQLAlchemyDialect.POSTGRESQL.value):
-            insert_func = postgresql.insert
-        else:
-            raise ValueError('Unsupported database dialect')
-        insert_stmnt = insert_func(config_yaml_table).values(key=key,
-                                                             value=config_str)
-        do_update_stmt = insert_stmnt.on_conflict_do_update(
-            index_elements=[config_yaml_table.c.key],
-            set_={config_yaml_table.c.value: config_str})
-        session.execute(do_update_stmt)
-        session.commit()
-
-
 class ConfigContext:
 
     def __init__(self,
@@ -586,7 +550,6 @@ def _reload_config_from_internal_file(internal_config_path: str) -> None:
 
 
 def _reload_config_as_server() -> None:
-    global _SQLALCHEMY_ENGINE
     # Reset the global variables, to avoid using stale values.
     _set_loaded_config(config_utils.Config())
     _set_loaded_config_path(None)
@@ -607,16 +570,33 @@ def _reload_config_as_server() -> None:
             'if db config is specified, no other config is allowed')
 
     if db_url:
-        if _SQLALCHEMY_ENGINE is None:
-            _SQLALCHEMY_ENGINE = sqlalchemy.create_engine(db_url)
-            create_table()
-        db_config = _get_config_yaml_from_db(API_SERVER_CONFIG_KEY)
-        if db_config:
-            if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
-                logger.debug(f'Config loaded from db:\n'
-                             f'{common_utils.dump_yaml_str(dict(db_config))}')
-            server_config = overlay_skypilot_config(server_config, db_config)
-
+        with _DB_USE_LOCK:
+            sqlalchemy_engine = sqlalchemy.create_engine(db_url,
+                                                         poolclass=NullPool)
+            Base.metadata.create_all(bind=sqlalchemy_engine)
+
+            def _get_config_yaml_from_db(
+                    key: str) -> Optional[config_utils.Config]:
+                assert sqlalchemy_engine is not None
+                with orm.Session(sqlalchemy_engine) as session:
+                    row = session.query(config_yaml_table).filter_by(
+                        key=key).first()
+                if row:
+                    db_config = config_utils.Config(yaml.safe_load(row.value))
+                    db_config.pop_nested(('db',), None)
+                    return db_config
+                return None
+
+            db_config = _get_config_yaml_from_db(API_SERVER_CONFIG_KEY)
+            if db_config:
+                if sky_logging.logging_enabled(logger, sky_logging.DEBUG):
+                    logger.debug(
+                        f'Config loaded from db:\n'
+                        f'{common_utils.dump_yaml_str(dict(db_config))}')
+                server_config = overlay_skypilot_config(server_config,
+                                                        db_config)
+            # Close the engine to avoid connection leaks
+            sqlalchemy_engine.dispose()
     _set_loaded_config(server_config)
     _set_loaded_config_path(server_config_path)
 
@@ -876,9 +856,38 @@ def update_api_server_config_no_lock(config: config_utils.Config) -> None:
             new_db_url = config.get_nested(('db',), None)
             if new_db_url and new_db_url != existing_db_url:
                 raise ValueError('Cannot change db url while server is running')
-            logger.debug('saving api_server config to db')
-            _set_config_yaml_to_db(API_SERVER_CONFIG_KEY, config)
-            db_updated = True
+            with _DB_USE_LOCK:
+                sqlalchemy_engine = sqlalchemy.create_engine(existing_db_url,
+                                                             poolclass=NullPool)
+                Base.metadata.create_all(bind=sqlalchemy_engine)
+
+                def _set_config_yaml_to_db(key: str,
+                                           config: config_utils.Config):
+                    assert sqlalchemy_engine is not None
+                    config.pop_nested(('db',), None)
+                    config_str = common_utils.dump_yaml_str(dict(config))
+                    with orm.Session(sqlalchemy_engine) as session:
+                        if (sqlalchemy_engine.dialect.name ==
+                                db_utils.SQLAlchemyDialect.SQLITE.value):
+                            insert_func = sqlite.insert
+                        elif (sqlalchemy_engine.dialect.name ==
+                              db_utils.SQLAlchemyDialect.POSTGRESQL.value):
+                            insert_func = postgresql.insert
+                        else:
+                            raise ValueError('Unsupported database dialect')
+                        insert_stmnt = insert_func(config_yaml_table).values(
+                            key=key, value=config_str)
+                        do_update_stmt = insert_stmnt.on_conflict_do_update(
+                            index_elements=[config_yaml_table.c.key],
+                            set_={config_yaml_table.c.value: config_str})
+                        session.execute(do_update_stmt)
+                        session.commit()
+
+                logger.debug('saving api_server config to db')
+                _set_config_yaml_to_db(API_SERVER_CONFIG_KEY, config)
+                db_updated = True
+                # Close the engine to avoid connection leaks
+                sqlalchemy_engine.dispose()
 
     if not db_updated:
         # save to the local file (PVC in Kubernetes, local file otherwise)

sky/task.py

@@ -1512,7 +1512,7 @@ class Task:
                 d[k] = v
         return d
 
-    def to_yaml_config(self, redact_secrets: bool = True) -> Dict[str, Any]:
+    def to_yaml_config(self, redact_secrets: bool = False) -> Dict[str, Any]:
         """Returns a yaml-style dict representation of the task.
 
         INTERNAL: this method is internal-facing.

sky/utils/accelerator_registry.py

@@ -1,6 +1,6 @@
 """Accelerator registry."""
 import typing
-from typing import Optional
+from typing import List, Optional
 
 from sky import catalog
 from sky.utils import rich_utils
@@ -35,6 +35,7 @@ if typing.TYPE_CHECKING:
 # Use a cached version of accelerators to cloud mapping, so that we don't have
 # to download and read the catalog file for every cloud locally.
 _accelerator_df = catalog.common.read_catalog('common/accelerators.csv')
+_memory_df = catalog.common.read_catalog('common/metadata.csv')
 
 # List of non-GPU accelerators that are supported by our backend for job queue
 # scheduling.
@@ -45,6 +46,32 @@ _SCHEDULABLE_NON_GPU_ACCELERATORS = [
 ]
 
 
+def get_devices_by_memory(memory: float,
+                          plus: bool = False,
+                          manufacturer: Optional[str] = None) -> List[str]:
+    """Returns a list of device names that meet the memory and manufacturer
+       requirements.
+
+    Args:
+        memory: The minimum memory size in GB.
+        plus: If True, returns devices with memory >= memory, otherwise returns
+              devices with memory == memory.
+        manufacturer: The manufacturer of the GPU.
+    """
+
+    # Filter by memory requirements
+    if plus:
+        df = _memory_df[_memory_df['MemoryGB'] >= memory]
+    else:
+        df = _memory_df[_memory_df['MemoryGB'] == memory]
+
+    # Filter by manufacturer if specified
+    if manufacturer is not None:
+        df = df[df['Manufacturer'].str.lower() == manufacturer.lower()]
+
+    return df['GPU'].tolist()
+
+
 def is_schedulable_non_gpu_accelerator(accelerator_name: str) -> bool:
     """Returns if this accelerator is a 'schedulable' non-GPU accelerator."""
     for name in _SCHEDULABLE_NON_GPU_ACCELERATORS:

sky/utils/dag_utils.py

@@ -147,11 +147,13 @@ def load_chain_dag_from_yaml_str(
     return _load_chain_dag(configs, env_overrides, secrets_overrides)
 
 
-def dump_chain_dag_to_yaml_str(dag: dag_lib.Dag) -> str:
+def dump_chain_dag_to_yaml_str(dag: dag_lib.Dag,
+                               redact_secrets: bool = False) -> str:
     """Dumps a chain DAG to a YAML string.
 
     Args:
         dag: the DAG to dump.
+        redact_secrets: whether to redact secrets in the YAML string.
 
     Returns:
         The YAML string.
@@ -159,7 +161,7 @@ def dump_chain_dag_to_yaml_str(dag: dag_lib.Dag) -> str:
     assert dag.is_chain(), dag
     configs = [{'name': dag.name}]
     for task in dag.tasks:
-        configs.append(task.to_yaml_config())
+        configs.append(task.to_yaml_config(redact_secrets=redact_secrets))
     return common_utils.dump_yaml_str(configs)
 
 

sky/utils/schemas.py

@@ -311,6 +311,9 @@ def _get_single_resources_schema():
                     }
                 }
             },
+            '_no_missing_accel_warnings': {
+                'type': 'boolean'
+            },
             'image_id': {
                 'anyOf': [{
                     'type': 'string',

sky/utils/ux_utils.py

@@ -12,6 +12,7 @@ import colorama
 from sky import sky_logging
 from sky.skylet import constants
 from sky.utils import common_utils
+from sky.utils import env_options
 from sky.utils import rich_console_utils
 
 if typing.TYPE_CHECKING:
@@ -57,10 +58,14 @@ def print_exception_no_traceback():
             if error():
                 raise ValueError('...')
     """
-    original_tracelimit = getattr(sys, 'tracebacklimit', 1000)
-    sys.tracebacklimit = 0
-    yield
-    sys.tracebacklimit = original_tracelimit
+    if env_options.Options.SHOW_DEBUG_INFO.get():
+        # When SKYPILOT_DEBUG is set, show the full traceback
+        yield
+    else:
+        original_tracelimit = getattr(sys, 'tracebacklimit', 1000)
+        sys.tracebacklimit = 0
+        yield
+        sys.tracebacklimit = original_tracelimit
 
 
 @contextlib.contextmanager

{skypilot_nightly-1.0.0.dev20250630.dist-info → skypilot_nightly-1.0.0.dev20250702.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: skypilot-nightly
-Version: 1.0.0.dev20250630
+Version: 1.0.0.dev20250702
 Summary: SkyPilot: Run AI on Any Infra — Unified, Faster, Cheaper.
 Author: SkyPilot Team
 License: Apache 2.0
@@ -53,6 +53,7 @@ Requires-Dist: casbin
 Requires-Dist: sqlalchemy_adapter
 Requires-Dist: prometheus_client>=0.8.0
 Requires-Dist: passlib
+Requires-Dist: pyjwt
 Provides-Extra: aws
 Requires-Dist: awscli>=1.27.10; extra == "aws"
 Requires-Dist: botocore>=1.29.10; extra == "aws"
@@ -126,6 +127,7 @@ Provides-Extra: server
 Requires-Dist: casbin; extra == "server"
 Requires-Dist: sqlalchemy_adapter; extra == "server"
 Requires-Dist: passlib; extra == "server"
+Requires-Dist: pyjwt; extra == "server"
 Provides-Extra: all
 Requires-Dist: awscli>=1.27.10; extra == "all"
 Requires-Dist: botocore>=1.29.10; extra == "all"
@@ -178,6 +180,7 @@ Requires-Dist: colorama<0.4.5; extra == "all"
 Requires-Dist: casbin; extra == "all"
 Requires-Dist: sqlalchemy_adapter; extra == "all"
 Requires-Dist: passlib; extra == "all"
+Requires-Dist: pyjwt; extra == "all"
 Dynamic: author
 Dynamic: classifier
 Dynamic: description