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

sky/utils/annotations.py

@@ -1,13 +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')
@@ -29,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.
@@ -50,13 +145,40 @@ 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
+
+
+def ttl_cache(scope: Literal['global', 'request'], *ttl_cache_args,
+              **ttl_cache_kwargs) -> Callable:
+    """TTLCache decorator for functions.
+
+    This decorator allows us to track which functions need to be reloaded for a
+    new request using the scope argument.
+    """
+
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
+        if scope == 'global':
+            return cachetools.cached(
+                cachetools.TTLCache(*ttl_cache_args, **ttl_cache_kwargs))(func)
+        else:
+            cached_func = cachetools.cached(
+                cachetools.TTLCache(*ttl_cache_args, **ttl_cache_kwargs))(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/asyncio_utils.py

@@ -0,0 +1,78 @@
+"""Asyncio utilities."""
+
+import asyncio
+import functools
+from typing import Set
+
+_background_tasks: Set[asyncio.Task] = set()
+
+
+def shield(func):
+    """Shield the decorated async function from cancellation.
+
+    If the outer coroutine is cancelled, the inner decorated function
+    will be protected from cancellation by asyncio.shield(). And we will
+    maintain a reference to the the inner task to avoid it get GCed before
+    it is done.
+
+    For example, filelock.AsyncFileLock is not cancellation safe. The
+    following code:
+
+        async def fn_with_lock():
+            async with filelock.AsyncFileLock('lock'):
+                await asyncio.sleep(1)
+
+    is equivalent to:
+
+        # The lock may leak if the cancellation happens in
+        # lock.acquire() or lock.release()
+        async def fn_with_lock():
+            lock = filelock.AsyncFileLock('lock')
+            await lock.acquire()
+            try:
+                await asyncio.sleep(1)
+            finally:
+                await lock.release()
+
+    Shilding the function ensures there is no cancellation will happen in the
+    function, thus the lock will be released properly:
+
+        @shield
+        async def fn_with_lock()
+
+    Note that the resource acquisition and release should usually be protected
+    in one @shield block but not separately, e.g.:
+
+        lock = filelock.AsyncFileLock('lock')
+
+        @shield
+        async def acquire():
+            await lock.acquire()
+
+        @shield
+        async def release():
+            await lock.release()
+
+        async def fn_with_lock():
+            await acquire()
+            try:
+                do_something()
+            finally:
+                await release()
+
+    The above code is not safe because if `fn_with_lock` is cancelled,
+    `acquire()` and `release()` will be executed in the background
+    concurrently and causes race conditions.
+    """
+
+    @functools.wraps(func)
+    async def async_wrapper(*args, **kwargs):
+        task = asyncio.create_task(func(*args, **kwargs))
+        try:
+            return await asyncio.shield(task)
+        except asyncio.CancelledError:
+            _background_tasks.add(task)
+            task.add_done_callback(lambda _: _background_tasks.discard(task))
+            raise
+
+    return async_wrapper

sky/utils/atomic.py

@@ -1,4 +1,4 @@
-"""Atomic structures and utilties."""
+"""Atomic structures and utilities."""
 
 import threading
 

sky/utils/auth_utils.py

@@ -0,0 +1,153 @@
+"""Utils for managing SkyPilot SSH key pairs."""
+
+import functools
+import os
+from typing import Tuple
+
+import filelock
+
+from sky import global_user_state
+from sky import sky_logging
+from sky.utils import common_utils
+
+logger = sky_logging.init_logger(__name__)
+
+MAX_TRIALS = 64
+# TODO(zhwu): Support user specified key pair.
+# We intentionally not have the ssh key pair to be stored in
+# ~/.sky/api_server/clients, i.e. sky.server.common.API_SERVER_CLIENT_DIR,
+# because ssh key pair need to persist across API server restarts, while
+# the former dir is ephemeral.
+_SSH_KEY_PATH_PREFIX = '~/.sky/clients/{user_hash}/ssh'
+
+
+def get_ssh_key_and_lock_path(user_hash: str) -> Tuple[str, str, str]:
+    user_ssh_key_prefix = _SSH_KEY_PATH_PREFIX.format(user_hash=user_hash)
+
+    os.makedirs(os.path.expanduser(user_ssh_key_prefix),
+                exist_ok=True,
+                mode=0o700)
+    private_key_path = os.path.join(user_ssh_key_prefix, 'sky-key')
+    public_key_path = os.path.join(user_ssh_key_prefix, 'sky-key.pub')
+    lock_path = os.path.join(user_ssh_key_prefix, '.__internal-sky-key.lock')
+    return private_key_path, public_key_path, lock_path
+
+
+def _generate_rsa_key_pair() -> Tuple[str, str]:
+    # Keep the import of the cryptography local to avoid expensive
+    # third-party imports when not needed.
+    # pylint: disable=import-outside-toplevel
+    from cryptography.hazmat.backends import default_backend
+    from cryptography.hazmat.primitives import serialization
+    from cryptography.hazmat.primitives.asymmetric import rsa
+
+    key = rsa.generate_private_key(backend=default_backend(),
+                                   public_exponent=65537,
+                                   key_size=2048)
+
+    private_key = key.private_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PrivateFormat.TraditionalOpenSSL,
+        encryption_algorithm=serialization.NoEncryption()).decode(
+            'utf-8').strip()
+
+    public_key = key.public_key().public_bytes(
+        serialization.Encoding.OpenSSH,
+        serialization.PublicFormat.OpenSSH).decode('utf-8').strip()
+
+    return public_key, private_key
+
+
+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)
+    os.makedirs(key_dir, exist_ok=True, mode=0o700)
+
+    with open(
+            private_key_path,
+            'w',
+            encoding='utf-8',
+            opener=functools.partial(os.open, mode=0o600),
+    ) as f:
+        f.write(private_key)
+
+    with open(public_key_path,
+              'w',
+              encoding='utf-8',
+              opener=functools.partial(os.open, mode=0o644)) as f:
+        f.write(public_key)
+
+
+def get_or_generate_keys() -> Tuple[str, str]:
+    """Returns the absolute private and public key paths."""
+    user_hash = common_utils.get_user_hash()
+    private_key_path, public_key_path, lock_path = get_ssh_key_and_lock_path(
+        user_hash)
+    private_key_path = os.path.expanduser(private_key_path)
+    public_key_path = os.path.expanduser(public_key_path)
+    lock_path = os.path.expanduser(lock_path)
+
+    lock_dir = os.path.dirname(lock_path)
+    # 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
+    # backend_utils.SSHConfigHelper
+    os.makedirs(lock_dir, exist_ok=True, mode=0o700)
+    with filelock.FileLock(lock_path, timeout=10):
+        if not os.path.exists(private_key_path):
+            ssh_public_key, ssh_private_key, exists = (
+                global_user_state.get_ssh_keys(user_hash))
+            if not exists:
+                ssh_public_key, ssh_private_key = _generate_rsa_key_pair()
+                global_user_state.set_ssh_keys(user_hash, ssh_public_key,
+                                               ssh_private_key)
+            _save_key_pair(private_key_path, public_key_path, ssh_private_key,
+                           ssh_public_key)
+    assert os.path.exists(public_key_path), (
+        'Private key found, but associated public key '
+        f'{public_key_path} does not exist.')
+    return private_key_path, public_key_path
+
+
+def create_ssh_key_files_from_db(private_key_path: str) -> bool:
+    """Creates the ssh key files from the database.
+
+    Returns:
+        True if the ssh key files are created successfully, False otherwise.
+    """
+    # Assume private key path is in the format of
+    # ~/.sky/clients/<user_hash>/ssh/sky-key
+    separated_path = os.path.normpath(private_key_path).split(os.path.sep)
+    assert separated_path[-1] == 'sky-key'
+    assert separated_path[-2] == 'ssh'
+    user_hash = separated_path[-3]
+
+    private_key_path_generated, public_key_path, lock_path = (
+        get_ssh_key_and_lock_path(user_hash))
+    assert private_key_path == os.path.expanduser(private_key_path_generated), (
+        f'Private key path {private_key_path} does not '
+        'match the generated path '
+        f'{os.path.expanduser(private_key_path_generated)}')
+    private_key_path = os.path.expanduser(private_key_path)
+    public_key_path = os.path.expanduser(public_key_path)
+    lock_path = os.path.expanduser(lock_path)
+    lock_dir = os.path.dirname(lock_path)
+
+    if os.path.exists(private_key_path) and os.path.exists(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
+    # backend_utils.SSHConfigHelper
+    os.makedirs(lock_dir, exist_ok=True, mode=0o700)
+    with filelock.FileLock(lock_path, timeout=10):
+        if not os.path.exists(private_key_path):
+            ssh_public_key, ssh_private_key, exists = (
+                global_user_state.get_ssh_keys(user_hash))
+            if not exists:
+                logger.debug(f'SSH keys not found for user {user_hash}')
+                return False
+            _save_key_pair(private_key_path, public_key_path, ssh_private_key,
+                           ssh_public_key)
+    assert os.path.exists(public_key_path), (
+        'Private key found, but associated public key '
+        f'{public_key_path} does not exist.')
+    return True

sky/utils/cli_utils/status_utils.py

@@ -11,6 +11,7 @@ from sky.utils import common_utils
 from sky.utils import log_utils
 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
@@ -105,11 +106,9 @@ def show_status_table(cluster_records: List[responses.StatusResponse],
 
     if query_clusters:
         cluster_names = {record['name'] for record in cluster_records}
-        not_found_clusters = [
-            repr(cluster)
-            for cluster in query_clusters
-            if cluster not in cluster_names
-        ]
+        not_found_clusters = ux_utils.get_non_matched_query(
+            query_clusters, cluster_names)
+        not_found_clusters = [repr(cluster) for cluster in not_found_clusters]
         if not_found_clusters:
             cluster_str = 'Cluster'
             if len(not_found_clusters) > 1:
@@ -283,8 +282,14 @@ def _get_resources(cluster_record: _ClusterRecord,
             if resources_str_full is not None:
                 resources_str = resources_str_full
         if resources_str is None:
-            resources_str = resources_utils.get_readable_resources_repr(
-                handle, simplify=truncate)
+            resources_str_simple, resources_str_full = (
+                resources_utils.get_readable_resources_repr(
+                    handle, simplified_only=truncate))
+            if truncate:
+                resources_str = resources_str_simple
+            else:
+                assert resources_str_full is not None
+                resources_str = resources_str_full
 
         return resources_str
     return '-'

sky/utils/cluster_utils.py

@@ -144,6 +144,9 @@ class SSHConfigHelper(object):
             username = docker_user
 
         key_path = cls.generate_local_key_file(cluster_name, auth_config)
+        # Keep the unexpanded path for SSH config (with ~)
+        key_path_for_config = key_path
+        # Expand the path for internal operations that need absolute path
         key_path = os.path.expanduser(key_path)
         sky_autogen_comment = ('# Added by sky (use `sky stop/down '
                                f'{cluster_name}` to remove)')
@@ -190,11 +193,29 @@ class SSHConfigHelper(object):
         proxy_command = auth_config.get('ssh_proxy_command', None)
 
         docker_proxy_command_generator = None
+        proxy_command_for_nodes = proxy_command
         if docker_user is not None:
-            docker_proxy_command_generator = lambda ip, port: ' '.join(
-                ['ssh'] + command_runner.ssh_options_list(
-                    key_path, ssh_control_name=None, port=port) +
-                ['-W', '%h:%p', f'{auth_config["ssh_user"]}@{ip}'])
+
+            def _docker_proxy_cmd(ip: str, port: int) -> str:
+                inner_proxy = proxy_command
+                inner_port = port or 22
+                if inner_proxy is not None:
+                    inner_proxy = inner_proxy.replace('%h', ip)
+                    inner_proxy = inner_proxy.replace('%p', str(inner_port))
+                return ' '.join(['ssh'] + command_runner.ssh_options_list(
+                    key_path,
+                    ssh_control_name=None,
+                    ssh_proxy_command=inner_proxy,
+                    port=inner_port,
+                    # ProxyCommand (ssh -W) is a forwarding tunnel, not an
+                    # interactive session. ControlMaster would cache these
+                    # processes, causing them to hang and block subsequent
+                    # connections. Each ProxyCommand should be ephemeral.
+                    disable_control_master=True
+                ) + ['-W', '%h:%p', f'{auth_config["ssh_user"]}@{ip}'])
+
+            docker_proxy_command_generator = _docker_proxy_cmd
+            proxy_command_for_nodes = None
 
         codegen = ''
         # Add the nodes to the codegen
@@ -208,8 +229,9 @@ class SSHConfigHelper(object):
             node_name = cluster_name if i == 0 else cluster_name + f'-worker{i}'
             # TODO(romilb): Update port number when k8s supports multinode
             codegen += cls._get_generated_config(
-                sky_autogen_comment, node_name, ip, username, key_path,
-                proxy_command, port, docker_proxy_command) + '\n'
+                sky_autogen_comment, node_name, ip, username,
+                key_path_for_config, proxy_command_for_nodes, port,
+                docker_proxy_command) + '\n'
 
         cluster_config_path = os.path.expanduser(
             cls.ssh_cluster_path.format(cluster_name))