skypilot-nightly 1.0.0.dev20250607__py3-none-any.whl → 1.0.0.dev20250609__py3-none-any.whl

sky/users/rbac.py

@@ -5,6 +5,8 @@ from typing import Dict, List
 
 from sky import sky_logging
 from sky import skypilot_config
+from sky.skylet import constants
+from sky.workspaces import utils as workspaces_utils
 
 logger = sky_logging.init_logger(__name__)
 
@@ -84,3 +86,27 @@ def get_role_permissions(
             }
         }
     return config_permissions
+
+
+def get_workspace_policy_permissions() -> Dict[str, List[str]]:
+    """Get workspace policy permissions from config.
+
+    Returns:
+        A dictionary of workspace policy permissions.
+        Example:
+        {
+            'workspace1': ['user1-id', 'user2-id'],
+            'workspace2': ['user3-id', 'user4-id']
+            'default': ['*']
+        }
+    """
+    current_workspaces = skypilot_config.get_nested(('workspaces',),
+                                                    default_value={})
+    if constants.SKYPILOT_DEFAULT_WORKSPACE not in current_workspaces:
+        current_workspaces[constants.SKYPILOT_DEFAULT_WORKSPACE] = {}
+    workspaces_to_policy = {}
+    for workspace_name, workspace_config in current_workspaces.items():
+        users = workspaces_utils.get_workspace_users(workspace_config)
+        workspaces_to_policy[workspace_name] = users
+    logger.debug(f'Workspace policy permissions: {workspaces_to_policy}')
+    return workspaces_to_policy

sky/users/server.py

@@ -1,6 +1,5 @@
 """REST API for workspace management."""
 
-import hashlib
 from typing import Any, Dict, List
 
 import fastapi
@@ -8,16 +7,15 @@ import fastapi
 from sky import global_user_state
 from sky import sky_logging
 from sky.server.requests import payloads
+from sky.skylet import constants
 from sky.users import permission
 from sky.users import rbac
-from sky.utils import common_utils
+from sky.utils import common
 
 logger = sky_logging.init_logger(__name__)
 
 router = fastapi.APIRouter()
 
-permission_service = permission.PermissionService()
-
 
 @router.get('')
 async def users() -> List[Dict[str, Any]]:
@@ -25,7 +23,7 @@ async def users() -> List[Dict[str, Any]]:
     all_users = []
     user_list = global_user_state.get_all_users()
     for user in user_list:
-        user_roles = permission_service.get_user_roles(user.id)
+        user_roles = permission.permission_service.get_user_roles(user.id)
         all_users.append({
             'id': user.id,
             'name': user.name,
@@ -39,13 +37,11 @@ async def get_current_user_role(request: fastapi.Request):
     """Get current user's role."""
     # TODO(hailong): is there a reliable way to get the user
     # hash for the request without 'X-Auth-Request-Email' header?
-    if 'X-Auth-Request-Email' not in request.headers:
+    auth_user = request.state.auth_user
+    if auth_user is None:
         return {'name': '', 'role': rbac.RoleName.ADMIN.value}
-    user_name = request.headers['X-Auth-Request-Email']
-    user_hash = hashlib.md5(
-        user_name.encode()).hexdigest()[:common_utils.USER_HASH_LENGTH]
-    user_roles = permission_service.get_user_roles(user_hash)
-    return {'name': user_name, 'role': user_roles[0] if user_roles else ''}
+    user_roles = permission.permission_service.get_user_roles(auth_user.id)
+    return {'name': auth_user.name, 'role': user_roles[0] if user_roles else ''}
 
 
 @router.post('/update')
@@ -58,9 +54,14 @@ async def user_update(user_update_body: payloads.UserUpdateBody) -> None:
         raise fastapi.HTTPException(status_code=400,
                                     detail=f'Invalid role: {role}')
     user_info = global_user_state.get_user(user_id)
-    if not user_info.name:
+    if user_info is None:
         raise fastapi.HTTPException(status_code=400,
                                     detail=f'User {user_id} does not exist')
+    # Disallow updating roles for the internal users.
+    if user_info.id in [common.SERVER_ID, constants.SKYPILOT_SYSTEM_USER_ID]:
+        raise fastapi.HTTPException(status_code=400,
+                                    detail=f'Cannot update role for internal '
+                                    f'API server user {user_info.name}')
 
     # Update user role in casbin policy
-    permission_service.update_role(user_id, role)
+    permission.permission_service.update_role(user_info.id, role)

sky/utils/common.py

@@ -5,6 +5,7 @@ import enum
 import os
 from typing import Generator
 
+from sky import models
 from sky.skylet import constants
 from sky.utils import common_utils
 
@@ -25,10 +26,13 @@ JOB_CONTROLLER_NAME: str = f'{JOB_CONTROLLER_PREFIX}{SERVER_ID}'
 
 
 @contextlib.contextmanager
-def with_server_user_hash() -> Generator[None, None, None]:
+def with_server_user() -> Generator[None, None, None]:
     """Temporarily set the user hash to common.SERVER_ID."""
     old_env_user_hash = os.getenv(constants.USER_ID_ENV_VAR)
+    # TODO(zhwu): once we have fully moved our code to use `get_current_user()`
+    # instead of `common_utils.get_user_hash()`, we can remove the env override.
     os.environ[constants.USER_ID_ENV_VAR] = SERVER_ID
+    common_utils.set_current_user(models.User.get_current_user())
     try:
         yield
     finally:
@@ -36,6 +40,7 @@ def with_server_user_hash() -> Generator[None, None, None]:
             os.environ[constants.USER_ID_ENV_VAR] = old_env_user_hash
         else:
             os.environ.pop(constants.USER_ID_ENV_VAR)
+        common_utils.set_current_user(models.User.get_current_user())
 
 
 class StatusRefreshMode(enum.Enum):

sky/utils/common_utils.py

@@ -20,6 +20,7 @@ import uuid
 import jsonschema
 
 from sky import exceptions
+from sky import models
 from sky import sky_logging
 from sky.adaptors import common as adaptors_common
 from sky.skylet import constants
@@ -256,11 +257,13 @@ class Backoff:
 _current_command: Optional[str] = None
 _current_client_entrypoint: Optional[str] = None
 _using_remote_api_server: Optional[bool] = None
+_current_user: Optional['models.User'] = None
 
 
-def set_client_status(client_entrypoint: Optional[str],
-                      client_command: Optional[str],
-                      using_remote_api_server: bool):
+def set_request_context(client_entrypoint: Optional[str],
+                        client_command: Optional[str],
+                        using_remote_api_server: bool,
+                        user: Optional['models.User']):
     """Override the current client entrypoint and command.
 
     This is useful when we are on the SkyPilot API server side and we have a
@@ -269,9 +272,11 @@ def set_client_status(client_entrypoint: Optional[str],
     global _current_command
     global _current_client_entrypoint
     global _using_remote_api_server
+    global _current_user
     _current_command = client_command
     _current_client_entrypoint = client_entrypoint
     _using_remote_api_server = using_remote_api_server
+    _current_user = user
 
 
 def get_current_command() -> str:
@@ -286,6 +291,19 @@ def get_current_command() -> str:
     return get_pretty_entrypoint_cmd()
 
 
+def get_current_user() -> 'models.User':
+    """Returns the current user."""
+    if _current_user is not None:
+        return _current_user
+    return models.User.get_current_user()
+
+
+def set_current_user(user: 'models.User'):
+    """Sets the current user."""
+    global _current_user
+    _current_user = user
+
+
 def get_current_client_entrypoint(server_entrypoint: str) -> str:
     """Returns the current client entrypoint.
 

sky/utils/schemas.py

@@ -1249,6 +1249,15 @@ def get_config_schema():
             'properties': {
                 # Explicit definition for GCP allows both project_id and
                 # disabled
+                'private': {
+                    'type': 'boolean',
+                },
+                'allowed_users': {
+                    'type': 'array',
+                    'items': {
+                        'type': 'string',
+                    },
+                },
                 'gcp': {
                     'type': 'object',
                     'properties': {

sky/workspaces/core.py

@@ -1,20 +1,24 @@
 """Workspace management core."""
 
 import concurrent.futures
-from typing import Any, Callable, Dict
+from typing import Any, Callable, Dict, List
 
 import filelock
 
 from sky import check as sky_check
 from sky import exceptions
 from sky import global_user_state
+from sky import models
 from sky import sky_logging
 from sky import skypilot_config
 from sky.skylet import constants
 from sky.usage import usage_lib
+from sky.users import permission
+from sky.utils import annotations
 from sky.utils import common_utils
 from sky.utils import config_utils
 from sky.utils import schemas
+from sky.workspaces import utils as workspaces_utils
 
 logger = sky_logging.init_logger(__name__)
 
@@ -28,10 +32,7 @@ _WORKSPACE_CONFIG_LOCK_TIMEOUT_SECONDS = 60
 
 def get_workspaces() -> Dict[str, Any]:
     """Returns the workspace config."""
-    workspaces = skypilot_config.get_nested(('workspaces',), default_value={})
-    if constants.SKYPILOT_DEFAULT_WORKSPACE not in workspaces:
-        workspaces[constants.SKYPILOT_DEFAULT_WORKSPACE] = {}
-    return workspaces
+    return workspaces_for_user(common_utils.get_current_user().id)
 
 
 def _update_workspaces_config(
@@ -160,7 +161,7 @@ def _check_workspaces_have_no_active_resources(
                 f'{len(workspace_clusters)} active cluster(s): {cluster_list}')
 
         if workspace_active_jobs:
-            job_names = [job['job_id'] for job in workspace_active_jobs]
+            job_names = [str(job['job_id']) for job in workspace_active_jobs]
             job_list = ', '.join(job_names)
             workspace_errors.append(
                 f'{len(workspace_active_jobs)} active managed job(s): '
@@ -223,14 +224,19 @@ def update_workspace(workspace_name: str, config: Dict[str,
         FileNotFoundError: If the config file cannot be found.
         PermissionError: If the config file cannot be written.
     """
+    _validate_workspace_config(workspace_name, config)
+
     # Check for active clusters and managed jobs in the workspace
+    # TODO(zhwu): we should allow the edits that only contain changes to
+    # allowed_users or private.
     _check_workspace_has_no_active_resources(workspace_name, 'update')
 
-    _validate_workspace_config(workspace_name, config)
-
     def update_workspace_fn(workspaces: Dict[str, Any]) -> None:
         """Function to update workspace inside the lock."""
         workspaces[workspace_name] = config
+        users = workspaces_utils.get_workspace_users(config)
+        permission_service = permission.permission_service
+        permission_service.update_workspace_policy(workspace_name, users)
 
     # Use the internal helper function to save
     result = _update_workspaces_config(update_workspace_fn)
@@ -275,6 +281,10 @@ def create_workspace(workspace_name: str, config: Dict[str,
             raise ValueError(f'Workspace {workspace_name!r} already exists. '
                              'Use update instead.')
         workspaces[workspace_name] = config
+        # Add policy for the workspace and allowed users
+        users = workspaces_utils.get_workspace_users(config)
+        permission_service = permission.permission_service
+        permission_service.add_workspace_policy(workspace_name, users)
 
     # Use the internal helper function to save
     result = _update_workspaces_config(create_workspace_fn)
@@ -324,6 +334,8 @@ def delete_workspace(workspace_name: str) -> Dict[str, Any]:
         if workspace_name not in workspaces:
             raise ValueError(f'Workspace {workspace_name!r} does not exist.')
         del workspaces[workspace_name]
+        permission_service = permission.permission_service
+        permission_service.remove_workspace_policy(workspace_name)
 
     # Use the internal helper function to save
     return _update_workspaces_config(delete_workspace_fn)
@@ -369,6 +381,9 @@ def update_config(config: Dict[str, Any]) -> Dict[str, Any]:
 
     # Check for API server changes and validate them
     current_config = skypilot_config.to_dict()
+    # If there is no changes to the config, we can return early
+    if current_config == config:
+        return config
 
     current_endpoint = current_config.get('api_server', {}).get('endpoint')
     new_endpoint = config.get('api_server', {}).get('endpoint')
@@ -382,15 +397,27 @@ def update_config(config: Dict[str, Any]) -> Dict[str, Any]:
 
     # Collect all workspaces that need to be checked for active resources
     workspaces_to_check = []
+    workspaces_to_check_policy: Dict[str, Dict[str, List[str]]] = {
+        'add': {},
+        'update': {},
+        'delete': {}
+    }
 
     # Check each workspace that is being modified
     for workspace_name, new_workspace_config in new_workspaces.items():
+        if workspace_name not in current_workspaces:
+            users = workspaces_utils.get_workspace_users(new_workspace_config)
+            workspaces_to_check_policy['add'][workspace_name] = users
+            continue
+
         current_workspace_config = current_workspaces.get(workspace_name, {})
 
         # If workspace configuration is changing, validate and mark for checking
         if current_workspace_config != new_workspace_config:
             _validate_workspace_config(workspace_name, new_workspace_config)
             workspaces_to_check.append((workspace_name, 'update'))
+            users = workspaces_utils.get_workspace_users(new_workspace_config)
+            workspaces_to_check_policy['update'][workspace_name] = users
 
     # Check for workspace deletions
     for workspace_name in current_workspaces:
@@ -400,6 +427,7 @@ def update_config(config: Dict[str, Any]) -> Dict[str, Any]:
                 raise ValueError(f'Cannot delete the default workspace '
                                  f'{constants.SKYPILOT_DEFAULT_WORKSPACE!r}.')
             workspaces_to_check.append((workspace_name, 'delete'))
+            workspaces_to_check_policy['delete'][workspace_name] = ['*']
 
     # Check all workspaces for active resources in one efficient call
     _check_workspaces_have_no_active_resources(workspaces_to_check)
@@ -412,6 +440,18 @@ def update_config(config: Dict[str, Any]) -> Dict[str, Any]:
             # Convert to config_utils.Config and save
             config_obj = config_utils.Config.from_dict(config)
             skypilot_config.update_api_server_config_no_lock(config_obj)
+            permission_service = permission.permission_service
+            for operation, workspaces in workspaces_to_check_policy.items():
+                for workspace_name, users in workspaces.items():
+                    if operation == 'add':
+                        permission_service.add_workspace_policy(
+                            workspace_name, users)
+                    elif operation == 'update':
+                        permission_service.update_workspace_policy(
+                            workspace_name, users)
+                    elif operation == 'delete':
+                        permission_service.remove_workspace_policy(
+                            workspace_name)
     except filelock.Timeout as e:
         raise RuntimeError(
             f'Failed to update configuration due to a timeout '
@@ -429,3 +469,55 @@ def update_config(config: Dict[str, Any]) -> Dict[str, Any]:
         # Don't fail the update if the check fails, just warn
 
     return config
+
+
+def reject_request_for_unauthorized_workspace(user: models.User) -> None:
+    """Rejects a request that has no permission to access active workspace.
+
+    Args:
+        user: The user making the request.
+
+    Raises:
+        PermissionDeniedError: If the user does not have permission to access
+            the active workspace.
+    """
+    active_workspace = skypilot_config.get_active_workspace()
+    if not permission.permission_service.check_workspace_permission(
+            user.id, active_workspace):
+        raise exceptions.PermissionDeniedError(
+            f'User {user.name} ({user.id}) does not have '
+            f'permission to access workspace {active_workspace!r}')
+
+
+def is_workspace_private(workspace_config: Dict[str, Any]) -> bool:
+    """Check if a workspace is private.
+
+    Args:
+        workspace_config: The workspace configuration dictionary.
+
+    Returns:
+        True if the workspace is private, False if it's public.
+    """
+    return workspace_config.get('private', False)
+
+
+@annotations.lru_cache(scope='request', maxsize=1)
+def workspaces_for_user(user_id: str) -> Dict[str, Any]:
+    """Returns the workspaces that the user has access to.
+
+    Args:
+        user_id: The user id to check.
+
+    Returns:
+        A map from workspace name to workspace configuration.
+    """
+    workspaces = skypilot_config.get_nested(('workspaces',), default_value={})
+    if constants.SKYPILOT_DEFAULT_WORKSPACE not in workspaces:
+        workspaces[constants.SKYPILOT_DEFAULT_WORKSPACE] = {}
+    user_workspaces = {}
+
+    for workspace_name, workspace_config in workspaces.items():
+        if permission.permission_service.check_workspace_permission(
+                user_id, workspace_name):
+            user_workspaces[workspace_name] = workspace_config
+    return user_workspaces

sky/workspaces/server.py

@@ -14,10 +14,18 @@ router = fastapi.APIRouter()
 # pylint: disable=redefined-builtin
 async def get(request: fastapi.Request) -> None:
     """Gets workspace config on the server."""
+    # Have to manually inject user info into the request body because the
+    # request body is not available in the GET endpoint.
+    auth_user = request.state.auth_user
+    auth_user_env_vars_kwargs = {
+        'env_vars': auth_user.to_env_vars()
+    } if auth_user else {}
+    request_body = payloads.RequestBody(**auth_user_env_vars_kwargs)
+
     executor.schedule_request(
         request_id=request.state.request_id,
         request_name='workspaces.get',
-        request_body=payloads.RequestBody(),
+        request_body=request_body,
         func=core.get_workspaces,
         schedule_type=api_requests.ScheduleType.SHORT,
     )
@@ -65,10 +73,15 @@ async def delete(request: fastapi.Request,
 @router.get('/config')
 async def get_config(request: fastapi.Request) -> None:
     """Gets the entire SkyPilot configuration."""
+    auth_user = request.state.auth_user
+    auth_user_env_vars_kwargs = {
+        'env_vars': auth_user.to_env_vars()
+    } if auth_user else {}
+    get_config_body = payloads.GetConfigBody(**auth_user_env_vars_kwargs)
     executor.schedule_request(
         request_id=request.state.request_id,
         request_name='workspaces.get_config',
-        request_body=payloads.GetConfigBody(),
+        request_body=get_config_body,
         func=core.get_config,
         schedule_type=api_requests.ScheduleType.SHORT,
     )

sky/workspaces/utils.py

@@ -0,0 +1,56 @@
+"""Utils for workspaces."""
+import collections
+from typing import Any, Dict, List
+
+from sky import global_user_state
+from sky import sky_logging
+
+logger = sky_logging.init_logger(__name__)
+
+
+def get_workspace_users(workspace_config: Dict[str, Any]) -> List[str]:
+    """Get the users that should have access to a workspace.
+
+    workspace_config is a dict with the following keys:
+    - private: bool
+    - allowed_users: list of user names or IDs
+
+    This function will automatically resolve the user names to IDs.
+
+    Args:
+        workspace_config: The configuration of the workspace.
+
+    Returns:
+        List of user IDs that should have access to the workspace.
+        For private workspaces, returns specific user IDs.
+        For public workspaces, returns ['*'] to indicate all users.
+    """
+    if workspace_config.get('private', False):
+        user_ids = []
+        workspace_user_name_or_ids = workspace_config.get('allowed_users', [])
+        all_users = global_user_state.get_all_users()
+        all_user_ids = {user.id for user in all_users}
+        all_user_map = collections.defaultdict(list)
+        for user in all_users:
+            all_user_map[user.name].append(user.id)
+
+        # Resolve user names to IDs
+        for user_name_or_id in workspace_user_name_or_ids:
+            if user_name_or_id in all_user_ids:
+                user_ids.append(user_name_or_id)
+            elif user_name_or_id in all_user_map:
+                if len(all_user_map[user_name_or_id]) > 1:
+                    user_ids_str = ', '.join(all_user_map[user_name_or_id])
+                    raise ValueError(
+                        f'User {user_name_or_id!r} has multiple IDs: '
+                        f'{user_ids_str}. Please specify the user '
+                        f'ID instead.')
+                user_ids.append(all_user_map[user_name_or_id][0])
+            else:
+                logger.warning(
+                    f'User {user_name_or_id!r} not found in all users')
+                continue
+        return user_ids
+    else:
+        # Public workspace - return '*' to indicate all users should have access
+        return ['*']

{skypilot_nightly-1.0.0.dev20250607.dist-info → skypilot_nightly-1.0.0.dev20250609.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: skypilot-nightly
-Version: 1.0.0.dev20250607
+Version: 1.0.0.dev20250609
 Summary: SkyPilot: Run AI on Any Infra — Unified, Faster, Cheaper.
 Author: SkyPilot Team
 License: Apache 2.0