skypilot-nightly 1.0.0.dev20251203__py3-none-any.whl → 1.0.0.dev20260112__py3-none-any.whl

sky/users/permission.py

@@ -3,6 +3,7 @@ import contextlib
 import hashlib
 import logging
 import os
+import threading
 from typing import Generator, List, Optional
 
 import casbin
@@ -36,16 +37,23 @@ class PermissionService:
 
     def __init__(self):
         self.enforcer: Optional[casbin.Enforcer] = None
+        self._lock = threading.Lock()
 
-    def _lazy_initialize(self):
+    def initialize(self):
+        self._lazy_initialize(full_initialize=True)
+
+    def _lazy_initialize(self, full_initialize: bool = False):
         if self.enforcer is not None:
             return
-        with _policy_lock():
+        with self._lock:
+            if self.enforcer is not None:
+                return
             global _enforcer_instance
             if _enforcer_instance is None:
                 engine = global_user_state.initialize_and_get_db()
-                db_utils.add_all_tables_to_db_sqlalchemy(
-                    sqlalchemy_adapter.Base.metadata, engine)
+                if full_initialize:
+                    db_utils.add_all_tables_to_db_sqlalchemy(
+                        sqlalchemy_adapter.Base.metadata, engine)
                 adapter = sqlalchemy_adapter.Adapter(
                     engine, db_class=sqlalchemy_adapter.CasbinRule)
                 model_path = os.path.join(os.path.dirname(__file__),
@@ -56,8 +64,10 @@ class PermissionService:
                 # 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()
+                if full_initialize:
+                    with _policy_lock():
+                        self._maybe_initialize_policies()
+                        self._maybe_initialize_basic_auth_user()
             else:
                 assert _enforcer_instance is not None
                 self.enforcer = _enforcer_instance.enforcer
@@ -69,6 +79,26 @@ class PermissionService:
             '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)
@@ -92,26 +122,29 @@ class PermissionService:
     def _maybe_initialize_policies(self) -> None:
         """Initialize policies if they don't already exist."""
         logger.debug(f'Initializing policies in process: {os.getpid()}')
-        self._load_policy_no_lock()
 
         policy_updated = False
 
         # Check if policies are already initialized by looking for existing
         # permission policies in the enforcer
         enforcer = self._ensure_enforcer()
-        existing_policies = enforcer.get_policy()
+        # Convert existing policies to set of tuples for O(1) lookups
+        existing_policies = {tuple(p) for p in 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[
-                    'permissions']:
+            if permissions.get('permissions'
+                              ) and 'blocklist' in permissions['permissions']:
                 blocklist = permissions['permissions']['blocklist']
                 for item in blocklist:
                     expected_policies.append(
-                        [role, item['path'], item['method']])
+                        (role, item['path'], item['method']))
 
         # Add workspace policy
         workspace_policy_permissions = rbac.get_workspace_policy_permissions()
@@ -120,50 +153,50 @@ class PermissionService:
 
         for workspace_name, users in workspace_policy_permissions.items():
             for user in users:
-                expected_policies.append([user, workspace_name, '*'])
-                logger.debug(f'Expected workspace policy: user={user}, '
-                             f'workspace={workspace_name}')
-
-        # 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)
-            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']
-                        logger.debug(f'Adding role policy: role={role}, '
-                                     f'path={path}, method={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}')
-                    enforcer.add_policy(user, workspace_name, '*')
-                    policy_updated = True
-            logger.debug('Policies initialized successfully')
-        else:
-            logger.debug('Policies already exist, skipping initialization')
+                expected_policies.append((user, workspace_name, '*'))
+        # Check if all expected policies already exist and find missing ones
+        missing_policies = [
+            p for p in expected_policies if p not in existing_policies
+        ]
+        # Find policies to remove
+        expected_policies_set = set(expected_policies)
+        redundant_policies = [
+            p for p in existing_policies if p not in expected_policies_set
+        ]
+        if missing_policies:
+            # Add missing policies
+            logger.debug(f'Found {len(missing_policies)} missing policies, '
+                         'initializing...')
+            for p in missing_policies:
+                logger.debug(f'Adding policy: {p}')
+                enforcer.add_policy(*p)
+                policy_updated = True
+            logger.debug('Missing policies added successfully')
+
+        if redundant_policies:
+            # Remove redundant policies
+            logger.debug(f'Found {len(redundant_policies)} redundant policies, '
+                         'cleaning up...')
+            for p in redundant_policies:
+                logger.debug(f'Removing policy: {p}')
+                enforcer.remove_policy(*p)
+                policy_updated = True
+            logger.debug('Redundant policies removed successfully')
+
+        if not missing_policies and not redundant_policies:
+            logger.debug('Policies already in sync, skipping initialization')
 
         # Always ensure users have default roles (this is idempotent)
+        # Get users who already have roles (g policies) to avoid redundant calls
+        users_with_roles = {tuple(g)[0] for g in enforcer.get_grouping_policy()}
         all_users = global_user_state.get_all_users()
         for existing_user in all_users:
-            user_added = self._add_user_if_not_exists_no_lock(existing_user.id)
-            policy_updated = policy_updated or user_added
+            if str(existing_user.id) not in users_with_roles:
+                logger.debug(f'Adding role for user: {existing_user.name}'
+                             f'({existing_user.id})')
+                user_added = self._add_user_if_not_exists_no_lock(
+                    existing_user.id)
+                policy_updated = policy_updated or user_added
 
         if policy_updated:
             enforcer.save_policy()

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/annotations.py

@@ -1,14 +1,20 @@
 """Annotations for public APIs."""
 
 import functools
-from typing import Callable, Literal, TypeVar
+import threading
+import time
+from typing import Callable, List, Literal, TypeVar
+import weakref
 
 import cachetools
 from typing_extensions import ParamSpec
 
 # Whether the current process is a SkyPilot API server process.
 is_on_api_server = True
-_FUNCTIONS_NEED_RELOAD_CACHE = []
+_FUNCTIONS_NEED_RELOAD_CACHE_LOCK = threading.Lock()
+# Caches can be thread-local, use weakref to avoid blocking the GC when the
+# thread is destroyed.
+_FUNCTIONS_NEED_RELOAD_CACHE: List[weakref.ReferenceType] = []
 
 T = TypeVar('T')
 P = ParamSpec('P')
@@ -30,6 +36,94 @@ def client_api(func: Callable[P, T]) -> Callable[P, T]:
     return wrapper
 
 
+def _register_functions_need_reload_cache(func: Callable) -> Callable:
+    """Register a cachefunction that needs to be reloaded for a new request.
+
+    The function will be registered as a weak reference to avoid blocking GC.
+    """
+    assert hasattr(func, 'cache_clear'), f'{func.__name__} is not cacheable'
+    wrapped_fn = func
+    try:
+        func_ref = weakref.ref(func)
+    except TypeError:
+        # The function might be not weakrefable (e.g. functools.lru_cache),
+        # wrap it in this case.
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+
+        wrapper.cache_clear = func.cache_clear  # type: ignore[attr-defined]
+        func_ref = weakref.ref(wrapper)
+        wrapped_fn = wrapper
+    with _FUNCTIONS_NEED_RELOAD_CACHE_LOCK:
+        _FUNCTIONS_NEED_RELOAD_CACHE.append(func_ref)
+    return wrapped_fn
+
+
+class ThreadLocalTTLCache(threading.local):
+    """Thread-local storage for _thread_local_lru_cache decorator."""
+
+    def __init__(self, func, maxsize: int, ttl: int):
+        super().__init__()
+        self.func = func
+        self.maxsize = maxsize
+        self.ttl = ttl
+
+    def get_cache(self):
+        if not hasattr(self, 'cache'):
+            self.cache = ttl_cache(scope='request',
+                                   maxsize=self.maxsize,
+                                   ttl=self.ttl,
+                                   timer=time.time)(self.func)
+        return self.cache
+
+    def __del__(self):
+        if hasattr(self, 'cache'):
+            self.cache.cache_clear()
+            self.cache = None
+
+
+def thread_local_ttl_cache(maxsize=32, ttl=60 * 55):
+    """Thread-local TTL cache decorator.
+
+    Args:
+        maxsize: Maximum size of the cache.
+        ttl: Time to live for the cache in seconds.
+             Default is 55 minutes, a bit less than 1 hour
+             default lifetime of an STS token.
+    """
+
+    def decorator(func):
+        # Create thread-local storage for the LRU cache
+        local_cache = ThreadLocalTTLCache(func, maxsize, ttl)
+
+        # We can't apply the lru_cache here, because this runs at import time
+        # so we will always have the main thread's cache.
+
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            # We are within the actual function call, which may be on a thread,
+            # so local_cache.cache will return the correct thread-local cache,
+            # which we can now apply and immediately call.
+            return local_cache.get_cache()(*args, **kwargs)
+
+        def cache_info():
+            # Note that this will only give the cache info for the current
+            # thread's cache.
+            return local_cache.get_cache().cache_info()
+
+        def cache_clear():
+            # Note that this will only clear the cache for the current thread.
+            local_cache.get_cache().cache_clear()
+
+        wrapper.cache_info = cache_info  # type: ignore[attr-defined]
+        wrapper.cache_clear = cache_clear  # type: ignore[attr-defined]
+
+        return wrapper
+
+    return decorator
+
+
 def lru_cache(scope: Literal['global', 'request'], *lru_cache_args,
               **lru_cache_kwargs) -> Callable:
     """LRU cache decorator for functions.
@@ -51,8 +145,7 @@ def lru_cache(scope: Literal['global', 'request'], *lru_cache_args,
         else:
             cached_func = functools.lru_cache(*lru_cache_args,
                                               **lru_cache_kwargs)(func)
-            _FUNCTIONS_NEED_RELOAD_CACHE.append(cached_func)
-            return cached_func
+            return _register_functions_need_reload_cache(cached_func)
 
     return decorator
 
@@ -72,13 +165,20 @@ def ttl_cache(scope: Literal['global', 'request'], *ttl_cache_args,
         else:
             cached_func = cachetools.cached(
                 cachetools.TTLCache(*ttl_cache_args, **ttl_cache_kwargs))(func)
-            _FUNCTIONS_NEED_RELOAD_CACHE.append(cached_func)
-            return cached_func
+            return _register_functions_need_reload_cache(cached_func)
 
     return decorator
 
 
 def clear_request_level_cache():
     """Clear the request-level cache."""
-    for func in _FUNCTIONS_NEED_RELOAD_CACHE:
-        func.cache_clear()
+    alive_entries = []
+    with _FUNCTIONS_NEED_RELOAD_CACHE_LOCK:
+        for entry in _FUNCTIONS_NEED_RELOAD_CACHE:
+            func = entry()
+            if func is None:
+                # Has been GC'ed, drop the reference.
+                continue
+            func.cache_clear()
+            alive_entries.append(entry)
+        _FUNCTIONS_NEED_RELOAD_CACHE[:] = alive_entries

sky/utils/auth_utils.py

@@ -58,6 +58,34 @@ def _generate_rsa_key_pair() -> Tuple[str, str]:
     return public_key, private_key
 
 
+def _ensure_key_permissions(private_key_path: str,
+                            public_key_path: str) -> None:
+    """Ensure SSH key files and parent directory have correct permissions.
+
+    This is necessary because external factors (e.g., Kubernetes fsGroup,
+    volume mounts, umask) can modify file permissions after creation.
+    SSH requires private keys to have strict permissions (0600) and the
+    parent directory to not be group/world writable (0700).
+
+    This function is best-effort and will not raise exceptions if permission
+    changes fail (e.g., due to permission denied or read-only filesystem).
+    """
+
+    def _safe_chmod(path: str, mode: int) -> None:
+        """Attempt to chmod, logging warning on failure."""
+        try:
+            if os.path.exists(path):
+                os.chmod(path, mode)
+        except OSError as e:
+            logger.warning(f'Failed to set permissions on {path}: {e}')
+
+    # Ensure parent directory has correct permissions (0700)
+    key_dir = os.path.dirname(private_key_path)
+    _safe_chmod(key_dir, 0o700)
+    _safe_chmod(private_key_path, 0o600)
+    _safe_chmod(public_key_path, 0o644)
+
+
 def _save_key_pair(private_key_path: str, public_key_path: str,
                    private_key: str, public_key: str) -> None:
     key_dir = os.path.dirname(private_key_path)
@@ -77,6 +105,11 @@ def _save_key_pair(private_key_path: str, public_key_path: str,
               opener=functools.partial(os.open, mode=0o644)) as f:
         f.write(public_key)
 
+    # Explicitly set permissions to ensure they are correct regardless of
+    # umask or pre-existing file permissions. The opener's mode parameter
+    # only applies when creating new files, and is still subject to umask.
+    _ensure_key_permissions(private_key_path, public_key_path)
+
 
 def get_or_generate_keys() -> Tuple[str, str]:
     """Returns the absolute private and public key paths."""
@@ -105,6 +138,9 @@ def get_or_generate_keys() -> Tuple[str, str]:
     assert os.path.exists(public_key_path), (
         'Private key found, but associated public key '
         f'{public_key_path} does not exist.')
+    # Ensure correct permissions every time, as external factors (e.g.,
+    # Kubernetes fsGroup) can modify them after creation.
+    _ensure_key_permissions(private_key_path, public_key_path)
     return private_key_path, public_key_path
 
 
@@ -133,6 +169,9 @@ def create_ssh_key_files_from_db(private_key_path: str) -> bool:
     lock_dir = os.path.dirname(lock_path)
 
     if os.path.exists(private_key_path) and os.path.exists(public_key_path):
+        # Ensure correct permissions every time, as external factors (e.g.,
+        # Kubernetes fsGroup) can modify them after creation.
+        _ensure_key_permissions(private_key_path, public_key_path)
         return True
     # We should have the folder ~/.sky/generated/ssh to have 0o700 permission,
     # as the ssh configs will be written to this folder as well in
@@ -150,4 +189,7 @@ def create_ssh_key_files_from_db(private_key_path: str) -> bool:
     assert os.path.exists(public_key_path), (
         'Private key found, but associated public key '
         f'{public_key_path} does not exist.')
+    # Ensure correct permissions every time, as external factors (e.g.,
+    # Kubernetes fsGroup) can modify them after creation.
+    _ensure_key_permissions(private_key_path, public_key_path)
     return True

sky/utils/cli_utils/status_utils.py

@@ -13,9 +13,6 @@ from sky.utils import resources_utils
 from sky.utils import status_lib
 from sky.utils import ux_utils
 
-if typing.TYPE_CHECKING:
-    from sky.provision.kubernetes import utils as kubernetes_utils
-
 if typing.TYPE_CHECKING:
     from sky.provision.kubernetes import utils as kubernetes_utils
 
@@ -225,8 +222,25 @@ def show_cost_report_table(cluster_records: List[_ClusterCostReportRecord],
 # exist in those cases.
 _get_name = (lambda cluster_record, _: cluster_record['name'])
 _get_user_hash = (lambda cluster_record, _: cluster_record['user_hash'])
-_get_user_name = (
-    lambda cluster_record, _: cluster_record.get('user_name', '-'))
+
+
+def get_user_display_name(user_name: str, user_id: Optional[str] = None) -> str:
+    """ Appends SA to the user name if the user is a service account. """
+    if user_id and user_id.lower().startswith('sa-'):
+        return f'{user_name} (SA)'
+    return user_name
+
+
+def _get_user_name(cluster_record: _ClusterRecord,
+                   truncate: bool = True) -> str:
+    del truncate
+    user_name = cluster_record.get('user_name', '-')
+    if user_name == '-':
+        return user_name
+    user_hash = cluster_record.get('user_hash')
+    return get_user_display_name(user_name, user_hash)
+
+
 _get_launched = (lambda cluster_record, _: log_utils.readable_time_duration(
     cluster_record['launched_at']))
 _get_duration = (lambda cluster_record, _: log_utils.readable_time_duration(

sky/utils/cluster_utils.py

@@ -46,7 +46,8 @@ class SSHConfigHelper(object):
     ssh_cluster_key_path = constants.SKY_USER_FILE_PATH + '/ssh-keys/{}.key'
 
     @classmethod
-    def _get_generated_config(cls, autogen_comment: str, host_name: str,
+    def _get_generated_config(cls, autogen_comment: str,
+                              cluster_name_on_cloud: str, host_name: str,
                               ip: str, username: str, ssh_key_path: str,
                               proxy_command: Optional[str], port: int,
                               docker_proxy_command: Optional[str]):
@@ -79,6 +80,7 @@ class SSHConfigHelper(object):
               UserKnownHostsFile=/dev/null
               GlobalKnownHostsFile=/dev/null
               Port {port}
+              SetEnv {constants.SKY_CLUSTER_NAME_ENV_VAR_KEY}={cluster_name_on_cloud}
               {proxy}
             """.rstrip())
         codegen = codegen + '\n'
@@ -111,6 +113,7 @@ class SSHConfigHelper(object):
     def add_cluster(
         cls,
         cluster_name: str,
+        cluster_name_on_cloud: str,
         ips: List[str],
         auth_config: Dict[str, str],
         ports: List[int],
@@ -135,6 +138,7 @@ class SSHConfigHelper(object):
             ports: List of port numbers for SSH corresponding to ips
             docker_user: If not None, use this user to ssh into the docker
             ssh_user: Override the ssh_user in auth_config
+            cluster_name_on_cloud: The cluster name as it appears in the cloud.
         """
         if ssh_user is None:
             username = auth_config['ssh_user']
@@ -227,10 +231,13 @@ class SSHConfigHelper(object):
                 ip = 'localhost'
                 port = constants.DEFAULT_DOCKER_PORT
             node_name = cluster_name if i == 0 else cluster_name + f'-worker{i}'
+            node_proxy_command = proxy_command_for_nodes
+            if node_proxy_command is not None:
+                node_proxy_command = node_proxy_command.replace('%w', str(i))
             # TODO(romilb): Update port number when k8s supports multinode
             codegen += cls._get_generated_config(
-                sky_autogen_comment, node_name, ip, username,
-                key_path_for_config, proxy_command_for_nodes, port,
+                sky_autogen_comment, cluster_name_on_cloud, node_name, ip,
+                username, key_path_for_config, node_proxy_command, port,
                 docker_proxy_command) + '\n'
 
         cluster_config_path = os.path.expanduser(