skypilot-nightly 1.0.0.dev20251009__py3-none-any.whl → 1.0.0.dev20251107__py3-none-any.whl

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 outter 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/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

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

@@ -14,6 +14,7 @@ from sky import exceptions
 from sky import sky_logging
 from sky.skylet import constants
 from sky.skylet import log_lib
+from sky.utils import auth_utils
 from sky.utils import common_utils
 from sky.utils import context_utils
 from sky.utils import control_master_utils
@@ -649,6 +650,8 @@ 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)
         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}.')
@@ -961,8 +964,16 @@ class KubernetesCommandRunner(CommandRunner):
         kubectl_args = [
             '--pod-running-timeout', f'{connect_timeout}s', '-n', self.namespace
         ]
+        # The same logic to either set `--context` to the k8s context where
+        # the sky cluster is hosted, or `--kubeconfig` to /dev/null for
+        # in-cluster k8s is used below in the `run()` method.
         if self.context:
             kubectl_args += ['--context', self.context]
+        # If context is none, it means the cluster is hosted on in-cluster k8s.
+        # In this case, we need to set KUBECONFIG to /dev/null to avoid looking
+        # for the cluster in whatever active context is set in the kubeconfig.
+        else:
+            kubectl_args += ['--kubeconfig', '/dev/null']
         local_port, remote_port = port_forward[0]
         local_port_str = f'{local_port}' if local_port is not None else ''
 

sky/utils/common.py

@@ -31,7 +31,7 @@ JOB_CONTROLLER_NAME: str
 def refresh_server_id() -> None:
     """Refresh the server id.
 
-    This function is used to ensure the server id is read from the authorative
+    This function is used to ensure the server id is read from the authoritative
     source.
     """
     global SERVER_ID
@@ -42,6 +42,8 @@ def refresh_server_id() -> None:
     JOB_CONTROLLER_NAME = f'{JOB_CONTROLLER_PREFIX}{SERVER_ID}'
 
 
+# TODO(kevin): Remove this side effect and have callers call
+# refresh_server_id() explicitly as needed.
 refresh_server_id()
 
 

sky/utils/common_utils.py

@@ -265,13 +265,16 @@ def get_global_job_id(job_timestamp: str,
 
 class Backoff:
     """Exponential backoff with jittering."""
-    MULTIPLIER = 1.6
     JITTER = 0.4
 
-    def __init__(self, initial_backoff: float = 5, max_backoff_factor: int = 5):
+    def __init__(self,
+                 initial_backoff: float = 5,
+                 max_backoff_factor: int = 5,
+                 multiplier: float = 1.6):
         self._initial = True
         self._backoff = 0.0
         self._initial_backoff = initial_backoff
+        self._multiplier = multiplier
         self._max_backoff = max_backoff_factor * self._initial_backoff
 
     # https://github.com/grpc/grpc/blob/2d4f3c56001cd1e1f85734b2f7c5ce5f2797c38a/doc/connection-backoff.md
@@ -283,7 +286,7 @@ class Backoff:
             self._initial = False
             self._backoff = min(self._initial_backoff, self._max_backoff)
         else:
-            self._backoff = min(self._backoff * self.MULTIPLIER,
+            self._backoff = min(self._backoff * self._multiplier,
                                 self._max_backoff)
         self._backoff += random.uniform(-self.JITTER * self._backoff,
                                         self.JITTER * self._backoff)
@@ -1112,7 +1115,7 @@ def release_memory():
         gc.collect()
         if sys.platform.startswith('linux'):
             # Will fail on musl (alpine), but at least it works on our
-            # offical docker images.
+            # official docker images.
             libc = ctypes.CDLL('libc.so.6')
             return libc.malloc_trim(0)
         return 0

sky/utils/context.py

@@ -5,13 +5,12 @@ from collections.abc import Mapping
 import contextvars
 import copy
 import functools
-import inspect
 import os
 import pathlib
 import subprocess
 import sys
-from typing import (Callable, Dict, Iterator, MutableMapping, Optional, TextIO,
-                    TYPE_CHECKING, TypeVar)
+from typing import (Any, Callable, Coroutine, Dict, Iterator, MutableMapping,
+                    Optional, TextIO, TYPE_CHECKING, TypeVar)
 
 from typing_extensions import ParamSpec
 
@@ -19,7 +18,7 @@ if TYPE_CHECKING:
     from sky.skypilot_config import ConfigContext
 
 
-class Context(object):
+class SkyPilotContext(object):
     """SkyPilot typed context vars for threads and coroutines.
 
     This is a wrapper around `contextvars.ContextVar` that provides a typed
@@ -114,7 +113,14 @@ class Context(object):
             self._log_file_handle.close()
             self._log_file_handle = None
 
-    def copy(self) -> 'Context':
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        del exc_type, exc_val, exc_tb
+        self.cleanup()
+
+    def copy(self) -> 'SkyPilotContext':
         """Create a copy of the context.
 
         Changes to the current context after this call will not affect the copy.
@@ -123,18 +129,18 @@ class Context(object):
         The new context will get an independent copy of the config context.
         Cancellation of the current context will not be propagated to the copy.
         """
-        new_context = Context()
+        new_context = SkyPilotContext()
         new_context.redirect_log(self._log_file)
         new_context.env_overrides = self.env_overrides.copy()
         new_context.config_context = copy.deepcopy(self.config_context)
         return new_context
 
 
-_CONTEXT = contextvars.ContextVar[Optional[Context]]('sky_context',
-                                                     default=None)
+_CONTEXT = contextvars.ContextVar[Optional[SkyPilotContext]]('sky_context',
+                                                             default=None)
 
 
-def get() -> Optional[Context]:
+def get() -> Optional[SkyPilotContext]:
     """Get the current SkyPilot context.
 
     If the context is not initialized, get() will return None. This helps
@@ -200,7 +206,7 @@ class ContextualEnviron(MutableMapping[str, str]):
 
     def __iter__(self) -> Iterator[str]:
 
-        def iter_from_context(ctx: Context) -> Iterator[str]:
+        def iter_from_context(ctx: SkyPilotContext) -> Iterator[str]:
             deleted_keys = set()
             for key, value in ctx.env_overrides.items():
                 if value is None:
@@ -311,56 +317,56 @@ def contextual(func: Callable[P, T]) -> Callable[P, T]:
     context that inherits the values from the existing context.
     """
 
+    def run_in_context(*args: P.args, **kwargs: P.kwargs) -> T:
+        # Within the new contextvars Context, set up the SkyPilotContext.
+        original_ctx = get()
+        with initialize(original_ctx):
+            return func(*args, **kwargs)
+
     @functools.wraps(func)
     def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+        # Create a copy of the current contextvars Context so that setting the
+        # SkyPilotContext does not affect the caller's context in async
+        # environments.
+        context = contextvars.copy_context()
+        return context.run(run_in_context, *args, **kwargs)
+
+    return wrapper
+
+
+def contextual_async(
+    func: Callable[P, Coroutine[Any, Any, T]]
+) -> Callable[P, Coroutine[Any, Any, T]]:
+    """Decorator to initialize a context before executing the function.
+
+    If a context is already initialized, this decorator will create a new
+    context that inherits the values from the existing context.
+    """
+
+    async def run_in_context(*args: P.args, **kwargs: P.kwargs) -> T:
+        # Within the new contextvars Context, set up the SkyPilotContext.
         original_ctx = get()
-        initialize(original_ctx)
-        ctx = get()
-        cleanup_after_await = False
-
-        def cleanup():
-            try:
-                if ctx is not None:
-                    ctx.cleanup()
-            finally:
-                # Note: _CONTEXT.reset() is not reliable - may fail with
-                # ValueError: <Token ... at ...> was created in a different
-                # Context
-                # We must make sure this happens because otherwise we may try to
-                # write to the wrong log.
-                _CONTEXT.set(original_ctx)
-
-        # There are two cases:
-        # 1. The function is synchronous (that is, return type is not awaitable)
-        #    In this case, we use a finally block to cleanup the context.
-        # 2. The function is asynchronous (that is, return type is awaitable)
-        #    In this case, we need to construct an async def wrapper and await
-        #    the value, then call the cleanup function in the finally block.
-
-        async def await_with_cleanup(awaitable):
-            try:
-                return await awaitable
-            finally:
-                cleanup()
-
-        try:
-            ret = func(*args, **kwargs)
-            if inspect.isawaitable(ret):
-                cleanup_after_await = True
-                return await_with_cleanup(ret)
-            else:
-                return ret
-        finally:
-            if not cleanup_after_await:
-                cleanup()
+        with initialize(original_ctx):
+            return await func(*args, **kwargs)
+
+    @functools.wraps(func)
+    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+        # Create a copy of the current contextvars Context so that setting the
+        # SkyPilotContext does not affect the caller's context in async
+        # environments.
+        context = contextvars.copy_context()
+        return await context.run(run_in_context, *args, **kwargs)
 
     return wrapper
 
 
-def initialize(base_context: Optional[Context] = None) -> None:
+def initialize(
+        base_context: Optional[SkyPilotContext] = None) -> SkyPilotContext:
     """Initialize the current SkyPilot context."""
-    new_context = base_context.copy() if base_context is not None else Context()
+    new_context = base_context.copy(
+    ) if base_context is not None else SkyPilotContext()
     _CONTEXT.set(new_context)
+    return new_context
 
 
 class _ContextualStream:

sky/utils/context_utils.py

@@ -1,5 +1,6 @@
 """Utilities for SkyPilot context."""
 import asyncio
+import concurrent.futures
 import contextvars
 import functools
 import io
@@ -7,6 +8,7 @@ import multiprocessing
 import os
 import subprocess
 import sys
+import time
 import typing
 from typing import Any, Callable, IO, Optional, Tuple, TypeVar
 
@@ -17,6 +19,9 @@ from sky.utils import context
 from sky.utils import subprocess_utils
 
 StreamHandler = Callable[[IO[Any], IO[Any]], str]
+PASSTHROUGH_FLUSH_INTERVAL_SECONDS = 0.5
+
+logger = sky_logging.init_logger(__name__)
 
 
 # TODO(aylei): call hijack_sys_attrs() proactivly in module init at server-side
@@ -43,6 +48,7 @@ def hijack_sys_attrs():
 
 def passthrough_stream_handler(in_stream: IO[Any], out_stream: IO[Any]) -> str:
     """Passthrough the stream from the process to the output stream"""
+    last_flush_time = time.time()
     wrapped = io.TextIOWrapper(in_stream,
                                encoding='utf-8',
                                newline='',
@@ -52,14 +58,23 @@ def passthrough_stream_handler(in_stream: IO[Any], out_stream: IO[Any]) -> str:
         line = wrapped.readline()
         if line:
             out_stream.write(line)
-            out_stream.flush()
+
+            # Flush based on timeout instead of on every line
+            current_time = time.time()
+            if (current_time - last_flush_time >=
+                    PASSTHROUGH_FLUSH_INTERVAL_SECONDS):
+                out_stream.flush()
+                last_flush_time = current_time
         else:
             break
+
+    # Final flush to ensure all data is written
+    out_stream.flush()
     return ''
 
 
 def pipe_and_wait_process(
-        ctx: context.Context,
+        ctx: context.SkyPilotContext,
         proc: subprocess.Popen,
         poll_interval: float = 0.5,
         cancel_callback: Optional[Callable[[], None]] = None,
@@ -112,7 +127,7 @@ def pipe_and_wait_process(
         return stdout, stderr
 
 
-def wait_process(ctx: context.Context,
+def wait_process(ctx: context.SkyPilotContext,
                  proc: subprocess.Popen,
                  poll_interval: float = 0.5,
                  cancel_callback: Optional[Callable[[], None]] = None):
@@ -191,14 +206,17 @@ def to_thread(func: Callable[P, T], /, *args: P.args,
 
     This is same as asyncio.to_thread added in python 3.9
     """
+    return to_thread_with_executor(None, func, *args, **kwargs)
+
+
+def to_thread_with_executor(executor: Optional[concurrent.futures.Executor],
+                            func: Callable[P, T], /, *args: P.args,
+                            **kwargs: P.kwargs) -> 'asyncio.Future[T]':
+    """Asynchronously run function *func* in a separate thread with
+    a custom executor."""
+
     loop = asyncio.get_running_loop()
-    # This is critical to pass the current coroutine context to the new thread
     pyctx = contextvars.copy_context()
-    func_call: Callable[..., T] = functools.partial(
-        # partial deletes arguments type and thus can't figure out the return
-        # type of pyctx.run
-        pyctx.run,  # type: ignore
-        func,
-        *args,
-        **kwargs)
-    return loop.run_in_executor(None, func_call)
+    func_call: Callable[..., T] = functools.partial(pyctx.run, func, *args,
+                                                    **kwargs)
+    return loop.run_in_executor(executor, func_call)