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

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

@@ -63,6 +63,22 @@ def _ssh_control_path(ssh_control_filename: Optional[str]) -> Optional[str]:
     return path
 
 
+def _is_skypilot_managed_key(key_path: str) -> bool:
+    """Check if SSH key follows SkyPilot's managed key format.
+
+    SkyPilot-managed keys follow the pattern: ~/.sky/clients/<hash>/ssh/sky-key
+    External keys (like ~/.ssh/id_rsa) do not follow this pattern.
+
+    Args:
+        key_path: Path to the SSH private key.
+
+    Returns:
+        True if the key follows SkyPilot's managed format, False otherwise.
+    """
+    parts = os.path.normpath(key_path).split(os.path.sep)
+    return len(parts) >= 2 and parts[-1] == 'sky-key' and parts[-2] == 'ssh'
+
+
 # Disable sudo for root user. This is useful when the command is running in a
 # docker container, i.e. image_id is a docker image.
 ALIAS_SUDO_TO_EMPTY_FOR_ROOT_CMD = (
@@ -603,7 +619,7 @@ class SSHCommandRunner(CommandRunner):
         self,
         node: Tuple[str, int],
         ssh_user: str,
-        ssh_private_key: str,
+        ssh_private_key: Optional[str],
         ssh_control_name: Optional[str] = '__default__',
         ssh_proxy_command: Optional[str] = None,
         docker_user: Optional[str] = None,
@@ -613,7 +629,7 @@ class SSHCommandRunner(CommandRunner):
         """Initialize SSHCommandRunner.
 
         Example Usage:
-            runner = SSHCommandRunner(ip, ssh_user, ssh_private_key)
+            runner = SSHCommandRunner((ip, port), ssh_user, ssh_private_key)
             runner.run('ls -l', mode=SshMode.NON_INTERACTIVE)
             runner.rsync(source, target, up=True)
 
@@ -650,8 +666,17 @@ class SSHCommandRunner(CommandRunner):
         self.disable_control_master = (
             disable_control_master or
             control_master_utils.should_disable_control_master())
-        # ensure the ssh key files are created from the database
-        auth_utils.create_ssh_key_files_from_db(ssh_private_key)
+        # Ensure SSH key is available. For SkyPilot-managed keys, create from
+        # database. For external keys (e.g., Slurm clusters), verify existence.
+        if ssh_private_key is not None and _is_skypilot_managed_key(
+                ssh_private_key):
+            auth_utils.create_ssh_key_files_from_db(ssh_private_key)
+        elif ssh_private_key is not None:
+            # Externally managed key - just verify it exists
+            expanded_key_path = os.path.expanduser(ssh_private_key)
+            if not os.path.exists(expanded_key_path):
+                raise FileNotFoundError(
+                    f'SSH private key not found: {expanded_key_path}')
         if docker_user is not None:
             assert port is None or port == 22, (
                 f'port must be None or 22 for docker_user, got {port}.')
@@ -867,6 +892,7 @@ class SSHCommandRunner(CommandRunner):
         log_path: str = os.devnull,
         stream_logs: bool = True,
         max_retry: int = 1,
+        get_remote_home_dir: Callable[[], str] = lambda: '~',
     ) -> None:
         """Uses 'rsync' to sync 'source' to 'target'.
 
@@ -879,6 +905,8 @@ class SSHCommandRunner(CommandRunner):
             stream_logs: Stream logs to the stdout/stderr.
             max_retry: The maximum number of retries for the rsync command.
               This value should be non-negative.
+            get_remote_home_dir: A callable that returns the remote home
+              directory. Defaults to '~'.
 
         Raises:
             exceptions.CommandError: rsync command failed.
@@ -903,7 +931,8 @@ class SSHCommandRunner(CommandRunner):
                     rsh_option=rsh_option,
                     log_path=log_path,
                     stream_logs=stream_logs,
-                    max_retry=max_retry)
+                    max_retry=max_retry,
+                    get_remote_home_dir=get_remote_home_dir)
 
 
 class KubernetesCommandRunner(CommandRunner):
@@ -1247,3 +1276,166 @@ class LocalProcessCommandRunner(CommandRunner):
                     log_path=log_path,
                     stream_logs=stream_logs,
                     max_retry=max_retry)
+
+
+class SlurmCommandRunner(SSHCommandRunner):
+    """Runner for Slurm commands.
+
+    SlurmCommandRunner sends commands over an SSH connection through the Slurm
+    controller, to the virtual instances.
+    """
+
+    def __init__(
+        self,
+        node: Tuple[str, int],
+        ssh_user: str,
+        ssh_private_key: Optional[str],
+        *,
+        sky_dir: str,
+        skypilot_runtime_dir: str,
+        job_id: str,
+        slurm_node: str,
+        **kwargs,
+    ):
+        """Initialize SlurmCommandRunner.
+
+        Example Usage:
+            runner = SlurmCommandRunner(
+                (ip, port),
+                ssh_user,
+                ssh_private_key,
+                sky_dir=sky_dir,
+                skypilot_runtime_dir=skypilot_runtime_dir,
+                job_id=job_id,
+                slurm_node=slurm_node)
+            runner.run('ls -l', mode=SshMode.NON_INTERACTIVE)
+            runner.rsync(source, target, up=True)
+
+        Args:
+            node: (ip, port) The IP address and port of the remote machine
+              (login node).
+            ssh_user: SSH username.
+            ssh_private_key: Path to SSH private key.
+            sky_dir: The private directory for the SkyPilot cluster on the
+              Slurm cluster.
+            skypilot_runtime_dir: The directory for the SkyPilot runtime
+              on the Slurm cluster.
+            job_id: The Slurm job ID for this instance.
+            slurm_node: The Slurm node hostname for this instance
+              (compute node).
+            **kwargs: Additional arguments forwarded to SSHCommandRunner
+              (e.g., ssh_proxy_command).
+        """
+        super().__init__(node, ssh_user, ssh_private_key, **kwargs)
+        self.sky_dir = sky_dir
+        self.skypilot_runtime_dir = skypilot_runtime_dir
+        self.job_id = job_id
+        self.slurm_node = slurm_node
+
+        # Build a chained ProxyCommand that goes through the login node to reach
+        # the compute node where the job is running.
+
+        # First, build SSH options to reach the login node, using the user's
+        # existing proxy command if provided.
+        proxy_ssh_options = ' '.join(
+            ssh_options_list(self.ssh_private_key,
+                             None,
+                             ssh_proxy_command=self._ssh_proxy_command,
+                             port=self.port,
+                             disable_control_master=True))
+        login_node_proxy_command = (f'ssh {proxy_ssh_options} '
+                                    f'-W %h:%p {self.ssh_user}@{self.ip}')
+
+        # Update the proxy command to be the login node proxy, which will
+        # be used by super().run() to reach the compute node.
+        self._ssh_proxy_command = login_node_proxy_command
+        # Update self.ip to target the compute node.
+        self.ip = slurm_node
+        # Assume the compute node's SSH port is 22.
+        # TODO(kevin): Make this configurable if needed.
+        self.port = 22
+
+    def rsync(
+        self,
+        source: str,
+        target: str,
+        *,
+        up: bool,
+        log_path: str = os.devnull,
+        stream_logs: bool = True,
+        max_retry: int = 1,
+    ) -> None:
+        """Rsyncs files directly to the Slurm compute node,
+        by proxying through the Slurm login node.
+
+        For Slurm, files need to be accessible by compute nodes where jobs
+        execute via srun. This means either it has to be on the compute node's
+        local filesystem, or on a shared filesystem.
+        """
+        # TODO(kevin): We can probably optimize this to skip the proxying
+        # if the target dir is in a shared filesystem, since it will
+        # be accessible by the compute node.
+
+        # Build SSH options for rsync using the ProxyCommand set up in __init__
+        # to reach the compute node through the login node.
+        ssh_options = ' '.join(
+            ssh_options_list(
+                # Assume nothing and rely on default SSH behavior when -i is
+                # not specified.
+                None,
+                None,
+                ssh_proxy_command=self._ssh_proxy_command,
+                disable_control_master=True))
+        rsh_option = f'ssh {ssh_options}'
+
+        self._rsync(
+            source,
+            target,
+            # Compute node
+            node_destination=f'{self.ssh_user}@{self.slurm_node}',
+            up=up,
+            rsh_option=rsh_option,
+            log_path=log_path,
+            stream_logs=stream_logs,
+            max_retry=max_retry,
+            get_remote_home_dir=lambda: self.sky_dir)
+
+    @timeline.event
+    @context_utils.cancellation_guard
+    def run(self, cmd: Union[str, List[str]],
+            **kwargs) -> Union[int, Tuple[int, str, str]]:
+        """Run Slurm-supported user commands over an SSH connection.
+
+        Args:
+            cmd: The Slurm-supported user command to run.
+
+        Returns:
+            returncode
+            or
+            A tuple of (returncode, stdout, stderr).
+        """
+        # Override $HOME so that each SkyPilot cluster's state is isolated
+        # from one another. We rely on the assumption that ~ is exclusively
+        # used by a cluster, and in Slurm that is not the case, as $HOME
+        # could be part of a shared filesystem.
+        # And similarly for SKY_RUNTIME_DIR. See constants.\
+        # SKY_RUNTIME_DIR_ENV_VAR_KEY for more details.
+        #
+        # SSH directly to the compute node instead of using srun.
+        # This avoids Slurm's proctrack/cgroup which kills all processes
+        # when the job step ends (including child processes launched as
+        # a separate process group), breaking background process spawning
+        # (e.g., JobScheduler._run_job which uses launch_new_process_tree).
+        # Note: proctrack/cgroup is enabled by default on Nebius'
+        # Managed Soperator.
+        cmd = (
+            f'export {constants.SKY_RUNTIME_DIR_ENV_VAR_KEY}='
+            f'"{self.skypilot_runtime_dir}" && '
+            # Set the uv cache directory to /tmp/uv_cache_$(id -u) to speed up
+            # package installation while avoiding permission conflicts when
+            # multiple users share the same host. Otherwise it defaults to
+            # ~/.cache/uv.
+            f'export UV_CACHE_DIR=/tmp/uv_cache_$(id -u) && '
+            f'cd {self.sky_dir} && export HOME=$(pwd) && {cmd}')
+
+        return super().run(cmd, **kwargs)

sky/utils/command_runner.pyi

@@ -6,7 +6,7 @@ determine the return type based on the value of require_outputs.
 """
 import enum
 import typing
-from typing import Any, Iterable, List, Optional, Tuple, Union
+from typing import Any, Callable, Iterable, List, Optional, Tuple, Union
 
 from typing_extensions import Literal
 
@@ -130,7 +130,7 @@ class SSHCommandRunner(CommandRunner):
     ip: str
     port: int
     ssh_user: str
-    ssh_private_key: str
+    ssh_private_key: Optional[str]
     ssh_control_name: Optional[str]
     docker_user: str
     disable_control_master: Optional[bool]
@@ -140,7 +140,7 @@ class SSHCommandRunner(CommandRunner):
         self,
         node: Tuple[str, int],
         ssh_user: str,
-        ssh_private_key: str,
+        ssh_private_key: Optional[str],
         ssh_control_name: Optional[str] = ...,
         ssh_proxy_command: Optional[str] = ...,
         docker_user: Optional[str] = ...,
@@ -216,7 +216,8 @@ class SSHCommandRunner(CommandRunner):
               up: bool,
               log_path: str = ...,
               stream_logs: bool = ...,
-              max_retry: int = ...) -> None:
+              max_retry: int = ...,
+              get_remote_home_dir: Callable[[], str] = ...) -> None:
         ...
 
     def port_forward_command(
@@ -306,6 +307,28 @@ class KubernetesCommandRunner(CommandRunner):
         ...
 
 
+class SlurmCommandRunner(SSHCommandRunner):
+    """Runner for Slurm commands."""
+    sky_dir: str
+    skypilot_runtime_dir: str
+    job_id: str
+    slurm_node: str
+
+    def __init__(
+        self,
+        node: Tuple[str, int],
+        ssh_user: str,
+        ssh_private_key: Optional[str],
+        *,
+        sky_dir: str,
+        skypilot_runtime_dir: str,
+        job_id: str,
+        slurm_node: str,
+        **kwargs,
+    ) -> None:
+        ...
+
+
 class LocalProcessCommandRunner(CommandRunner):
 
     def __init__(self) -> None:

sky/utils/common_utils.py

@@ -300,6 +300,7 @@ _current_user: Optional['models.User'] = None
 _current_request_id: Optional[str] = None
 
 
+# TODO(aylei,hailong): request context should be contextual
 def set_request_context(client_entrypoint: Optional[str],
                         client_command: Optional[str],
                         using_remote_api_server: bool,
@@ -341,19 +342,32 @@ def get_current_command() -> str:
 
 
 def get_current_user() -> 'models.User':
-    """Returns the current user."""
+    """Returns the user in current server session."""
     if _current_user is not None:
         return _current_user
     return models.User.get_current_user()
 
 
 def get_current_user_name() -> str:
-    """Returns the current user name."""
+    """Returns the user name in current server session."""
     name = get_current_user().name
     assert name is not None
     return name
 
 
+def get_local_user_name() -> str:
+    """Returns the user name in local environment.
+
+    This is for backward compatibility where anonymous access is implicitly
+    allowed when no authentication method at server-side is configured and
+    the username from client environment variable will be used to identify the
+    user.
+    """
+    name = os.getenv(constants.USER_ENV_VAR, getpass.getuser())
+    assert name is not None
+    return name
+
+
 def set_current_user(user: 'models.User'):
     """Sets the current user."""
     global _current_user
@@ -724,7 +738,8 @@ def find_free_port(start_port: int) -> int:
             try:
                 s.bind(('', port))
                 return port
-            except OSError:
+            except OSError as e:
+                logger.debug(f'Error binding port {port}: {e}')
                 pass
     raise OSError('No free ports available.')