Repository: skypilot-org/skypilot
Pull request #5582 (2025-05-14 08:14:13)

[API server] handle logs request in coroutine
close https://github.com/skypilot-org/skypilot/issues/4767

Reopen this PR to address backward compatibility issue

The review and smoke test history can be found in the original PR: https://github.com/skypilot-org/skypilot/pull/5366

commit to fix the compatibility issue: https://github.com/skypilot-org/skypilot/pull/5582/commits/e1691afe897694fbc8298d0538ef9281f8b184d7


Patch:
diff --git a/.github/workflows/release-publish.yml b/.github/workflows/release-publish.yml
index e0c369760..a6a12f64e 100644
--- a/.github/workflows/release-publish.yml
+++ b/.github/workflows/release-publish.yml
@@ -116,6 +116,15 @@ jobs:
       package_name: 'skypilot'
       version: ${{ needs.extract-version.outputs.version }}
     secrets: inherit
+  
+  publish-helm:
+    needs: [extract-version, publish-to-pypi, trigger-docker-build]
+    uses: ./.github/workflows/publish-helm.yml
+    with:
+      version: ${{ needs.extract-version.outputs.version }}
+      commit_message: 'Update Helm charts for skypilot version ${{ needs.extract-version.outputs.version }}'
+      package_name: 'skypilot'
+    secrets: inherit
 
   cleanup-branches:
     needs: extract-version
diff --git a/docs/requirements-docs.txt b/docs/requirements-docs.txt
index 6566296f7..8ad29b501 100644
--- a/docs/requirements-docs.txt
+++ b/docs/requirements-docs.txt
@@ -19,3 +19,4 @@ sphinxcontrib-htmlhelp==2.0.4
 sphinxcontrib-jsmath==1.0.1
 sphinxcontrib-qthelp==1.0.6
 sphinxcontrib-serializinghtml==1.1.9
+sphinx-tabs==3.4.7
diff --git a/docs/source/conf.py b/docs/source/conf.py
index bbfcb05f9..9799a522b 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -36,6 +36,7 @@ extensions = [
     'sphinx.ext.viewcode',
     'sphinx_autodoc_typehints',
     'sphinx_click',
+    'sphinx_tabs.tabs',
     'sphinx_copybutton',
     "sphinx_togglebutton",  # Used for collapsible sections in examples.
     'sphinxcontrib.googleanalytics',
diff --git a/docs/source/examples/frameworks/aws_efa.md b/docs/source/examples/frameworks/aws_efa.md
new file mode 120000
index 000000000..fe45a9043
--- /dev/null
+++ b/docs/source/examples/frameworks/aws_efa.md
@@ -0,0 +1,1 @@
+../../generated-examples/aws_efa.md
\ No newline at end of file
diff --git a/docs/source/examples/frameworks/gcp_gpu_direct_tcpx.md b/docs/source/examples/frameworks/gcp_gpu_direct_tcpx.md
new file mode 120000
index 000000000..cab896d31
--- /dev/null
+++ b/docs/source/examples/frameworks/gcp_gpu_direct_tcpx.md
@@ -0,0 +1,1 @@
+../../generated-examples/gcp_gpu_direct_tcpx.md
\ No newline at end of file
diff --git a/docs/source/examples/frameworks/index.rst b/docs/source/examples/frameworks/index.rst
index f8b0592ef..a1793959c 100644
--- a/docs/source/examples/frameworks/index.rst
+++ b/docs/source/examples/frameworks/index.rst
@@ -5,9 +5,11 @@ Frameworks
    :maxdepth: 1
 
    Airflow <airflow>
+   AWS EFA <aws_efa>
    Cross-cloud data transfer <https://nebius.com/blog/posts/bulk-object-storage-s3-data-migration-with-skypilot>
    DVC <dvc>
    GCP DWS/Kueue <https://gke-ai-labs.dev/docs/tutorials/skypilot/resource-management-using-kueue/>
+   GCP GPUDirect-TCPX <gcp_gpu_direct_tcpx>
    Jupyter <jupyter>
    MLFlow <https://nebius.com/blog/posts/orchestrating-llm-fine-tuning-k8s-skypilot-mlflow>
    MPI <mpi>
diff --git a/docs/source/examples/index.rst b/docs/source/examples/index.rst
index 39546f9fd..3ccfc3d2c 100644
--- a/docs/source/examples/index.rst
+++ b/docs/source/examples/index.rst
@@ -45,7 +45,7 @@ We welcome contributions from the community. Instructions below:
    4. Go to the appropriate subdir under Examples and add your example.
 
       - For example, if you want to add to "Examples / Serving":
-      - ``cd docs/source/examples/serving; ln -s ../generated-examples/my-llm.md``.
+      - ``cd docs/source/examples/serving; ln -s ../../generated-examples/my-llm.md``.
       - Add a heading for your example to the TOC in the index.rst file.
       - ``git add index.rst my-llm.md``
 
diff --git a/docs/source/reference/training-guide.rst b/docs/source/reference/training-guide.rst
index 0e0795ffe..b37ceaf64 100644
--- a/docs/source/reference/training-guide.rst
+++ b/docs/source/reference/training-guide.rst
@@ -37,6 +37,46 @@ Example configuration:
       A100-80GB:1  
       H100:1  
 
+Use high-performance networking
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. tabs::
+
+    .. tab:: AWS EFA
+
+        AWS Elastic Fabric Adapter (EFA) is a network interface similar to Nvidia Infiniband that enables users to run applications requiring high levels of inter-node communications at scale on AWS. You can enable EFA on AWS HyperPod/EKS clusters with a simple additional setting in your SkyPilot YAML.
+
+        Example configuration:
+
+        .. code-block:: yaml
+
+          config:
+            kubernetes:
+              pod_config:
+                spec:
+                  containers:
+                    - resources:
+                      limits:
+                        vpc.amazonaws.com/efa: 4
+                      requests:
+                        vpc.amazonaws.com/efa: 4
+
+        See `EFA example <https://github.com/skypilot-org/skypilot/tree/master/examples/aws_efa>`_ for more details.
+
+    .. tab:: GCP GPUDirect-TCPX
+
+        `GPUDirect-TCPX <https://cloud.google.com/compute/docs/gpus/gpudirect>`_ is a high-performance networking technology that enables direct communication between GPUs and network interfaces for `a3-highgpu-8g` or `a3-edgegpu-8g` VMs. You can enable it with the following additional setting in your SkyPilot YAML.
+
+        Example configuration:
+
+        .. code-block:: yaml
+
+          config:
+            gcp:
+              enable_gpu_direct: true
+
+        See `GPUDirect-TCPX example <https://github.com/skypilot-org/skypilot/tree/master/examples/gcp_gpu_direct_tcpx>`_ for more details.
+
 Use ``disk_tier: best``
 ~~~~~~~~~~~~~~~~~~~~~~~
 Fast storage is critical for loading and storing data and model checkpoints.
diff --git a/examples/aws_efa/README.md b/examples/aws_efa/README.md
index f342b1326..4813227c3 100644
--- a/examples/aws_efa/README.md
+++ b/examples/aws_efa/README.md
@@ -41,7 +41,7 @@ hyperpod-i-0da69b9076c7ff6a4   ml.p4d.24xlarge   8     4
 
 ## Access HyperPod and run distributed job with SkyPilot
 
-To accesss HyperPod and run distributed job with SkyPilot, see the SkyPilot [HyperPod example](../hyperpod-eks/).
+To access HyperPod and run distributed job with SkyPilot, see the SkyPilot [HyperPod example](https://github.com/skypilot-org/skypilot/blob/master/examples/hyperpod-eks).
 
 ### Adding EFA configurations in SkyPilot YAML
 
@@ -101,11 +101,11 @@ Check the following table for the GPU and EFA count mapping for AWS instance typ
 | g6e.48xlarge  | L40S:8   | 4    |
 
 
-Update the EFA number in the [nccl_efa.yaml](nccl_efa.yaml) for the GPUs you used.
+Update the EFA number in the [`nccl_efa.yaml`](https://github.com/skypilot-org/skypilot/blob/master/examples/aws_efa/nccl_efa.yaml) for the GPUs you use.
 
 ## Running NCCL test with EFA using SkyPilot
 
-Check the [nccl_efa.yaml](nccl_efa.yaml) for the complete SkyPilot cluster yaml cofnigurations.
+Check the [`nccl_efa.yaml`](https://github.com/skypilot-org/skypilot/blob/master/examples/aws_efa/nccl_efa.yaml) for the complete SkyPilot cluster yaml configurations.
 
 The `image_id` provides the environment setup for [NCCL](https://developer.nvidia.com/nccl) (NVIDIA Collective Communications Library) and EFA (Elastic Fabric Adapter).
 
@@ -121,7 +121,7 @@ SkyPilot will:
 3. Execute the NCCL performance test with EFA networking
 4. Output performance metrics showing the benefits of EFA for distributed training
 
-> [!NOTE]
+> **NOTE:**
 > We can turn off EFA with `nccl_efa.yaml` by passing an env:
 > ```bash
 > sky launch -c efa --env USE_EFA=false nccl_efa.yaml
diff --git a/examples/gcp_gpu_direct_tcpx/README.md b/examples/gcp_gpu_direct_tcpx/README.md
index 2cba5f8c5..1c51ec171 100644
--- a/examples/gcp_gpu_direct_tcpx/README.md
+++ b/examples/gcp_gpu_direct_tcpx/README.md
@@ -32,7 +32,7 @@ With `enable_gpu_direct: true`, SkyPilot automatically:
 
 ## Running NCCL Tests with GPUDirect-TCPX
 
-The complete configuration is available in `gpu_direct_tcpx.yaml`. The configuration includes:
+The complete configuration is available in [`gpu_direct_tcpx.yaml`](https://github.com/skypilot-org/skypilot/blob/master/examples/gcp_gpu_direct_tcpx/gpu_direct_tcpx.yaml). The configuration includes:
 - `image_id`: Pre-configured environment for NCCL testing
 - `instance_type`: Set to `a3-highgpu-8g`
 
@@ -43,7 +43,7 @@ sky launch -c tcpx gpu_direct_tcpx.yaml
 ```
 
 SkyPilot will:
-1. Deploy a GCP cluster with GPUDirect-TCPX-enabled nodes
+1. Deploy a GCP cluster with GPUDirect-TCPX enabled nodes
 2. Execute NCCL performance tests
 3. Output detailed performance metrics
 
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index 0ec8664df..0463fc274 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -40,6 +40,7 @@ from sky.utils import cluster_utils
 from sky.utils import command_runner
 from sky.utils import common
 from sky.utils import common_utils
+from sky.utils import context_utils
 from sky.utils import controller_utils
 from sky.utils import env_options
 from sky.utils import registry
@@ -2204,6 +2205,7 @@ def refresh_cluster_record(
 
 
 @timeline.event
+@context_utils.cancellation_guard
 def refresh_cluster_status_handle(
     cluster_name: str,
     *,
@@ -2253,6 +2255,7 @@ def check_cluster_available(
     ...
 
 
+@context_utils.cancellation_guard
 def check_cluster_available(
     cluster_name: str,
     *,
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index e920b53bf..374ef434c 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -61,6 +61,7 @@ from sky.utils import cluster_utils
 from sky.utils import command_runner
 from sky.utils import common
 from sky.utils import common_utils
+from sky.utils import context_utils
 from sky.utils import controller_utils
 from sky.utils import env_options
 from sky.utils import log_utils
@@ -274,6 +275,7 @@ class RayCodeGen:
         ray_address = 'auto'
         self._code = [
             textwrap.dedent(f"""\
+            import functools
             import getpass
             import hashlib
             import io
@@ -363,6 +365,7 @@ class RayCodeGen:
             # by ray.remote. This should be removed once we have a better way to
             # specify dependencies for ray.
             inspect.getsource(log_lib._ProcessingArgs),  # pylint: disable=protected-access
+            inspect.getsource(log_lib._get_context),  # pylint: disable=protected-access
             inspect.getsource(log_lib._handle_io_stream),  # pylint: disable=protected-access
             inspect.getsource(log_lib.process_subprocess_stream),
             inspect.getsource(log_lib.run_with_log),
@@ -2415,6 +2418,7 @@ class CloudVmRayResourceHandle(backends.backend.ResourceHandle):
             internal_external_ips[1:], key=lambda x: x[1])
         self.stable_internal_external_ips = stable_internal_external_ips
 
+    @context_utils.cancellation_guard
     @annotations.lru_cache(scope='global')
     @timeline.event
     def get_command_runners(self,
@@ -3842,6 +3846,7 @@ class CloudVmRayBackend(backends.Backend['CloudVmRayResourceHandle']):
         subprocess_utils.run_in_parallel(_rsync_down, parallel_args)
         return dict(zip(job_ids, local_log_dirs))
 
+    @context_utils.cancellation_guard
     def tail_logs(self,
                   handle: CloudVmRayResourceHandle,
                   job_id: Optional[int],
@@ -4559,6 +4564,7 @@ class CloudVmRayBackend(backends.Backend['CloudVmRayResourceHandle']):
     # TODO(zhwu): Refactor this to a CommandRunner class, so different backends
     # can support its own command runner.
     @timeline.event
+    @context_utils.cancellation_guard
     def run_on_head(
         self,
         handle: CloudVmRayResourceHandle,
diff --git a/sky/global_user_state.py b/sky/global_user_state.py
index e1c9ce21d..6dc1bb6ad 100644
--- a/sky/global_user_state.py
+++ b/sky/global_user_state.py
@@ -19,6 +19,7 @@ import uuid
 from sky import models
 from sky import sky_logging
 from sky.utils import common_utils
+from sky.utils import context_utils
 from sky.utils import db_utils
 from sky.utils import registry
 from sky.utils import status_lib
@@ -671,6 +672,7 @@ def _load_storage_mounts_metadata(
     return pickle.loads(record_storage_mounts_metadata)
 
 
+@context_utils.cancellation_guard
 def get_cluster_from_name(
         cluster_name: Optional[str]) -> Optional[Dict[str, Any]]:
     rows = _DB.cursor.execute(
diff --git a/sky/optimizer.py b/sky/optimizer.py
index 3c2f2c320..c5f5b6ba1 100644
--- a/sky/optimizer.py
+++ b/sky/optimizer.py
@@ -1195,10 +1195,12 @@ def _check_specified_clouds(dag: 'dag_lib.Dag') -> None:
             all_clouds_specified.add(cloud_str)
 
         # Explicitly check again to update the enabled cloud list.
-        sky_check.check_capability(sky_cloud.CloudCapability.COMPUTE,
-                                   quiet=True,
-                                   clouds=list(clouds_need_recheck -
-                                               global_disabled_clouds))
+        clouds_to_check_again = list(clouds_need_recheck -
+                                     global_disabled_clouds)
+        if len(clouds_to_check_again) > 0:
+            sky_check.check_capability(sky_cloud.CloudCapability.COMPUTE,
+                                       quiet=True,
+                                       clouds=clouds_to_check_again)
         enabled_clouds = sky_check.get_cached_enabled_clouds_or_refresh(
             capability=sky_cloud.CloudCapability.COMPUTE,
             raise_if_no_cloud_access=True)
diff --git a/sky/server/requests/executor.py b/sky/server/requests/executor.py
index 5317856bd..2cd019dd4 100644
--- a/sky/server/requests/executor.py
+++ b/sky/server/requests/executor.py
@@ -18,7 +18,10 @@ The number of the workers is determined by the system resources.
 
 See the [README.md](../README.md) for detailed architecture of the executor.
 """
+import asyncio
 import contextlib
+import contextvars
+import functools
 import multiprocessing
 import os
 import queue as queue_lib
@@ -47,6 +50,7 @@ from sky.server.requests.queues import mp_queue
 from sky.skylet import constants
 from sky.utils import annotations
 from sky.utils import common_utils
+from sky.utils import context
 from sky.utils import subprocess_utils
 from sky.utils import timeline
 
@@ -60,7 +64,6 @@ else:
     from typing_extensions import ParamSpec
 
 P = ParamSpec('P')
-
 logger = sky_logging.init_logger(__name__)
 
 # On macOS, the default start method for multiprocessing is 'fork', which
@@ -341,6 +344,114 @@ def _request_execution_wrapper(request_id: str,
             logger.info(f'Request {request_id} finished')
 
 
+async def execute_request_coroutine(request: api_requests.Request):
+    """Execute a request in current event loop.
+
+    Similar to _request_execution_wrapper, but executed as coroutine in current
+    event loop. This is designed for executing tasks that are not CPU
+    intensive, e.g. sky logs.
+    """
+    ctx = context.get()
+    if ctx is None:
+        raise ValueError('Context is not initialized')
+    logger.info(f'Executing request {request.request_id} in coroutine')
+    func = request.entrypoint
+    request_body = request.request_body
+    with api_requests.update_request(request.request_id) as request_task:
+        request_task.status = api_requests.RequestStatus.RUNNING
+    # Redirect stdout and stderr to the request log path.
+    original_output = ctx.redirect_log(request.log_path)
+    # Override environment variables that backs env_options.Options
+    # TODO(aylei): compared to process executor, running task in coroutine has
+    # two issues to fix:
+    # 1. skypilot config is not contextual
+    # 2. envs that read directly from os.environ are not contextual
+    ctx.override_envs(request_body.env_vars)
+    loop = asyncio.get_running_loop()
+    pyctx = contextvars.copy_context()
+    func_call = functools.partial(pyctx.run, func, **request_body.to_kwargs())
+    fut: asyncio.Future = loop.run_in_executor(None, func_call)
+
+    async def poll_task(request_id: str) -> bool:
+        request = api_requests.get_request(request_id)
+        if request is None:
+            raise RuntimeError('Request not found')
+
+        if request.status == api_requests.RequestStatus.CANCELLED:
+            ctx.cancel()
+            return True
+
+        if fut.done():
+            try:
+                result = await fut
+                api_requests.set_request_succeeded(request_id, result)
+            except asyncio.CancelledError:
+                # The task is cancelled by ctx.cancel(), where the status
+                # should already be set to CANCELLED.
+                pass
+            except Exception as e:  # pylint: disable=broad-except
+                ctx.redirect_log(original_output)
+                api_requests.set_request_failed(request_id, e)
+                logger.error(f'Request {request_id} failed due to '
+                             f'{common_utils.format_exception(e)}')
+            return True
+        return False
+
+    try:
+        while True:
+            res = await poll_task(request.request_id)
+            if res:
+                break
+            await asyncio.sleep(0.5)
+    except asyncio.CancelledError:
+        # Current coroutine is cancelled due to client disconnect, set the
+        # request status for consistency.
+        api_requests.set_request_cancelled(request.request_id)
+        pass
+    # pylint: disable=broad-except
+    except (Exception, KeyboardInterrupt, SystemExit) as e:
+        # Handle any other error
+        ctx.redirect_log(original_output)
+        ctx.cancel()
+        api_requests.set_request_failed(request.request_id, e)
+        logger.error(f'Request {request.request_id} interrupted due to '
+                     f'unhandled exception: {common_utils.format_exception(e)}')
+        raise
+
+
+def prepare_request(
+    request_id: str,
+    request_name: str,
+    request_body: payloads.RequestBody,
+    func: Callable[P, Any],
+    request_cluster_name: Optional[str] = None,
+    schedule_type: api_requests.ScheduleType = (api_requests.ScheduleType.LONG),
+    is_skypilot_system: bool = False,
+) -> api_requests.Request:
+    """Prepare a request for execution."""
+    user_id = request_body.env_vars[constants.USER_ID_ENV_VAR]
+    if is_skypilot_system:
+        user_id = server_constants.SKYPILOT_SYSTEM_USER_ID
+        global_user_state.add_or_update_user(
+            models.User(id=user_id, name=user_id))
+    request = api_requests.Request(request_id=request_id,
+                                   name=server_constants.REQUEST_NAME_PREFIX +
+                                   request_name,
+                                   entrypoint=func,
+                                   request_body=request_body,
+                                   status=api_requests.RequestStatus.PENDING,
+                                   created_at=time.time(),
+                                   schedule_type=schedule_type,
+                                   user_id=user_id,
+                                   cluster_name=request_cluster_name)
+
+    if not api_requests.create_if_not_exists(request):
+        raise RuntimeError(f'Request {request_id} already exists.')
+
+    request.log_path.touch()
+    return request
+
+
 def schedule_request(
         request_id: str,
         request_name: str,
@@ -372,27 +483,8 @@ def schedule_request(
             The precondition is waited asynchronously and does not block the
             caller.
     """
-    user_id = request_body.env_vars[constants.USER_ID_ENV_VAR]
-    if is_skypilot_system:
-        user_id = server_constants.SKYPILOT_SYSTEM_USER_ID
-        global_user_state.add_or_update_user(
-            models.User(id=user_id, name=user_id))
-    request = api_requests.Request(request_id=request_id,
-                                   name=server_constants.REQUEST_NAME_PREFIX +
-                                   request_name,
-                                   entrypoint=func,
-                                   request_body=request_body,
-                                   status=api_requests.RequestStatus.PENDING,
-                                   created_at=time.time(),
-                                   schedule_type=schedule_type,
-                                   user_id=user_id,
-                                   cluster_name=request_cluster_name)
-
-    if not api_requests.create_if_not_exists(request):
-        logger.debug(f'Request {request_id} already exists.')
-        return
-
-    request.log_path.touch()
+    prepare_request(request_id, request_name, request_body, func,
+                    request_cluster_name, schedule_type, is_skypilot_system)
 
     def enqueue():
         input_tuple = (request_id, ignore_return_value)
diff --git a/sky/server/requests/requests.py b/sky/server/requests/requests.py
index d417e0295..cce2c0138 100644
--- a/sky/server/requests/requests.py
+++ b/sky/server/requests/requests.py
@@ -606,3 +606,18 @@ def set_request_failed(request_id: str, e: BaseException) -> None:
         assert request_task is not None, request_id
         request_task.status = RequestStatus.FAILED
         request_task.set_error(e)
+
+
+def set_request_succeeded(request_id: str, result: Any) -> None:
+    """Set a request to succeeded and populate the result."""
+    with update_request(request_id) as request_task:
+        assert request_task is not None, request_id
+        request_task.status = RequestStatus.SUCCEEDED
+        request_task.set_return_value(result)
+
+
+def set_request_cancelled(request_id: str) -> None:
+    """Set a request to cancelled."""
+    with update_request(request_id) as request_task:
+        assert request_task is not None, request_id
+        request_task.status = RequestStatus.CANCELLED
diff --git a/sky/server/server.py b/sky/server/server.py
index 7c0f0fc06..395a675bd 100644
--- a/sky/server/server.py
+++ b/sky/server/server.py
@@ -47,6 +47,7 @@ from sky.usage import usage_lib
 from sky.utils import admin_policy_utils
 from sky.utils import common as common_lib
 from sky.utils import common_utils
+from sky.utils import context
 from sky.utils import dag_utils
 from sky.utils import env_options
 from sky.utils import status_lib
@@ -673,24 +674,28 @@ async def logs(
     # TODO(zhwu): This should wait for the request on the cluster, e.g., async
     # launch, to finish, so that a user does not need to manually pull the
     # request status.
-    executor.schedule_request(
+    # Only initialize the context in logs handler to limit the scope of this
+    # experimental change.
+    # TODO(aylei): init in lifespan() to enable SkyPilot context in all APIs.
+    context.initialize()
+    request_task = executor.prepare_request(
         request_id=request.state.request_id,
         request_name='logs',
         request_body=cluster_job_body,
         func=core.tail_logs,
-        # TODO(aylei): We have tail logs scheduled as SHORT request, because it
-        # should be responsive. However, it can be long running if the user's
-        # job keeps running, and we should avoid it taking the SHORT worker.
         schedule_type=requests_lib.ScheduleType.SHORT,
-        request_cluster_name=cluster_job_body.cluster_name,
     )
+    task = asyncio.create_task(executor.execute_request_coroutine(request_task))
 
-    request_task = requests_lib.get_request(request.state.request_id)
+    def cancel_task():
+        task.cancel()
 
+    # Cancel the task after the request is done or client disconnects
+    background_tasks.add_task(cancel_task)
     # TODO(zhwu): This makes viewing logs in browser impossible. We should adopt
     # the same approach as /stream.
     return stream_utils.stream_response(
-        request_id=request_task.request_id,
+        request_id=request.state.request_id,
         logs_path=request_task.log_path,
         background_tasks=background_tasks,
     )
diff --git a/sky/server/uvicorn.py b/sky/server/uvicorn.py
index 5ad7db485..e99213ea6 100644
--- a/sky/server/uvicorn.py
+++ b/sky/server/uvicorn.py
@@ -3,6 +3,7 @@
 This module is a wrapper around uvicorn to customize the behavior of the
 server.
 """
+import functools
 import os
 import threading
 from typing import Optional
@@ -10,6 +11,7 @@ from typing import Optional
 import uvicorn
 from uvicorn.supervisors import multiprocess
 
+from sky.utils import context_utils
 from sky.utils import subprocess_utils
 
 
@@ -21,19 +23,27 @@ def run(config: uvicorn.Config):
         # guard by an exception.
         raise ValueError('Reload is not supported yet.')
     server = uvicorn.Server(config=config)
+    run_server_process = functools.partial(_run_server_process, server)
     try:
         if config.workers is not None and config.workers > 1:
             sock = config.bind_socket()
-            SlowStartMultiprocess(config, target=server.run,
+            SlowStartMultiprocess(config,
+                                  target=run_server_process,
                                   sockets=[sock]).run()
         else:
-            server.run()
+            run_server_process()
     finally:
         # Copied from unvicorn.run()
         if config.uds and os.path.exists(config.uds):
             os.remove(config.uds)
 
 
+def _run_server_process(server: uvicorn.Server, *args, **kwargs):
+    """Run the server process with contextually aware."""
+    context_utils.hijack_sys_attrs()
+    server.run(*args, **kwargs)
+
+
 class SlowStartMultiprocess(multiprocess.Multiprocess):
     """Uvicorn Multiprocess wrapper with slow start.
 
diff --git a/sky/sky_logging.py b/sky/sky_logging.py
index 7021ec9bc..f38147dff 100644
--- a/sky/sky_logging.py
+++ b/sky/sky_logging.py
@@ -10,6 +10,7 @@ import threading
 import colorama
 
 from sky.skylet import constants
+from sky.utils import context
 from sky.utils import env_options
 from sky.utils import rich_utils
 
@@ -47,6 +48,43 @@ class NewLineFormatter(logging.Formatter):
         return msg
 
 
+class EnvAwareHandler(rich_utils.RichSafeStreamHandler):
+    """A handler that awares environment variables.
+
+    This handler dynamically reflects the log level from environment variables.
+    """
+
+    def __init__(self, stream=None, level=logging.NOTSET, sensitive=False):
+        super().__init__(stream)
+        self.level = level
+        self._sensitive = sensitive
+
+    @property
+    def level(self):
+        # Only refresh log level if we are in a context, since the log level
+        # has already been reloaded eagerly in multi-processing. Refresh again
+        # is a no-op and can be avoided.
+        # TODO(aylei): unify the mechanism for coroutine context and
+        # multi-processing.
+        if context.get() is not None:
+            if self._sensitive:
+                # For sensitive logger, suppress debug log despite the
+                # SKYPILOT_DEBUG env var if SUPPRESS_SENSITIVE_LOG is set
+                if env_options.Options.SUPPRESS_SENSITIVE_LOG.get():
+                    return logging.INFO
+            if env_options.Options.SHOW_DEBUG_INFO.get():
+                return logging.DEBUG
+            else:
+                return self._level
+        else:
+            return self._level
+
+    @level.setter
+    def level(self, level):
+        # pylint: disable=protected-access
+        self._level = logging._checkLevel(level)
+
+
 _root_logger = logging.getLogger('sky')
 _default_handler = None
 _logging_config = threading.local()
@@ -67,7 +105,7 @@ def _setup_logger():
     _root_logger.setLevel(logging.DEBUG)
     global _default_handler
     if _default_handler is None:
-        _default_handler = rich_utils.RichSafeStreamHandler(sys.stdout)
+        _default_handler = EnvAwareHandler(sys.stdout)
         _default_handler.flush = sys.stdout.flush  # type: ignore
         if env_options.Options.SHOW_DEBUG_INFO.get():
             _default_handler.setLevel(logging.DEBUG)
@@ -87,7 +125,7 @@ def _setup_logger():
         # for certain loggers.
         for logger_name in _SENSITIVE_LOGGER:
             logger = logging.getLogger(logger_name)
-            handler_to_logger = rich_utils.RichSafeStreamHandler(sys.stdout)
+            handler_to_logger = EnvAwareHandler(sys.stdout, sensitive=True)
             handler_to_logger.flush = sys.stdout.flush  # type: ignore
             logger.addHandler(handler_to_logger)
             logger.setLevel(logging.INFO)
diff --git a/sky/skylet/log_lib.py b/sky/skylet/log_lib.py
index 24df608cd..878ee26a4 100644
--- a/sky/skylet/log_lib.py
+++ b/sky/skylet/log_lib.py
@@ -4,6 +4,7 @@ This is a remote utility module that provides logging functionality.
 """
 import collections
 import copy
+import functools
 import io
 import multiprocessing.pool
 import os
@@ -21,6 +22,8 @@ import colorama
 from sky import sky_logging
 from sky.skylet import constants
 from sky.skylet import job_lib
+from sky.utils import context
+from sky.utils import context_utils
 from sky.utils import log_utils
 from sky.utils import subprocess_utils
 from sky.utils import ux_utils
@@ -59,6 +62,16 @@ class _ProcessingArgs:
         self.streaming_prefix = streaming_prefix
 
 
+def _get_context():
+    # TODO(aylei): remove this after we drop the backward-compatibility for
+    # 0.9.x in 0.12.0
+    # Keep backward-compatibility for the old version of SkyPilot runtimes.
+    if 'context' in globals():
+        return context.get()
+    else:
+        return None
+
+
 def _handle_io_stream(io_stream, out_stream, args: _ProcessingArgs):
     """Process the stream of a process."""
     out_io = io.TextIOWrapper(io_stream,
@@ -77,6 +90,9 @@ def _handle_io_stream(io_stream, out_stream, args: _ProcessingArgs):
     with open(args.log_path, 'a', encoding='utf-8') as fout:
         with line_processor:
             while True:
+                ctx = _get_context()
+                if ctx is not None and ctx.is_canceled():
+                    return
                 line = out_io.readline()
                 if not line:
                     break
@@ -111,26 +127,24 @@ def _handle_io_stream(io_stream, out_stream, args: _ProcessingArgs):
     return ''.join(out)
 
 
-def process_subprocess_stream(proc, args: _ProcessingArgs) -> Tuple[str, str]:
-    """Redirect the process's filtered stdout/stderr to both stream and file"""
+def process_subprocess_stream(proc, stdout_stream_handler,
+                              stderr_stream_handler) -> Tuple[str, str]:
+    """Process the stream of a process in threads, blocking."""
     if proc.stderr is not None:
         # Asyncio does not work as the output processing can be executed in a
         # different thread.
         # selectors is possible to handle the multiplexing of stdout/stderr,
         # but it introduces buffering making the output not streaming.
         with multiprocessing.pool.ThreadPool(processes=1) as pool:
-            err_args = copy.copy(args)
-            err_args.line_processor = None
-            stderr_fut = pool.apply_async(_handle_io_stream,
-                                          args=(proc.stderr, sys.stderr,
-                                                err_args))
+            stderr_fut = pool.apply_async(stderr_stream_handler,
+                                          args=(proc.stderr, sys.stderr))
             # Do not launch a thread for stdout as the rich.status does not
             # work in a thread, which is used in
             # log_utils.RayUpLineProcessor.
-            stdout = _handle_io_stream(proc.stdout, sys.stdout, args)
+            stdout = stdout_stream_handler(proc.stdout, sys.stdout)
             stderr = stderr_fut.get()
     else:
-        stdout = _handle_io_stream(proc.stdout, sys.stdout, args)
+        stdout = stdout_stream_handler(proc.stdout, sys.stdout)
         stderr = ''
     return stdout, stderr
 
@@ -176,7 +190,12 @@ def run_with_log(
     # Redirect stderr to stdout when using ray, to preserve the order of
     # stdout and stderr.
     stdout_arg = stderr_arg = None
-    if process_stream:
+    ctx = _get_context()
+    if process_stream or ctx is not None:
+        # Capture stdout/stderr of the subprocess if:
+        # 1. Post-processing is needed (process_stream=True)
+        # 2. Potential contextual handling is needed (ctx is not None)
+        # TODO(aylei): can we always capture the stdout/stderr?
         stdout_arg = subprocess.PIPE
         stderr_arg = subprocess.PIPE if not with_ray else subprocess.STDOUT
     # Use stdin=subprocess.DEVNULL by default, as allowing inputs will mess up
@@ -197,6 +216,8 @@ def run_with_log(
             subprocess_utils.kill_process_daemon(proc.pid)
             stdout = ''
             stderr = ''
+            stdout_stream_handler = None
+            stderr_stream_handler = None
 
             if process_stream:
                 if skip_lines is None:
@@ -223,7 +244,35 @@ def run_with_log(
                     replace_crlf=with_ray,
                     streaming_prefix=streaming_prefix,
                 )
-                stdout, stderr = process_subprocess_stream(proc, args)
+                stdout_stream_handler = functools.partial(
+                    _handle_io_stream,
+                    args=args,
+                )
+                if proc.stderr is not None:
+                    err_args = copy.copy(args)
+                    err_args.line_processor = None
+                    stderr_stream_handler = functools.partial(
+                        _handle_io_stream,
+                        args=err_args,
+                    )
+            if ctx is not None:
+                # When runs in a coroutine, always process the subprocess
+                # stream to:
+                # 1. handle context cancellation
+                # 2. redirect subprocess stdout/stderr to the contextual
+                #    stdout/stderr of current coroutine.
+                stdout, stderr = context_utils.pipe_and_wait_process(
+                    ctx,
+                    proc,
+                    cancel_callback=subprocess_utils.kill_children_processes,
+                    stdout_stream_handler=stdout_stream_handler,
+                    stderr_stream_handler=stderr_stream_handler)
+            elif process_stream:
+                # When runs in a process, only process subprocess stream if
+                # necessary to avoid unnecessary stream handling overhead.
+                stdout, stderr = process_subprocess_stream(
+                    proc, stdout_stream_handler, stderr_stream_handler)
+            # Ensure returncode is set.
             proc.wait()
             if require_outputs:
                 return proc.returncode, stdout, stderr
diff --git a/sky/skylet/log_lib.pyi b/sky/skylet/log_lib.pyi
index c7028e121..204867f98 100644
--- a/sky/skylet/log_lib.pyi
+++ b/sky/skylet/log_lib.pyi
@@ -11,6 +11,7 @@ from typing_extensions import Literal
 from sky import sky_logging as sky_logging
 from sky.skylet import constants as constants
 from sky.skylet import job_lib as job_lib
+from sky.utils import context
 from sky.utils import log_utils as log_utils
 
 SKY_LOG_WAITING_GAP_SECONDS: int = ...
@@ -41,6 +42,10 @@ class _ProcessingArgs:
         ...
 
 
+def _get_context() -> Optional[context.Context]:
+    ...
+
+
 def _handle_io_stream(io_stream, out_stream, args: _ProcessingArgs):
     ...
 
diff --git a/sky/utils/command_runner.py b/sky/utils/command_runner.py
index d1bdd9ab6..b6e2bbb9e 100644
--- a/sky/utils/command_runner.py
+++ b/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]],
diff --git a/sky/utils/context.py b/sky/utils/context.py
new file mode 100644
index 000000000..9376e88de
--- /dev/null
+++ b/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)
diff --git a/sky/utils/context_utils.py b/sky/utils/context_utils.py
new file mode 100644
index 000000000..842b86ce1
--- /dev/null
+++ b/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)
diff --git a/sky/utils/rich_utils.py b/sky/utils/rich_utils.py
index 606df16aa..53c759e17 100644
--- a/sky/utils/rich_utils.py
+++ b/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()
 
 
diff --git a/sky/utils/subprocess_utils.py b/sky/utils/subprocess_utils.py
index 992b5d591..554d7cc65 100644
--- a/sky/utils/subprocess_utils.py
+++ b/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


Test patch:
diff --git a/tests/unit_tests/test_sky/server/test_server.py b/tests/unit_tests/test_sky/server/test_server.py
index ec2110944..db681b4a8 100644
--- a/tests/unit_tests/test_sky/server/test_server.py
+++ b/tests/unit_tests/test_sky/server/test_server.py
@@ -2,6 +2,8 @@
 
 import argparse
 import asyncio
+import threading
+import time
 from unittest import mock
 
 import fastapi
@@ -10,6 +12,7 @@ import uvicorn
 
 from sky.server import server
 from sky.utils import common_utils
+from sky.utils import context
 
 
 @mock.patch('uvicorn.run')
@@ -135,7 +138,6 @@ async def test_validate():
 
     def slow_validate(*args, **kwargs):
         # Simulate slow validation
-        import time
         time.sleep(0.1)
         validation_complete.set()
 
@@ -155,3 +157,57 @@ async def test_validate():
         # Wait for validation to complete
         await validation_task
         assert validation_complete.is_set()
+
+
+@pytest.mark.asyncio
+async def test_logs():
+    """Test the logs endpoint."""
+    mock_cluster_job_body = mock.MagicMock()
+    mock_cluster_job_body.cluster_name = 'test-cluster'
+    mock_background_tasks = mock.MagicMock()
+
+    # Create an event to track when logs streaming starts
+    streaming_started = threading.Event()
+
+    # Mock the stream_response function
+    def mock_stream_response(*args, **kwargs):
+        streaming_started.set()
+        return fastapi.responses.StreamingResponse(
+            content=iter([]),  # Empty iterator for testing
+            media_type='text/plain')
+
+    def slow_execute(*args, **kwargs):
+        # Simulate slow execution
+        time.sleep(1)
+
+    with mock.patch('sky.server.requests.executor.prepare_request') as mock_prepare, \
+         mock.patch('sky.server.requests.executor.execute_request_coroutine',
+                   side_effect=slow_execute) as mock_execute, \
+         mock.patch('sky.server.stream_utils.stream_response',
+                   side_effect=mock_stream_response) as mock_stream:
+
+        # Mock prepare_request to return a request task
+        mock_request_task = mock.MagicMock()
+        mock_request_task.log_path = '/tmp/test.log'
+        mock_prepare.return_value = mock_request_task
+
+        # Start logs endpoint in background
+        logs_task = asyncio.create_task(
+            server.logs(mock.MagicMock(), mock_cluster_job_body,
+                        mock_background_tasks))
+
+        # Execute should be run in background and does not block streaming start
+        streaming_started.wait(timeout=0.1)
+
+        # Verify the response was created
+        response = await logs_task
+        assert isinstance(response, fastapi.responses.StreamingResponse)
+        assert response.media_type == 'text/plain'
+
+        # Verify the executor calls
+        mock_prepare.assert_called_once()
+        mock_execute.assert_called_once_with(mock_request_task)
+        mock_stream.assert_called_once_with(
+            request_id=mock.ANY,
+            logs_path=mock_request_task.log_path,
+            background_tasks=mock_background_tasks)
diff --git a/tests/unit_tests/test_sky/test_sky_logging.py b/tests/unit_tests/test_sky/test_sky_logging.py
new file mode 100644
index 000000000..e18556f6f
--- /dev/null
+++ b/tests/unit_tests/test_sky/test_sky_logging.py
@@ -0,0 +1,123 @@
+"""Unit tests for sky.sky_logging module."""
+
+import io
+import logging
+from unittest import mock
+
+import pytest
+
+from sky import sky_logging
+from sky.utils import context
+from sky.utils import env_options
+
+
+@pytest.fixture
+def stream():
+    """Fixture to provide a string IO stream."""
+    return io.StringIO()
+
+
+@pytest.fixture
+def handler(stream):
+    """Fixture to provide an EnvAwareHandler instance."""
+    return sky_logging.EnvAwareHandler(stream=stream, level=logging.INFO)
+
+
+@pytest.fixture
+def sensitive_handler(stream):
+    """Fixture to provide a sensitive EnvAwareHandler instance."""
+    return sky_logging.EnvAwareHandler(stream=stream,
+                                       level=logging.INFO,
+                                       sensitive=True)
+
+
+@mock.patch('sky.utils.context.get', return_value=None)
+def test_handler_default_level(mock_get, handler, monkeypatch):
+    """Test handler's default log level behavior."""
+    monkeypatch.setattr(env_options.Options.SHOW_DEBUG_INFO, 'get',
+                        lambda: False)
+    assert handler.level == logging.INFO
+    handler.level = logging.DEBUG
+    assert handler.level == logging.DEBUG
+
+
+@mock.patch('sky.utils.context.get', return_value=context.Context())
+def test_handler_with_debug_env(mock_get, handler, monkeypatch):
+    """Test handler respects SHOW_DEBUG_INFO env var."""
+    # Mock the env var to enable debug
+    monkeypatch.setattr(env_options.Options.SHOW_DEBUG_INFO, 'get',
+                        lambda: True)
+    assert handler.level == logging.DEBUG
+
+    # Mock the env var to disable debug
+    monkeypatch.setattr(env_options.Options.SHOW_DEBUG_INFO, 'get',
+                        lambda: False)
+    assert handler.level == logging.INFO
+
+
+@mock.patch('sky.utils.context.get', return_value=None)
+def test_handler_without_context(mock_get, handler, monkeypatch):
+    """Test handler behavior when no context is available."""
+    # Even with debug env var set, should return original level
+    monkeypatch.setattr(env_options.Options.SHOW_DEBUG_INFO, 'get',
+                        lambda: True)
+    assert handler.level == logging.INFO
+
+
+@mock.patch('sky.utils.context.get', return_value=context.Context())
+def test_sensitive_handler(mock_get, sensitive_handler, monkeypatch):
+    """Test sensitive handler's log level behavior."""
+    # Test when sensitive logs are suppressed
+    monkeypatch.setattr(env_options.Options.SUPPRESS_SENSITIVE_LOG, 'get',
+                        lambda: True)
+    monkeypatch.setattr(env_options.Options.SHOW_DEBUG_INFO, 'get',
+                        lambda: True)
+    assert sensitive_handler.level == logging.INFO
+
+    # Test when sensitive logs are not suppressed
+    monkeypatch.setattr(env_options.Options.SUPPRESS_SENSITIVE_LOG, 'get',
+                        lambda: False)
+    assert sensitive_handler.level == logging.DEBUG
+
+
+@mock.patch('sky.utils.context.get', return_value=context.Context())
+def test_handler_invalid_level(mock_get, handler, monkeypatch):
+    """Test setting invalid log levels."""
+    monkeypatch.setattr(env_options.Options.SHOW_DEBUG_INFO, 'get',
+                        lambda: False)
+    with pytest.raises(ValueError):
+        handler.level = 'INVALID_LEVEL'
+
+    # Test valid level setting
+    handler.level = logging.WARNING
+    assert handler._level == logging.WARNING
+
+
+@mock.patch('sky.utils.context.get', return_value=None)
+def test_handler_with_context_override(mock_get, handler, monkeypatch):
+    """Test setting invalid log levels."""
+    assert handler.level == logging.INFO
+    ctx = context.Context()
+    ctx.override_envs({'SKYPILOT_DEBUG': '1'})
+    mock_get.return_value = ctx
+    assert handler.level == logging.DEBUG
+
+
+@mock.patch('sky.utils.context.get', return_value=None)
+def test_handler_output(mock_get, handler, stream, monkeypatch):
+    """Sanity check."""
+    monkeypatch.setattr(env_options.Options.SHOW_DEBUG_INFO, 'get',
+                        lambda: False)
+    logger = logging.getLogger('test')
+    logger.addHandler(handler)
+    logger.setLevel(logging.DEBUG)
+
+    logger.debug('Debug message')
+    assert stream.getvalue() == ''
+
+    logger.info('Info message')
+    assert 'Info message' in stream.getvalue()
+
+    handler.level = logging.DEBUG
+    logger.debug('Debug message 2')
+    assert 'Debug message 2' in stream.getvalue()
diff --git a/tests/unit_tests/test_sky/utils/test_context_utils.py b/tests/unit_tests/test_sky/utils/test_context_utils.py
new file mode 100644
index 000000000..027d317e1
--- /dev/null
+++ b/tests/unit_tests/test_sky/utils/test_context_utils.py
@@ -0,0 +1,22 @@
+"""Unit tests for sky.utils.context_utils module."""
+from typing import Optional, Union
+
+from sky.utils import context_utils
+
+
+@context_utils.cancellation_guard
+def original_function(arg1: int, arg2: str) -> Optional[Union[int, str]]:
+    return None
+
+
+def test_cancellation_guard_perserves_typecheck():
+    # Verify that the decorated function has the same signature
+    assert original_function.__name__ == 'original_function'
+    assert original_function.__annotations__ == {
+        'arg1': int,
+        'arg2': str,
+        'return': Optional[Union[int, str]]
+    }
+
+    # Verify that the decorated function can be called with the same signature
+    assert original_function(1, 'test') is None
diff --git a/tests/unit_tests/test_sky/utils/text_context.py b/tests/unit_tests/test_sky/utils/text_context.py
new file mode 100644
index 000000000..a8ba88058
--- /dev/null
+++ b/tests/unit_tests/test_sky/utils/text_context.py
@@ -0,0 +1,270 @@
+"""Unit tests for sky.utils.context module."""
+
+import asyncio
+import os
+import pathlib
+import tempfile
+
+import pytest
+
+from sky.utils import context
+
+
+@pytest.fixture
+def ctx():
+    """Fixture to provide a fresh context for each test."""
+    context.initialize()
+    yield context.get()
+    # Cleanup
+    if context.get()._log_file_handle is not None:
+        context.get()._log_file_handle.close()
+
+
+def test_context_initialization():
+    """Test context initialization."""
+    assert context.get() is None
+    context.initialize()
+    assert isinstance(context.get(), context.Context)
+
+
+def test_context_cancellation(ctx):
+    """Test context cancellation."""
+    assert not ctx.is_canceled()
+    ctx.cancel()
+    assert ctx.is_canceled()
+
+
+def test_log_redirection(ctx):
+    """Test log file redirection."""
+    with tempfile.NamedTemporaryFile(mode='w', delete=False) as f:
+        log_path = pathlib.Path(f.name)
+
+    try:
+        # Test setting log file
+        old_log = ctx.redirect_log(log_path)
+        assert old_log is None
+        assert ctx._log_file == log_path
+        assert ctx._log_file_handle is not None
+
+        # Test removing log file
+        old_log = ctx.redirect_log(None)
+        assert old_log == log_path
+        assert ctx._log_file is None
+        assert ctx._log_file_handle is None
+    finally:
+        # Cleanup
+        if log_path.exists():
+            log_path.unlink()
+
+
+def test_env_overrides(ctx):
+    """Test environment variable overrides."""
+    # Setup env overrides
+    ctx.override_envs({'TEST_VAR': 'overridden'})
+    assert 'TEST_VAR' in ctx.env_overrides
+    assert ctx.env_overrides['TEST_VAR'] == 'overridden'
+
+    # Test adding more overrides
+    ctx.override_envs({'TEST_VAR2': 'overridden2'})
+    assert 'TEST_VAR2' in ctx.env_overrides
+    assert ctx.env_overrides['TEST_VAR2'] == 'overridden2'
+    # Original override should still be present
+    assert ctx.env_overrides['TEST_VAR'] == 'overridden'
+
+
+def test_contextual_environ_getitem(ctx):
+    """Test ContextualEnviron __getitem__ functionality."""
+    # Create a ContextualEnviron with a simple dict as base environ
+    base_environ = {'BASE_VAR': 'base_value'}
+    env = context.ContextualEnviron(base_environ)
+
+    # Test getting a value from base environ
+    assert env['BASE_VAR'] == 'base_value'
+
+    # Test getting a value that doesn't exist should raise KeyError
+    with pytest.raises(KeyError):
+        _ = env['NON_EXISTENT']
+
+    # Test override with context
+    ctx.override_envs({'BASE_VAR': 'overridden_value'})
+    assert env['BASE_VAR'] == 'overridden_value'
+
+    # Test adding a new override
+    ctx.override_envs({'NEW_VAR': 'new_value'})
+    assert env['NEW_VAR'] == 'new_value'
+
+
+def test_contextual_environ_iter(ctx):
+    """Test ContextualEnviron __iter__ functionality."""
+    # Create a ContextualEnviron with a simple dict as base environ
+    base_environ = {'BASE_VAR1': 'value1', 'BASE_VAR2': 'value2'}
+    env = context.ContextualEnviron(base_environ)
+
+    # Without overrides, should iterate base environ
+    keys = set(env)
+    assert keys == {'BASE_VAR1', 'BASE_VAR2'}
+
+    # With overrides, should include both base and overridden vars
+    ctx.override_envs({'NEW_VAR': 'new_value', 'BASE_VAR1': 'override'})
+    keys = set(env)
+    assert keys == {'BASE_VAR1', 'BASE_VAR2', 'NEW_VAR'}
+
+
+def test_contextual_environ_set_del(ctx):
+    """Test ContextualEnviron set and delete operations."""
+    # Create a ContextualEnviron with a simple dict as base environ
+    base_environ = {'BASE_VAR': 'base_value'}
+    env = context.ContextualEnviron(base_environ)
+
+    # Set a new value in the base environ
+    env['NEW_VAR'] = 'direct_value'
+    assert base_environ['NEW_VAR'] == 'direct_value'
+
+    # Delete a value from base environ
+    del env['BASE_VAR']
+    assert 'BASE_VAR' not in base_environ
+
+    # Context overrides should take precedence over base
+    ctx.override_envs({'NEW_VAR': 'override_value'})
+    assert env['NEW_VAR'] == 'override_value'
+    assert base_environ['NEW_VAR'] == 'direct_value'
+
+
+def test_contextual_environ_methods(ctx):
+    """Test other ContextualEnviron methods."""
+    # Create a ContextualEnviron with a simple dict as base environ
+    base_environ = {'BASE_VAR': 'base_value'}
+    env = context.ContextualEnviron(base_environ)
+    ctx.override_envs({'BASE_VAR': 'overridden_value'})
+
+    copied = env.copy()
+    assert copied == {'BASE_VAR': 'overridden_value'}
+    assert copied is not base_environ
+
+    # Test setdefault
+    assert env.setdefault('BASE_VAR', 'default') == 'base_value'
+    assert env.setdefault('NEW_DEFAULT', 'default_value') == 'default_value'
+    assert base_environ['NEW_DEFAULT'] == 'default_value'
+
+
+def test_contextual_environ_subprocess(ctx):
+    """Test ContextualEnviron behavior with subprocesses."""
+    import subprocess
+    import sys
+
+    # Save original os.environ to restore it later
+    original_environ = os.environ
+
+    try:
+        # Create a ContextualEnviron and replace os.environ
+        base_environ = dict(os.environ)
+        contextual_env = context.ContextualEnviron(base_environ)
+        os.environ = contextual_env
+
+        # Set an environment variable override in the context
+        test_var_name = 'SKYPILOT_TEST_SUBPROCESS_VAR'
+        ctx.override_envs({test_var_name: 'initial_value'})
+
+        assert os.getenv(test_var_name) == 'initial_value'
+
+        print_envs = (f'import os; '
+                      f'import time; '
+                      f'print(os.getenv(\'{test_var_name}\', \'NOT_FOUND\')); '
+                      'time.sleep(1); '
+                      f'print(os.getenv(\'{test_var_name}\', \'NOT_FOUND\'))')
+
+        # Run subprocess and capture its output
+        proc = context.Popen([sys.executable, '-c', print_envs],
+                             stdout=subprocess.PIPE,
+                             stderr=subprocess.PIPE,
+                             text=True)
+        ctx.override_envs({test_var_name: 'updated_value'})
+        proc.wait()
+
+        output = proc.stdout.readlines()
+        assert len(output) == 2
+        # The subprocess should see the context-overridden env vars at spawn time
+        assert output[0].strip() == 'initial_value'
+        # The subprocess should not see the updated value after spawn
+        assert output[1].strip() == 'initial_value'
+
+        # Verify the environment variable is updated in current process
+        assert os.getenv(test_var_name) == 'updated_value'
+    finally:
+        # Restore original os.environ
+        os.environ = original_environ
+
+
+def test_output_stream(ctx):
+    """Test output stream selection."""
+    import sys
+
+    # Default should return fallback
+    assert ctx.output_stream(sys.stdout) == sys.stdout
+
+    # With log file should return log file handle
+    with tempfile.NamedTemporaryFile(mode='w', delete=False) as f:
+        log_path = pathlib.Path(f.name)
+    try:
+        ctx.redirect_log(log_path)
+        assert ctx.output_stream(sys.stdout) == ctx._log_file_handle
+    finally:
+        if log_path.exists():
+            log_path.unlink()
+
+
+def test_contextual_streams(ctx):
+    """Test contextual stdout/stderr."""
+    import sys
+    stdout = context.Stdout()
+    stderr = context.Stderr()
+
+    # Without log file should use original streams
+    assert stdout._active_stream() == sys.stdout
+    assert stderr._active_stream() == sys.stderr
+
+    # With log file should use log file handle
+    with tempfile.NamedTemporaryFile(mode='w', delete=False) as f:
+        log_path = pathlib.Path(f.name)
+    try:
+        ctx.redirect_log(log_path)
+        assert stdout._active_stream() == ctx._log_file_handle
+        assert stderr._active_stream() == ctx._log_file_handle
+    finally:
+        if log_path.exists():
+            log_path.unlink()
+
+
+@pytest.mark.asyncio
+async def test_async_context_isolation():
+    """Test context isolation between coroutines."""
+    context.initialize()
+    ctx1 = context.get()
+
+    async def other_coro():
+        context.initialize()
+        return context.get()
+
+    ctx2 = await other_coro()
+    assert ctx1 is not ctx2
+
+
+@pytest.mark.asyncio
+async def test_async_cancellation():
+    """Test context cancellation in async environment."""
+    context.initialize()
+    ctx = context.get()
+
+    cancel_called = False
+
+    async def check_cancel():
+        nonlocal cancel_called
+        while not ctx.is_canceled():
+            await asyncio.sleep(0.1)
+        cancel_called = True
+
+    task = asyncio.create_task(check_cancel())
+    await asyncio.sleep(0.2)
+    ctx.cancel()
+    await asyncio.wait_for(task, timeout=1.0)
+    assert cancel_called