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

sky/utils/context.py

@@ -0,0 +1,402 @@
+"""SkyPilot context for threads and coroutines."""
+
+import asyncio
+from collections.abc import Mapping
+import contextvars
+import copy
+import functools
+import os
+import pathlib
+import subprocess
+import sys
+from typing import (Any, Callable, Coroutine, Dict, Iterator, MutableMapping,
+                    Optional, TextIO, TYPE_CHECKING, TypeVar)
+
+from typing_extensions import ParamSpec
+
+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
+    interface for the SkyPilot specific context variables that can be accessed
+    at any layer of the call stack. ContextVar is coroutine local, an empty
+    Context will be initialized for each coroutine when it is created.
+
+    Adding a new context variable for a new feature is as simple as:
+    1. Add a new instance variable to the Context class.
+    2. (Optional) Add new accessor methods if the variable should be protected.
+
+    To propagate the context to a new thread/coroutine, use
+    `contextvars.copy_context()`.
+
+    Example:
+        import asyncio
+        import contextvars
+        import time
+        from sky.utils import context
+
+        def sync_task():
+            while True:
+                if context.get().is_canceled():
+                    break
+                time.sleep(1)
+
+        async def fastapi_handler():
+            # context.initialize() has been called in lifespan
+            ctx = contextvars.copy_context()
+            # asyncio.to_thread copies current context implicitly
+            task = asyncio.to_thread(sync_task)
+            # Or explicitly:
+            # loop = asyncio.get_running_loop()
+            # ctx = contextvars.copy_context()
+            # task = loop.run_in_executor(None, ctx.run, sync_task)
+            await asyncio.sleep(1)
+            context.get().cancel()
+            await task
+    """
+
+    def __init__(self):
+        self._canceled = asyncio.Event()
+        self._log_file = None
+        self._log_file_handle = None
+        self.env_overrides = {}
+        self.config_context = None
+
+    def cancel(self):
+        """Cancel the context."""
+        self._canceled.set()
+
+    def is_canceled(self):
+        """Check if the context is canceled."""
+        return self._canceled.is_set()
+
+    def redirect_log(
+            self, log_file: Optional[pathlib.Path]) -> Optional[pathlib.Path]:
+        """Redirect the stdout and stderr of current context to a file.
+
+        Args:
+            log_file: The log file to redirect to. If None, the stdout and
+            stderr will be restored to the original streams.
+
+        Returns:
+            The old log file, or None if the stdout and stderr were not
+            redirected.
+        """
+        original_log_file = self._log_file
+        original_log_handle = self._log_file_handle
+        if log_file is None:
+            self._log_file_handle = None
+        else:
+            self._log_file_handle = open(log_file, 'a', encoding='utf-8')
+        self._log_file = log_file
+        if original_log_handle is not None:
+            original_log_handle.close()
+        return original_log_file
+
+    def output_stream(self, fallback: TextIO) -> TextIO:
+        if self._log_file_handle is None:
+            return fallback
+        else:
+            return self._log_file_handle
+
+    def override_envs(self, envs: Dict[str, str]):
+        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[Optional[SkyPilotContext]]('sky_context',
+                                                             default=None)
+
+
+def get() -> Optional[SkyPilotContext]:
+    """Get the current SkyPilot context.
+
+    If the context is not initialized, get() will return None. This helps
+    sync code to check whether it runs in a cancellable context and avoid
+    polling the cancellation event if it is not.
+    """
+    return _CONTEXT.get()
+
+
+class ContextualEnviron(MutableMapping[str, str]):
+    """Environment variables wrapper with contextual overrides.
+
+    An instance of ContextualEnviron will typically be used to replace
+    os.environ to make the envron access of current process contextual
+    aware.
+
+    Behavior of spawning a subprocess:
+    - The contexual 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
+      environment variables at the time of the spawn, that is, it will not
+      see the updates to the environment variables after the spawn. Also,
+      os.environ of the subprocess will not be a ContextualEnviron unless
+      the subprocess hijacks os.environ explicitly.
+    - Optionally, context.Popen() can be used to automatically pass
+      os.environ with overrides to subprocess.
+
+
+    Example:
+    1. Parent process:
+       # Hijack os.environ to be a ContextualEnviron
+       os.environ = ContextualEnviron(os.environ)
+       ctx = context.get()
+       ctx.override_envs({'FOO': 'BAR1'})
+       proc = subprocess.Popen(..., env=os.environ)
+       # Or use context.Popen instead
+       # proc = context.Popen(...)
+       ctx.override_envs({'FOO': 'BAR2'})
+    2. Subprocess:
+       assert os.environ['FOO'] == 'BAR1'
+       ctx = context.get()
+       # Override the contextual env var in the subprocess does not take
+       # effect since the os.environ is not hijacked.
+       ctx.override_envs({'FOO': 'BAR3'})
+       assert os.environ['FOO'] == 'BAR1'
+    """
+
+    def __init__(self, environ: 'os._Environ[str]') -> None:
+        self._environ = environ
+
+    def __getitem__(self, key: str) -> str:
+        ctx = get()
+        if ctx is not None:
+            if key in ctx.env_overrides:
+                value = ctx.env_overrides[key]
+                # None is used to indicate that the key is deleted in the
+                # context.
+                if value is None:
+                    raise KeyError(key)
+                return value
+        return self._environ[key]
+
+    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)
+                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) -> int:
+        return len(dict(self))
+
+    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: str) -> None:
+        ctx = get()
+        if ctx is not 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) -> 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) -> Dict[str, str]:
+        copied = self._environ.copy()
+        ctx = get()
+        if ctx is not None:
+            for key in ctx.env_overrides:
+                if ctx.env_overrides[key] is None:
+                    copied.pop(key)
+                else:
+                    copied[key] = ctx.env_overrides[key]
+        return copied
+
+    def setdefault(self, key: str, default: str) -> str:
+        return self._environ.setdefault(key, default)
+
+    def __ior__(self, other):
+        if not isinstance(other, Mapping):
+            return NotImplemented
+        self.update(other)
+        return self
+
+    def __or__(self, other):
+        if not isinstance(other, Mapping):
+            return NotImplemented
+        new = dict(self)
+        new.update(other)
+        return new
+
+    def __ror__(self, other):
+        if not isinstance(other, Mapping):
+            return NotImplemented
+        new = dict(other)
+        new.update(self)
+        return new
+
+
+class Popen(subprocess.Popen):
+
+    def __init__(self, *args, **kwargs):
+        env = kwargs.pop('env', None)
+        if env is None:
+            # 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)
+
+
+P = ParamSpec('P')
+T = TypeVar('T')
+
+
+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 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)
+    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[SkyPilotContext] = None) -> SkyPilotContext:
+    """Initialize the current SkyPilot context."""
+    new_context = base_context.copy(
+    ) if base_context is not None else SkyPilotContext()
+    _CONTEXT.set(new_context)
+    return new_context
+
+
+class _ContextualStream:
+    """A base class for streams that are contextually aware.
+
+    This class implements the TextIO interface via __getattr__ to delegate
+    attribute access to the original or contextual stream.
+    """
+    _original_stream: TextIO
+
+    def __init__(self, original_stream: TextIO):
+        self._original_stream = original_stream
+
+    def __getattr__(self, attr: str):
+        return getattr(self._active_stream(), attr)
+
+    def _active_stream(self) -> TextIO:
+        ctx = get()
+        if ctx is None:
+            return self._original_stream
+        return ctx.output_stream(self._original_stream)
+
+
+class Stdout(_ContextualStream):
+
+    def __init__(self):
+        super().__init__(sys.stdout)
+
+
+class Stderr(_ContextualStream):
+
+    def __init__(self):
+        super().__init__(sys.stderr)

sky/utils/context_utils.py

@@ -0,0 +1,222 @@
+"""Utilities for SkyPilot context."""
+import asyncio
+import concurrent.futures
+import contextvars
+import functools
+import io
+import multiprocessing
+import os
+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
+# once we have context widely adopted.
+def hijack_sys_attrs():
+    """hijack system attributes to be context aware
+
+    This function should be called at the very beginning of the processes
+    that might use sky.utils.context.
+    """
+    # Modify stdout and stderr of unvicorn process to be contextually aware,
+    # use setattr to bypass the TextIO type check.
+    setattr(sys, 'stdout', context.Stdout())
+    setattr(sys, 'stderr', context.Stderr())
+    # Reload logger to apply latest stdout and stderr.
+    sky_logging.reload_logger()
+    # Hijack os.environ with ContextualEnviron to make env variables
+    # contextually aware.
+    setattr(os, 'environ', context.ContextualEnviron(os.environ))
+    # Hijack subprocess.Popen to pass the contextual environ to subprocess
+    # by default.
+    setattr(subprocess, 'Popen', context.Popen)
+
+
+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='',
+                               errors='replace',
+                               write_through=True)
+    while True:
+        line = wrapped.readline()
+        if line:
+            out_stream.write(line)
+
+            # 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.SkyPilotContext,
+        proc: subprocess.Popen,
+        poll_interval: float = 0.5,
+        cancel_callback: Optional[Callable[[], None]] = None,
+        stdout_stream_handler: Optional[StreamHandler] = None,
+        stderr_stream_handler: Optional[StreamHandler] = None
+) -> Tuple[str, str]:
+    """Wait for the process to finish or cancel it if the context is cancelled.
+
+    Args:
+        proc: The process to wait for.
+        poll_interval: The interval to poll the process.
+        cancel_callback: The callback to call if the context is cancelled.
+        stdout_stream_handler: An optional handler to handle the stdout stream,
+            if None, the stdout stream will be passed through.
+        stderr_stream_handler: An optional handler to handle the stderr stream,
+            if None, the stderr stream will be passed through.
+    """
+
+    if stdout_stream_handler is None:
+        stdout_stream_handler = passthrough_stream_handler
+    if stderr_stream_handler is None:
+        stderr_stream_handler = passthrough_stream_handler
+
+    # Threads are lazily created, so no harm if stderr is None
+    with multiprocessing.pool.ThreadPool(processes=2) as pool:
+        # Context will be lost in the new thread, capture current output stream
+        # and pass it to the new thread directly.
+        stdout_fut = pool.apply_async(
+            stdout_stream_handler, (proc.stdout, ctx.output_stream(sys.stdout)))
+        stderr_fut = None
+        if proc.stderr is not None:
+            stderr_fut = pool.apply_async(
+                stderr_stream_handler,
+                (proc.stderr, ctx.output_stream(sys.stderr)))
+        try:
+            wait_process(ctx,
+                         proc,
+                         poll_interval=poll_interval,
+                         cancel_callback=cancel_callback)
+        finally:
+            # Wait for the stream handler threads to exit when process is done
+            # or cancelled
+            stdout_fut.wait()
+            if stderr_fut is not None:
+                stderr_fut.wait()
+        stdout = stdout_fut.get()
+        stderr = ''
+        if stderr_fut is not None:
+            stderr = stderr_fut.get()
+        return stdout, stderr
+
+
+def wait_process(ctx: context.SkyPilotContext,
+                 proc: subprocess.Popen,
+                 poll_interval: float = 0.5,
+                 cancel_callback: Optional[Callable[[], None]] = None):
+    """Wait for the process to finish or cancel it if the context is cancelled.
+
+    Args:
+        proc: The process to wait for.
+        poll_interval: The interval to poll the process.
+        cancel_callback: The callback to call if the context is cancelled.
+    """
+    while True:
+        if ctx.is_canceled():
+            if cancel_callback is not None:
+                cancel_callback()
+            # Kill the process despite the caller's callback, the utility
+            # function gracefully handles the case where the process is
+            # already terminated.
+            # 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)
+        except subprocess.TimeoutExpired:
+            pass
+        else:
+            # Process exited
+            break
+
+
+F = TypeVar('F', bound=Callable[..., Any])
+
+
+def cancellation_guard(func: F) -> F:
+    """Decorator to make a synchronous function cancellable via context.
+
+    Guards the function execution by checking context.is_canceled() before
+    executing the function and raises asyncio.CancelledError if the context
+    is already cancelled.
+
+    This basically mimics the behavior of asyncio, which checks coroutine
+    cancelled in await call.
+
+    Args:
+        func: The function to be decorated.
+
+    Returns:
+        The wrapped function that checks cancellation before execution.
+
+    Raises:
+        asyncio.CancelledError: If the context is cancelled before execution.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        ctx = context.get()
+        if ctx is not None and ctx.is_canceled():
+            raise asyncio.CancelledError(
+                f'Function {func.__name__} cancelled before execution')
+        return func(*args, **kwargs)
+
+    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: 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()
+    pyctx = contextvars.copy_context()
+    func_call: Callable[..., T] = functools.partial(pyctx.run, func, *args,
+                                                    **kwargs)
+    return loop.run_in_executor(executor, func_call)