skypilot-nightly 1.0.0.dev20250922__py3-none-any.whl → 1.0.0.dev20250925__py3-none-any.whl

sky/server/requests/preconditions.py

@@ -146,10 +146,9 @@ class ClusterStartCompletePrecondition(Precondition):
         self.cluster_name = cluster_name
 
     async def check(self) -> Tuple[bool, Optional[str]]:
-        cluster_record = global_user_state.get_cluster_from_name(
+        cluster_status = global_user_state.get_status_from_cluster_name(
             self.cluster_name)
-        if (cluster_record and
-                cluster_record['status'] is status_lib.ClusterStatus.UP):
+        if cluster_status is status_lib.ClusterStatus.UP:
             # Shortcut for started clusters, ignore cluster not found
             # since the cluster record might not yet be created by the
             # launch task.

sky/server/rest.py

@@ -9,6 +9,7 @@ import typing
 from typing import Any, Callable, cast, Optional, TypeVar
 
 import colorama
+import urllib3.exceptions
 
 from sky import exceptions
 from sky import sky_logging
@@ -53,6 +54,7 @@ _session.headers[constants.VERSION_HEADER] = (
 _transient_errors = [
     requests.exceptions.RequestException,
     ConnectionError,
+    urllib3.exceptions.HTTPError,
 ]
 
 

sky/server/server.py

@@ -445,6 +445,22 @@ async def loop_lag_monitor(loop: asyncio.AbstractEventLoop,
     loop.call_at(target, tick)
 
 
+def schedule_on_boot_check():
+    try:
+        executor.schedule_request(
+            request_id='skypilot-server-on-boot-check',
+            request_name='check',
+            request_body=payloads.CheckBody(),
+            func=sky_check.check,
+            schedule_type=requests_lib.ScheduleType.SHORT,
+            is_skypilot_system=True,
+        )
+    except exceptions.RequestAlreadyExistsError:
+        # Lifespan will be executed in each uvicorn worker process, we
+        # can safely ignore the error if the task is already scheduled.
+        logger.debug('Request skypilot-server-on-boot-check already exists.')
+
+
 @contextlib.asynccontextmanager
 async def lifespan(app: fastapi.FastAPI):  # pylint: disable=redefined-outer-name
     """FastAPI lifespan context manager."""
@@ -469,6 +485,7 @@ async def lifespan(app: fastapi.FastAPI):  # pylint: disable=redefined-outer-nam
             # Lifespan will be executed in each uvicorn worker process, we
             # can safely ignore the error if the task is already scheduled.
             logger.debug(f'Request {event.id} already exists.')
+    schedule_on_boot_check()
     asyncio.create_task(cleanup_upload_ids())
     if metrics_utils.METRICS_ENABLED:
         # Start monitoring the event loop lag in each server worker
@@ -1216,19 +1233,8 @@ async def logs(
         schedule_type=requests_lib.ScheduleType.SHORT,
         request_cluster_name=cluster_job_body.cluster_name,
     )
-    task = asyncio.create_task(executor.execute_request_coroutine(request_task))
-
-    async def cancel_task():
-        try:
-            logger.info('Client disconnected for request: '
-                        f'{request.state.request_id}')
-            task.cancel()
-            await task
-        except asyncio.CancelledError:
-            pass
-
-    # Cancel the task after the request is done or client disconnects
-    background_tasks.add_task(cancel_task)
+    task = executor.execute_request_in_coroutine(request_task)
+    background_tasks.add_task(task.cancel)
     # TODO(zhwu): This makes viewing logs in browser impossible. We should adopt
     # the same approach as /stream.
     return stream_utils.stream_response(
@@ -1354,10 +1360,12 @@ def provision_logs(cluster_body: payloads.ClusterNameBody,
     effective_tail = None if tail is None or tail <= 0 else tail
 
     return fastapi.responses.StreamingResponse(
-        content=stream_utils.log_streamer(None,
-                                          log_path,
-                                          tail=effective_tail,
-                                          follow=follow),
+        content=stream_utils.log_streamer(
+            None,
+            log_path,
+            tail=effective_tail,
+            follow=follow,
+            cluster_name=cluster_body.cluster_name),
         media_type='text/plain',
         headers={
             'Cache-Control': 'no-cache, no-transform',
@@ -1419,12 +1427,13 @@ async def local_up(request: fastapi.Request,
 
 
 @app.post('/local_down')
-async def local_down(request: fastapi.Request) -> None:
+async def local_down(request: fastapi.Request,
+                     local_down_body: payloads.LocalDownBody) -> None:
     """Tears down the Kubernetes cluster started by local_up."""
     executor.schedule_request(
         request_id=request.state.request_id,
         request_name='local_down',
-        request_body=payloads.RequestBody(),
+        request_body=local_down_body,
         func=core.local_down,
         schedule_type=requests_lib.ScheduleType.LONG,
     )

sky/server/stream_utils.py

@@ -8,10 +8,12 @@ from typing import AsyncGenerator, Deque, List, Optional
 import aiofiles
 import fastapi
 
+from sky import global_user_state
 from sky import sky_logging
 from sky.server.requests import requests as requests_lib
 from sky.utils import message_utils
 from sky.utils import rich_utils
+from sky.utils import status_lib
 
 logger = sky_logging.init_logger(__name__)
 
@@ -22,6 +24,7 @@ logger = sky_logging.init_logger(__name__)
 _BUFFER_SIZE = 8 * 1024  # 8KB
 _BUFFER_TIMEOUT = 0.02  # 20ms
 _HEARTBEAT_INTERVAL = 30
+_CLUSTER_STATUS_INTERVAL = 1
 
 
 async def _yield_log_file_with_payloads_skipped(
@@ -37,11 +40,13 @@ async def _yield_log_file_with_payloads_skipped(
         yield line_str
 
 
-async def log_streamer(request_id: Optional[str],
-                       log_path: pathlib.Path,
-                       plain_logs: bool = False,
-                       tail: Optional[int] = None,
-                       follow: bool = True) -> AsyncGenerator[str, None]:
+async def log_streamer(
+        request_id: Optional[str],
+        log_path: pathlib.Path,
+        plain_logs: bool = False,
+        tail: Optional[int] = None,
+        follow: bool = True,
+        cluster_name: Optional[str] = None) -> AsyncGenerator[str, None]:
     """Streams the logs of a request.
 
     Args:
@@ -51,6 +56,8 @@ async def log_streamer(request_id: Optional[str],
         plain_logs: Whether to show plain logs.
         tail: The number of lines to tail. If None, tail the whole file.
         follow: Whether to follow the log file.
+        cluster_name: The cluster name to check status for provision logs.
+            If provided and cluster status is UP, streaming will terminate.
     """
 
     if request_id is not None:
@@ -104,15 +111,17 @@ async def log_streamer(request_id: Optional[str],
 
     async with aiofiles.open(log_path, 'rb') as f:
         async for chunk in _tail_log_file(f, request_id, plain_logs, tail,
-                                          follow):
+                                          follow, cluster_name):
             yield chunk
 
 
-async def _tail_log_file(f: aiofiles.threadpool.binary.AsyncBufferedReader,
-                         request_id: Optional[str] = None,
-                         plain_logs: bool = False,
-                         tail: Optional[int] = None,
-                         follow: bool = True) -> AsyncGenerator[str, None]:
+async def _tail_log_file(
+        f: aiofiles.threadpool.binary.AsyncBufferedReader,
+        request_id: Optional[str] = None,
+        plain_logs: bool = False,
+        tail: Optional[int] = None,
+        follow: bool = True,
+        cluster_name: Optional[str] = None) -> AsyncGenerator[str, None]:
     """Tail the opened log file, buffer the lines and flush in chunks."""
 
     if tail is not None:
@@ -128,6 +137,7 @@ async def _tail_log_file(f: aiofiles.threadpool.binary.AsyncBufferedReader,
             yield line_str
 
     last_heartbeat_time = asyncio.get_event_loop().time()
+    last_cluster_status_check_time = asyncio.get_event_loop().time()
 
     # Buffer the lines in memory and flush them in chunks to improve log
     # tailing throughput.
@@ -176,7 +186,19 @@ async def _tail_log_file(f: aiofiles.threadpool.binary.AsyncBufferedReader,
                     break
             if not follow:
                 break
-
+            # Provision logs pass in cluster_name, check cluster status
+            # periodically to see if provisioning is done. We only
+            # check once a second to avoid overloading the DB.
+            check_status = (current_time - last_cluster_status_check_time
+                           ) >= _CLUSTER_STATUS_INTERVAL
+            if cluster_name is not None and check_status:
+                cluster_record = await (
+                    global_user_state.get_status_from_cluster_name_async(
+                        cluster_name))
+                if (cluster_record is None or
+                        cluster_record != status_lib.ClusterStatus.INIT):
+                    break
+                last_cluster_status_check_time = current_time
             if current_time - last_heartbeat_time >= _HEARTBEAT_INTERVAL:
                 # Currently just used to keep the connection busy, refer to
                 # https://github.com/skypilot-org/skypilot/issues/5750 for

sky/setup_files/dependencies.py

@@ -112,6 +112,7 @@ server_dependencies = [
     GRPC,
     PROTOBUF,
     'aiosqlite',
+    'greenlet',
 ]
 
 local_ray = [
@@ -192,7 +193,9 @@ extras_require: Dict[str, List[str]] = {
     'remote': remote,
     # For the container registry auth api. Reference:
     # https://github.com/runpod/runpod-python/releases/tag/1.6.1
-    'runpod': ['runpod>=1.6.1'],
+    # RunPod needs a TOML parser to read ~/.runpod/config.toml. On Python 3.11+
+    # stdlib provides tomllib; on lower versions we depend on tomli explicitly.
+    'runpod': ['runpod>=1.6.1', 'tomli; python_version < "3.11"'],
     'fluidstack': [],  # No dependencies needed for fluidstack
     'cudo': ['cudo-compute>=0.1.10'],
     'paperspace': [],  # No dependencies needed for paperspace

sky/setup_files/setup.py

@@ -148,47 +148,47 @@ if os.path.exists(readme_filepath):
     long_description = io.open(readme_filepath, 'r', encoding='utf-8').read()
     long_description = parse_readme(long_description)
 
-atexit.register(revert_commit_hash)
-replace_commit_hash()
-
-setuptools.setup(
-    # NOTE: this affects the package.whl wheel name. When changing this (if
-    # ever), you must grep for '.whl' and change all corresponding wheel paths
-    # (templates/*.j2 and wheel_utils.py).
-    name='skypilot-nightly',
-    version=find_version(),
-    packages=setuptools.find_packages(),
-    author='SkyPilot Team',
-    license='Apache 2.0',
-    readme='README.md',
-    description='SkyPilot: Run AI on Any Infra — Unified, Faster, Cheaper.',
-    long_description=long_description,
-    long_description_content_type='text/markdown',
-    setup_requires=['wheel'],
-    requires_python='>=3.7',
-    install_requires=dependencies['install_requires'],
-    extras_require=dependencies['extras_require'],
-    entry_points={
-        'console_scripts': ['sky = sky.cli:cli'],
-    },
-    include_package_data=True,
-    classifiers=[
-        'Programming Language :: Python :: 3.7',
-        'Programming Language :: Python :: 3.8',
-        'Programming Language :: Python :: 3.9',
-        'Programming Language :: Python :: 3.10',
-        'Programming Language :: Python :: 3.11',
-        'Programming Language :: Python :: 3.12',
-        'Programming Language :: Python :: 3.13',
-        'License :: OSI Approved :: Apache Software License',
-        'Operating System :: OS Independent',
-        'Topic :: Software Development :: Libraries :: Python Modules',
-        'Topic :: System :: Distributed Computing',
-    ],
-    project_urls={
-        'Homepage': 'https://github.com/skypilot-org/skypilot',
-        'Issues': 'https://github.com/skypilot-org/skypilot/issues',
-        'Discussion': 'https://github.com/skypilot-org/skypilot/discussions',
-        'Documentation': 'https://docs.skypilot.co/',
-    },
-)
+if __name__ == '__main__':
+    atexit.register(revert_commit_hash)
+    replace_commit_hash()
+    setuptools.setup(
+        # NOTE: this affects the package.whl wheel name. When changing this (if
+        # ever), you must grep for '.whl' and change all corresponding wheel paths
+        # (templates/*.j2 and wheel_utils.py).
+        name='skypilot-nightly',
+        version=find_version(),
+        packages=setuptools.find_packages(),
+        author='SkyPilot Team',
+        license='Apache 2.0',
+        readme='README.md',
+        description='SkyPilot: Run AI on Any Infra — Unified, Faster, Cheaper.',
+        long_description=long_description,
+        long_description_content_type='text/markdown',
+        setup_requires=['wheel'],
+        requires_python='>=3.7',
+        install_requires=dependencies['install_requires'],
+        extras_require=dependencies['extras_require'],
+        entry_points={
+            'console_scripts': ['sky = sky.cli:cli'],
+        },
+        include_package_data=True,
+        classifiers=[
+            'Programming Language :: Python :: 3.7',
+            'Programming Language :: Python :: 3.8',
+            'Programming Language :: Python :: 3.9',
+            'Programming Language :: Python :: 3.10',
+            'Programming Language :: Python :: 3.11',
+            'Programming Language :: Python :: 3.12',
+            'Programming Language :: Python :: 3.13',
+            'License :: OSI Approved :: Apache Software License',
+            'Operating System :: OS Independent',
+            'Topic :: Software Development :: Libraries :: Python Modules',
+            'Topic :: System :: Distributed Computing',
+        ],
+        project_urls={
+            'Homepage': 'https://github.com/skypilot-org/skypilot',
+            'Issues': 'https://github.com/skypilot-org/skypilot/issues',
+            'Discussion': 'https://github.com/skypilot-org/skypilot/discussions',
+            'Documentation': 'https://docs.skypilot.co/',
+        },
+    )

sky/templates/kubernetes-ray.yml.j2

@@ -510,6 +510,16 @@ available_node_types:
               valueFrom:
                 fieldRef:
                   fieldPath: metadata.labels['ray-node-type']
+            - name: SKYPILOT_POD_CPU_CORE_LIMIT
+              valueFrom:
+                resourceFieldRef:
+                  containerName: ray-node
+                  resource: requests.cpu
+            - name: SKYPILOT_POD_MEMORY_BYTES_LIMIT
+              valueFrom:
+                resourceFieldRef:
+                  containerName: ray-node
+                  resource: requests.memory
             {% for key, value in k8s_env_vars.items() if k8s_env_vars is not none %}
             - name: {{ key }}
               value: {{ value }}
@@ -630,13 +640,6 @@ available_node_types:
           command: ["/bin/bash", "-c", "--"]
           args:
             - |
-              # For backwards compatibility, we put a marker file in the pod
-              # to indicate that the pod is running with the changes introduced
-              # in project nimbus: https://github.com/skypilot-org/skypilot/pull/4393
-              # TODO: Remove this marker file and it's usage in setup_commands
-              # after v0.10.0 release.
-              touch /tmp/skypilot_is_nimbus
-
               # Helper function to conditionally use sudo
               # TODO(zhwu): consolidate the two prefix_cmd and sudo replacements
               prefix_cmd() { if [ $(id -u) -ne 0 ]; then echo "sudo"; else echo ""; fi; }
@@ -1333,18 +1336,16 @@ setup_commands:
     # Wait for SSH setup to complete before proceeding
     if [ -f /tmp/apt_ssh_setup_started ]; then
       echo "=== Logs for asynchronous SSH setup ===";
-      [ -f /tmp/apt_ssh_setup_complete ] && cat /tmp/${STEPS[0]}.log ||
-      { tail -f -n +1 /tmp/${STEPS[0]}.log & TAIL_PID=$!; echo "Tail PID: $TAIL_PID"; until [ -f /tmp/apt_ssh_setup_complete ] || [ -f /tmp/${STEPS[0]}.failed ]; do sleep 0.5; done; kill $TAIL_PID || true; };
+      ([ -f /tmp/apt_ssh_setup_complete ]|| [ -f /tmp/${STEPS[0]}.failed ]) && cat /tmp/${STEPS[0]}.log ||
+      { tail -f -n +1 /tmp/${STEPS[0]}.log & TAIL_PID=$!; echo "Tail PID: $TAIL_PID"; sleep 0.5; until [ -f /tmp/apt_ssh_setup_complete ] || [ -f /tmp/${STEPS[0]}.failed ]; do sleep 0.5; done; kill $TAIL_PID || true; };
       [ -f /tmp/${STEPS[0]}.failed ] && { echo "Error: ${STEPS[0]} failed. Exiting."; exit 1; } || true;
     fi
 
     echo "=== Logs for asynchronous ray and skypilot installation ===";
-    if [ -f /tmp/skypilot_is_nimbus ]; then
-      echo "=== Logs for asynchronous ray and skypilot installation ===";
-      [ -f /tmp/ray_skypilot_installation_complete ] && cat /tmp/${STEPS[1]}.log ||
-      { tail -f -n +1 /tmp/${STEPS[1]}.log & TAIL_PID=$!; echo "Tail PID: $TAIL_PID"; until [ -f /tmp/ray_skypilot_installation_complete ] || [ -f /tmp/${STEPS[1]}.failed ]; do sleep 0.5; done; kill $TAIL_PID || true; };
-      [ -f /tmp/${STEPS[1]}.failed ] && { echo "Error: ${STEPS[1]} failed. Exiting."; exit 1; } || true;
-    fi
+    ([ -f /tmp/ray_skypilot_installation_complete ]|| [ -f /tmp/${STEPS[1]}.failed ]) && cat /tmp/${STEPS[1]}.log ||
+    { tail -f -n +1 /tmp/${STEPS[1]}.log & TAIL_PID=$!; echo "Tail PID: $TAIL_PID"; sleep 0.5; until [ -f /tmp/ray_skypilot_installation_complete ] || [ -f /tmp/${STEPS[1]}.failed ]; do sleep 0.5; done; kill $TAIL_PID || true; };
+    [ -f /tmp/${STEPS[1]}.failed ] && { echo "Error: ${STEPS[1]} failed. Exiting."; exit 1; } || true;
+
     end_epoch=$(date +%s);
     echo "=== Ray and skypilot dependencies installation completed in $(($end_epoch - $start_epoch)) secs ===";
     start_epoch=$(date +%s);

sky/usage/usage_lib.py

@@ -14,6 +14,7 @@ from typing_extensions import ParamSpec
 
 import sky
 from sky import sky_logging
+from sky import skypilot_config
 from sky.adaptors import common as adaptors_common
 from sky.usage import constants
 from sky.utils import common_utils
@@ -167,6 +168,7 @@ class UsageMessageToReport(MessageToReport):
         self.runtimes: Dict[str, float] = {}  # update_runtime
         self.exception: Optional[str] = None  # entrypoint_context
         self.stacktrace: Optional[str] = None  # entrypoint_context
+        self.skypilot_config: Optional[Dict[str, Any]] = None
 
         # Whether API server is deployed remotely.
         self.using_remote_api_server: bool = (
@@ -177,6 +179,7 @@ class UsageMessageToReport(MessageToReport):
             self.client_entrypoint = common_utils.get_current_client_entrypoint(
                 msg)
         self.entrypoint = msg
+        self.skypilot_config = dict(skypilot_config.to_dict())
 
     def set_internal(self):
         self.internal = True

sky/utils/cli_utils/status_utils.py

@@ -11,6 +11,7 @@ from sky.utils import common_utils
 from sky.utils import log_utils
 from sky.utils import resources_utils
 from sky.utils import status_lib
+from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
     from sky.provision.kubernetes import utils as kubernetes_utils
@@ -105,11 +106,9 @@ def show_status_table(cluster_records: List[responses.StatusResponse],
 
     if query_clusters:
         cluster_names = {record['name'] for record in cluster_records}
-        not_found_clusters = [
-            repr(cluster)
-            for cluster in query_clusters
-            if cluster not in cluster_names
-        ]
+        not_found_clusters = ux_utils.get_non_matched_query(
+            query_clusters, cluster_names)
+        not_found_clusters = [repr(cluster) for cluster in not_found_clusters]
         if not_found_clusters:
             cluster_str = 'Cluster'
             if len(not_found_clusters) > 1:

sky/utils/context.py

@@ -2,15 +2,21 @@
 
 import asyncio
 from collections.abc import Mapping
-from collections.abc import MutableMapping
 import contextvars
+import copy
 import functools
+import inspect
 import os
 import pathlib
 import subprocess
 import sys
-import typing
-from typing import Any, Callable, Dict, Optional, TextIO, TypeVar
+from typing import (Callable, Dict, Iterator, MutableMapping, Optional, TextIO,
+                    TYPE_CHECKING, TypeVar)
+
+from typing_extensions import ParamSpec
+
+if TYPE_CHECKING:
+    from sky.skypilot_config import ConfigContext
 
 
 class Context(object):
@@ -88,7 +94,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,8 +108,30 @@ 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 copy(self) -> 'Context':
+        """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 = Context()
+        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[Context]]('sky_context',
+                                                     default=None)
 
 
 def get() -> Optional[Context]:
@@ -116,7 +144,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
@@ -155,10 +183,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,10 +198,10 @@ 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: Context) -> Iterator[str]:
+            deleted_keys = set()
             for key, value in ctx.env_overrides.items():
                 if value is None:
                     deleted_keys.add(key)
@@ -182,20 +210,24 @@ class ContextualEnviron(MutableMapping):
                 # 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:
@@ -211,10 +243,13 @@ class ContextualEnviron(MutableMapping):
         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 +260,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):
@@ -260,27 +295,67 @@ class Popen(subprocess.Popen):
         super().__init__(*args, env=env, **kwargs)
 
 
-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.
     """
 
     @functools.wraps(func)
-    def wrapper(*args, **kwargs):
-        initialize()
-        return func(*args, **kwargs)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+        original_ctx = get()
+        initialize(original_ctx)
+        ctx = get()
+        cleanup_after_await = False
+
+        def cleanup():
+            try:
+                if ctx is not None:
+                    ctx.cleanup()
+            finally:
+                # Note: _CONTEXT.reset() is not reliable - may fail with
+                # ValueError: <Token ... at ...> was created in a different
+                # Context
+                # We must make sure this happens because otherwise we may try to
+                # write to the wrong log.
+                _CONTEXT.set(original_ctx)
+
+        # There are two cases:
+        # 1. The function is synchronous (that is, return type is not awaitable)
+        #    In this case, we use a finally block to cleanup the context.
+        # 2. The function is asynchronous (that is, return type is awaitable)
+        #    In this case, we need to construct an async def wrapper and await
+        #    the value, then call the cleanup function in the finally block.
+
+        async def await_with_cleanup(awaitable):
+            try:
+                return await awaitable
+            finally:
+                cleanup()
+
+        try:
+            ret = func(*args, **kwargs)
+            if inspect.isawaitable(ret):
+                cleanup_after_await = True
+                return await_with_cleanup(ret)
+            else:
+                return ret
+        finally:
+            if not cleanup_after_await:
+                cleanup()
 
-    return typing.cast(F, wrapper)
+    return wrapper
 
 
-def initialize():
+def initialize(base_context: Optional[Context] = None) -> None:
     """Initialize the current SkyPilot context."""
-    _CONTEXT.set(Context())
+    new_context = base_context.copy() if base_context is not None else Context()
+    _CONTEXT.set(new_context)
 
 
 class _ContextualStream: