skypilot-nightly 1.0.0.dev20250603__py3-none-any.whl → 1.0.0.dev20250605__py3-none-any.whl

sky/skypilot_config.py

@@ -55,7 +55,7 @@ import os
 import tempfile
 import threading
 import typing
-from typing import Any, Dict, Iterator, List, Optional, Tuple
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
 import filelock
 
@@ -163,11 +163,23 @@ def _set_loaded_config(config: config_utils.Config) -> None:
     _get_config_context().config = config
 
 
-def _get_loaded_config_path() -> Optional[str]:
-    return _get_config_context().config_path
+def _get_loaded_config_path() -> List[Optional[str]]:
+    serialized = _get_config_context().config_path
+    if not serialized:
+        return []
+    return json.loads(serialized)
+
+
+def _set_loaded_config_path(
+        path: Optional[Union[str, List[Optional[str]]]]) -> None:
+    if not path:
+        _get_config_context().config_path = None
+    if isinstance(path, str):
+        path = [path]
+    _get_config_context().config_path = json.dumps(path)
 
 
-def _set_loaded_config_path(path: Optional[str]) -> None:
+def _set_loaded_config_path_serialized(path: Optional[str]) -> None:
     _get_config_context().config_path = path
 
 
@@ -184,9 +196,14 @@ def get_user_config_path() -> str:
     return _GLOBAL_CONFIG_PATH
 
 
-def get_user_config() -> config_utils.Config:
-    """Returns the user config."""
-    # find the user config file
+def _get_config_from_path(path: Optional[str]) -> config_utils.Config:
+    if path is None:
+        return config_utils.Config()
+    return parse_and_validate_config_file(path)
+
+
+def _resolve_user_config_path() -> Optional[str]:
+    # find the user config file path, None if not resolved.
     user_config_path = _get_config_file_path(ENV_VAR_GLOBAL_CONFIG)
     if user_config_path:
         logger.debug('using user config file specified by '
@@ -203,16 +220,17 @@ def get_user_config() -> config_utils.Config:
         user_config_path = get_user_config_path()
         logger.debug(f'using default user config file: {user_config_path}')
         user_config_path = os.path.expanduser(user_config_path)
-
-    # load the user config file
     if os.path.exists(user_config_path):
-        user_config = parse_and_validate_config_file(user_config_path)
-    else:
-        user_config = config_utils.Config()
-    return user_config
+        return user_config_path
+    return None
 
 
-def _get_project_config() -> config_utils.Config:
+def get_user_config() -> config_utils.Config:
+    """Returns the user config."""
+    return _get_config_from_path(_resolve_user_config_path())
+
+
+def _resolve_project_config_path() -> Optional[str]:
     # find the project config file
     project_config_path = _get_config_file_path(ENV_VAR_PROJECT_CONFIG)
     if project_config_path:
@@ -231,17 +249,17 @@ def _get_project_config() -> config_utils.Config:
             f'using default project config file: {_PROJECT_CONFIG_PATH}')
         project_config_path = _PROJECT_CONFIG_PATH
         project_config_path = os.path.expanduser(project_config_path)
-
-    # load the project config file
     if os.path.exists(project_config_path):
-        project_config = parse_and_validate_config_file(project_config_path)
-    else:
-        project_config = config_utils.Config()
-    return project_config
+        return project_config_path
+    return None
 
 
-def get_server_config() -> config_utils.Config:
-    """Returns the server config."""
+def _get_project_config() -> config_utils.Config:
+    """Returns the project config."""
+    return _get_config_from_path(_resolve_project_config_path())
+
+
+def _resolve_server_config_path() -> Optional[str]:
     # find the server config file
     server_config_path = _get_config_file_path(ENV_VAR_GLOBAL_CONFIG)
     if server_config_path:
@@ -259,13 +277,14 @@ def get_server_config() -> config_utils.Config:
         server_config_path = _GLOBAL_CONFIG_PATH
         logger.debug(f'using default server config file: {server_config_path}')
         server_config_path = os.path.expanduser(server_config_path)
-
-    # load the server config file
     if os.path.exists(server_config_path):
-        server_config = parse_and_validate_config_file(server_config_path)
-    else:
-        server_config = config_utils.Config()
-    return server_config
+        return server_config_path
+    return None
+
+
+def get_server_config() -> config_utils.Config:
+    """Returns the server config."""
+    return _get_config_from_path(_resolve_server_config_path())
 
 
 def get_nested(keys: Tuple[str, ...],
@@ -487,9 +506,11 @@ def _reload_config_from_internal_file(internal_config_path: str) -> None:
 def _reload_config_as_server() -> None:
     # Reset the global variables, to avoid using stale values.
     _set_loaded_config(config_utils.Config())
+    _set_loaded_config_path(None)
 
     overrides: List[config_utils.Config] = []
-    server_config = get_server_config()
+    server_config_path = _resolve_server_config_path()
+    server_config = _get_config_from_path(server_config_path)
     if server_config:
         overrides.append(server_config)
 
@@ -503,17 +524,21 @@ def _reload_config_as_server() -> None:
             f'server config: \n'
             f'{common_utils.dump_yaml_str(dict(overlaid_server_config))}')
     _set_loaded_config(overlaid_server_config)
+    _set_loaded_config_path(server_config_path)
 
 
 def _reload_config_as_client() -> None:
     # Reset the global variables, to avoid using stale values.
     _set_loaded_config(config_utils.Config())
+    _set_loaded_config_path(None)
 
     overrides: List[config_utils.Config] = []
-    user_config = get_user_config()
+    user_config_path = _resolve_user_config_path()
+    user_config = _get_config_from_path(user_config_path)
     if user_config:
         overrides.append(user_config)
-    project_config = _get_project_config()
+    project_config_path = _resolve_project_config_path()
+    project_config = _get_config_from_path(project_config_path)
     if project_config:
         overrides.append(project_config)
 
@@ -527,14 +552,26 @@ def _reload_config_as_client() -> None:
             f'client config (before task and CLI overrides): \n'
             f'{common_utils.dump_yaml_str(dict(overlaid_client_config))}')
     _set_loaded_config(overlaid_client_config)
+    _set_loaded_config_path([user_config_path, project_config_path])
 
 
 def loaded_config_path() -> Optional[str]:
-    """Returns the path to the loaded config file, or
-    '<overridden>' if the config is overridden."""
-    if _is_config_overridden():
-        return '<overridden>'
-    return _get_loaded_config_path()
+    """Returns the path to the loaded config file, or '<overridden>' if the
+    config is overridden."""
+    path = [p for p in set(_get_loaded_config_path()) if p is not None]
+    if len(path) == 0:
+        return '<overridden>' if _is_config_overridden() else None
+    if len(path) == 1:
+        return path[0]
+
+    header = 'overridden' if _is_config_overridden() else 'merged'
+    path_str = ', '.join(p for p in path if p is not None)
+    return f'<{header} ({path_str})>'
+
+
+def loaded_config_path_serialized() -> Optional[str]:
+    """Returns the json serialized config path list"""
+    return _get_config_context().config_path
 
 
 # Load on import, synchronization is guaranteed by python interpreter.
@@ -548,7 +585,9 @@ def loaded() -> bool:
 
 @contextlib.contextmanager
 def override_skypilot_config(
-        override_configs: Optional[Dict[str, Any]]) -> Iterator[None]:
+        override_configs: Optional[Dict[str, Any]],
+        override_config_path_serialized: Optional[str] = None
+) -> Iterator[None]:
     """Overrides the user configurations."""
     # TODO(SKY-1215): allow admin user to extend the disallowed keys or specify
     # allowed keys.
@@ -557,7 +596,13 @@ def override_skypilot_config(
         yield
         return
     original_config = _get_loaded_config()
+    original_config_path = loaded_config_path_serialized()
     override_configs = config_utils.Config(override_configs)
+    if override_config_path_serialized is None:
+        override_config_path = []
+    else:
+        override_config_path = json.loads(override_config_path_serialized)
+
     disallowed_diff_keys = []
     for key in constants.SKIPPED_CLIENT_OVERRIDE_KEYS:
         value = override_configs.pop_nested(key, default_value=None)
@@ -602,6 +647,8 @@ def override_skypilot_config(
             skip_none=False)
         _set_config_overridden(True)
         _set_loaded_config(config)
+        _set_loaded_config_path(_get_loaded_config_path() +
+                                override_config_path)
         yield
     except exceptions.InvalidSkyPilotConfigError as e:
         with ux_utils.print_exception_no_traceback():
@@ -616,6 +663,7 @@ def override_skypilot_config(
     finally:
         _set_loaded_config(original_config)
         _set_config_overridden(False)
+        _set_loaded_config_path_serialized(original_config_path)
 
 
 @contextlib.contextmanager
@@ -628,6 +676,7 @@ def replace_skypilot_config(new_configs: config_utils.Config) -> Iterator[None]:
        sky_utils.context for more details.
     """
     original_config = _get_loaded_config()
+    original_config_path = loaded_config_path_serialized()
     original_env_var = os.environ.get(ENV_VAR_SKYPILOT_CONFIG)
     if new_configs != original_config:
         # Modify the global config of current process or context
@@ -642,9 +691,11 @@ def replace_skypilot_config(new_configs: config_utils.Config) -> Iterator[None]:
         # Note that this code modifies os.environ directly because it
         # will be hijacked to be context-aware if a context is active.
         os.environ[ENV_VAR_SKYPILOT_CONFIG] = temp_file.name
+        _set_loaded_config_path(temp_file.name)
         yield
         # Restore the original config and env var.
         _set_loaded_config(original_config)
+        _set_loaded_config_path_serialized(original_config_path)
         if original_env_var:
             os.environ[ENV_VAR_SKYPILOT_CONFIG] = original_env_var
         else:
@@ -705,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/usage/usage_lib.py

@@ -205,8 +205,8 @@ class UsageMessageToReport(MessageToReport):
                     logger.debug('Multiple accelerators are not supported: '
                                  f'{resources.accelerators}.')
                 self.task_accelerators = list(resources.accelerators.keys())[0]
-                self.task_num_accelerators = resources.accelerators[
-                    self.task_accelerators]
+                self.task_num_accelerators = int(
+                    resources.accelerators[self.task_accelerators])
             else:
                 self.task_accelerators = None
                 self.task_num_accelerators = None
@@ -245,7 +245,8 @@ class UsageMessageToReport(MessageToReport):
                 logger.debug('Multiple accelerators are not supported: '
                              f'{resources.accelerators}.')
             self.accelerators = list(resources.accelerators.keys())[0]
-            self.num_accelerators = resources.accelerators[self.accelerators]
+            self.num_accelerators = int(
+                resources.accelerators[self.accelerators])
         else:
             self.accelerators = None
             self.num_accelerators = None

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

sky/users/server.py

@@ -0,0 +1,66 @@
+"""REST API for workspace management."""
+
+import hashlib
+from typing import Any, Dict, List
+
+import fastapi
+
+from sky import global_user_state
+from sky import sky_logging
+from sky.server.requests import payloads
+from sky.users import permission
+from sky.users import rbac
+from sky.utils import common_utils
+
+logger = sky_logging.init_logger(__name__)
+
+router = fastapi.APIRouter()
+
+permission_service = permission.PermissionService()
+
+
+@router.get('')
+async def users() -> List[Dict[str, Any]]:
+    """Gets all users."""
+    all_users = []
+    user_list = global_user_state.get_all_users()
+    for user in user_list:
+        user_roles = permission_service.get_user_roles(user.id)
+        all_users.append({
+            'id': user.id,
+            'name': user.name,
+            'role': user_roles[0] if user_roles else ''
+        })
+    return all_users
+
+
+@router.get('/role')
+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:
+        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 ''}
+
+
+@router.post('/update')
+async def user_update(user_update_body: payloads.UserUpdateBody) -> None:
+    """Updates the user role."""
+    user_id = user_update_body.user_id
+    role = user_update_body.role
+    supported_roles = rbac.get_supported_roles()
+    if role not in supported_roles:
+        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:
+        raise fastapi.HTTPException(status_code=400,
+                                    detail=f'User {user_id} does not exist')
+
+    # Update user role in casbin policy
+    permission_service.update_role(user_id, role)