skypilot-nightly 1.0.0.dev20250604__py3-none-any.whl → 1.0.0.dev20250606__py3-none-any.whl

sky/serve/load_balancer.py

@@ -2,7 +2,8 @@
 import asyncio
 import logging
 import threading
-from typing import Dict, Optional, Union
+import traceback
+from typing import Dict, List, Optional, Union
 
 import aiohttp
 import fastapi
@@ -69,6 +70,48 @@ class SkyServeLoadBalancer:
         # updating it from _sync_with_controller.
         self._client_pool_lock: threading.Lock = threading.Lock()
 
+    async def _sync_with_controller_once(self) -> List[asyncio.Task]:
+        close_client_tasks = []
+        async with aiohttp.ClientSession() as session:
+            try:
+                # Send request information
+                async with session.post(
+                        self._controller_url + '/controller/load_balancer_sync',
+                        json={
+                            'request_aggregator':
+                                self._request_aggregator.to_dict()
+                        },
+                        timeout=aiohttp.ClientTimeout(5),
+                ) as response:
+                    # Clean up after reporting request info to avoid OOM.
+                    self._request_aggregator.clear()
+                    response.raise_for_status()
+                    response_json = await response.json()
+                    ready_replica_urls = response_json.get(
+                        'ready_replica_urls', [])
+            except (aiohttp.ClientError, asyncio.TimeoutError) as e:
+                logger.error(f'An error occurred when syncing with '
+                             f'the controller: {e}'
+                             f'\nTraceback: {traceback.format_exc()}')
+            else:
+                logger.info(f'Available Replica URLs: {ready_replica_urls}')
+                with self._client_pool_lock:
+                    self._load_balancing_policy.set_ready_replicas(
+                        ready_replica_urls)
+                    for replica_url in ready_replica_urls:
+                        if replica_url not in self._client_pool:
+                            self._client_pool[replica_url] = httpx.AsyncClient(
+                                base_url=replica_url)
+                    urls_to_close = set(
+                        self._client_pool.keys()) - set(ready_replica_urls)
+                    client_to_close = []
+                    for replica_url in urls_to_close:
+                        client_to_close.append(
+                            self._client_pool.pop(replica_url))
+                for client in client_to_close:
+                    close_client_tasks.append(client.aclose())
+        return close_client_tasks
+
     async def _sync_with_controller(self):
         """Sync with controller periodically.
 
@@ -82,49 +125,16 @@ class SkyServeLoadBalancer:
         await asyncio.sleep(5)
 
         while True:
-            close_client_tasks = []
-            async with aiohttp.ClientSession() as session:
-                try:
-                    # Send request information
-                    async with session.post(
-                            self._controller_url +
-                            '/controller/load_balancer_sync',
-                            json={
-                                'request_aggregator':
-                                    self._request_aggregator.to_dict()
-                            },
-                            timeout=aiohttp.ClientTimeout(5),
-                    ) as response:
-                        # Clean up after reporting request info to avoid OOM.
-                        self._request_aggregator.clear()
-                        response.raise_for_status()
-                        response_json = await response.json()
-                        ready_replica_urls = response_json.get(
-                            'ready_replica_urls', [])
-                except aiohttp.ClientError as e:
-                    logger.error('An error occurred when syncing with '
-                                 f'the controller: {e}')
-                else:
-                    logger.info(f'Available Replica URLs: {ready_replica_urls}')
-                    with self._client_pool_lock:
-                        self._load_balancing_policy.set_ready_replicas(
-                            ready_replica_urls)
-                        for replica_url in ready_replica_urls:
-                            if replica_url not in self._client_pool:
-                                self._client_pool[replica_url] = (
-                                    httpx.AsyncClient(base_url=replica_url))
-                        urls_to_close = set(
-                            self._client_pool.keys()) - set(ready_replica_urls)
-                        client_to_close = []
-                        for replica_url in urls_to_close:
-                            client_to_close.append(
-                                self._client_pool.pop(replica_url))
-                    for client in client_to_close:
-                        close_client_tasks.append(client.aclose())
-
-            await asyncio.sleep(constants.LB_CONTROLLER_SYNC_INTERVAL_SECONDS)
-            # Await those tasks after the interval to avoid blocking.
-            await asyncio.gather(*close_client_tasks)
+            try:
+                close_client_tasks = await self._sync_with_controller_once()
+                await asyncio.sleep(
+                    constants.LB_CONTROLLER_SYNC_INTERVAL_SECONDS)
+                # Await those tasks after the interval to avoid blocking.
+                await asyncio.gather(*close_client_tasks)
+            except Exception as e:  # pylint: disable=broad-except
+                logger.error(f'An error occurred when syncing with '
+                             f'the controller: {e}'
+                             f'\nTraceback: {traceback.format_exc()}')
 
     async def _proxy_request_to(
         self, url: str, request: fastapi.Request
@@ -168,7 +178,8 @@ class SkyServeLoadBalancer:
                 background=background.BackgroundTask(background_func))
         except (httpx.RequestError, httpx.HTTPStatusError) as e:
             logger.error(f'Error when proxy request to {url}: '
-                         f'{common_utils.format_exception(e)}')
+                         f'{common_utils.format_exception(e)}'
+                         f'\nTraceback: {traceback.format_exc()}')
             return e
 
     async def _proxy_with_retries(

sky/server/constants.py

@@ -7,7 +7,7 @@ from sky.skylet import constants
 # API server version, whenever there is a change in API server that requires a
 # restart of the local API server or error out when the client does not match
 # the server version.
-API_VERSION = '8'
+API_VERSION = '9'
 
 # Prefix for API request names.
 REQUEST_NAME_PREFIX = 'sky.'
@@ -25,8 +25,10 @@ API_SERVER_REQUEST_DB_PATH = '~/.sky/api_server/requests.db'
 CLUSTER_REFRESH_DAEMON_INTERVAL_SECONDS = 60
 
 # Environment variable for a file path to the API cookie file.
+# Keep in sync with websocket_proxy.py
 API_COOKIE_FILE_ENV_VAR = f'{constants.SKYPILOT_ENV_VAR_PREFIX}API_COOKIE_FILE'
 # Default file if unset.
+# Keep in sync with websocket_proxy.py
 API_COOKIE_FILE_DEFAULT_LOCATION = '~/.sky/cookies.txt'
 
 # The path to the dashboard build output

sky/server/requests/payloads.py

@@ -196,8 +196,10 @@ class LaunchBody(RequestBody):
     task: str
     cluster_name: str
     retry_until_up: bool = False
+    # TODO(aylei): remove this field in v0.12.0
     idle_minutes_to_autostop: Optional[int] = None
     dryrun: bool = False
+    # TODO(aylei): remove this field in v0.12.0
     down: bool = False
     backend: Optional[str] = None
     optimize_target: common_lib.OptimizeTarget = common_lib.OptimizeTarget.COST
@@ -331,6 +333,12 @@ class ClusterJobsDownloadLogsBody(RequestBody):
     local_dir: str = constants.SKY_LOGS_DIRECTORY
 
 
+class UserUpdateBody(RequestBody):
+    """The request body for the user update endpoint."""
+    user_id: str
+    role: str
+
+
 class DownloadBody(RequestBody):
     """The request body for the download endpoint."""
     folder_paths: List[str]
@@ -375,6 +383,7 @@ class JobsQueueBody(RequestBody):
     refresh: bool = False
     skip_finished: bool = False
     all_users: bool = False
+    job_ids: Optional[List[int]] = None
 
 
 class JobsCancelBody(RequestBody):

sky/server/server.py

@@ -49,6 +49,7 @@ from sky.server.requests import preconditions
 from sky.server.requests import requests as requests_lib
 from sky.skylet import constants
 from sky.usage import usage_lib
+from sky.users import server as users_rest
 from sky.utils import admin_policy_utils
 from sky.utils import common as common_lib
 from sky.utils import common_utils
@@ -100,6 +101,27 @@ logger = sky_logging.init_logger(__name__)
 # response will block other requests from being processed.
 
 
+class RBACMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
+    """Middleware to handle RBAC."""
+
+    async def dispatch(self, request: fastapi.Request, call_next):
+        if request.url.path.startswith('/dashboard/'):
+            return await call_next(request)
+
+        auth_user = _get_auth_user_header(request)
+        if auth_user is None:
+            return await call_next(request)
+
+        permission_service = users_rest.permission_service
+        # Check the role permission
+        if permission_service.check_permission(auth_user.id, request.url.path,
+                                               request.method):
+            return fastapi.responses.JSONResponse(
+                status_code=403, content={'detail': 'Forbidden'})
+
+        return await call_next(request)
+
+
 class RequestIDMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
     """Middleware to add a request ID to each request."""
 
@@ -130,7 +152,10 @@ class AuthProxyMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
 
         # Add user to database if auth_user is present
         if auth_user is not None:
-            global_user_state.add_or_update_user(auth_user)
+            newly_added = global_user_state.add_or_update_user(auth_user)
+            if newly_added:
+                users_rest.permission_service.add_user_if_not_exists(
+                    auth_user.id)
 
         body = await request.body()
         if auth_user and body:
@@ -244,11 +269,13 @@ class PathCleanMiddleware(starlette.middleware.base.BaseHTTPMiddleware):
             parent = pathlib.Path('/dashboard')
             request_path = pathlib.Path(posixpath.normpath(request.url.path))
             if not _is_relative_to(request_path, parent):
-                raise fastapi.HTTPException(status_code=403, detail='Forbidden')
+                return fastapi.responses.JSONResponse(
+                    status_code=403, content={'detail': 'Forbidden'})
         return await call_next(request)
 
 
 app = fastapi.FastAPI(prefix='/api/v1', debug=True, lifespan=lifespan)
+app.add_middleware(RBACMiddleware)
 app.add_middleware(InternalDashboardPrefixMiddleware)
 app.add_middleware(PathCleanMiddleware)
 app.add_middleware(CacheControlStaticMiddleware)
@@ -266,6 +293,7 @@ app.add_middleware(AuthProxyMiddleware)
 app.add_middleware(RequestIDMiddleware)
 app.include_router(jobs_rest.router, prefix='/jobs', tags=['jobs'])
 app.include_router(serve_rest.router, prefix='/serve', tags=['serve'])
+app.include_router(users_rest.router, prefix='/users', tags=['users'])
 app.include_router(workspaces_rest.router,
                    prefix='/workspaces',
                    tags=['workspaces'])
@@ -835,13 +863,6 @@ async def logs(
     )
 
 
-@app.get('/users')
-async def users() -> List[Dict[str, Any]]:
-    """Gets all users."""
-    user_list = global_user_state.get_all_users()
-    return [user.to_dict() for user in user_list]
-
-
 @app.post('/download_logs')
 async def download_logs(
         request: fastapi.Request,

sky/setup_files/MANIFEST.in

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

sky/setup_files/dependencies.py

@@ -58,6 +58,8 @@ install_requires = [
     'setproctitle',
     'sqlalchemy',
     'psycopg2-binary',
+    'casbin',
+    'sqlalchemy_adapter',
 ]
 
 local_ray = [

sky/skylet/constants.py

@@ -379,7 +379,7 @@ OVERRIDEABLE_CONFIG_KEYS_IN_TASK: List[Tuple[str, ...]] = [
 SKIPPED_CLIENT_OVERRIDE_KEYS: List[Tuple[str, ...]] = [('admin_policy',),
                                                        ('api_server',),
                                                        ('allowed_clouds',),
-                                                       ('workspaces',)]
+                                                       ('workspaces',), ('db',)]
 
 # Constants for Azure blob storage
 WAIT_FOR_STORAGE_ACCOUNT_CREATION = 60
@@ -409,6 +409,12 @@ ENV_VAR_IS_SKYPILOT_SERVER = 'IS_SKYPILOT_SERVER'
 
 SKYPILOT_DEFAULT_WORKSPACE = 'default'
 
-# Experimental - may be deprecated in the future without notice.
-SKYPILOT_API_SERVER_DB_URL_ENV_VAR: str = (
-    f'{SKYPILOT_ENV_VAR_PREFIX}API_SERVER_DB_URL')
+# BEGIN constants used for service catalog.
+HOSTED_CATALOG_DIR_URL = 'https://raw.githubusercontent.com/skypilot-org/skypilot-catalog/master/catalogs'  # pylint: disable=line-too-long
+HOSTED_CATALOG_DIR_URL_S3_MIRROR = 'https://skypilot-catalog.s3.us-east-1.amazonaws.com/catalogs'  # pylint: disable=line-too-long
+CATALOG_SCHEMA_VERSION = 'v7'
+CATALOG_DIR = '~/.sky/catalogs'
+ALL_CLOUDS = ('aws', 'azure', 'gcp', 'ibm', 'lambda', 'scp', 'oci',
+              'kubernetes', 'runpod', 'vast', 'vsphere', 'cudo', 'fluidstack',
+              'paperspace', 'do', 'nebius', 'ssh')
+# END constants used for service catalog.

sky/skypilot_config.py

@@ -756,13 +756,15 @@ def apply_cli_config(cli_config: Optional[List[str]]) -> Dict[str, Any]:
     return parsed_config
 
 
-def update_config_no_lock(config: config_utils.Config) -> None:
+def update_api_server_config_no_lock(config: config_utils.Config) -> None:
     """Dumps the new config to a file and syncs to ConfigMap if in Kubernetes.
 
     Args:
         config: The config to save and sync.
     """
-    global_config_path = os.path.expanduser(get_user_config_path())
+    global_config_path = _resolve_server_config_path()
+    if global_config_path is None:
+        global_config_path = get_user_config_path()
 
     # Always save to the local file (PVC in Kubernetes, local file otherwise)
     common_utils.dump_yaml(global_config_path, dict(config))

sky/templates/websocket_proxy.py

@@ -21,11 +21,21 @@ from websockets.asyncio.client import connect
 
 BUFFER_SIZE = 2**16  # 64KB
 
+# Environment variable for a file path to the API cookie file.
+# Keep in sync with server/constants.py
+API_COOKIE_FILE_ENV_VAR = 'SKYPILOT_API_COOKIE_FILE'
+# Default file if unset.
+# Keep in sync with server/constants.py
+API_COOKIE_FILE_DEFAULT_LOCATION = '~/.sky/cookies.txt'
+
 
 def _get_cookie_header(url: str) -> Dict[str, str]:
     """Extract Cookie header value from a cookie jar for a specific URL"""
-    cookie_path = os.environ.get('SKYPILOT_API_COOKIE_FILE')
+    cookie_path = os.environ.get(API_COOKIE_FILE_ENV_VAR)
     if cookie_path is None:
+        cookie_path = API_COOKIE_FILE_DEFAULT_LOCATION
+    cookie_path = os.path.expanduser(cookie_path)
+    if not os.path.exists(cookie_path):
         return {}
 
     request = Request(url)

sky/users/model.conf

@@ -0,0 +1,15 @@
+# rbac_model.conf
+[request_definition]
+r = sub, obj, act
+
+[policy_definition]
+p = sub, obj, act
+
+[role_definition]
+g = _, _
+
+[policy_effect]
+e = some(where (p.eft == allow))
+
+[matchers]
+m = g(r.sub, p.sub) && r.obj == p.obj && r.act == p.act

sky/users/permission.py

@@ -0,0 +1,178 @@
+"""Permission service for SkyPilot API Server."""
+import contextlib
+import logging
+import os
+import threading
+from typing import List
+
+import casbin
+import filelock
+import sqlalchemy_adapter
+
+from sky import global_user_state
+from sky import sky_logging
+from sky.users import rbac
+
+logger = sky_logging.init_logger(__name__)
+
+# Filelocks for the policy update.
+POLICY_UPDATE_LOCK_PATH = os.path.expanduser('~/.sky/.policy_update.lock')
+POLICY_UPDATE_LOCK_TIMEOUT_SECONDS = 20
+
+_enforcer_instance = None
+_lock = threading.Lock()
+
+
+class PermissionService:
+    """Permission service for SkyPilot API Server."""
+
+    def __init__(self):
+        global _enforcer_instance
+        if _enforcer_instance is None:
+            # For different threads, we share the same enforcer instance.
+            with _lock:
+                if _enforcer_instance is None:
+                    _enforcer_instance = self
+                    engine = global_user_state.SQLALCHEMY_ENGINE
+                    adapter = sqlalchemy_adapter.Adapter(engine)
+                    model_path = os.path.join(os.path.dirname(__file__),
+                                              'model.conf')
+                    enforcer = casbin.Enforcer(model_path, adapter)
+                    logging.getLogger('casbin.policy').setLevel(
+                        sky_logging.ERROR)
+                    logging.getLogger('casbin.role').setLevel(sky_logging.ERROR)
+                    self.enforcer = enforcer
+        else:
+            self.enforcer = _enforcer_instance.enforcer
+        self._maybe_initialize_policies()
+
+    def _maybe_initialize_policies(self):
+        """Initialize policies if they don't already exist."""
+        logger.debug(f'Initializing policies in process: {os.getpid()}')
+
+        # Check if policies are already initialized by looking for existing
+        # permission policies in the enforcer
+        existing_policies = self.enforcer.get_policy()
+
+        # If we already have policies for the expected roles, skip
+        # initialization
+        role_permissions = rbac.get_role_permissions()
+        expected_policies = []
+        for role, permissions in role_permissions.items():
+            if permissions['permissions'] and 'blocklist' in permissions[
+                    'permissions']:
+                blocklist = permissions['permissions']['blocklist']
+                for item in blocklist:
+                    expected_policies.append(
+                        [role, item['path'], item['method']])
+
+        # Check if all expected policies already exist
+        policies_exist = all(
+            any(policy == expected
+                for policy in existing_policies)
+            for expected in expected_policies)
+
+        if not policies_exist:
+            # Only clear and reinitialize if policies don't exist or are
+            # incomplete
+            logger.debug('Policies not found or incomplete, initializing...')
+            # Only clear p policies (permission policies),
+            # keep g policies (role policies)
+            self.enforcer.remove_filtered_policy(0)
+            for role, permissions in role_permissions.items():
+                if permissions['permissions'] and 'blocklist' in permissions[
+                        'permissions']:
+                    blocklist = permissions['permissions']['blocklist']
+                    for item in blocklist:
+                        path = item['path']
+                        method = item['method']
+                        self.enforcer.add_policy(role, path, method)
+            self.enforcer.save_policy()
+        else:
+            logger.debug('Policies already exist, skipping initialization')
+
+        # Always ensure users have default roles (this is idempotent)
+        all_users = global_user_state.get_all_users()
+        for user in all_users:
+            self.add_user_if_not_exists(user.id)
+
+    def add_user_if_not_exists(self, user: str) -> None:
+        """Add user role relationship."""
+        with _policy_lock():
+            user_roles = self.enforcer.get_roles_for_user(user)
+            if not user_roles:
+                logger.info(f'User {user} has no roles, adding'
+                            f' default role {rbac.get_default_role()}')
+                self.enforcer.add_grouping_policy(user, rbac.get_default_role())
+                self.enforcer.save_policy()
+
+    def update_role(self, user: str, new_role: str):
+        """Update user role relationship."""
+        with _policy_lock():
+            # Get current roles
+            self._load_policy_no_lock()
+            # Avoid calling get_user_roles, as it will require the lock.
+            current_roles = self.enforcer.get_roles_for_user(user)
+            if not current_roles:
+                logger.warning(f'User {user} has no roles')
+            else:
+                # TODO(hailong): how to handle multiple roles?
+                current_role = current_roles[0]
+                if current_role == new_role:
+                    logger.info(f'User {user} already has role {new_role}')
+                    return
+                self.enforcer.remove_grouping_policy(user, current_role)
+
+            # Update user role
+            self.enforcer.add_grouping_policy(user, new_role)
+            self.enforcer.save_policy()
+
+    def get_user_roles(self, user: str) -> List[str]:
+        """Get all roles for a user.
+
+        This method returns all roles that the user has, including inherited
+        roles. For example, if a user has role 'admin' and 'admin' inherits
+        from 'user', this method will return ['admin', 'user'].
+
+        Args:
+            user: The user ID to get roles for.
+
+        Returns:
+            A list of role names that the user has.
+        """
+        self._load_policy()
+        return self.enforcer.get_roles_for_user(user)
+
+    def check_permission(self, user: str, path: str, method: str) -> bool:
+        """Check permission."""
+        # We intentionally don't load the policy here, as it is a hot path, and
+        # we don't support updating the policy.
+        # We don't hold the lock for checking permission, as it is read only and
+        # it is a hot path in every request. It is ok to have a stale policy,
+        # as long as it is eventually consistent.
+        # self._load_policy_no_lock()
+        return self.enforcer.enforce(user, path, method)
+
+    def _load_policy_no_lock(self):
+        """Load policy from storage."""
+        self.enforcer.load_policy()
+
+    def _load_policy(self):
+        """Load policy from storage with lock."""
+        with _policy_lock():
+            self._load_policy_no_lock()
+
+
+@contextlib.contextmanager
+def _policy_lock():
+    """Context manager for policy update lock."""
+    try:
+        with filelock.FileLock(POLICY_UPDATE_LOCK_PATH,
+                               POLICY_UPDATE_LOCK_TIMEOUT_SECONDS):
+            yield
+    except filelock.Timeout as e:
+        raise RuntimeError(f'Failed to load policy due to a timeout '
+                           f'when trying to acquire the lock at '
+                           f'{POLICY_UPDATE_LOCK_PATH}. '
+                           'Please try again or manually remove the lock '
+                           f'file if you believe it is stale.') from e

sky/users/rbac.py

@@ -0,0 +1,86 @@
+"""RBAC (Role-Based Access Control) functionality for SkyPilot API Server."""
+
+import enum
+from typing import Dict, List
+
+from sky import sky_logging
+from sky import skypilot_config
+
+logger = sky_logging.init_logger(__name__)
+
+# Default user blocklist for user role
+# Cannot access workspace CUD operations
+_DEFAULT_USER_BLOCKLIST = [{
+    'path': '/workspaces/config',
+    'method': 'POST'
+}, {
+    'path': '/workspaces/update',
+    'method': 'POST'
+}, {
+    'path': '/workspaces/create',
+    'method': 'POST'
+}, {
+    'path': '/workspaces/delete',
+    'method': 'POST'
+}, {
+    'path': '/users/update',
+    'method': 'POST'
+}]
+
+
+# Define roles
+class RoleName(str, enum.Enum):
+    ADMIN = 'admin'
+    USER = 'user'
+
+
+def get_supported_roles() -> List[str]:
+    return [role_name.value for role_name in RoleName]
+
+
+def get_default_role() -> str:
+    return skypilot_config.get_nested(('rbac', 'default_role'),
+                                      default_value=RoleName.ADMIN.value)
+
+
+def get_role_permissions(
+) -> Dict[str, Dict[str, Dict[str, List[Dict[str, str]]]]]:
+    """Get all role permissions from config.
+
+    Returns:
+        Dictionary containing all roles and their permissions configuration.
+        Example:
+        {
+            'admin': {
+                'permissions': {
+                    'blocklist': []
+                }
+            },
+            'user': {
+                'permissions': {
+                    'blocklist': [
+                        {'path': '/workspaces/config', 'method': 'POST'},
+                        {'path': '/workspaces/update', 'method': 'POST'}
+                    ]
+                }
+            }
+        }
+    """
+    # Get all roles from the config
+    config_permissions = skypilot_config.get_nested(('rbac', 'roles'),
+                                                    default_value={})
+    supported_roles = get_supported_roles()
+    for role, permissions in config_permissions.items():
+        role_name = role.lower()
+        if role_name not in supported_roles:
+            logger.warning(f'Invalid role: {role_name}')
+            continue
+        config_permissions[role_name] = permissions
+    # Add default roles if not present
+    if 'user' not in config_permissions:
+        config_permissions['user'] = {
+            'permissions': {
+                'blocklist': _DEFAULT_USER_BLOCKLIST
+            }
+        }
+    return config_permissions