skypilot-nightly 1.0.0.dev20250513__py3-none-any.whl → 1.0.0.dev20250514__py3-none-any.whl

sky/utils/context.py

@@ -1,264 +0,0 @@
-"""SkyPilot context for threads and coroutines."""
-
-import asyncio
-from collections.abc import Mapping
-from collections.abc import MutableMapping
-import contextvars
-import os
-import pathlib
-import subprocess
-import sys
-from typing import Dict, Optional, TextIO
-
-
-class Context(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 intialized 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 = {}
-
-    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_file 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
-
-
-_CONTEXT = contextvars.ContextVar('sky_context', default=None)
-
-
-def get() -> Optional[Context]:
-    """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):
-    """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):
-        self._environ = environ
-
-    def __getitem__(self, key):
-        ctx = get()
-        if ctx is not None:
-            if key in ctx.env_overrides:
-                return ctx.env_overrides[key]
-        return self._environ[key]
-
-    def __iter__(self):
-        ctx = get()
-        if ctx is not None:
-            for key in ctx.env_overrides:
-                yield key
-            for key in self._environ:
-                # Deduplicate the keys
-                if key not in ctx.env_overrides:
-                    yield key
-        else:
-            return self._environ.__iter__()
-
-    def __len__(self):
-        return len(dict(self))
-
-    def __setitem__(self, key, value):
-        return self._environ.__setitem__(key, value)
-
-    def __delitem__(self, key):
-        return self._environ.__delitem__(key)
-
-    def __repr__(self):
-        return self._environ.__repr__()
-
-    def copy(self):
-        copied = self._environ.copy()
-        ctx = get()
-        if ctx is not None:
-            copied.update(ctx.env_overrides)
-        return copied
-
-    def setdefault(self, key, default=None):
-        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:
-            env = os.environ
-        super().__init__(*args, env=env, **kwargs)
-
-
-def initialize():
-    """Initialize the current SkyPilot context."""
-    _CONTEXT.set(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

@@ -1,172 +0,0 @@
-"""Utilities for SkyPilot context."""
-import asyncio
-import functools
-import io
-import multiprocessing
-import os
-import subprocess
-import sys
-import typing
-from typing import Any, Callable, IO, Optional, Tuple, TypeVar
-
-from sky import sky_logging
-from sky.utils import context
-from sky.utils import subprocess_utils
-
-StreamHandler = Callable[[IO[Any], IO[Any]], str]
-
-
-# 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"""
-    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)
-            out_stream.flush()
-        else:
-            break
-    return ''
-
-
-def pipe_and_wait_process(
-        ctx: context.Context,
-        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.Context,
-                 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.
-            subprocess_utils.kill_process_with_grace_period(proc)
-            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)