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

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)