skypilot-nightly 1.0.0.dev20251210__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
@@ -112,14 +122,14 @@ 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()
@@ -129,12 +139,12 @@ class PermissionService:
         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()
@@ -143,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/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(