skypilot-nightly 1.0.0.dev20250509__py3-none-any.whl → 1.0.0.dev20251107__py3-none-any.whl

sky/users/token_service.py

@@ -0,0 +1,218 @@
+"""JWT-based service account token management for SkyPilot."""
+
+import contextlib
+import datetime
+import hashlib
+import os
+import secrets
+import threading
+from typing import Any, Dict, Generator, Optional
+
+import filelock
+import jwt
+
+from sky import global_user_state
+from sky import sky_logging
+
+logger = sky_logging.init_logger(__name__)
+
+# JWT Configuration
+JWT_ALGORITHM = 'HS256'
+JWT_ISSUER = 'sky'  # Shortened for compact tokens
+JWT_SECRET_DB_KEY = 'jwt_secret'
+
+# File lock for JWT secret initialization
+JWT_SECRET_LOCK_PATH = os.path.expanduser('~/.sky/.jwt_secret_init.lock')
+JWT_SECRET_LOCK_TIMEOUT_SECONDS = 20
+
+
+@contextlib.contextmanager
+def _jwt_secret_lock() -> Generator[None, None, None]:
+    """Context manager for JWT secret initialization lock."""
+    try:
+        with filelock.FileLock(JWT_SECRET_LOCK_PATH,
+                               JWT_SECRET_LOCK_TIMEOUT_SECONDS):
+            yield
+    except filelock.Timeout as e:
+        raise RuntimeError(f'Failed to initialize JWT secret due to a timeout '
+                           f'when trying to acquire the lock at '
+                           f'{JWT_SECRET_LOCK_PATH}. '
+                           'Please try again or manually remove the lock '
+                           f'file if you believe it is stale.') from e
+
+
+class TokenService:
+    """Service for managing JWT-based service account tokens."""
+
+    def __init__(self):
+        self.secret_key = None
+        self.init_lock = threading.Lock()
+
+    def _lazy_initialize(self):
+        if self.secret_key is not None:
+            return
+        with self.init_lock:
+            if self.secret_key is not None:
+                return
+            self.secret_key = self._get_or_generate_secret()
+
+    def _get_or_generate_secret(self) -> str:
+        """Get JWT secret from database or generate a new one."""
+
+        def _get_secret_from_db():
+            try:
+                db_secret = global_user_state.get_system_config(
+                    JWT_SECRET_DB_KEY)
+                if db_secret:
+                    logger.debug('Retrieved existing JWT secret from database')
+                    return db_secret
+            except Exception as e:  # pylint: disable=broad-except
+                logger.debug(f'Failed to get JWT secret from database: {e}')
+            return None
+
+        # Try to get from database (persistent across deployments)
+        token_from_db = _get_secret_from_db()
+        if token_from_db:
+            return token_from_db
+
+        with _jwt_secret_lock():
+            token_from_db = _get_secret_from_db()
+            if token_from_db:
+                return token_from_db
+            # Generate a new secret and store in database
+            new_secret = secrets.token_urlsafe(64)
+            try:
+                global_user_state.set_system_config(JWT_SECRET_DB_KEY,
+                                                    new_secret)
+                logger.info(
+                    'Generated new JWT secret and stored in database. '
+                    'This secret will persist across API server restarts.')
+            except Exception as e:  # pylint: disable=broad-except
+                logger.warning(
+                    f'Failed to store new JWT secret in database: {e}. '
+                    f'Using in-memory secret (tokens will not persist '
+                    f'across restarts).')
+
+            return new_secret
+
+    def create_token(self,
+                     creator_user_id: str,
+                     service_account_user_id: str,
+                     token_name: str,
+                     expires_in_days: Optional[int] = None) -> Dict[str, Any]:
+        """Create a new JWT service account token.
+
+        Args:
+            creator_user_id: The creator's user hash
+            service_account_user_id: The service account's own user ID
+            token_name: Descriptive name for the token
+            expires_in_days: Optional expiration in days
+
+        Returns:
+            Dict containing token info including the JWT token
+        """
+        self._lazy_initialize()
+        now = datetime.datetime.now(datetime.timezone.utc)
+        token_id = secrets.token_urlsafe(12)  # Shorter ID for JWT
+
+        # Build minimal JWT payload with single-character field names for
+        # compactness
+        payload = {
+            'i': JWT_ISSUER,  # Issuer (use constant)
+            't': int(now.timestamp()),  # Issued at (shortened from 'iat')
+            # Service account user ID (shortened from 'sub')
+            'u': service_account_user_id,
+            'k': token_id,  # Token ID (shortened from 'token_id')
+            'y': 'sa',  # Type: service account (shortened from 'type')
+        }
+
+        # Add expiration if specified
+        expires_at = None
+        if expires_in_days:
+            exp_time = now + datetime.timedelta(days=expires_in_days)
+            payload['e'] = int(
+                exp_time.timestamp())  # Expiration (shortened from 'exp')
+            expires_at = int(exp_time.timestamp())
+
+        # Generate JWT
+        jwt_token = jwt.encode(payload,
+                               self.secret_key,
+                               algorithm=JWT_ALGORITHM)
+
+        # Create token with SkyPilot prefix
+        full_token = f'sky_{jwt_token}'
+
+        # Generate hash for database storage (we still hash the full token)
+        token_hash = hashlib.sha256(full_token.encode()).hexdigest()
+
+        return {
+            'token_id': token_id,
+            'token': full_token,
+            'token_hash': token_hash,
+            'creator_user_id': creator_user_id,
+            'service_account_user_id': service_account_user_id,
+            'token_name': token_name,
+            'created_at': int(now.timestamp()),
+            'expires_at': expires_at,
+        }
+
+    def verify_token(self, token: str) -> Optional[Dict[str, Any]]:
+        """Verify and decode a JWT token.
+
+        Args:
+            token: The full token (with sky_ prefix)
+
+        Returns:
+            Decoded token payload or None if invalid
+        """
+        self._lazy_initialize()
+        if not token.startswith('sky_'):
+            return None
+
+        # Remove the sky_ prefix
+        jwt_token = token[4:]
+
+        try:
+            # Decode and verify JWT (without issuer verification)
+            payload = jwt.decode(jwt_token,
+                                 self.secret_key,
+                                 algorithms=[JWT_ALGORITHM])
+
+            # Manually verify issuer using our shortened field name
+            token_issuer = payload.get('i')
+            if token_issuer != JWT_ISSUER:
+                logger.warning(f'Invalid token issuer: {token_issuer}')
+                return None
+
+            # Verify token type
+            token_type = payload.get('y')
+            if token_type != 'sa':
+                logger.warning(f'Invalid token type: {token_type}')
+                return None
+
+            # Convert shortened field names back to standard names for
+            # compatibility
+            normalized_payload = {
+                'iss': payload.get('i'),  # issuer
+                'iat': payload.get('t'),  # issued at
+                'sub': payload.get('u'),  # subject (service account user ID)
+                'token_id': payload.get('k'),  # token ID
+                'type': 'service_account',  # expand shortened type
+            }
+
+            # Add expiration if present
+            if 'e' in payload:
+                normalized_payload['exp'] = payload['e']
+
+            return normalized_payload
+
+        except jwt.ExpiredSignatureError:
+            logger.warning('Token has expired')
+            return None
+        except jwt.InvalidTokenError as e:
+            logger.warning(f'Invalid token: {e}')
+            return None
+
+
+# Singleton instance
+token_service = TokenService()

sky/utils/accelerator_registry.py

@@ -1,8 +1,8 @@
 """Accelerator registry."""
 import typing
-from typing import Optional
+from typing import List, Optional
 
-from sky.clouds import service_catalog
+from sky import catalog
 from sky.utils import rich_utils
 from sky.utils import ux_utils
 
@@ -34,7 +34,8 @@ 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 = service_catalog.common.read_catalog('common/accelerators.csv')
+_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:
@@ -80,10 +107,12 @@ def canonicalize_accelerator_name(accelerator: str,
     if not names and cloud_str in ['Kubernetes', None]:
         with rich_utils.safe_status(
                 ux_utils.spinner_message('Listing accelerators on Kubernetes')):
-            searched = service_catalog.list_accelerators(
+            # Only search for Kubernetes to reduce the lookup cost.
+            # For other clouds, the catalog has been searched in previous steps.
+            searched = catalog.list_accelerators(
                 name_filter=accelerator,
                 case_sensitive=False,
-                clouds=cloud_str,
+                clouds='Kubernetes',
             )
         names = list(searched.keys())
         if accelerator in names:

sky/utils/admin_policy_utils.py

@@ -1,9 +1,10 @@
 """Admin policy utils."""
+import contextlib
 import copy
 import importlib
-import os
-import tempfile
-from typing import Optional, Tuple, Union
+import typing
+from typing import Iterator, Optional, Tuple, Union
+import urllib.parse
 
 import colorama
 
@@ -13,25 +14,45 @@ from sky import exceptions
 from sky import sky_logging
 from sky import skypilot_config
 from sky import task as task_lib
+from sky.server.requests import request_names
 from sky.utils import common_utils
 from sky.utils import config_utils
 from sky.utils import ux_utils
 
 logger = sky_logging.init_logger(__name__)
 
+if typing.TYPE_CHECKING:
+    from sky import models
 
-def _get_policy_cls(
-        policy: Optional[str]) -> Optional[admin_policy.AdminPolicy]:
+
+def _is_url(policy_string: str) -> bool:
+    """Check if the policy string is a URL."""
+    try:
+        parsed = urllib.parse.urlparse(policy_string)
+        return parsed.scheme in ('http', 'https')
+    except Exception:  # pylint: disable=broad-except
+        return False
+
+
+def _get_policy_impl(
+        policy_location: Optional[str]
+) -> Optional[admin_policy.PolicyInterface]:
     """Gets admin-defined policy."""
-    if policy is None:
+    if policy_location is None:
         return None
+
+    if _is_url(policy_location):
+        # Use the built-in URL policy class when an URL is specified.
+        return admin_policy.RestfulAdminPolicy(policy_location)
+
+    # Handle module path format
     try:
-        module_path, class_name = policy.rsplit('.', 1)
+        module_path, class_name = policy_location.rsplit('.', 1)
         module = importlib.import_module(module_path)
     except ImportError as e:
         with ux_utils.print_exception_no_traceback():
             raise ImportError(
-                f'Failed to import policy module: {policy}. '
+                f'Failed to import policy module: {policy_location}. '
                 'Please check if the module is installed in your Python '
                 'environment.') from e
 
@@ -43,19 +64,48 @@ def _get_policy_cls(
                 f'Could not find {class_name} class in module {module_path}. '
                 'Please check with your policy admin for details.') from e
 
-    # Check if the module implements the AdminPolicy interface.
+    # Currently we only allow users to define subclass of AdminPolicy
+    # instead of inheriting from PolicyInterface or PolicyTemplate.
     if not issubclass(policy_cls, admin_policy.AdminPolicy):
         with ux_utils.print_exception_no_traceback():
             raise ValueError(
-                f'Policy class {policy!r} does not implement the AdminPolicy '
-                'interface. Please check with your policy admin for details.')
-    return policy_cls
+                f'Policy class {policy_cls!r} does not implement the '
+                'AdminPolicy interface. Please check with your policy admin '
+                'for details.')
+    return policy_cls()
+
+
+@contextlib.contextmanager
+def apply_and_use_config_in_current_request(
+    entrypoint: Union['dag_lib.Dag', 'task_lib.Task'],
+    request_name: request_names.AdminPolicyRequestName,
+    request_options: Optional[admin_policy.RequestOptions] = None,
+    at_client_side: bool = False,
+) -> Iterator['dag_lib.Dag']:
+    """Applies an admin policy and override SkyPilot config for current request
+
+    This is a helper function of `apply()` that applies an admin policy and
+    overrides the SkyPilot config for the current request as a context manager.
+    The original SkyPilot config will be restored when the context manager is
+    exited.
+
+    Refer to `apply()` for more details.
+    """
+    original_config = skypilot_config.to_dict()
+    dag, mutated_config = apply(entrypoint, request_name, request_options,
+                                at_client_side)
+    if mutated_config != original_config:
+        with skypilot_config.replace_skypilot_config(mutated_config):
+            yield dag
+    else:
+        yield dag
 
 
 def apply(
     entrypoint: Union['dag_lib.Dag', 'task_lib.Task'],
-    use_mutated_config_in_current_request: bool = True,
+    request_name: request_names.AdminPolicyRequestName,
     request_options: Optional[admin_policy.RequestOptions] = None,
+    at_client_side: bool = False,
 ) -> Tuple['dag_lib.Dag', config_utils.Config]:
     """Applies an admin policy (if registered) to a DAG or a task.
 
@@ -79,29 +129,41 @@ def apply(
     else:
         dag = entrypoint
 
-    policy = skypilot_config.get_nested(('admin_policy',), None)
-    policy_cls = _get_policy_cls(policy)
-    if policy_cls is None:
+    policy_location = skypilot_config.get_nested(('admin_policy',), None)
+    policy = _get_policy_impl(policy_location)
+    if policy is None:
         return dag, skypilot_config.to_dict()
 
-    logger.info(f'Applying policy: {policy}')
-    original_config = skypilot_config.to_dict()
-    config = copy.deepcopy(original_config)
+    user = None
+    if at_client_side:
+        logger.info(f'Applying client admin policy: {policy}')
+    else:
+        # When being called by the server, the middleware has set the
+        # current user and this information is available at this point.
+        user = common_utils.get_current_user()
+        logger.info(f'Applying server admin policy: {policy}')
+    config = copy.deepcopy(skypilot_config.to_dict())
     mutated_dag = dag_lib.Dag()
     mutated_dag.name = dag.name
 
     mutated_config = None
     for task in dag.tasks:
-        user_request = admin_policy.UserRequest(task, config, request_options)
+        user_request = admin_policy.UserRequest(task, config, request_name,
+                                                request_options, at_client_side,
+                                                user)
         try:
-            mutated_user_request = policy_cls.validate_and_mutate(user_request)
+            mutated_user_request = policy.apply(user_request)
+        # Avoid duplicate exception wrapping.
+        except exceptions.UserRequestRejectedByPolicy as e:
+            with ux_utils.print_exception_no_traceback():
+                raise e
         except Exception as e:  # pylint: disable=broad-except
             with ux_utils.print_exception_no_traceback():
                 raise exceptions.UserRequestRejectedByPolicy(
                     f'{colorama.Fore.RED}User request rejected by policy '
                     f'{policy!r}{colorama.Fore.RESET}: '
                     f'{common_utils.format_exception(e, use_bracket=True)}'
-                ) from e
+                ) from None
         if mutated_config is None:
             mutated_config = mutated_user_request.skypilot_config
         else:
@@ -126,22 +188,6 @@ def apply(
         mutated_dag.graph.add_edge(mutated_dag.tasks[u_idx],
                                    mutated_dag.tasks[v_idx])
 
-    if (use_mutated_config_in_current_request and
-            original_config != mutated_config):
-        with tempfile.NamedTemporaryFile(
-                delete=False,
-                mode='w',
-                prefix='policy-mutated-skypilot-config-',
-                suffix='.yaml') as temp_file:
-
-            common_utils.dump_yaml(temp_file.name, dict(**mutated_config))
-            os.environ[skypilot_config.ENV_VAR_SKYPILOT_CONFIG] = temp_file.name
-            logger.debug(f'Updated SkyPilot config: {temp_file.name}')
-            # TODO(zhwu): This is not a clean way to update the SkyPilot config,
-            # because we are resetting the global context for a single DAG,
-            # which is conceptually weird.
-            importlib.reload(skypilot_config)
-
     logger.debug(f'Mutated user request: {mutated_user_request}')
     mutated_dag.policy_applied = True
     return mutated_dag, mutated_config

sky/utils/annotations.py

@@ -1,14 +1,19 @@
 """Annotations for public APIs."""
 
 import functools
-from typing import Callable, Literal
+from typing import Callable, Literal, TypeVar
+
+from typing_extensions import ParamSpec
 
 # Whether the current process is a SkyPilot API server process.
 is_on_api_server = True
-FUNCTIONS_NEED_RELOAD_CACHE = []
+_FUNCTIONS_NEED_RELOAD_CACHE = []
+
+T = TypeVar('T')
+P = ParamSpec('P')
 
 
-def client_api(func):
+def client_api(func: Callable[P, T]) -> Callable[P, T]:
     """Mark a function as a client-side API.
 
     Code invoked by server-side functions will find annotations.is_on_api_server
@@ -38,14 +43,20 @@ def lru_cache(scope: Literal['global', 'request'], *lru_cache_args,
         lru_cache_kwargs: Keyword arguments for functools.lru_cache.
     """
 
-    def decorator(func: Callable) -> Callable:
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
         if scope == 'global':
             return functools.lru_cache(*lru_cache_args,
                                        **lru_cache_kwargs)(func)
         else:
             cached_func = functools.lru_cache(*lru_cache_args,
                                               **lru_cache_kwargs)(func)
-            FUNCTIONS_NEED_RELOAD_CACHE.append(cached_func)
+            _FUNCTIONS_NEED_RELOAD_CACHE.append(cached_func)
             return cached_func
 
     return decorator
+
+
+def clear_request_level_cache():
+    """Clear the request-level cache."""
+    for func in _FUNCTIONS_NEED_RELOAD_CACHE:
+        func.cache_clear()

sky/utils/asyncio_utils.py

@@ -0,0 +1,78 @@
+"""Asyncio utilities."""
+
+import asyncio
+import functools
+from typing import Set
+
+_background_tasks: Set[asyncio.Task] = set()
+
+
+def shield(func):
+    """Shield the decorated async function from cancellation.
+
+    If the outter coroutine is cancelled, the inner decorated function
+    will be protected from cancellation by asyncio.shield(). And we will
+    maintain a reference to the the inner task to avoid it get GCed before
+    it is done.
+
+    For example, filelock.AsyncFileLock is not cancellation safe. The
+    following code:
+
+        async def fn_with_lock():
+            async with filelock.AsyncFileLock('lock'):
+                await asyncio.sleep(1)
+
+    is equivalent to:
+
+        # The lock may leak if the cancellation happens in
+        # lock.acquire() or lock.release()
+        async def fn_with_lock():
+            lock = filelock.AsyncFileLock('lock')
+            await lock.acquire()
+            try:
+                await asyncio.sleep(1)
+            finally:
+                await lock.release()
+
+    Shilding the function ensures there is no cancellation will happen in the
+    function, thus the lock will be released properly:
+
+        @shield
+        async def fn_with_lock()
+
+    Note that the resource acquisition and release should usually be protected
+    in one @shield block but not separately, e.g.:
+
+        lock = filelock.AsyncFileLock('lock')
+
+        @shield
+        async def acquire():
+            await lock.acquire()
+
+        @shield
+        async def release():
+            await lock.release()
+
+        async def fn_with_lock():
+            await acquire()
+            try:
+                do_something()
+            finally:
+                await release()
+
+    The above code is not safe because if `fn_with_lock` is cancelled,
+    `acquire()` and `release()` will be executed in the background
+    concurrently and causes race conditions.
+    """
+
+    @functools.wraps(func)
+    async def async_wrapper(*args, **kwargs):
+        task = asyncio.create_task(func(*args, **kwargs))
+        try:
+            return await asyncio.shield(task)
+        except asyncio.CancelledError:
+            _background_tasks.add(task)
+            task.add_done_callback(lambda _: _background_tasks.discard(task))
+            raise
+
+    return async_wrapper