skypilot-nightly 1.0.0.dev20250905__py3-none-any.whl → 1.0.0.dev20251210__py3-none-any.whl

sky/usage/usage_lib.py

@@ -14,6 +14,7 @@ from typing_extensions import ParamSpec
 
 import sky
 from sky import sky_logging
+from sky import skypilot_config
 from sky.adaptors import common as adaptors_common
 from sky.usage import constants
 from sky.utils import common_utils
@@ -167,6 +168,7 @@ class UsageMessageToReport(MessageToReport):
         self.runtimes: Dict[str, float] = {}  # update_runtime
         self.exception: Optional[str] = None  # entrypoint_context
         self.stacktrace: Optional[str] = None  # entrypoint_context
+        self.skypilot_config: Optional[Dict[str, Any]] = None
 
         # Whether API server is deployed remotely.
         self.using_remote_api_server: bool = (
@@ -177,6 +179,7 @@ class UsageMessageToReport(MessageToReport):
             self.client_entrypoint = common_utils.get_current_client_entrypoint(
                 msg)
         self.entrypoint = msg
+        self.skypilot_config = dict(skypilot_config.to_dict())
 
     def set_internal(self):
         self.internal = True
@@ -313,21 +316,30 @@ class MessageCollection:
     """A collection of messages."""
 
     def __init__(self):
-        self._messages = {
+        self._messages: Dict[MessageType, MessageToReport] = {
             MessageType.USAGE: UsageMessageToReport(),
             MessageType.HEARTBEAT: HeartbeatMessageToReport()
         }
 
     @property
     def usage(self) -> UsageMessageToReport:
-        return self._messages[MessageType.USAGE]
+        msg = self._messages[MessageType.USAGE]
+        assert isinstance(msg, UsageMessageToReport)
+        return msg
 
     @property
     def heartbeat(self) -> HeartbeatMessageToReport:
-        return self._messages[MessageType.HEARTBEAT]
+        msg = self._messages[MessageType.HEARTBEAT]
+        assert isinstance(msg, HeartbeatMessageToReport)
+        return msg
 
     def reset(self, message_type: MessageType):
-        self._messages[message_type] = self._messages[message_type].__class__()
+        if message_type == MessageType.USAGE:
+            self._messages[message_type] = UsageMessageToReport()
+        elif message_type == MessageType.HEARTBEAT:
+            self._messages[message_type] = HeartbeatMessageToReport()
+        else:
+            raise ValueError(f'Unknown message type: {message_type}')
 
     def __getitem__(self, key):
         return self._messages[key]

sky/users/model.conf

@@ -12,4 +12,4 @@ g = _, _
 e = some(where (p.eft == allow))
 
 [matchers]
-m = (g(r.sub, p.sub)|| p.sub == '*') && r.obj == p.obj && r.act == p.act
+m = (g(r.sub, p.sub)|| p.sub == '*') && keyMatch2(r.obj, p.obj) && r.act == p.act

sky/users/permission.py

@@ -3,7 +3,7 @@ import contextlib
 import hashlib
 import logging
 import os
-from typing import Generator, List
+from typing import Generator, List, Optional
 
 import casbin
 import filelock
@@ -14,6 +14,7 @@ from sky import models
 from sky import sky_logging
 from sky.skylet import constants
 from sky.users import rbac
+from sky.utils import annotations
 from sky.utils import common_utils
 from sky.utils.db import db_utils
 
@@ -27,14 +28,14 @@ logger = sky_logging.init_logger(__name__)
 POLICY_UPDATE_LOCK_PATH = os.path.expanduser('~/.sky/.policy_update.lock')
 POLICY_UPDATE_LOCK_TIMEOUT_SECONDS = 20
 
-_enforcer_instance = None
+_enforcer_instance: Optional['PermissionService'] = None
 
 
 class PermissionService:
     """Permission service for SkyPilot API Server."""
 
     def __init__(self):
-        self.enforcer = None
+        self.enforcer: Optional[casbin.Enforcer] = None
 
     def _lazy_initialize(self):
         if self.enforcer is not None:
@@ -42,7 +43,6 @@ class PermissionService:
         with _policy_lock():
             global _enforcer_instance
             if _enforcer_instance is None:
-                _enforcer_instance = self
                 engine = global_user_state.initialize_and_get_db()
                 db_utils.add_all_tables_to_db_sqlalchemy(
                     sqlalchemy_adapter.Base.metadata, engine)
@@ -52,11 +52,43 @@ class PermissionService:
                                           'model.conf')
                 enforcer = casbin.Enforcer(model_path, adapter)
                 self.enforcer = enforcer
+                # Only set the enforcer instance once the enforcer
+                # is successfully initialized, if we change it and then fail
+                # we will set it to None and all subsequent calls will fail.
+                _enforcer_instance = self
                 self._maybe_initialize_policies()
                 self._maybe_initialize_basic_auth_user()
             else:
+                assert _enforcer_instance is not None
                 self.enforcer = _enforcer_instance.enforcer
 
+    def _ensure_enforcer(self) -> casbin.Enforcer:
+        """Ensure enforcer is initialized and return it."""
+        self._lazy_initialize()
+        assert self.enforcer is not None, (
+            'Enforcer should be initialized after _lazy_initialize()')
+        return self.enforcer
+
+    def _get_plugin_rbac_rules(self):
+        """Get RBAC rules from loaded plugins.
+
+        Returns:
+            Dictionary of plugin RBAC rules, or empty dict if plugins module
+            is not available or no rules are defined.
+        """
+        try:
+            # pylint: disable=import-outside-toplevel
+            from sky.server import plugins as server_plugins
+            return server_plugins.get_plugin_rbac_rules()
+        except ImportError:
+            # Plugin module not available (e.g., not running as server)
+            logger.debug(
+                'Plugin module not available, skipping plugin RBAC rules')
+            return {}
+        except Exception as e:  # pylint: disable=broad-except
+            logger.warning(f'Failed to get plugin RBAC rules: {e}')
+            return {}
+
     def _maybe_initialize_basic_auth_user(self) -> None:
         """Initialize basic auth user if it is enabled."""
         basic_auth = os.environ.get(constants.SKYPILOT_INITIAL_BASIC_AUTH)
@@ -72,9 +104,9 @@ class PermissionService:
                 return
             global_user_state.add_or_update_user(
                 models.User(id=user_hash, name=username, password=password))
-            self.enforcer.add_grouping_policy(user_hash,
-                                              rbac.RoleName.ADMIN.value)
-            self.enforcer.save_policy()
+            enforcer = self._ensure_enforcer()
+            enforcer.add_grouping_policy(user_hash, rbac.RoleName.ADMIN.value)
+            enforcer.save_policy()
             logger.info(f'Basic auth user {username} initialized')
 
     def _maybe_initialize_policies(self) -> None:
@@ -86,11 +118,15 @@ class PermissionService:
 
         # Check if policies are already initialized by looking for existing
         # permission policies in the enforcer
-        existing_policies = self.enforcer.get_policy()
+        enforcer = self._ensure_enforcer()
+        existing_policies = enforcer.get_policy()
+
+        # Get plugin RBAC rules dynamically
+        plugin_rules = self._get_plugin_rbac_rules()
 
         # If we already have policies for the expected roles, skip
         # initialization
-        role_permissions = rbac.get_role_permissions()
+        role_permissions = rbac.get_role_permissions(plugin_rules=plugin_rules)
         expected_policies = []
         for role, permissions in role_permissions.items():
             if permissions['permissions'] and 'blocklist' in permissions[
@@ -123,7 +159,7 @@ class PermissionService:
             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)
+            enforcer.remove_filtered_policy(0)
             for role, permissions in role_permissions.items():
                 if permissions['permissions'] and 'blocklist' in permissions[
                         'permissions']:
@@ -133,14 +169,14 @@ class PermissionService:
                         method = item['method']
                         logger.debug(f'Adding role policy: role={role}, '
                                      f'path={path}, method={method}')
-                        self.enforcer.add_policy(role, path, method)
+                        enforcer.add_policy(role, path, method)
                         policy_updated = True
 
             for workspace_name, users in workspace_policy_permissions.items():
                 for user in users:
                     logger.debug(f'Initializing workspace policy: user={user}, '
                                  f'workspace={workspace_name}')
-                    self.enforcer.add_policy(user, workspace_name, '*')
+                    enforcer.add_policy(user, workspace_name, '*')
                     policy_updated = True
             logger.debug('Policies initialized successfully')
         else:
@@ -153,7 +189,7 @@ class PermissionService:
             policy_updated = policy_updated or user_added
 
         if policy_updated:
-            self.enforcer.save_policy()
+            enforcer.save_policy()
 
     def add_user_if_not_exists(self, user_id: str) -> None:
         """Add user role relationship."""
@@ -167,34 +203,35 @@ class PermissionService:
         Returns:
             True if the user was added, False otherwise.
         """
-        user_roles = self.enforcer.get_roles_for_user(user_id)
+        enforcer = self._ensure_enforcer()
+        user_roles = enforcer.get_roles_for_user(user_id)
         if not user_roles:
-            self.enforcer.add_grouping_policy(user_id, rbac.get_default_role())
+            enforcer.add_grouping_policy(user_id, rbac.get_default_role())
             return True
         return False
 
     def delete_user(self, user_id: str) -> None:
         """Delete user role relationship."""
-        self._lazy_initialize()
         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_id)
+            enforcer = self._ensure_enforcer()
+            current_roles = enforcer.get_roles_for_user(user_id)
             if not current_roles:
                 logger.debug(f'User {user_id} has no roles')
                 return
-            self.enforcer.remove_grouping_policy(user_id, current_roles[0])
-            self.enforcer.save_policy()
+            enforcer.remove_grouping_policy(user_id, current_roles[0])
+            enforcer.save_policy()
 
     def update_role(self, user_id: str, new_role: str) -> None:
         """Update user role relationship."""
-        self._lazy_initialize()
         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_id)
+            enforcer = self._ensure_enforcer()
+            current_roles = enforcer.get_roles_for_user(user_id)
             if not current_roles:
                 logger.debug(f'User {user_id} has no roles')
             else:
@@ -203,11 +240,11 @@ class PermissionService:
                 if current_role == new_role:
                     logger.debug(f'User {user_id} already has role {new_role}')
                     return
-                self.enforcer.remove_grouping_policy(user_id, current_role)
+                enforcer.remove_grouping_policy(user_id, current_role)
 
             # Update user role
-            self.enforcer.add_grouping_policy(user_id, new_role)
-            self.enforcer.save_policy()
+            enforcer.add_grouping_policy(user_id, new_role)
+            enforcer.save_policy()
 
     def get_user_roles(self, user_id: str) -> List[str]:
         """Get all roles for a user.
@@ -222,15 +259,15 @@ class PermissionService:
         Returns:
             A list of role names that the user has.
         """
-        self._lazy_initialize()
         self._load_policy_no_lock()
-        return self.enforcer.get_roles_for_user(user_id)
+        enforcer = self._ensure_enforcer()
+        return enforcer.get_roles_for_user(user_id)
 
     def get_users_for_role(self, role: str) -> List[str]:
         """Get all users for a role."""
-        self._lazy_initialize()
         self._load_policy_no_lock()
-        return self.enforcer.get_users_for_role(role)
+        enforcer = self._ensure_enforcer()
+        return enforcer.get_users_for_role(role)
 
     def check_endpoint_permission(self, user_id: str, path: str,
                                   method: str) -> bool:
@@ -241,19 +278,22 @@ class PermissionService:
         # 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()
-        self._lazy_initialize()
-        return self.enforcer.enforce(user_id, path, method)
+        enforcer = self._ensure_enforcer()
+        return enforcer.enforce(user_id, path, method)
 
     def _load_policy_no_lock(self):
         """Load policy from storage."""
-        self.enforcer.load_policy()
+        enforcer = self._ensure_enforcer()
+        enforcer.load_policy()
 
     def load_policy(self):
         """Load policy from storage with lock."""
-        self._lazy_initialize()
         with _policy_lock():
             self._load_policy_no_lock()
 
+    # Right now, not a lot of users are using multiple workspaces,
+    # so 5 should be more than enough.
+    @annotations.lru_cache(scope='request', maxsize=5)
     def check_workspace_permission(self, user_id: str,
                                    workspace_name: str) -> bool:
         """Check workspace permission.
@@ -266,7 +306,6 @@ class PermissionService:
         For public workspaces, the permission is granted via a wildcard policy
         ('*').
         """
-        self._lazy_initialize()
         if os.getenv(constants.ENV_VAR_IS_SKYPILOT_SERVER) is None:
             # When it is not on API server, we allow all users to access all
             # workspaces, as the workspace check has been done on API server.
@@ -279,7 +318,8 @@ class PermissionService:
         # r.act == p.act
         # This means if there's a policy ('*', workspace_name, '*'), it will
         # match any user
-        result = self.enforcer.enforce(user_id, workspace_name, '*')
+        enforcer = self._ensure_enforcer()
+        result = enforcer.enforce(user_id, workspace_name, '*')
         logger.debug(f'Workspace permission check: user={user_id}, '
                      f'workspace={workspace_name}, result={result}')
         return result
@@ -323,13 +363,13 @@ class PermissionService:
                    For public workspaces, this should be ['*'].
                    For private workspaces, this should be specific user IDs.
         """
-        self._lazy_initialize()
         with _policy_lock():
+            enforcer = self._ensure_enforcer()
             for user in users:
                 logger.debug(f'Adding workspace policy: user={user}, '
                              f'workspace={workspace_name}')
-                self.enforcer.add_policy(user, workspace_name, '*')
-            self.enforcer.save_policy()
+                enforcer.add_policy(user, workspace_name, '*')
+            enforcer.save_policy()
 
     def update_workspace_policy(self, workspace_name: str,
                                 users: List[str]) -> None:
@@ -341,24 +381,24 @@ class PermissionService:
                    For public workspaces, this should be ['*'].
                    For private workspaces, this should be specific user IDs.
         """
-        self._lazy_initialize()
         with _policy_lock():
             self._load_policy_no_lock()
+            enforcer = self._ensure_enforcer()
             # Remove all existing policies for this workspace
-            self.enforcer.remove_filtered_policy(1, workspace_name)
+            enforcer.remove_filtered_policy(1, workspace_name)
             # Add new policies
             for user in users:
                 logger.debug(f'Updating workspace policy: user={user}, '
                              f'workspace={workspace_name}')
-                self.enforcer.add_policy(user, workspace_name, '*')
-            self.enforcer.save_policy()
+                enforcer.add_policy(user, workspace_name, '*')
+            enforcer.save_policy()
 
     def remove_workspace_policy(self, workspace_name: str) -> None:
         """Remove workspace policy."""
-        self._lazy_initialize()
         with _policy_lock():
-            self.enforcer.remove_filtered_policy(1, workspace_name)
-            self.enforcer.save_policy()
+            enforcer = self._ensure_enforcer()
+            enforcer.remove_filtered_policy(1, workspace_name)
+            enforcer.save_policy()
 
 
 @contextlib.contextmanager

sky/users/rbac.py

@@ -1,7 +1,7 @@
 """RBAC (Role-Based Access Control) functionality for SkyPilot API Server."""
 
 import enum
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 from sky import sky_logging
 from sky import skypilot_config
@@ -55,8 +55,13 @@ def get_default_role() -> str:
 
 
 def get_role_permissions(
+    plugin_rules: Optional[Dict[str, List[Dict[str, str]]]] = None
 ) -> Dict[str, Dict[str, Dict[str, List[Dict[str, str]]]]]:
-    """Get all role permissions from config.
+    """Get all role permissions from config and plugins.
+
+    Args:
+        plugin_rules: Optional dictionary of plugin RBAC rules to merge.
+                     Format: {'user': [{'path': '...', 'method': '...'}]}
 
     Returns:
         Dictionary containing all roles and their permissions configuration.
@@ -91,9 +96,32 @@ def get_role_permissions(
     if 'user' not in config_permissions:
         config_permissions['user'] = {
             'permissions': {
-                'blocklist': _DEFAULT_USER_BLOCKLIST
+                'blocklist': _DEFAULT_USER_BLOCKLIST.copy()
             }
         }
+
+    # Merge plugin rules into the appropriate roles
+    if plugin_rules:
+        for role, rules in plugin_rules.items():
+            if role not in supported_roles:
+                logger.warning(f'Plugin specified invalid role: {role}')
+                continue
+            if role not in config_permissions:
+                config_permissions[role] = {'permissions': {'blocklist': []}}
+            if 'permissions' not in config_permissions[role]:
+                config_permissions[role]['permissions'] = {'blocklist': []}
+            if 'blocklist' not in config_permissions[role]['permissions']:
+                config_permissions[role]['permissions']['blocklist'] = []
+
+            # Merge plugin rules, avoiding duplicates
+            existing_rules = config_permissions[role]['permissions'][
+                'blocklist']
+            for rule in rules:
+                if rule not in existing_rules:
+                    existing_rules.append(rule)
+                    logger.debug(f'Added plugin RBAC rule for {role}: '
+                                 f'{rule["method"]} {rule["path"]}')
+
     return config_permissions
 
 

sky/utils/accelerator_registry.py

@@ -3,6 +3,7 @@ import typing
 from typing import List, Optional
 
 from sky import catalog
+from sky.catalog import common as catalog_common
 from sky.utils import rich_utils
 from sky.utils import ux_utils
 
@@ -34,8 +35,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 = catalog.common.read_catalog('common/accelerators.csv')
-_memory_df = catalog.common.read_catalog('common/metadata.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.
@@ -107,10 +108,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')):
+            # 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

@@ -2,8 +2,9 @@
 import contextlib
 import copy
 import importlib
+import typing
 from typing import Iterator, Optional, Tuple, Union
-import urllib.parse
+from urllib import parse as urlparse
 
 import colorama
 
@@ -13,17 +14,21 @@ 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 _is_url(policy_string: str) -> bool:
     """Check if the policy string is a URL."""
     try:
-        parsed = urllib.parse.urlparse(policy_string)
+        parsed = urlparse.urlparse(policy_string)
         return parsed.scheme in ('http', 'https')
     except Exception:  # pylint: disable=broad-except
         return False
@@ -73,6 +78,7 @@ def _get_policy_impl(
 @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']:
@@ -86,7 +92,8 @@ def apply_and_use_config_in_current_request(
     Refer to `apply()` for more details.
     """
     original_config = skypilot_config.to_dict()
-    dag, mutated_config = apply(entrypoint, request_options, at_client_side)
+    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
@@ -96,6 +103,7 @@ def apply_and_use_config_in_current_request(
 
 def apply(
     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,
 ) -> Tuple['dag_lib.Dag', config_utils.Config]:
@@ -126,9 +134,13 @@ def apply(
     if policy is None:
         return dag, skypilot_config.to_dict()
 
+    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()
@@ -136,8 +148,9 @@ def apply(
 
     mutated_config = None
     for task in dag.tasks:
-        user_request = admin_policy.UserRequest(task, config, request_options,
-                                                at_client_side)
+        user_request = admin_policy.UserRequest(task, config, request_name,
+                                                request_options, at_client_side,
+                                                user)
         try:
             mutated_user_request = policy.apply(user_request)
         # Avoid duplicate exception wrapping.