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

sky/utils/command_runner.py

@@ -11,6 +11,7 @@ from sky import sky_logging
 from sky.skylet import constants
 from sky.skylet import log_lib
 from sky.utils import common_utils
+from sky.utils import context_utils
 from sky.utils import control_master_utils
 from sky.utils import subprocess_utils
 from sky.utils import timeline
@@ -574,6 +575,7 @@ class SSHCommandRunner(CommandRunner):
                                      shell=True)
 
     @timeline.event
+    @context_utils.cancellation_guard
     def run(
             self,
             cmd: Union[str, List[str]],
@@ -779,6 +781,7 @@ class KubernetesCommandRunner(CommandRunner):
         return kubectl_cmd
 
     @timeline.event
+    @context_utils.cancellation_guard
     def run(
             self,
             cmd: Union[str, List[str]],

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/controller_utils.py

@@ -6,7 +6,7 @@ import getpass
 import os
 import tempfile
 import typing
-from typing import Any, Dict, Iterable, List, Optional, Set, Tuple
+from typing import Any, Dict, Iterable, List, Optional, Set
 import uuid
 
 import colorama
@@ -517,6 +517,30 @@ def get_controller_resources(
         if custom_controller_resources_config is not None:
             controller_resources_config_copied.update(
                 custom_controller_resources_config)
+        # Compatibility with the old way of specifying the controller autostop
+        # config. TODO(cooperc): Remove this before 0.12.0.
+        custom_controller_autostop_config = skypilot_config.get_nested(
+            (controller.value.controller_type, 'controller', 'autostop'), None)
+        if custom_controller_autostop_config is not None:
+            logger.warning(
+                f'{colorama.Fore.YELLOW}Warning: Config value '
+                f'`{controller.value.controller_type}.controller.autostop` '
+                'is deprecated. Please use '
+                f'`{controller.value.controller_type}.controller.resources.'
+                f'autostop` instead.{colorama.Style.RESET_ALL}')
+            # Only set the autostop config if it is not already specified.
+            if controller_resources_config_copied.get('autostop') is None:
+                controller_resources_config_copied['autostop'] = (
+                    custom_controller_autostop_config)
+            else:
+                logger.warning(f'{colorama.Fore.YELLOW}Ignoring the old '
+                               'config, since it is already specified in '
+                               f'resources.{colorama.Style.RESET_ALL}')
+    # Set the default autostop config for the controller, if not already
+    # specified.
+    if controller_resources_config_copied.get('autostop') is None:
+        controller_resources_config_copied['autostop'] = (
+            controller.value.default_autostop_config)
 
     try:
         controller_resources = resources.Resources.from_yaml_config(
@@ -547,7 +571,10 @@ def get_controller_resources(
     if controller_record is not None:
         handle = controller_record.get('handle', None)
         if handle is not None:
-            controller_resources_to_use = handle.launched_resources
+            # Use the existing resources, but override the autostop config with
+            # the one currently specified in the config.
+            controller_resources_to_use = handle.launched_resources.copy(
+                autostop=controller_resources_config_copied.get('autostop'))
 
     # If the controller and replicas are from the same cloud (and region/zone),
     # it should provide better connectivity. We will let the controller choose
@@ -608,8 +635,9 @@ def get_controller_resources(
     controller_zone = controller_resources_to_use.zone
 
     # Filter clouds if controller_resources_to_use.cloud is specified.
-    filtered_clouds = ({controller_cloud} if controller_cloud is not None else
-                       requested_clouds_with_region_zone.keys())
+    filtered_clouds: Set[str] = {controller_cloud
+                                } if controller_cloud is not None else set(
+                                    requested_clouds_with_region_zone.keys())
 
     # Filter regions and zones and construct the result.
     result: Set[resources.Resources] = set()
@@ -618,15 +646,17 @@ def get_controller_resources(
                                                         {None: {None}})
 
         # Filter regions if controller_resources_to_use.region is specified.
-        filtered_regions = ({controller_region} if controller_region is not None
-                            else regions.keys())
+        filtered_regions: Set[Optional[str]] = ({
+            controller_region
+        } if controller_region is not None else set(regions.keys()))
 
         for region in filtered_regions:
             zones = regions.get(region, {None})
 
             # Filter zones if controller_resources_to_use.zone is specified.
-            filtered_zones = ({controller_zone}
-                              if controller_zone is not None else zones)
+            filtered_zones: Set[Optional[str]] = ({
+                controller_zone
+            } if controller_zone is not None else set(zones))
 
             # Create combinations of cloud, region, and zone.
             for zone in filtered_zones:
@@ -641,40 +671,6 @@ def get_controller_resources(
     return result
 
 
-def get_controller_autostop_config(
-        controller: Controllers) -> Tuple[Optional[int], bool]:
-    """Get the autostop config for the controller.
-
-    Returns:
-      A tuple of (idle_minutes_to_autostop, down), which correspond to the
-      values passed to execution.launch().
-    """
-    controller_autostop_config_copied: Dict[str, Any] = copy.copy(
-        controller.value.default_autostop_config)
-    if skypilot_config.loaded():
-        custom_controller_autostop_config = skypilot_config.get_nested(
-            (controller.value.controller_type, 'controller', 'autostop'), None)
-        if custom_controller_autostop_config is False:
-            # Disabled with `autostop: false` in config.
-            # To indicate autostop is disabled, we return None for
-            # idle_minutes_to_autostop.
-            return None, False
-        elif custom_controller_autostop_config is True:
-            # Enabled with default values. There is no change in behavior, but
-            # this is included by for completeness, since `False` is valid.
-            pass
-        elif custom_controller_autostop_config is not None:
-            # We have specific config values.
-            # Override the controller autostop config with the ones specified in
-            # the config.
-            assert isinstance(custom_controller_autostop_config, dict)
-            controller_autostop_config_copied.update(
-                custom_controller_autostop_config)
-
-    return (controller_autostop_config_copied['idle_minutes'],
-            controller_autostop_config_copied['down'])
-
-
 def _setup_proxy_command_on_controller(
         controller_launched_cloud: 'clouds.Cloud',
         user_config: Dict[str, Any]) -> config_utils.Config:
@@ -703,7 +699,7 @@ def _setup_proxy_command_on_controller(
     # NOTE: suppose that we have a controller in old VPC, then user
     # changes 'vpc_name' in the config and does a 'job launch' /
     # 'serve up'. In general, the old controller may not successfully
-    # launch the job in the new VPC. This happens if the two VPCs don’t
+    # launch the job in the new VPC. This happens if the two VPCs don't
     # have peering set up. Like other places in the code, we assume
     # properly setting up networking is user's responsibilities.
     # TODO(zongheng): consider adding a basic check that checks

sky/utils/dag_utils.py

@@ -1,6 +1,6 @@
 """Utilities for loading and dumping DAGs from/to YAML files."""
 import copy
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from sky import dag as dag_lib
 from sky import sky_logging
@@ -195,7 +195,9 @@ def fill_default_config_in_dag_for_job_launch(dag: dag_lib.Dag) -> None:
         assert default_strategy is not None
         for resources in list(task_.resources):
             original_job_recovery = resources.job_recovery
-            job_recovery = {'strategy': default_strategy}
+            job_recovery: Dict[str, Optional[Union[str, int]]] = {
+                'strategy': default_strategy
+            }
             if isinstance(original_job_recovery, str):
                 job_recovery['strategy'] = original_job_recovery
             elif isinstance(original_job_recovery, dict):

sky/utils/resources_utils.py

@@ -140,10 +140,12 @@ def simplify_ports(ports: List[str]) -> List[str]:
 def format_resource(resource: 'resources_lib.Resources',
                     simplify: bool = False) -> str:
     if simplify:
+        resource = resource.assert_launchable()
         cloud = resource.cloud
         if resource.accelerators is None:
             vcpu, _ = cloud.get_vcpus_mem_from_instance_type(
                 resource.instance_type)
+            assert vcpu is not None, 'vCPU must be specified'
             hardware = f'vCPU={int(vcpu)}'
         else:
             hardware = f'{resource.accelerators}'
@@ -248,6 +250,7 @@ def make_launchables_for_valid_region_zones(
     launchables = []
     regions = launchable_resources.get_valid_regions_for_launchable()
     for region in regions:
+        assert launchable_resources.cloud is not None, 'Cloud must be specified'
         optimize_by_zone = (override_optimize_by_zone or
                             launchable_resources.cloud.optimize_by_zone())
         # It is possible that we force the optimize_by_zone but some clouds