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

sky/utils/context.py

@@ -0,0 +1,264 @@
+"""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

@@ -0,0 +1,172 @@
+"""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)

sky/utils/rich_utils.py

@@ -1,13 +1,15 @@
 """Rich status spinner utils."""
 import contextlib
+import contextvars
 import enum
 import logging
 import threading
 import typing
-from typing import Dict, Iterator, Optional, Tuple, Union
+from typing import Callable, Iterator, Optional, Tuple, Union
 
 from sky.adaptors import common as adaptors_common
 from sky.utils import annotations
+from sky.utils import context
 from sky.utils import message_utils
 from sky.utils import rich_console_utils
 
@@ -18,11 +20,31 @@ else:
     requests = adaptors_common.LazyImport('requests')
     rich_console = adaptors_common.LazyImport('rich.console')
 
-_statuses: Dict[str, Optional[Union['EncodedStatus',
-                                    'rich_console.Status']]] = {
-                                        'server': None,
-                                        'client': None,
-                                    }
+GeneralStatus = Union['rich_console.Status', 'EncodedStatus']
+
+_client_status: Optional[GeneralStatus] = None
+_server_status: contextvars.ContextVar[
+    Optional[GeneralStatus]] = contextvars.ContextVar('server_status',
+                                                      default=None)
+
+
+def _get_client_status() -> Optional[GeneralStatus]:
+    return _client_status
+
+
+def _get_server_status() -> Optional[GeneralStatus]:
+    return _server_status.get()
+
+
+def _set_client_status(status: Optional[GeneralStatus]):
+    global _client_status
+    _client_status = status
+
+
+def _set_server_status(status: Optional[GeneralStatus]):
+    _server_status.set(status)
+
+
 _status_nesting_level = 0
 
 _logging_lock = threading.RLock()
@@ -128,20 +150,22 @@ class _NoOpConsoleStatus:
 class _RevertibleStatus:
     """A wrapper for status that can revert to previous message after exit."""
 
-    def __init__(self, message: str, status_type: str):
+    def __init__(self, message: str, get_status_fn: Callable[[], GeneralStatus],
+                 set_status_fn: Callable[[Optional[GeneralStatus]], None]):
         self.previous_message = None
-        self.status_type = status_type
-        status = _statuses[status_type]
+        self.get_status_fn = get_status_fn
+        self.set_status_fn = set_status_fn
+        status = self.get_status_fn()
         if status is not None:
             self.previous_message = status.status
         self.message = message
 
     def __enter__(self):
         global _status_nesting_level
-        _statuses[self.status_type].update(self.message)
+        self.get_status_fn().update(self.message)
         _status_nesting_level += 1
-        _statuses[self.status_type].__enter__()
-        return _statuses[self.status_type]
+        self.get_status_fn().__enter__()
+        return self.get_status_fn()
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         # We use the same lock with the `safe_logger` to avoid the following 2
@@ -160,32 +184,48 @@ class _RevertibleStatus:
             _status_nesting_level -= 1
             if _status_nesting_level <= 0:
                 _status_nesting_level = 0
-                if _statuses[self.status_type] is not None:
-                    _statuses[self.status_type].__exit__(
-                        exc_type, exc_val, exc_tb)
-                    _statuses[self.status_type] = None
+                if self.get_status_fn() is not None:
+                    self.get_status_fn().__exit__(exc_type, exc_val, exc_tb)
+                    self.set_status_fn(None)
             else:
-                _statuses[self.status_type].update(self.previous_message)
+                self.get_status_fn().update(self.previous_message)
 
     def update(self, *args, **kwargs):
-        _statuses[self.status_type].update(*args, **kwargs)
+        self.get_status_fn().update(*args, **kwargs)
 
     def stop(self):
-        _statuses[self.status_type].stop()
+        self.get_status_fn().stop()
 
     def start(self):
-        _statuses[self.status_type].start()
+        self.get_status_fn().start()
+
+
+def _is_thread_safe() -> bool:
+    """Check if the current status context is thread-safe.
+
+    We are thread-safe if we are on the main thread or the server_status is
+    context-local, i.e. an async context has been initialized.
+    """
+    return (threading.current_thread() is threading.main_thread() or
+            context.get() is not None)
 
 
 def safe_status(msg: str) -> Union['rich_console.Status', _NoOpConsoleStatus]:
-    """A wrapper for multi-threaded console.status."""
+    """A wrapper for multi-threaded server-side console.status.
+
+    This function will encode rich status with control codes and output the
+    encoded string to stdout. Client-side decode control codes from server
+    output and update the rich status. This function is safe to be called in
+    async/multi-threaded context.
+
+    See also: :func:`client_status`, :class:`EncodedStatus`.
+    """
     from sky import sky_logging  # pylint: disable=import-outside-toplevel
-    if (annotations.is_on_api_server and
-            threading.current_thread() is threading.main_thread() and
+    if (annotations.is_on_api_server and _is_thread_safe() and
             not sky_logging.is_silent()):
-        if _statuses['server'] is None:
-            _statuses['server'] = EncodedStatus(msg)
-        return _RevertibleStatus(msg, 'server')
+        if _get_server_status() is None:
+            _set_server_status(EncodedStatus(msg))
+        return _RevertibleStatus(msg, _get_server_status, _set_server_status)
     return _NoOpConsoleStatus()
 
 
@@ -196,22 +236,26 @@ def stop_safe_status():
     stream logs from user program and do not want it to interfere with the
     spinner display.
     """
-    if (threading.current_thread() is threading.main_thread() and
-            _statuses['server'] is not None):
-        _statuses['server'].stop()
+    if _is_thread_safe():
+        return
+    server_status = _get_server_status()
+    if server_status is not None:
+        server_status.stop()
 
 
 def force_update_status(msg: str):
     """Update the status message even if sky_logging.is_silent() is true."""
-    if (threading.current_thread() is threading.main_thread() and
-            _statuses['server'] is not None):
-        _statuses['server'].update(msg)
+    if not _is_thread_safe():
+        return
+    server_status = _get_server_status()
+    if server_status is not None:
+        server_status.update(msg)
 
 
 @contextlib.contextmanager
 def safe_logger():
     with _logging_lock:
-        client_status_obj = _statuses['client']
+        client_status_obj = _get_client_status()
 
         client_status_live = (client_status_obj is not None and
                               client_status_obj._live.is_started)  # pylint: disable=protected-access
@@ -230,13 +274,13 @@ class RichSafeStreamHandler(logging.StreamHandler):
 
 
 def client_status(msg: str) -> Union['rich_console.Status', _NoOpConsoleStatus]:
-    """A wrapper for multi-threaded console.status."""
+    """A wrapper for multi-threaded client-side console.status."""
     from sky import sky_logging  # pylint: disable=import-outside-toplevel
     if (threading.current_thread() is threading.main_thread() and
             not sky_logging.is_silent()):
-        if _statuses['client'] is None:
-            _statuses['client'] = rich_console_utils.get_console().status(msg)
-        return _RevertibleStatus(msg, 'client')
+        if _get_client_status() is None:
+            _set_client_status(rich_console_utils.get_console().status(msg))
+        return _RevertibleStatus(msg, _get_client_status, _set_client_status)
     return _NoOpConsoleStatus()
 
 

sky/utils/schemas.py

@@ -837,6 +837,12 @@ def get_config_schema():
                 'enable_gvnic': {
                     'type': 'boolean'
                 },
+                'enable_gpu_direct': {
+                    'type': 'boolean'
+                },
+                'placement_policy': {
+                    'type': 'string',
+                },
                 'vpc_name': {
                     'oneOf': [
                         {
@@ -966,7 +972,9 @@ def get_config_schema():
         'nebius': {
             'type': 'object',
             'required': [],
-            'properties': {},
+            'properties': {
+                **_NETWORK_CONFIG_SCHEMA,
+            },
             'additionalProperties': {
                 'type': 'object',
                 'required': [],

sky/utils/subprocess_utils.py

@@ -208,8 +208,11 @@ def kill_children_processes(parent_pids: Optional[Union[
             kill_process_with_grace_period(child, force=force)
 
 
-def kill_process_with_grace_period(proc: Union[multiprocessing.Process,
-                                               psutil.Process],
+GenericProcess = Union[multiprocessing.Process, psutil.Process,
+                       subprocess.Popen]
+
+
+def kill_process_with_grace_period(proc: GenericProcess,
                                    force: bool = False,
                                    grace_period: int = 10) -> None:
     """Kill a process with SIGTERM and wait for it to exit.
@@ -223,6 +226,9 @@ def kill_process_with_grace_period(proc: Union[multiprocessing.Process,
     if isinstance(proc, psutil.Process):
         alive = proc.is_running
         wait = proc.wait
+    elif isinstance(proc, subprocess.Popen):
+        alive = lambda: proc.poll() is None
+        wait = proc.wait
     else:
         alive = proc.is_alive
         wait = proc.join

{skypilot_nightly-1.0.0.dev20250510.dist-info → skypilot_nightly-1.0.0.dev20250513.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: skypilot-nightly
-Version: 1.0.0.dev20250510
+Version: 1.0.0.dev20250513
 Summary: SkyPilot: Run AI on Any Infra — Unified, Faster, Cheaper.
 Author: SkyPilot Team
 License: Apache 2.0