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

sky/utils/common_utils.py

@@ -1,8 +1,10 @@
 """Utils shared between all of sky"""
 
+import ctypes
 import difflib
 import enum
 import functools
+import gc
 import getpass
 import hashlib
 import inspect
@@ -263,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
@@ -281,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)
@@ -295,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,
@@ -336,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
@@ -719,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.')
 
@@ -994,7 +1014,17 @@ def get_mem_size_gb() -> float:
         except ValueError as e:
             with ux_utils.print_exception_no_traceback():
                 raise ValueError(
-                    f'Failed to parse the memory size from {mem_size}') from e
+                    f'Failed to parse the memory size from {mem_size} (GB)'
+                ) from e
+    mem_size = os.getenv('SKYPILOT_POD_MEMORY_BYTES_LIMIT')
+    if mem_size is not None:
+        try:
+            return float(mem_size) / (1024**3)
+        except ValueError as e:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    f'Failed to parse the memory size from {mem_size} (bytes)'
+                ) from e
     return _mem_size_gb()
 
 
@@ -1090,3 +1120,21 @@ def removeprefix(string: str, prefix: str) -> str:
     if string.startswith(prefix):
         return string[len(prefix):]
     return string
+
+
+def release_memory():
+    """Release the process memory"""
+    # Do the best effort to release the python heap and let malloc_trim
+    # be more efficient.
+    try:
+        gc.collect()
+        if sys.platform.startswith('linux'):
+            # Will fail on musl (alpine), but at least it works on our
+            # official docker images.
+            libc = ctypes.CDLL('libc.so.6')
+            return libc.malloc_trim(0)
+        return 0
+    except Exception as e:  # pylint: disable=broad-except
+        logger.error(f'Failed to release memory: '
+                     f'{format_exception(e)}')
+        return 0

sky/utils/config_utils.py

@@ -272,7 +272,7 @@ def get_cloud_config_value_from_dict(
     """
     input_config = Config(dict_config)
     region_key = None
-    if cloud == 'kubernetes':
+    if cloud in ('kubernetes', 'ssh'):
         region_key = 'context_configs'
     elif cloud in _REGION_CONFIG_CLOUDS:
         region_key = 'region_configs'
@@ -283,19 +283,6 @@ def get_cloud_config_value_from_dict(
             keys=(cloud, region_key, region) + keys,
             default_value=None,
             override_configs=override_configs)
-        if not per_context_config and cloud in _REGION_CONFIG_CLOUDS:
-            # TODO (kyuds): Backward compatibility, remove after 0.11.0.
-            per_context_config = input_config.get_nested(
-                keys=(cloud, region) + keys,
-                default_value=None,
-                override_configs=override_configs)
-            if per_context_config is not None:
-                logger.info(
-                    f'{cloud} configuration is using the legacy format. \n'
-                    'This format will be deprecated after 0.11.0, refer to '
-                    '`https://docs.skypilot.co/en/latest/reference/config.html` '  # pylint: disable=line-too-long
-                    'for the new format. Please use `region_configs` to specify region specific configuration.'
-                )
     # if no override found for specified region
     general_config = input_config.get_nested(keys=(cloud,) + keys,
                                              default_value=default_value,

sky/utils/context.py

@@ -2,18 +2,23 @@
 
 import asyncio
 from collections.abc import Mapping
-from collections.abc import MutableMapping
 import contextvars
+import copy
 import functools
 import os
 import pathlib
 import subprocess
 import sys
-import typing
-from typing import Any, Callable, Dict, Optional, TextIO, TypeVar
+from typing import (Any, Callable, Coroutine, Dict, Iterator, MutableMapping,
+                    Optional, TextIO, TYPE_CHECKING, TypeVar)
 
+from typing_extensions import ParamSpec
 
-class Context(object):
+if TYPE_CHECKING:
+    from sky.skypilot_config import ConfigContext
+
+
+class SkyPilotContext(object):
     """SkyPilot typed context vars for threads and coroutines.
 
     This is a wrapper around `contextvars.ContextVar` that provides a typed
@@ -88,7 +93,7 @@ class Context(object):
         else:
             self._log_file_handle = open(log_file, 'a', encoding='utf-8')
         self._log_file = log_file
-        if original_log_file is not None:
+        if original_log_handle is not None:
             original_log_handle.close()
         return original_log_file
 
@@ -102,11 +107,40 @@ class Context(object):
         for k, v in envs.items():
             self.env_overrides[k] = v
 
+    def cleanup(self):
+        """Clean up the context."""
+        if self._log_file_handle is not None:
+            self._log_file_handle.close()
+            self._log_file_handle = None
+
+    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.
+        The new context will get its own handle/fd for the log file.
+        The new context will get an independent copy of the env var overrides.
+        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 = 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('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
@@ -116,7 +150,7 @@ def get() -> Optional[Context]:
     return _CONTEXT.get()
 
 
-class ContextualEnviron(MutableMapping):
+class ContextualEnviron(MutableMapping[str, str]):
     """Environment variables wrapper with contextual overrides.
 
     An instance of ContextualEnviron will typically be used to replace
@@ -124,7 +158,7 @@ class ContextualEnviron(MutableMapping):
     aware.
 
     Behavior of spawning a subprocess:
-    - The contexual overrides will not be applied to the subprocess by
+    - The contextual overrides will not be applied to the subprocess by
       default.
     - When using env=os.environ to pass the environment variables to the
       subprocess explicitly. The subprocess will inherit the contextual
@@ -155,10 +189,10 @@ class ContextualEnviron(MutableMapping):
        assert os.environ['FOO'] == 'BAR1'
     """
 
-    def __init__(self, environ):
+    def __init__(self, environ: 'os._Environ[str]') -> None:
         self._environ = environ
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: str) -> str:
         ctx = get()
         if ctx is not None:
             if key in ctx.env_overrides:
@@ -170,51 +204,63 @@ class ContextualEnviron(MutableMapping):
                 return value
         return self._environ[key]
 
-    def __iter__(self):
-        ctx = get()
-        deleted_keys = set()
-        if ctx is not None:
+    def __iter__(self) -> 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:
                     deleted_keys.add(key)
-                yield key
+                else:
+                    yield key
             for key in self._environ:
                 # Deduplicate the keys
                 if key not in ctx.env_overrides and key not in deleted_keys:
                     yield key
+
+        ctx = get()
+        if ctx is not None:
+            return iter_from_context(ctx)
         else:
             return self._environ.__iter__()
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(dict(self))
 
-    def __setitem__(self, key, value):
+    def __setitem__(self, key: str, value: str) -> None:
         ctx = get()
         if ctx is not None:
             ctx.env_overrides[key] = value
         else:
             self._environ.__setitem__(key, value)
 
-    def __delitem__(self, key):
+    def __delitem__(self, key: str) -> None:
         ctx = get()
         if ctx is not None:
-            if key in ctx.env_overrides:
-                del ctx.env_overrides[key]
-            elif key in self._environ:
-                # If the key is not set in the context but set in the environ
-                # of the process, we mark it as deleted in the context by
-                # setting the value to None.
+            if key in self._environ:
+                # If the key is set in the environ of the process, we mark it as
+                # deleted in the context by setting the value to None.
+                # Note: we must do this even if it was also set in the context,
+                # since it could be set in both, and deleting should delete it
+                # from both.
                 ctx.env_overrides[key] = None
+            elif key in ctx.env_overrides:
+                # If the key is set in the context, but not the original
+                # environ, we can just delete the override.
+                del ctx.env_overrides[key]
             else:
                 # The key is not set in the context nor the process.
                 raise KeyError(key)
         else:
             self._environ.__delitem__(key)
 
-    def __repr__(self):
-        return self._environ.__repr__()
+    def __repr__(self) -> str:
+        # Adapted from os._Environ.__repr__
+        formatted_items = ', '.join(
+            f'{key!r}: {value!r}' for key, value in self.items())
+        return f'ctx_environ({{{formatted_items}}})'
 
-    def copy(self):
+    def copy(self) -> Dict[str, str]:
         copied = self._environ.copy()
         ctx = get()
         if ctx is not None:
@@ -225,7 +271,7 @@ class ContextualEnviron(MutableMapping):
                     copied[key] = ctx.env_overrides[key]
         return copied
 
-    def setdefault(self, key, default=None):
+    def setdefault(self, key: str, default: str) -> str:
         return self._environ.setdefault(key, default)
 
     def __ior__(self, other):
@@ -257,30 +303,71 @@ class Popen(subprocess.Popen):
             # Pass a copy of current context.environ to avoid race condition
             # when the context is updated after the Popen is created.
             env = os.environ.copy()
-        super().__init__(*args, env=env, **kwargs)
+        super().__init__(*args, env=env,
+                         **kwargs)  # type: ignore[call-overload]
 
 
-F = TypeVar('F', bound=Callable[..., Any])
+P = ParamSpec('P')
+T = TypeVar('T')
 
 
-def contextual(func: F) -> F:
+def contextual(func: Callable[P, T]) -> Callable[P, T]:
     """Decorator to initialize a context before executing the function.
 
-    If a context is already initialized, this decorator will reset the context,
-    i.e. all contextual variables set previously will be cleared.
+    If a context is already initialized, this decorator will create a new
+    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()
+        with initialize(original_ctx):
+            return await func(*args, **kwargs)
+
     @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        initialize()
-        return func(*args, **kwargs)
+    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 typing.cast(F, wrapper)
+    return wrapper
 
 
-def initialize():
+def initialize(
+        base_context: Optional[SkyPilotContext] = None) -> SkyPilotContext:
     """Initialize the current SkyPilot context."""
-    _CONTEXT.set(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,20 +1,27 @@
 """Utilities for SkyPilot context."""
 import asyncio
+import concurrent.futures
 import contextvars
 import functools
-import io
 import multiprocessing
 import os
+import select
 import subprocess
 import sys
+import time
 import typing
 from typing import Any, Callable, IO, Optional, Tuple, TypeVar
 
+from typing_extensions import ParamSpec
+
 from sky import sky_logging
 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
@@ -41,23 +48,53 @@ 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"""
-    wrapped = io.TextIOWrapper(in_stream,
-                               encoding='utf-8',
-                               newline='',
-                               errors='replace',
-                               write_through=True)
+    last_flush_time = time.time()
+    has_unflushed_content = False
+
+    # Use poll() with timeout instead of readline() to avoid blocking.
+    # readline() blocks until a newline is available, which can take minutes
+    # for tasks that emit logs infrequently (e.g. jupyter lab server).
+    # While readline() is blocked, the timing code never executes, so buffered
+    # logs never get flushed. poll() with timeout allows us to periodically
+    # flush even when no new data is available, ensuring logs appear promptly.
+    fd = in_stream.fileno()
+    poller = select.poll()
+    poller.register(fd, select.POLLIN)
+
+    # Timeout in milliseconds for poll()
+    poll_timeout_ms = int(PASSTHROUGH_FLUSH_INTERVAL_SECONDS * 1000)
+
     while True:
-        line = wrapped.readline()
-        if line:
-            out_stream.write(line)
+        # Poll with timeout - returns when data available or timeout
+        events = poller.poll(poll_timeout_ms)
+
+        current_time = time.time()
+
+        if events:
+            # Data is available, read a chunk
+            chunk = os.read(fd, 4096)  # Read up to 4KB
+            if not chunk:
+                break  # EOF
+            out_stream.write(chunk.decode('utf-8', errors='replace'))
+            has_unflushed_content = True
+
+        # Flush only if we have unflushed content and timeout reached
+        if (has_unflushed_content and current_time - last_flush_time >=
+                PASSTHROUGH_FLUSH_INTERVAL_SECONDS):
             out_stream.flush()
-        else:
-            break
+            last_flush_time = current_time
+            has_unflushed_content = False
+
+    poller.unregister(fd)
+    # Final flush to ensure all data is written
+    if has_unflushed_content:
+        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,
@@ -110,7 +147,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):
@@ -128,7 +165,11 @@ def wait_process(ctx: context.Context,
             # Kill the process despite the caller's callback, the utility
             # function gracefully handles the case where the process is
             # already terminated.
-            subprocess_utils.kill_process_with_grace_period(proc)
+            # Bash script typically does not forward SIGTERM to childs, thus
+            # cannot be killed gracefully, shorten the grace period for faster
+            # termination.
+            subprocess_utils.kill_process_with_grace_period(proc,
+                                                            grace_period=1)
             raise asyncio.CancelledError()
         try:
             proc.wait(poll_interval)
@@ -173,15 +214,29 @@ def cancellation_guard(func: F) -> F:
     return typing.cast(F, wrapper)
 
 
+P = ParamSpec('P')
+T = TypeVar('T')
+
+
 # TODO(aylei): replace this with asyncio.to_thread once we drop support for
 # python 3.8
-def to_thread(func, /, *args, **kwargs):
+def to_thread(func: Callable[P, T], /, *args: P.args,
+              **kwargs: P.kwargs) -> 'asyncio.Future[T]':
     """Asynchronously run function *func* in a separate thread.
 
     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 = functools.partial(pyctx.run, 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)