Repository: skypilot-org/skypilot
Pull request #5559 (2025-05-12 22:42:31)

[docker] update docker login creds when relaunching cluster
<!-- Describe the changes in this PR -->
If you `sky launch` with new docker credentials on an existing cluster, the cluster resources will now be updated with the new credentials.

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [ ] Code formatting: install pre-commit (auto-check on commit) or `bash format.sh`
- [ ] Any manual or new tests for this PR (please specify below)
  - [ ] Launch, then change credentials, cluster YAML updated
  - [ ] Cluster fails to start, retry with fixed creds
  - [ ] launch jobs on live cluster with credentials but no docker creds in the launched task
  - [ ] stopped cluster relaunch without creds
- [ ] All smoke tests: `/smoke-test` (CI) or `pytest tests/test_smoke.py` (local)
- [ ] Relevant individual tests: `/smoke-test -k docker` (CI)
- [ ] Backward compatibility: `/quicktest-core` (CI) or `pytest tests/smoke_tests/test_backward_compat.py` (local)

<!-- CI commands (/-prefixed) can only be triggered by repo members -->



Patch:
diff --git a/charts/skypilot/values.yaml b/charts/skypilot/values.yaml
index 0a8f6a2d2..bbafb68d4 100644
--- a/charts/skypilot/values.yaml
+++ b/charts/skypilot/values.yaml
@@ -1,5 +1,5 @@
 apiService:
-  image: berkeleyskypilot/skypilot-nightly:latest
+  image: berkeleyskypilot/skypilot:0.9.2
   # Number of replicas to deploy - replicas > 1 is not well tested, and requires
   # a PVC that supports ReadWriteMany (see accessMode in storage section below).
   replicas: 1
diff --git a/docs/source/cloud-setup/cloud-permissions/gcp.rst b/docs/source/cloud-setup/cloud-permissions/gcp.rst
index ce366117f..40e42ee6d 100644
--- a/docs/source/cloud-setup/cloud-permissions/gcp.rst
+++ b/docs/source/cloud-setup/cloud-permissions/gcp.rst
@@ -339,6 +339,163 @@ To do so, you can use SkyPilot's global config file ``~/.sky/config.yaml`` to sp
 
 The ``gcp.ssh_proxy_command`` field is optional. If SkyPilot is run on a machine that can directly access the internal IPs of the instances, it can be omitted. Otherwise, it should be set to a command that can be used to proxy SSH connections to the internal IPs of the instances.
 
+Here are two ways to enable SkyPilot's access those private hosts:
+
+* **Option 1** – :ref:`SSH proxy / bastion host <gcp-proxy-host>`
+* **Option 2** – :ref:`IAP tunneling (Identity‑Aware Proxy) <gcp-iap-tunnel>`
+  This is often simpler because you do **not** need to run or secure your own
+  bastion; Google handles the tunnel at the edge.
+
+Choose whichever fits your environment.
+
+.. _gcp-proxy-host:
+
+Option 1: Simple single‑proxy example
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+If you have a "jump host" that's directly accessible from the machine running SkyPilot, tell SkyPilot how to pipe SSH through it:
+
+.. code-block:: yaml
+
+    gcp:
+      use_internal_ips: true
+      vpc_name: my-vpc-name
+      ssh_proxy_command: ssh -W %h:%p -o StrictHostKeyChecking=no myself@my.proxy
+
+When **multiple regions** are involved you can map each region to its own proxy
+(just as before):
+
+.. code-block:: yaml
+
+    gcp:
+      use_internal_ips: true
+      vpc_name: my-vpc-name
+      ssh_proxy_command:
+        us-west1: ssh -W %h:%p -o StrictHostKeyChecking=no myself@west.proxy
+        us-east1: ssh -W %h:%p -o StrictHostKeyChecking=no myself@east.proxy
+
+
+.. _gcp-iap-tunnel:
+
+Option 2: IAP tunneling
+~~~~~~~~~~~~~~~~~~~~~~~
+`IAP TCP forwarding <https://cloud.google.com/iap/docs/using-tcp-forwarding>`__
+creates an authenticated tunnel from your workstation to any private IP inside
+a GCP VPC—no public IP, VPN, or self-managed bastion required. SkyPilot
+implements the `IP-based (destination-group) mode <https://cloud.google.com/iap/docs/tcp-by-host>`__; the per-instance mode is not supported yet.
+
+Setup
+^^^^^
+* Enable the API once per project:
+
+  .. code-block:: console
+
+     $ gcloud services enable iap.googleapis.com
+
+* This method requires **IAP tunnel permissions**.
+  Adding the built-in role ``roles/iap.tunnelResourceAccessor`` on the destination
+  group (or project-wide) is the simplest approach, in addition to the :ref:`other permissions required by SkyPilot <cloud-permissions-gcp>`.
+* A **Cloud NAT** gateway must exist in every region where you launch
+  internal-only VMs so they can reach the internet for package installs.
+  See :ref:`gcp-cloud-nat`.
+
+The following script can be used to set up IAP tunneling in a specific VPC region:
+
+.. dropdown:: Automated setup script (per region)
+
+    .. code-block:: bash
+
+        #!/usr/bin/env bash
+        # Set your region and VPC.
+        REGION=us-east1        # e.g. us-west1, europe-west4, ...
+        VPC=default            # change if using a custom VPC
+
+        #######################################################################
+        # 1. Cloud NAT (if none exists)                                       #
+        #######################################################################
+        nat_exists=false
+        for r in $(gcloud compute routers list --regions="$REGION" \
+                     --format='value(name)'); do
+          if gcloud compute routers nats list --router="$r" --region="$REGION" \
+                 --format='value(name)' | grep -q .; then
+            nat_exists=true
+          fi
+        done
+        if [ "$nat_exists" = false ]; then
+          gcloud compute routers create nat-router --network="$VPC" --region="$REGION"
+          gcloud compute routers nats create nat-config \
+            --router=nat-router --router-region="$REGION" \
+            --nat-all-subnet-ip-ranges --auto-allocate-nat-external-ips
+        fi
+
+        #######################################################################
+        # 2. IAP destination group covering the subnet CIDR                   #
+        #######################################################################
+        CIDR=$(gcloud compute networks subnets list --network="$VPC" \
+                 --regions="$REGION" --format='value(ipCidrRange)')
+        gcloud iap tcp dest-groups create "internal-vpc-$VPC" \
+            --region="$REGION" --ip-range-list="$CIDR"
+
+        #######################################################################
+        # 3. Show the SkyPilot config snippet                                 #
+        #######################################################################
+        echo "Update your config.yaml to add the new gcp.ssh_proxy_command.$REGION value:"
+        echo
+        cat <<EOF
+        gcp:
+          use_internal_ips: true
+          ssh_proxy_command:
+            $REGION: gcloud compute start-iap-tunnel %h %p --listen-on-stdin --region=$REGION --network=$VPC --dest-group=internal-vpc-$VPC
+        EOF
+
+        # Automatically do it
+        read -p 'Automatically update ~/.sky/config.yaml? [y/N] ' choice
+        if echo "$choice" | grep -xqE '[yY]'; then
+          yq -Yi ".gcp.use_internal_ips=true | .gcp.ssh_proxy_command[\"$REGION\"]=\"gcloud compute start-iap-tunnel %h %p --listen-on-stdin --region=$REGION --network=$VPC --dest-group=internal-vpc-$VPC\"" ~/.sky/config.yaml
+        else
+          echo 'Aborting.'
+        fi
+
+    Paste the printed snippet into your ``~/.sky/config.yaml`` (or let the script
+    modify it automatically).
+
+Once the destination group is created, make sure to set the ``ssh_proxy_command`` :ref:`in the SkyPilot config <config-yaml-gcp-ssh-proxy-command>`. For instance, if you have
+
+- destination group: ``internal-vpc-default``
+- region: ``us-east1``
+- VPC: ``default``
+
+you could use this configuration:
+
+.. code-block:: yaml
+
+    gcp:
+      use_internal_ips: true
+      ssh_proxy_command:
+        us-east1: gcloud compute start-iap-tunnel %h %p --listen-on-stdin --region=us-east1 --network=default --dest-group=internal-vpc-default
+
+
+How it works
+^^^^^^^^^^^^
+1. ``sky launch`` creates VMs **without external IPs**.
+2. When connecting to the instance, ``ssh`` uses the **IAP proxy command** shown above.
+   ``gcloud compute start-iap-tunnel`` transparently
+   forwards the connection through Google’s edge to the VM’s private address.
+3. SSH packets travel over the authenticated IAP
+   tunnel—no public ingress is ever exposed.
+
+Troubleshooting
+^^^^^^^^^^^^^^^
+* *Permission denied / 403* when opening the tunnel → the caller is missing
+  ``roles/iap.tunnelResourceAccessor`` (or an equivalent custom role) on the
+  destination group.
+* VM cannot reach the Internet → ensure the Cloud NAT gateway is in **the same
+  region** as the subnet and that ``--nat-all-subnet-ip-ranges`` is enabled.
+* Multiple subnets or VPCs → create one destination group per
+  *region × CIDR* and map each region to the correct group in
+  ``ssh_proxy_command``.
+
+
+.. _gcp-cloud-nat:
 
 Cloud NAT setup
 ~~~~~~~~~~~~~~~~
diff --git a/docs/source/getting-started/installation.rst b/docs/source/getting-started/installation.rst
index 4fd3d4ce8..c14c1718a 100644
--- a/docs/source/getting-started/installation.rst
+++ b/docs/source/getting-started/installation.rst
@@ -229,6 +229,9 @@ SkyPilot can run workloads on on-prem or cloud-hosted Kubernetes clusters
 
 See :ref:`SkyPilot on Kubernetes <kubernetes-overview>` for more.
 
+.. tip::
+   If you do not have access to a Kubernetes cluster, you can :ref:`deploy a local Kubernetes cluster on your laptop <kubernetes-setup-kind>` with ``sky local up``.
+
 .. _aws-installation:
 
 AWS
diff --git a/docs/source/reference/api-server/api-server-admin-deploy.rst b/docs/source/reference/api-server/api-server-admin-deploy.rst
index 5615e8001..32276b85b 100644
--- a/docs/source/reference/api-server/api-server-admin-deploy.rst
+++ b/docs/source/reference/api-server/api-server-admin-deploy.rst
@@ -22,29 +22,9 @@ Prerequisites
 
     You can also deploy the API server on cloud VMs using an existing SkyPilot installation. See :ref:`sky-api-server-cloud-deploy`.
 
-Step 1: Create a namespace and add Helm repository
---------------------------------------------------
-
-The API server will be deployed in a namespace of your choice. You can either create the namespace manually:
-
-.. code-block:: bash
-
-    # The NAMESPACE variable will be used throughout the guide
-    NAMESPACE=skypilot
-    kubectl create namespace $NAMESPACE
-
-Or, let Helm create it automatically by adding the ``--create-namespace`` flag to the helm install command in Step 2.
-
-Next, add the SkyPilot Helm repository:
-
-.. code-block:: bash
-
-    helm repo add skypilot https://helm.skypilot.co
-    helm repo update
-
 .. _sky-api-server-helm-deploy-command:
 
-Step 2: Deploy the API server Helm chart
+Step 1: Deploy the API server Helm chart
 ----------------------------------------
 
 Install the SkyPilot Helm chart with the following command:
@@ -55,37 +35,40 @@ Install the SkyPilot Helm chart with the following command:
 
 .. code-block:: bash
 
+    # Ensure the helm repository is added and up to date
+    helm repo add skypilot https://helm.skypilot.co
+    helm repo update
+
     # The following variables will be used throughout the guide
+    # NAMESPACE is the namespace to deploy the API server in
+    NAMESPACE=skypilot
+    # RELEASE_NAME is the name of the helm release, must be unique within the namespace
     RELEASE_NAME=skypilot
+    # Replace with your username and password to configure the basic auth credentials for the API server
     WEB_USERNAME=skypilot
-    # Replace with your password to configure the password for the API server
     WEB_PASSWORD=yourpassword
     AUTH_STRING=$(htpasswd -nb $WEB_USERNAME $WEB_PASSWORD)
+    # Deploy the API server
     helm upgrade --install $RELEASE_NAME skypilot/skypilot-nightly --devel \
       --namespace $NAMESPACE \
       --create-namespace \
       --set ingress.authCredentials=$AUTH_STRING
 
-The ``--namespace`` flag specifies which namespace to deploy the API server in, and ``--create-namespace`` will create the namespace if it doesn't exist.
-
-To install a specific version, pass the ``--version`` flag to the ``helm upgrade`` command (e.g., ``--version 0.1.0``).
-
-.. tip::
-
-    The API server deployed will be configured to use the hosting Kubernetes cluster to launch tasks by default. Refer to :ref:`sky-api-server-configure-credentials` to configure credentials for more clouds and Kubernetes clusters.
+.. dropdown:: Flags explanation
 
-.. dropdown:: Use existing basic auth credentials
+    Here is a breakdown of the flags used in the command above:
 
-    If you already have a Kubernetes secret containing basic auth credentials, you can use it directly by setting ``ingress.authSecret`` instead of ``ingress.authCredentials``:
+    * ``upgrade --install``: Upgrade the API server if it already exists, or install it if it doesn't exist.
+    * ``--devel``: Use the latest development version of the SkyPilot helm chart. To use a specific version, pass the ``--version`` flag to the ``helm upgrade`` command (e.g., ``--version 0.1.0``).
+    * ``--namespace $NAMESPACE``: Specify the namespace to deploy the API server in.
+    * ``--create-namespace``: Create the namespace if it doesn't exist.
+    * :ref:`--set ingress.authCredentials=$AUTH_STRING <helm-values-ingress-authCredentials>`: Set the basic auth credentials for the API server.
 
-    .. code-block:: bash
+    For more details on the available configuration options, refer to :ref:`SkyPilot API Server Helm Chart Values <helm-values-spec>`.
 
-        helm upgrade --install skypilot skypilot/skypilot-nightly --devel \
-          --namespace $NAMESPACE \
-          --create-namespace \
-          --set ingress.authSecret=my-existing-auth-secret
+.. tip::
 
-    The secret must be in the same namespace as the API server and must contain a key named ``auth`` with the basic auth credentials in htpasswd format.
+    The API server deployed will be configured to use the hosting Kubernetes cluster to launch tasks by default. Refer to :ref:`sky-api-server-configure-credentials` to configure credentials for more clouds and Kubernetes clusters.
 
 After the API server is deployed, you can inspect the API server pod status with:
 
@@ -97,7 +80,7 @@ You should see the pod is initializing and finally becomes running and ready. If
 
 .. _sky-get-api-server-url:
 
-Step 3: Get the API server URL
+Step 2: Get the API server URL
 ------------------------------
 
 Once the API server is deployed, we can fetch the API server URL. We use nginx ingress to expose the API server.
@@ -160,7 +143,7 @@ Our default of using a NodePort service is the recommended way to expose the API
             To avoid frequent IP address changes on nodes by your cloud provider, you can attach a static IP address to your nodes (`instructions for GKE <https://cloud.google.com/compute/docs/ip-addresses/configure-static-external-ip-address>`_) and use it as the ``NODE_IP`` in the command above.
 
 
-Step 4: Test the API server
+Step 3: Test the API server
 ---------------------------
 
 Test the API server by curling the health endpoint:
@@ -478,7 +461,17 @@ To set the config file, pass ``--set-file apiService.config=path/to/your/config.
       --reuse-values \
       --set-file apiService.config=config.yaml
 
-You can also directly set config values in the ``values.yaml`` file.
+You can also directly set config values in the ``values.yaml`` file, e.g.:
+
+.. code-block:: yaml
+
+    apiService:
+      config: |
+        allowed_clouds:
+        - aws
+        - kubernetes
+
+To apply a new config, rerun ``helm upgrade`` with the updated ``values.yaml`` file.
 
 Setting an admin policy
 ^^^^^^^^^^^^^^^^^^^^^^^
diff --git a/docs/source/reference/api-server/api-server-tunning.rst b/docs/source/reference/api-server/api-server-tunning.rst
index b5f97f38e..742f00ad6 100644
--- a/docs/source/reference/api-server/api-server-tunning.rst
+++ b/docs/source/reference/api-server/api-server-tunning.rst
@@ -5,6 +5,8 @@ SkyPilot API Server Performance Best Practices
 
 This page describes performance best practices for remote SkyPilot API server in :ref:`team deployment <sky-api-server>`.
 
+.. _sky-api-server-resources-tuning:
+
 Tuning API server resources
 ---------------------------
 
diff --git a/docs/source/reference/api-server/api-server.rst b/docs/source/reference/api-server/api-server.rst
index f6cd5d21e..42fe4b43b 100644
--- a/docs/source/reference/api-server/api-server.rst
+++ b/docs/source/reference/api-server/api-server.rst
@@ -116,6 +116,7 @@ To verify that the API server is working, run ``sky api info``:
    Upgrading API Server <api-server-upgrade>
    Performance Best Practices <api-server-tunning>
    Troubleshooting <api-server-troubleshooting>
+   Helm Chart Reference <helm-values-spec>
 
 By default, each user connected to the API server will only see their own resources.
 
diff --git a/docs/source/reference/api-server/helm-values-spec.rst b/docs/source/reference/api-server/helm-values-spec.rst
new file mode 100644
index 000000000..79be55e93
--- /dev/null
+++ b/docs/source/reference/api-server/helm-values-spec.rst
@@ -0,0 +1,828 @@
+.. _helm-values-spec:
+
+SkyPilot API Server Helm Chart Values
+=====================================
+
+The SkyPilot API server helm chart provides typical `helm values <https://helm.sh/docs/chart_template_guide/values_files/>`_ as configuration entries. Configuration values can be passed in two ways when installing the chart:
+
+* ``--values`` (or ``-f``): Specify a YAML file with overrides.
+
+  .. code-block:: bash
+
+      cat <<EOF > values.yaml
+      apiService:
+        image: berkeleyskypilot/skypilot:0.9.1
+      EOF
+
+      helm install $RELEASE_NAME skypilot/skypilot-nightly --devel --values values.yaml
+
+* ``--set``: Specify overrides on the command line.
+
+  .. code-block:: bash
+      
+      helm install $RELEASE_NAME skypilot/skypilot-nightly --set apiService.image="berkeleyskypilot/skypilot:0.9.1"
+
+Values
+------
+
+Below is the available helm value keys and the default value of each key:
+
+..
+  Omitted values:
+  * storage.accessMode: accessMode other than ReadWriteOnce is not tested yet.
+
+.. parsed-literal::
+
+  :ref:`apiService <helm-values-apiService>`:
+    :ref:`image <helm-values-apiService-image>`: berkeleyskypilot/skypilot:0.9.1
+    :ref:`preDeployHook <helm-values-apiService-preDeployHook>`: \|-
+      # Run commands before deploying the API server, e.g. installing an admin
+      # policy. Remember to set the admin policy in the config section below.
+
+      echo "Pre-deploy hook"
+
+      # Uncomment the following lines to install the admin policy
+
+      # echo "Installing admin policy"
+      # pip install git+https://github.com/michaelvll/admin-policy-examples
+    :ref:`config <helm-values-apiService-config>`: null
+    :ref:`skipResourceCheck <helm-values-apiService-skipResourceCheck>`: false
+    :ref:`resources <helm-values-apiService-resources>`:
+      requests:
+        cpu: "4"
+        memory: "8Gi"
+      limits:
+        cpu: "4"
+        memory: "8Gi"
+    :ref:`skypilotDev <helm-values-apiService-skypilotDev>`: false
+
+  :ref:`storage <helm-values-storage>`:
+    :ref:`enabled <helm-values-storage-enabled>`: true
+    :ref:`storageClassName <helm-values-storage-storageClassName>`: ""
+    :ref:`size <helm-values-storage-size>`: 10Gi
+    :ref:`selector <helm-values-storage-selector>`: {}
+    :ref:`volumeName <helm-values-storage-volumeName>`: ""
+    :ref:`annotations <helm-values-storage-annotations>`: {}
+
+  :ref:`ingress <helm-values-ingress>`:
+    :ref:`enabled <helm-values-ingress-enabled>`: true
+    :ref:`authSecret <helm-values-ingress-authSecret>`: null
+    :ref:`authCredentials <helm-values-ingress-authCredentials>`: "username:$apr1$encrypted_password"
+    :ref:`path <helm-values-ingress-path>`: '/'
+
+  :ref:`ingress-nginx <helm-values-ingress-nginx>`:
+    :ref:`enabled <helm-values-ingress-nginx-enabled>`: true
+    :ref:`controller <helm-values-ingress-nginx-controller>`:
+      service:
+        type: LoadBalancer
+        annotations:
+          service.beta.kubernetes.io/aws-load-balancer-type: "nlb"
+          cloud.google.com/l4-rbs: "enabled"
+          service.beta.kubernetes.io/port_443_health-probe_protocol: "TCP"
+          service.beta.kubernetes.io/port_80_health-probe_protocol: "TCP"
+      config:
+        http-snippet: |
+          map $http_upgrade $connection_upgrade {
+              default upgrade;
+              ''      close;
+          }
+
+  :ref:`rbac <helm-values-rbac>`:
+    :ref:`create <helm-values-rbac-create>`: true
+    :ref:`serviceAccountName <helm-values-rbac-serviceAccountName>`: ""
+    :ref:`namespaceRules <helm-values-rbac-namespaceRules>`:
+      - apiGroups: [ "" ]
+        resources: [ "pods", "pods/status", "pods/exec", "pods/portforward" ]
+        verbs: [ "*" ]
+      - apiGroups: [ "" ]
+        resources: [ "services" ]
+        verbs: [ "*" ]
+      - apiGroups: [ "" ]
+        resources: [ "secrets" ]
+        verbs: [ "*" ]
+      - apiGroups: [ "" ]
+        resources: [ "events" ]
+        verbs: [ "get", "list", "watch" ]
+    :ref:`clusterRules <helm-values-rbac-clusterRules>`:
+      - apiGroups: [ "" ]
+        resources: [ "nodes" ]
+        verbs: [ "get", "list", "watch" ]
+      - apiGroups: [ "" ]
+        resources: [ "pods" ]
+        verbs: [ "get", "list", "watch" ]
+      - apiGroups: [ "node.k8s.io" ]
+        resources: [ "runtimeclasses" ]
+        verbs: [ "get", "list", "watch" ]
+      - apiGroups: [ "networking.k8s.io" ]
+        resources: [ "ingressclasses" ]
+        verbs: [ "get", "list", "watch" ]
+      - apiGroups: [""]
+        resources: ["services"]
+        verbs: ["list", "get"]
+    :ref:`manageRbacPolicies <helm-values-rbac-manageRbacPolicies>`: true
+    :ref:`manageSystemComponents <helm-values-rbac-manageSystemComponents>`: true
+
+  :ref:`kubernetesCredentials <helm-values-kubernetesCredentials>`:
+    :ref:`useApiServerCluster <helm-values-kubernetesCredentials-useApiServerCluster>`: true
+    :ref:`useKubeconfig <helm-values-kubernetesCredentials-useKubeconfig>`: false
+    :ref:`kubeconfigSecretName <helm-values-kubernetesCredentials-kubeconfigSecretName>`: kube-credentials
+    :ref:`inclusterNamespace <helm-values-kubernetesCredentials-inclusterNamespace>`: null
+
+  :ref:`awsCredentials <helm-values-awsCredentials>`:
+    :ref:`enabled <helm-values-awsCredentials-enabled>`: false
+    :ref:`awsSecretName <helm-values-awsCredentials-awsSecretName>`: aws-credentials
+    :ref:`accessKeyIdKeyName <helm-values-awsCredentials-accessKeyIdKeyName>`: aws_access_key_id
+    :ref:`secretAccessKeyKeyName <helm-values-awsCredentials-secretAccessKeyKeyName>`: aws_secret_access_key
+
+  :ref:`gcpCredentials <helm-values-gcpCredentials>`:
+    :ref:`enabled <helm-values-gcpCredentials-enabled>`: false
+    :ref:`projectId <helm-values-gcpCredentials-projectId>`: null
+    :ref:`gcpSecretName <helm-values-gcpCredentials-gcpSecretName>`: gcp-credentials
+
+  :ref:`podSecurityContext <helm-values-podSecurityContext>`: {}
+
+  :ref:`securityContext <helm-values-securityContext>`:
+    :ref:`capabilities <helm-values-securityContext-capabilities>`:
+      drop:
+      - ALL
+    :ref:`allowPrivilegeEscalation <helm-values-securityContext-allowPrivilegeEscalation>`: false
+
+  :ref:`runtimeClassName <helm-values-runtimeClassName>`: ""
+
+Fields
+----------
+
+.. _helm-values-apiService:
+
+``apiService``
+~~~~~~~~~~~~~~
+
+Configuration for the SkyPilot API server deployment.
+
+.. _helm-values-apiService-image:
+
+``apiService.image``
+^^^^^^^^^^^^^^^^^^^^
+
+Docker image to use for the API server.
+
+Default: ``"berkeleyskypilot/skypilot:0.9.1"``
+
+.. code-block:: yaml
+
+  apiService:
+    image: berkeleyskypilot/skypilot:0.9.1
+
+To use a nightly build, find the desired nightly version on `pypi <https://pypi.org/project/skypilot-nightly/#history>`_ and update the ``image`` value:
+
+.. code-block:: yaml
+
+  apiService:
+    # Replace 1.0.0.devYYYYMMDD with the desired nightly version
+    image: berkeleyskypilot/skypilot-nightly:1.0.0.devYYYYMMDD
+
+.. _helm-values-apiService-preDeployHook:
+
+``apiService.preDeployHook``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Commands to run before deploying the API server (e.g., install :ref:`admin policy <advanced-policy-config>`).
+
+Default: see the yaml below.
+
+.. code-block:: yaml
+
+  apiService:
+    preDeployHook: |-
+      # Run commands before deploying the API server, e.g. installing an admin
+      # policy. Remember to set the admin policy in the config section below.
+      echo "Pre-deploy hook"
+
+      # Uncomment the following lines to install the admin policy
+      # echo "Installing admin policy"
+      # pip install git+https://github.com/michaelvll/admin-policy-examples
+
+.. _helm-values-apiService-config:
+
+``apiService.config``
+^^^^^^^^^^^^^^^^^^^^^
+
+Content of the `SkyPilot config.yaml <https://docs.skypilot.co/en/latest/reference/config.html>`_ to set on the API server. Set to ``null`` to use an empty config. Refer to :ref:`setting the SkyPilot config <sky-api-server-config>` for more details.
+
+Default: ``null``
+
+.. code-block:: yaml
+
+  apiService:
+    config: |-
+      allowed_clouds:
+        - aws
+        - gcp
+
+.. _helm-values-apiService-skipResourceCheck:
+
+``apiService.skipResourceCheck``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Skip resource check for the API server (not recommended for production), refer to :ref:`tuning API server resources <sky-api-server-resources-tuning>` for more details.
+
+Default: ``false``
+
+.. code-block:: yaml
+
+  apiService:
+    skipResourceCheck: false
+
+.. _helm-values-apiService-resources:
+
+``apiService.resources``
+^^^^^^^^^^^^^^^^^^^^^^^^
+
+Resource requests and limits for the API server container. Refer to :ref:`tuning API server resources <sky-api-server-resources-tuning>` for how to tune the resources.
+
+Default: see the yaml below.
+
+.. code-block:: yaml
+
+  apiService:
+    resources:
+      requests:
+        cpu: "4"
+        memory: "8Gi"
+      limits:
+        cpu: "4"
+        memory: "8Gi"
+
+.. _helm-values-apiService-skypilotDev:
+
+``apiService.skypilotDev``
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Enable developer mode for SkyPilot.
+
+Default: ``false``
+
+.. code-block:: yaml
+
+  apiService:
+    skypilotDev: false
+
+.. _helm-values-storage:
+
+``storage``
+~~~~~~~~~~~
+
+.. _helm-values-storage-enabled:
+
+``storage.enabled``
+^^^^^^^^^^^^^^^^^^^
+
+Enable persistent storage for the API server, setting this to ``false`` is prone to data loss and should only be used for testing.
+
+Default: ``true``
+
+.. code-block:: yaml
+
+  storage:
+    enabled: true
+
+.. _helm-values-storage-storageClassName:
+
+``storage.storageClassName``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Storage class to use for the API server, leave empty to use the default storage class of the hosting Kubernetes cluster.
+
+Default: ``""``
+
+.. code-block:: yaml
+
+  storage:
+    storageClassName: gp2
+
+.. _helm-values-storage-size:
+
+``storage.size``
+^^^^^^^^^^^^^^^^
+
+Size of the persistent storage volume for the API server.
+
+Default: ``10Gi``
+
+.. code-block:: yaml
+
+  storage:
+    size: 10Gi
+
+.. _helm-values-storage-selector:
+
+``storage.selector``
+^^^^^^^^^^^^^^^^^^^^
+
+Selector for matching specific PersistentVolumes. Usually left empty.
+
+Default: ``{}``
+
+.. code-block:: yaml
+
+  storage:
+    selector: {}
+
+.. _helm-values-storage-volumeName:
+
+``storage.volumeName``
+^^^^^^^^^^^^^^^^^^^^^^
+
+Name of the PersistentVolume to bind to. Usually left empty to let Kubernetes select and bind the volume automatically.
+
+Default: ``""``
+
+.. code-block:: yaml
+
+  storage:
+    volumeName: ""
+
+.. _helm-values-storage-annotations:
+
+``storage.annotations``
+^^^^^^^^^^^^^^^^^^^^^^^
+
+Annotations to add to the PersistentVolumeClaim.
+
+Default: ``{}``
+
+.. code-block:: yaml
+
+  storage:
+    annotations: {}
+
+.. _helm-values-ingress:
+
+``ingress``
+~~~~~~~~~~~
+
+.. _helm-values-ingress-enabled:
+
+``ingress.enabled``
+^^^^^^^^^^^^^^^^^^^
+
+Enable ingress for the API server. Set to ``true`` to expose the API server via an ingress controller.
+
+Default: ``true``
+
+.. code-block:: yaml
+
+  ingress:
+    enabled: true
+
+.. _helm-values-ingress-authSecret:
+
+``ingress.authSecret``
+^^^^^^^^^^^^^^^^^^^^^^
+
+Name of the Kubernetes secret containing basic auth credentials for ingress. If not specified, a new secret will be created using ``authCredentials``.
+
+One of ``ingress.authSecret`` or ``ingress.authCredentials`` must be set.
+
+Default: ``null``
+
+.. code-block:: yaml
+
+  ingress:
+    authSecret: null
+
+.. _helm-values-ingress-authCredentials:
+
+``ingress.authCredentials``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Basic auth credentials in the format ``username:encrypted_password``. Used only if ``authSecret`` is not set.
+
+One of ``ingress.authSecret`` or ``ingress.authCredentials`` must be set.
+
+Default: ``"username:$apr1$encrypted_password"``
+
+.. code-block:: yaml
+
+  ingress:
+    authCredentials: "username:$apr1$encrypted_password"
+
+.. _helm-values-ingress-path:
+
+``ingress.path``
+^^^^^^^^^^^^^^^^
+
+The base path of the API server. You may use different paths to expose multiple API servers through a unified ingress controller.
+
+Default: ``'/'``
+
+.. code-block:: yaml
+
+  ingress:
+    path: '/'
+
+.. _helm-values-ingress-nginx:
+
+``ingress-nginx``
+~~~~~~~~~~~~~~~~~
+
+.. _helm-values-ingress-nginx-enabled:
+
+``ingress-nginx.enabled``
+^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Enable the ingress-nginx controller for the API server. If you have an existing ingress-nginx controller, you have to set this to ``false`` to avoid conflict.
+
+Default: ``true``
+
+.. code-block:: yaml
+
+  ingress-nginx:
+    enabled: true
+
+.. _helm-values-ingress-nginx-controller:
+
+``ingress-nginx.controller``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Fields under ``ingress-nginx.controller`` will be mapped to ``controller`` values for the ingress-nginx controller sub-chart. Refer to the `ingress-nginx chart documentation <https://artifacthub.io/packages/helm/ingress-nginx/ingress-nginx#values>`_ for more details.
+
+Default: see the yaml below.
+
+.. code-block:: yaml
+
+  ingress-nginx:
+    controller:
+      service:
+        # Service type of the ingress controller.
+        type: LoadBalancer
+        # Annotations for the ingress controller service.
+        annotations:
+          service.beta.kubernetes.io/aws-load-balancer-type: "nlb"
+          cloud.google.com/l4-rbs: "enabled"
+          service.beta.kubernetes.io/port_443_health-probe_protocol: "TCP"
+          service.beta.kubernetes.io/port_80_health-probe_protocol: "TCP"
+      config:
+        # Custom HTTP snippet to inject into the ingress-nginx configuration.
+        http-snippet: |
+          map $http_upgrade $connection_upgrade {
+              default upgrade;
+              ''      close;
+          }
+
+.. _helm-values-rbac:
+
+``rbac``
+~~~~~~~~
+
+.. _helm-values-rbac-create:
+
+``rbac.create``
+^^^^^^^^^^^^^^^
+
+Whether to create the service account and RBAC policies for the API server. If false, an external service account is expected.
+
+Default: ``true``
+
+.. code-block:: yaml
+
+  rbac:
+    create: true
+
+.. _helm-values-rbac-serviceAccountName:
+
+``rbac.serviceAccountName``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Name of the service account to use. Leave empty to let the chart generate one.
+
+Default: ``""``
+
+.. code-block:: yaml
+
+  rbac:
+    serviceAccountName: ""
+
+.. _helm-values-rbac-namespaceRules:
+
+``rbac.namespaceRules``
+^^^^^^^^^^^^^^^^^^^^^^^
+
+Namespace-scoped RBAC rules granted to the namespace where the SkyPilot tasks will be launched.
+
+.. note::
+
+  Modifying the rules may break functionalities of SkyPilot API server. Refer to :ref:`setting minimum permissions in helm deployment <minimum-permissions-in-helm>` for how to modify the rules based on your use case.
+
+Default: see the yaml below.
+
+.. code-block:: yaml
+
+  rbac:
+    namespaceRules:
+      - apiGroups: [ "" ]
+        resources: [ "pods", "pods/status", "pods/exec", "pods/portforward" ]
+        verbs: [ "*" ]
+      - apiGroups: [ "" ]
+        resources: [ "services" ]
+        verbs: [ "*" ]
+      - apiGroups: [ "" ]
+        resources: [ "secrets" ]
+        verbs: [ "*" ]
+      - apiGroups: [ "" ]
+        resources: [ "events" ]
+        verbs: [ "get", "list", "watch" ]
+
+.. _helm-values-rbac-clusterRules:
+
+``rbac.clusterRules``
+^^^^^^^^^^^^^^^^^^^^^^
+
+Cluster-scoped RBAC rules for the API server.
+
+.. note::
+
+  Modifying the rules may break functionalities of SkyPilot API server. Refer to :ref:`setting minimum permissions in helm deployment <minimum-permissions-in-helm>` for how to modify the rules based on your use case.
+
+Default: see the yaml below.
+
+.. code-block:: yaml
+
+  rbac:
+    clusterRules:
+      - apiGroups: [ "" ]
+        resources: [ "nodes" ]
+        verbs: [ "get", "list", "watch" ]
+      - apiGroups: [ "" ]
+        resources: [ "pods" ]
+        verbs: [ "get", "list", "watch" ]
+      - apiGroups: [ "node.k8s.io" ]
+        resources: [ "runtimeclasses" ]
+        verbs: [ "get", "list", "watch" ]
+      - apiGroups: [ "networking.k8s.io" ]
+        resources: [ "ingressclasses" ]
+        verbs: [ "get", "list", "watch" ]
+      - apiGroups: ["" ]
+        resources: ["services"]
+        verbs: ["list", "get"]
+
+.. _helm-values-rbac-manageRbacPolicies:
+
+``rbac.manageRbacPolicies``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Allow the API server to grant permissions to SkyPilot Pods and system components. Refer to :ref:`setting minimum permissions in helm deployment <minimum-permissions-in-helm>` for more details.
+
+Default: ``true``
+
+.. code-block:: yaml
+
+  rbac:
+    manageRbacPolicies: true
+
+.. _helm-values-rbac-manageSystemComponents:
+
+``rbac.manageSystemComponents``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Allow the API server to manage system components in the skypilot-system namespace. Required for object store mounting.
+
+Default: ``true``
+
+.. code-block:: yaml
+
+  rbac:
+    manageSystemComponents: true
+
+.. _helm-values-kubernetesCredentials:
+
+``kubernetesCredentials``
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. _helm-values-kubernetesCredentials-useApiServerCluster:
+
+``kubernetesCredentials.useApiServerCluster``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Enable using the API server's cluster for workloads.
+
+Default: ``true``
+
+.. code-block:: yaml
+
+  kubernetesCredentials:
+    useApiServerCluster: true
+
+.. _helm-values-kubernetesCredentials-useKubeconfig:
+
+``kubernetesCredentials.useKubeconfig``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Use the `kube-credentials` secret containing the kubeconfig to authenticate to Kubernetes.
+
+Default: ``false``
+
+.. code-block:: yaml
+
+  kubernetesCredentials:
+    useKubeconfig: false
+
+.. _helm-values-kubernetesCredentials-kubeconfigSecretName:
+
+``kubernetesCredentials.kubeconfigSecretName``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Name of the secret containing the kubeconfig file. Only used if useKubeconfig is true.
+
+Default: ``kube-credentials``
+
+.. code-block:: yaml
+
+  kubernetesCredentials:
+    kubeconfigSecretName: kube-credentials
+
+.. _helm-values-kubernetesCredentials-inclusterNamespace:
+
+``kubernetesCredentials.inclusterNamespace``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Namespace to use for in-cluster resources.
+
+Default: ``null``
+
+.. code-block:: yaml
+
+  kubernetesCredentials:
+    inclusterNamespace: null
+
+.. _helm-values-awsCredentials:
+
+``awsCredentials``
+~~~~~~~~~~~~~~~~~~
+
+.. _helm-values-awsCredentials-enabled:
+
+``awsCredentials.enabled``
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Enable AWS credentials for the API server.
+
+Default: ``false``
+
+.. code-block:: yaml
+
+  awsCredentials:
+    enabled: false
+
+.. _helm-values-awsCredentials-awsSecretName:
+
+``awsCredentials.awsSecretName``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Name of the secret containing the AWS credentials. Only used if enabled is true.
+
+Default: ``aws-credentials``
+
+.. code-block:: yaml
+
+  awsCredentials:
+    awsSecretName: aws-credentials
+
+.. _helm-values-awsCredentials-accessKeyIdKeyName:
+
+``awsCredentials.accessKeyIdKeyName``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Key name used to set AWS_ACCESS_KEY_ID.
+
+Default: ``aws_access_key_id``
+
+.. code-block:: yaml
+
+  awsCredentials:
+    accessKeyIdKeyName: aws_access_key_id
+
+.. _helm-values-awsCredentials-secretAccessKeyKeyName:
+
+``awsCredentials.secretAccessKeyKeyName``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Key name used to set AWS_SECRET_ACCESS_KEY.
+
+Default: ``aws_secret_access_key``
+
+.. code-block:: yaml
+
+  awsCredentials:
+    secretAccessKeyKeyName: aws_secret_access_key
+
+.. _helm-values-gcpCredentials:
+
+``gcpCredentials``
+~~~~~~~~~~~~~~~~~~
+
+.. _helm-values-gcpCredentials-enabled:
+
+``gcpCredentials.enabled``
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Enable GCP credentials for the API server.
+
+Default: ``false``
+
+.. code-block:: yaml
+
+  gcpCredentials:
+    enabled: false
+
+.. _helm-values-gcpCredentials-projectId:
+
+``gcpCredentials.projectId``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+GCP project ID. Only used if enabled is true.
+
+Default: ``null``
+
+.. code-block:: yaml
+
+  gcpCredentials:
+    projectId: null
+
+.. _helm-values-gcpCredentials-gcpSecretName:
+
+``gcpCredentials.gcpSecretName``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Name of the secret containing the GCP credentials. Only used if enabled is true.
+
+Default: ``gcp-credentials``
+
+.. code-block:: yaml
+
+  gcpCredentials:
+    gcpSecretName: gcp-credentials
+
+.. _helm-values-podSecurityContext:
+
+``podSecurityContext``
+~~~~~~~~~~~~~~~~~~~~~~
+
+Security context for the API server pod. Usually left empty to use defaults. Refer to `set the security context for Pod <https://kubernetes.io/docs/tasks/configure-pod-container/security-context/#set-the-security-context-for-a-pod>`_ for more details.
+
+Default: ``{}``
+
+.. code-block:: yaml
+
+  podSecurityContext:
+    runAsUser: 1000
+    runAsGroup: 3000
+    fsGroup: 2000
+
+.. _helm-values-securityContext:
+
+``securityContext``
+~~~~~~~~~~~~~~~~~~~
+
+.. _helm-values-securityContext-capabilities:
+
+``securityContext.capabilities``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Linux capabilities to drop for the API server container.
+
+Default: drop all capabilities.
+
+.. code-block:: yaml
+
+  securityContext:
+    capabilities:
+      drop:
+      - ALL
+
+.. _helm-values-securityContext-allowPrivilegeEscalation:
+
+``securityContext.allowPrivilegeEscalation``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Whether to allow privilege escalation in the API server container.
+
+Default: ``false``
+
+.. code-block:: yaml
+
+  securityContext:
+    allowPrivilegeEscalation: false
+
+.. _helm-values-runtimeClassName:
+
+``runtimeClassName``
+~~~~~~~~~~~~~~~~~~~~
+
+The runtime class to use for the API server pod. Usually left empty to use the default runtime class.
+
+Default: (empty)
+
+.. code-block:: yaml
+
+  runtimeClassName:
diff --git a/docs/source/reference/config.rst b/docs/source/reference/config.rst
index ca28a11c1..58bb7b660 100644
--- a/docs/source/reference/config.rst
+++ b/docs/source/reference/config.rst
@@ -105,6 +105,8 @@ Below is the configuration syntax and some example values. See detailed explanat
       provision_timeout: 900
     :ref:`remote_identity <config-yaml-gcp-remote-identity>`: LOCAL_CREDENTIALS
     :ref:`enable_gvnic <config-yaml-gcp-enable-gvnic>`: false
+    :ref:`enable_gpu_direct <config-yaml-gcp-enable-gpu-direct>`: false
+    :ref:`placement_policy <config-yaml-gcp-placement-policy>`: compact
 
   :ref:`azure <config-yaml-azure>`:
     :ref:`resource_group_vm <config-yaml-azure-resource-group-vm>`: user-resource-group-name
@@ -855,6 +857,26 @@ launched by SkyPilot. This is useful for improving network performance.
 
 Default: ``false``.
 
+.. _config-yaml-gcp-enable-gpu-direct:
+
+``gcp.enable_gpu_direct``
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Enable GPUDirect-TCPX, a high-performance networking technology that establishes direct communication between GPUs and network interfaces for `a3-highgpu-8g` or `a3-edgegpu-8g` instances launched by SkyPilot. When enabled, this configuration automatically activates the gVNIC network interface for optimal performance.
+
+Default: ``false``.
+
+.. _config-yaml-gcp-placement-policy:
+
+``gcp.placement_policy``
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+Placement policy for GCP instances. This setting controls how instances are physically placed within a data center to optimize performance and resource utilization.
+
+When `gcp.enable_gpu_direct` is enabled, the placement policy is automatically set to `compact` to ensure optimal communication performance. If `gcp.enable_gpu_direct` is disabled, no default placement policy is applied.
+
+Refer to the `GCP documentation <https://cloud.google.com/compute/docs/instances/placement-policies-overview>`_ for more information on placement policies.
+
 .. _config-yaml-azure:
 
 ``azure``
diff --git a/docs/source/reservations/existing-machines.rst b/docs/source/reservations/existing-machines.rst
index d760a444a..1a6c14db7 100644
--- a/docs/source/reservations/existing-machines.rst
+++ b/docs/source/reservations/existing-machines.rst
@@ -5,6 +5,11 @@ Deploy SkyPilot on existing machines
 
 This guide will help you deploy SkyPilot on your existing machines — whether they are on-premises or reserved instances on a cloud provider.
 
+
+.. tip::
+
+    To run SkyPilot on your local machine, use ``sky local up`` to :ref:`deploy a kubernetes cluster <kubernetes-setup-kind>` with `kind <https://kind.sigs.k8s.io/>`_.
+
 **Given a list of IP addresses and SSH credentials,**
 SkyPilot will install necessary dependencies on the remote machines and configure itself to run jobs and services on the cluster.
 
diff --git a/examples/gcp_gpu_direct_tcpx/README.md b/examples/gcp_gpu_direct_tcpx/README.md
new file mode 100644
index 000000000..2cba5f8c5
--- /dev/null
+++ b/examples/gcp_gpu_direct_tcpx/README.md
@@ -0,0 +1,109 @@
+# Using GCP GPUDirect-TCPX on A3 VM with SkyPilot
+
+[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. By bypassing the CPU and system memory, it significantly enhances network performance for A3 VMs, particularly for large data transfers.
+
+When deploying `a3-highgpu-8g` or `a3-edgegpu-8g` VMs, combining GPUDirect-TCPX with Google Virtual NIC (gVNIC) delivers optimal network performance with minimal latency between applications and the network infrastructure.
+
+This example demonstrates how to run [NCCL tests](https://github.com/NVIDIA/nccl-tests) on a GCP cluster with `a3-highgpu-8g` VMs, comparing performance with and without GPUDirect-TCPX enabled.
+
+## TL;DR: enable GPUDirect-TCPX with SkyPilot
+
+Enable GPUDirect-TCPX on GCP clusters with `a3-highgpu-8g` or `a3-edgegpu-8g` VMs by adding a single configuration parameter to your SkyPilot YAML:
+
+```yaml
+config:
+  gcp:
+    enable_gpu_direct: true
+```
+
+With `enable_gpu_direct: true`, SkyPilot automatically:
+
+1. Creates a dedicated network infrastructure:
+   - 1 management VPC
+   - 4 data VPCs
+   - Corresponding subnets for each VPC
+
+2. Provisions VMs with GPUDirect-TCPX support:
+   - Launches VMs with the specified instance type
+   - Uses GPUDirect-TCPX-compatible images
+   - Installs necessary GPU drivers
+   - Deploys the GPUDirect-TCPX Receive Data Path Manager service
+   - Configures NVIDIA Collective Communications Library (NCCL) and GPUDirect-TCPX plugin
+
+## Running NCCL Tests with GPUDirect-TCPX
+
+The complete configuration is available in `gpu_direct_tcpx.yaml`. The configuration includes:
+- `image_id`: Pre-configured environment for NCCL testing
+- `instance_type`: Set to `a3-highgpu-8g`
+
+To run the NCCL test with GPUDirect-TCPX:
+
+```bash
+sky launch -c tcpx gpu_direct_tcpx.yaml
+```
+
+SkyPilot will:
+1. Deploy a GCP cluster with GPUDirect-TCPX-enabled nodes
+2. Execute NCCL performance tests
+3. Output detailed performance metrics
+
+Successful GPUDirect-TCPX activation is confirmed by these log entries:
+
+```
+NCCL INFO NET/GPUDirectTCPX ver. 3.1.8.
+NCCL INFO NET/GPUDirectTCPX : GPUDirectTCPX enable: 1
+```
+
+> **Note:** To run tests without GPUDirect-TCPX, use:
+> ```bash
+> sky launch -c tcpx --env USE_GPU_DIRECT=false gpu_direct_tcpx.yaml
+> ```
+
+### Performance Benchmark Results
+
+We conducted performance comparisons using NCCL tests on a GCP cluster with 2x a3-highgpu-8g (2xH100:8) instances. The speed-up is calculated as:
+```
+Speed-up = busbw GPUDirect-TCPX (GB/s) / busbw Non-GPUDirect-TCPX (GB/s)
+```
+
+| Message Size | busbw GPUDirect-TCPX (GB/s) | busbw Non-GPUDirect-TCPX (GB/s) | Speed-up |
+|--------------|------------|-------------|--------------|
+| 8 B          | 0          | 0           | -            |
+| 16 B         | 0          | 0           | -            |
+| 32 B         | 0          | 0           | -            |
+| 64 B         | 0          | 0           | -            |
+| 128 B        | 0          | 0           | -            |
+| 256 B        | 0          | 0           | -            |
+| 512 B        | 0          | 0           | -            |
+| 1 KB         | 0.01       | 0.01        | 1 x          |
+| 2 KB         | 0.01       | 0.01        | 1 x          |
+| 4 KB         | 0.01       | 0.02        | 0.5 x        |
+| 8 KB         | 0.02       | 0.04        | 0.5 x        |
+| 16 KB        | 0.04       | 0.09        | 0.4 x        |
+| 32 KB        | 0.09       | 0.12        | 0.7 x        |
+| 64 KB        | 0.11       | 0.17        | 0.6 x        |
+| 128 KB       | 0.19       | 0.15        | 1.2 x        |
+| 256 KB       | 0.35       | 0.23        | 1.5 x        |
+| 512 KB       | 0.65       | 0.47        | 1.4 x        |
+| 1 MB         | 1.33       | 0.95        | 1.4 x        |
+| 2 MB         | 2.43       | 1.87        | 1.3 x        |
+| 4 MB         | 4.8        | 3.64        | 1.3 x        |
+| 8 MB         | 9.21       | 7.1         | 1.3 x        |
+| 16 MB        | 17.16      | 8.83        | 1.9 x        |
+| 32 MB        | 30.08      | 12.07       | 2.5 x        |
+| 64 MB        | 45.31      | 12.48       | 3.6 x        |
+| 128 MB       | 61.58      | 16.27       | 3.8 x        |
+| 256 MB       | 67.82      | 20.93       | 3.2 x        |
+| 512 MB       | 67.09      | 19.93       | 3.3 x        |
+| 1 GB         | 66.2       | 20.09       | 3.3 x        |
+| 2 GB         | 65.72      | 19.39       | 3.4 x        |
+
+### Key Performance Insights
+
+| Message Size Range | Performance Characteristics |
+|-------------------|----------------------------|
+| ≤ 128 KB          | Minimal benefit - GPUDirect-TCPX may introduce slight overhead for small messages, with comparable or lower bandwidth than non-GPUDirect mode |
+| 256 KB – 8 MB     | Moderate improvement - Speedup of 1.5–1.9×, with performance crossover point at 128–256 KB |
+| ≥ 16 MB           | Significant advantage - 2.5–3.8× speedup, with GPUDirect-TCPX maintaining 65–67 GB/s versus ~20 GB/s without it |
+
+GPUDirect-TCPX's direct GPU-to-NIC communication path eliminates CPU and system memory bottlenecks, delivering superior throughput for large-scale data transfers. This makes it particularly effective for distributed deep learning workloads and high-performance computing applications.
diff --git a/examples/gcp_gpu_direct_tcpx/gpu_direct_tcpx.yaml b/examples/gcp_gpu_direct_tcpx/gpu_direct_tcpx.yaml
new file mode 100644
index 000000000..23a205cf8
--- /dev/null
+++ b/examples/gcp_gpu_direct_tcpx/gpu_direct_tcpx.yaml
@@ -0,0 +1,126 @@
+name: nccl-gpu-direct-tcpx
+
+resources:
+  cloud: gcp
+  instance_type: a3-highgpu-8g
+  image_id: docker:us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx
+
+num_nodes: 2
+
+envs:
+  USE_GPU_DIRECT: "true"
+
+setup: |
+  # Check if the /usr/local/tcpx/lib64/libnccl.so.2 is
+  # present to ensure the user-data script has completed
+  while [ ! -f /usr/local/tcpx/lib64/libnccl.so.2 ]; do
+    echo "Waiting for user-data script to complete"
+    sleep 10
+  done
+  # Remount the directories with exec permissions
+  sudo mount -o remount,exec /usr/local/tcpx/lib64
+  sudo mount -o remount,exec /usr/local/nvidia/lib64
+  sudo mount -o remount,exec /usr/local/nvidia/bin
+
+
+run: |
+  if [ "${SKYPILOT_NODE_RANK}" == "0" ]; then
+    echo "Head node"
+
+    # Total number of processes, NP should be the total number of GPUs in the cluster
+    NP=$(($SKYPILOT_NUM_GPUS_PER_NODE * $SKYPILOT_NUM_NODES))
+
+    # Append :${SKYPILOT_NUM_GPUS_PER_NODE} to each IP as slots
+    nodes=""
+    for ip in $SKYPILOT_NODE_IPS; do
+      nodes="${nodes}${ip}:${SKYPILOT_NUM_GPUS_PER_NODE},"
+    done
+    nodes=${nodes::-1}
+    echo "All nodes: ${nodes}"
+
+    # Set environment variables
+    export PATH=/usr/local/nvidia/bin:/usr/local/cuda/bin:$PATH
+    export NCCL_SOCKET_IFNAME=eth0
+    export NCCL_CROSS_NIC=0
+    export NCCL_ALGO=Ring
+    export NCCL_PROTO=Simple
+    export NCCL_NSOCKS_PERTHREAD=4
+    export NCCL_SOCKET_NTHREADS=1
+    export NCCL_NET_GDR_LEVEL=PIX
+    export NCCL_DYNAMIC_CHUNK_SIZE=524288
+    export NCCL_P2P_PXN_LEVEL=0    
+    export NCCL_P2P_NET_CHUNKSIZE=524288
+    export NCCL_P2P_PCI_CHUNKSIZE=524288
+    export NCCL_P2P_NVL_CHUNKSIZE=1048576
+    export NCCL_BUFFSIZE=8388608
+    export NCCL_MAX_NCHANNELS=8
+    export NCCL_MIN_NCHANNELS=8
+    export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+    export NCCL_GPUDIRECTTCPX_SOCKET_IFNAME=eth1,eth2,eth3,eth4
+    export NCCL_GPUDIRECTTCPX_CTRL_DEV=eth0
+    export NCCL_GPUDIRECTTCPX_TX_BINDINGS="eth1:8-21,112-125;eth2:8-21,112-125;eth3:60-73,164-177;eth4:60-73,164-177"
+    export NCCL_GPUDIRECTTCPX_RX_BINDINGS="eth1:22-35,126-139;eth2:22-35,126-139;eth3:74-87,178-191;eth4:74-87,178-191"
+    export NCCL_GPUDIRECTTCPX_PROGRAM_FLOW_STEERING_WAIT_MICROS=50000
+    export NCCL_GPUDIRECTTCPX_UNIX_CLIENT_PREFIX="/run/tcpx"
+    export NCCL_GPUDIRECTTCPX_FORCE_ACK=0
+
+    if [ "${USE_GPU_DIRECT}" == "true" ]; then
+      export LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/tcpx/lib64
+    else
+      # Use the default NCCL library
+      export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:/usr/local/nvidia/lib64
+    fi
+    
+    mpirun \
+      --allow-run-as-root \
+      --tag-output \
+      -H $nodes \
+      -np $NP \
+      -N $SKYPILOT_NUM_GPUS_PER_NODE \
+      -x PATH \
+      -x LD_LIBRARY_PATH \
+      -x NCCL_SOCKET_IFNAME \
+      -x NCCL_CROSS_NIC \
+      -x NCCL_ALGO \
+      -x NCCL_PROTO \
+      -x NCCL_NSOCKS_PERTHREAD \
+      -x NCCL_SOCKET_NTHREADS \
+      -x NCCL_MAX_NCHANNELS \
+      -x NCCL_MIN_NCHANNELS \
+      -x NCCL_DYNAMIC_CHUNK_SIZE \
+      -x NCCL_P2P_NET_CHUNKSIZE \
+      -x NCCL_P2P_PCI_CHUNKSIZE \
+      -x NCCL_P2P_NVL_CHUNKSIZE \
+      -x NCCL_BUFFSIZE \
+      -x CUDA_VISIBLE_DEVICES \
+      -x NCCL_GPUDIRECTTCPX_SOCKET_IFNAME \
+      -x NCCL_GPUDIRECTTCPX_CTRL_DEV \
+      -x NCCL_GPUDIRECTTCPX_TX_BINDINGS \
+      -x NCCL_GPUDIRECTTCPX_RX_BINDINGS \
+      -x NCCL_GPUDIRECTTCPX_PROGRAM_FLOW_STEERING_WAIT_MICROS \
+      -x NCCL_GPUDIRECTTCPX_UNIX_CLIENT_PREFIX \
+      -x NCCL_GPUDIRECTTCPX_FORCE_ACK \
+      -x NCCL_NET_GDR_LEVEL \
+      -x NCCL_P2P_PXN_LEVEL \
+      -x NCCL_DEBUG=INFO -x NCCL_DEBUG_SUBSYS=ENV \
+      --mca btl tcp,self \
+      --mca btl_tcp_if_include eth0 \
+      --mca plm_rsh_args "-p 10022" \
+      /third_party/nccl-tests-mpi/build/all_reduce_perf \
+      -b 8 \
+      -e 2G \
+      -f 2 \
+      -g 1 \
+      -c 1 \
+      -w 5 \
+      -n 20
+  else
+    echo "Worker nodes"
+  fi
+
+config:
+  gcp:
+    enable_gpu_direct: true
+    managed_instance_group:
+      run_duration: 36000
+      provision_timeout: 900
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 11a2fc83a..e1564e933 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
@@ -301,6 +303,8 @@ class RayCodeGen:
             from sky.skylet import autostop_lib
             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
 
@@ -2415,6 +2419,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 +3847,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 +4565,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,
@@ -4657,6 +4664,7 @@ class CloudVmRayBackend(backends.Backend['CloudVmRayResourceHandle']):
             handle_before_refresh = record['handle']
             status_before_refresh = record['status']
 
+        handle: Optional[CloudVmRayResourceHandle]
         prev_cluster_status, handle = (status_before_refresh,
                                        handle_before_refresh)
 
@@ -4677,7 +4685,7 @@ class CloudVmRayBackend(backends.Backend['CloudVmRayResourceHandle']):
                 prev_cluster_status = None
                 handle = None
         # We should check the cluster_ever_up after refresh, because if the
-        # cluster is terminated (through console or auto-dwon), the record will
+        # cluster is terminated (through console or auto-down), the record will
         # become None and the cluster_ever_up should be considered as False.
         cluster_ever_up = record is not None and record['cluster_ever_up']
         prev_config_hash = record['config_hash'] if record is not None else None
@@ -4690,11 +4698,14 @@ class CloudVmRayBackend(backends.Backend['CloudVmRayResourceHandle']):
             self.check_resources_fit_cluster(handle, task)
             # Use the existing cluster.
             assert handle.launched_resources is not None, (cluster_name, handle)
+            # Take a random resource in order to get resource info that applies
+            # to all resources.
+            one_task_resource = list(task.resources)[0]
             # Assume resources share the same ports.
             for resource in task.resources:
-                assert resource.ports == list(task.resources)[0].ports
+                assert resource.ports == one_task_resource.ports
             requested_ports_set = resources_utils.port_ranges_to_set(
-                list(task.resources)[0].ports)
+                one_task_resource.ports)
             current_ports_set = resources_utils.port_ranges_to_set(
                 handle.launched_resources.ports)
             all_ports = resources_utils.port_set_to_ranges(current_ports_set |
@@ -4713,6 +4724,21 @@ class CloudVmRayBackend(backends.Backend['CloudVmRayResourceHandle']):
                             'a new cluster with the desired ports open.')
             if all_ports:
                 to_provision = to_provision.copy(ports=all_ports)
+            # Docker login should always be the same for all resources, since
+            # it's set from envs.
+            for resource in task.resources:
+                assert (resource.docker_login_config ==
+                        one_task_resource.docker_login_config), (
+                            resource.docker_login_config,
+                            one_task_resource.docker_login_config)
+            # If we have docker login config in the new task, override the
+            # existing resources to pick up new credentials. This allows the
+            # user to specify new or fixed credentials if the existing
+            # credentials are not working. If we don't do this, the credentials
+            # from the existing resources will always be reused.
+            if one_task_resource.docker_login_config is not None:
+                to_provision = to_provision.copy(
+                    _docker_login_config=one_task_resource.docker_login_config)
             return RetryingVmProvisioner.ToProvisionConfig(
                 cluster_name,
                 to_provision,
diff --git a/sky/clouds/gcp.py b/sky/clouds/gcp.py
index 1cfa23826..c22a6b5ce 100644
--- a/sky/clouds/gcp.py
+++ b/sky/clouds/gcp.py
@@ -17,6 +17,7 @@ from sky import skypilot_config
 from sky.adaptors import gcp
 from sky.clouds import service_catalog
 from sky.clouds.utils import gcp_utils
+from sky.provision.gcp import constants
 from sky.utils import annotations
 from sky.utils import common_utils
 from sky.utils import registry
@@ -112,6 +113,10 @@ _DEFAULT_CPU_IMAGE_ID = 'skypilot:custom-cpu-ubuntu-2204'
 # For GPU-related package version, see sky/clouds/service_catalog/images/provisioners/cuda.sh
 _DEFAULT_GPU_IMAGE_ID = 'skypilot:custom-gpu-ubuntu-2204'
 _DEFAULT_GPU_K80_IMAGE_ID = 'skypilot:k80-debian-10'
+# Use COS image with GPU Direct support.
+# Need to contact GCP support to build our own image for GPUDirect-TCPX support.
+# Refer to https://github.com/GoogleCloudPlatform/cluster-toolkit/blob/main/examples/machine-learning/a3-highgpu-8g/README.md#before-starting
+_DEFAULT_GPU_DIRECT_IMAGE_ID = 'skypilot:gpu-direct-cos'
 
 
 def _run_output(cmd):
@@ -488,6 +493,11 @@ class GCP(clouds.Cloud):
             'gcp_project_id': self.get_project_id(dryrun),
             **GCP._get_disk_specs(r.instance_type, _failover_disk_tier()),
         }
+        enable_gpu_direct = skypilot_config.get_nested(
+            ('gcp', 'enable_gpu_direct'),
+            False,
+            override_configs=resources.cluster_config_overrides)
+        resources_vars['enable_gpu_direct'] = enable_gpu_direct
         accelerators = r.accelerators
         if accelerators is not None:
             assert len(accelerators) == 1, r
@@ -520,14 +530,17 @@ class GCP(clouds.Cloud):
                     resources_vars['gpu'] = 'nvidia-tesla-{}'.format(
                         acc.lower())
                 resources_vars['gpu_count'] = acc_count
-                if acc == 'K80':
-                    # Though the image is called cu113, it actually has later
-                    # versions of CUDA as noted below.
-                    # CUDA driver version 470.57.02, CUDA Library 11.4
-                    image_id = _DEFAULT_GPU_K80_IMAGE_ID
+                if enable_gpu_direct:
+                    image_id = _DEFAULT_GPU_DIRECT_IMAGE_ID
                 else:
-                    # CUDA driver version 535.86.10, CUDA Library 12.2
-                    image_id = _DEFAULT_GPU_IMAGE_ID
+                    if acc == 'K80':
+                        # Though the image is called cu113, it actually has later
+                        # versions of CUDA as noted below.
+                        # CUDA driver version 470.57.02, CUDA Library 11.4
+                        image_id = _DEFAULT_GPU_K80_IMAGE_ID
+                    else:
+                        # CUDA driver version 535.86.10, CUDA Library 12.2
+                        image_id = _DEFAULT_GPU_IMAGE_ID
 
         if (resources.image_id is not None and
                 resources.extract_docker_image() is None):
@@ -580,7 +593,21 @@ class GCP(clouds.Cloud):
 
         # Add gVNIC from config
         resources_vars['enable_gvnic'] = skypilot_config.get_nested(
-            ('gcp', 'enable_gvnic'), False)
+            ('gcp', 'enable_gvnic'),
+            False,
+            override_configs=resources.cluster_config_overrides)
+        placement_policy = skypilot_config.get_nested(
+            ('gcp', 'placement_policy'),
+            None,
+            override_configs=resources.cluster_config_overrides)
+        resources_vars['user_data'] = None
+        if enable_gpu_direct:
+            resources_vars['user_data'] = constants.GPU_DIRECT_TCPX_USER_DATA
+            resources_vars[
+                'docker_run_options'] = constants.GPU_DIRECT_TCPX_SPECIFIC_OPTIONS
+            if placement_policy is None:
+                placement_policy = constants.COMPACT_GROUP_PLACEMENT_POLICY
+        resources_vars['placement_policy'] = placement_policy
 
         return resources_vars
 
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/provision/docker_utils.py b/sky/provision/docker_utils.py
index 007d2c41f..a23717b30 100644
--- a/sky/provision/docker_utils.py
+++ b/sky/provision/docker_utils.py
@@ -343,9 +343,12 @@ class DockerInitializer:
         # `mesg: ttyname failed: inappropriate ioctl for device`.
         # see https://www.educative.io/answers/error-mesg-ttyname-failed-inappropriate-ioctl-for-device  # pylint: disable=line-too-long
         port = constants.DEFAULT_DOCKER_PORT
+        # In case the port is already configured in the sshd_config file
+        # in some images, we delete it first and then append the new one.
         # pylint: disable=anomalous-backslash-in-string
         self._run(
-            f'sudo sed -i "s/#Port 22/Port {port}/" /etc/ssh/sshd_config;'
+            'sudo sed -i "/^Port .*/d" /etc/ssh/sshd_config;'
+            f'sudo echo "Port {port}" >> /etc/ssh/sshd_config;'
             'mkdir -p ~/.ssh;'
             'cat /tmp/host_ssh_authorized_keys >> ~/.ssh/authorized_keys;'
             'sudo service ssh start;'
diff --git a/sky/provision/gcp/config.py b/sky/provision/gcp/config.py
index 835645d78..9b1c4ecd4 100644
--- a/sky/provision/gcp/config.py
+++ b/sky/provision/gcp/config.py
@@ -75,6 +75,30 @@ def wait_for_compute_global_operation(project_name, operation, compute):
     return result
 
 
+def wait_for_compute_region_operation(project_name, region, operation, compute):
+    """Poll for region compute operation until finished."""
+    logger.info('wait_for_compute_region_operation: '
+                'Waiting for operation {} to finish...'.format(
+                    operation['name']))
+
+    for _ in range(constants.MAX_POLLS):
+        result = (compute.regionOperations().get(
+            project=project_name,
+            region=region,
+            operation=operation['name'],
+        ).execute())
+        if 'error' in result:
+            raise Exception(result['error'])
+
+        if result['status'] == 'DONE':
+            logger.info('wait_for_compute_region_operation: Operation done.')
+            break
+
+        time.sleep(constants.POLL_INTERVAL)
+
+    return result
+
+
 def _create_crm(gcp_credentials=None):
     return gcp.build('cloudresourcemanager',
                      'v1',
@@ -168,6 +192,7 @@ def bootstrap_instances(
     iam_role = _configure_iam_role(config, crm, iam)
     config.node_config.update(iam_role)
     config = _configure_subnet(region, cluster_name, config, compute)
+    config = _configure_placement_policy(region, cluster_name, config, compute)
 
     return config
 
@@ -660,6 +685,95 @@ def get_usable_vpc_and_subnet(
     return usable_vpc_name, usable_subnet
 
 
+def get_gpu_direct_usable_vpcs_and_subnets(
+    cluster_name: str,
+    region: str,
+    config: common.ProvisionConfig,
+    compute,
+) -> List[Tuple[str, 'google.cloud.compute_v1.types.compute.Subnetwork']]:
+    """Return a list of usable VPCs and subnets for GPU Direct."""
+    project_id = config.provider_config['project_id']
+    vpc_prefix = constants.SKYPILOT
+    cluster_prefix = cluster_name[:constants.CLUSTER_PREFIX_LENGTH]
+    vpc_subnet_pairs = []
+
+    # TODO(hailong): Determine the num_vpcs per different GPU Direct types
+    num_vpcs = constants.SKYPILOT_GPU_DIRECT_VPC_NUM
+
+    cidr_prefix = constants.SKYPILOT_GPU_DIRECT_VPC_CIDR_PREFIX
+    for i in range(num_vpcs):
+        if i == 0:
+            vpc_name = f'{vpc_prefix}-{cluster_prefix}-mgmt-net'
+        else:
+            vpc_name = f'{vpc_prefix}-{cluster_prefix}-data-net-{i}'
+        subnet_name = f'{vpc_name}-sub'
+        subnet_cidr_range = f'{cidr_prefix}.{i}.0/24'
+        # Check if VPC exists
+        vpc_list = _list_vpcnets(project_id, compute, filter=f'name={vpc_name}')
+        if not vpc_list:
+            body = constants.VPC_TEMPLATE.copy()
+            body['mtu'] = 8244
+            body['autoCreateSubnetworks'] = False
+            body['name'] = vpc_name
+            body['selfLink'] = body['selfLink'].format(PROJ_ID=project_id,
+                                                       VPC_NAME=vpc_name)
+            _create_vpcnet(project_id, compute, body)
+        # Check if subnet exists
+        subnets = _list_subnets(project_id, region, compute, network=vpc_name)
+        if not subnets:
+            _create_subnet(project_id, region, compute, vpc_name, subnet_name,
+                           subnet_cidr_range)
+            subnets = _list_subnets(project_id,
+                                    region,
+                                    compute,
+                                    network=vpc_name)
+        # Apply firewall rules
+        _create_rules(project_id, compute, constants.FIREWALL_RULES_TEMPLATE,
+                      vpc_name)
+        vpc_subnet_pairs.append((vpc_name, subnets[0]))
+    return vpc_subnet_pairs
+
+
+def _configure_placement_policy(region: str, cluster_name: str,
+                                config: common.ProvisionConfig, compute):
+    """Configure placement group for GPU Direct."""
+    node_config = config.node_config
+    project_id = config.provider_config['project_id']
+    group_placement_policy = config.provider_config.get('placement_policy',
+                                                        None)
+    # If the placement policy is not compact,
+    # or the managed instance group is specified,
+    # skip the placement policy creation.
+    # If placement policy is specified together with managed instance group,
+    # it will cause the following error:
+    # Reason: [{'code': 'UNSUPPORTED_OPERATION',
+    # 'message': 'Creating queued resource with
+    # resource policies is not supported.'}]
+    mig_configuration = config.provider_config.get('use_managed_instance_group',
+                                                   False)
+    if (group_placement_policy is None or group_placement_policy.lower() !=
+            constants.COMPACT_GROUP_PLACEMENT_POLICY or mig_configuration):
+        return config
+
+    cluster_prefix = cluster_name[:constants.CLUSTER_PREFIX_LENGTH]
+    policy_name = f'{cluster_prefix}-placement-policy'
+    resource_policy = {
+        'name': policy_name,
+        'groupPlacementPolicy': {
+            'collocation': constants.COLLOCATED_COLLOCATION,
+        }
+    }
+    # Try to get the placement policy first, if not found, create it
+    placement_policy = _get_placement_policy(project_id, region, compute,
+                                             policy_name)
+    if not placement_policy:
+        logger.info(f'Creating placement policy {policy_name}'
+                    f' for cluster {cluster_name}')
+        _create_placement_policy(project_id, region, compute, resource_policy)
+    node_config['resourcePolicies'] = [policy_name]
+    return config
+
+
 def _configure_subnet(region: str, cluster_name: str,
                       config: common.ProvisionConfig, compute):
     """Pick a reasonable subnet if not specified by the config."""
@@ -671,25 +785,54 @@ def _configure_subnet(region: str, cluster_name: str,
     if 'networkInterfaces' in node_config or 'networkConfig' in node_config:
         return config
 
-    # SkyPilot: make sure there's a usable VPC
-    _, default_subnet = get_usable_vpc_and_subnet(cluster_name, region, config,
-                                                  compute)
-
-    default_interfaces = [{
-        'subnetwork': default_subnet['selfLink'],
-        'accessConfigs': [{
-            'name': 'External NAT',
-            'type': 'ONE_TO_ONE_NAT',
-        }]
-    }]
-    # Add gVNIC if specified in config
+    default_interfaces = []
+    enable_gpu_direct = config.provider_config.get('enable_gpu_direct', False)
     enable_gvnic = config.provider_config.get('enable_gvnic', False)
-    if enable_gvnic:
-        default_interfaces[0]['nicType'] = 'gVNIC'
+    if enable_gpu_direct:
+        if not enable_gvnic:
+            logger.warning(
+                'Enable GPU Direct requires gvnic to be enabled, enabling gvnic'
+            )
+            config.provider_config['enable_gvnic'] = True
+            enable_gvnic = True
+        if 'machineType' not in node_config or node_config[
+                'machineType'] not in constants.GPU_DIRECT_TCPX_INSTANCE_TYPES:
+            raise ValueError(
+                'Enable GPU Direct requires machineType to be one of '
+                f'{constants.GPU_DIRECT_TCPX_INSTANCE_TYPES}')
+        logger.info(f'Enable GPU Direct for cluster {cluster_name} '
+                    f'with machineType {node_config["machineType"]}')
+        vpc_subnet_pairs = get_gpu_direct_usable_vpcs_and_subnets(
+            cluster_name, region, config, compute)
+        for _, subnet in vpc_subnet_pairs:
+            default_interfaces.append({
+                'subnetwork': subnet['selfLink'],
+                'accessConfigs': [{
+                    'name': 'External NAT',
+                    'type': 'ONE_TO_ONE_NAT',
+                }],
+                'nicType': 'gVNIC'
+            })
+    else:
+        # SkyPilot: make sure there's a usable VPC
+        _, default_subnet = get_usable_vpc_and_subnet(cluster_name, region,
+                                                      config, compute)
+
+        default_interfaces = [{
+            'subnetwork': default_subnet['selfLink'],
+            'accessConfigs': [{
+                'name': 'External NAT',
+                'type': 'ONE_TO_ONE_NAT',
+            }]
+        }]
+        # Add gVNIC if specified in config
+        if enable_gvnic:
+            default_interfaces[0]['nicType'] = 'gVNIC'
     enable_external_ips = _enable_external_ips(config)
     if not enable_external_ips:
         # Removing this key means the VM will not be assigned an external IP.
-        default_interfaces[0].pop('accessConfigs')
+        for interface in default_interfaces:
+            interface.pop('accessConfigs')
 
     # The not applicable key will be removed during node creation
 
@@ -840,3 +983,42 @@ def _add_iam_policy_binding(service_account, policy, crm, iam):
     ).execute())
 
     return result
+
+
+def _create_subnet(project_id: str, region: str, compute, vpc_name: str,
+                   subnet_name: str, ip_cidr_range: str):
+    body = {
+        'name': subnet_name,
+        'ipCidrRange': ip_cidr_range,
+        'network': f'projects/{project_id}/global/networks/{vpc_name}',
+        'region': region,
+    }
+    operation = compute.subnetworks().insert(project=project_id,
+                                             region=region,
+                                             body=body).execute()
+    response = wait_for_compute_region_operation(project_id, region, operation,
+                                                 compute)
+    return response
+
+
+def _create_placement_policy(project_id: str, region: str, compute,
+                             placement_policy: dict):
+    operation = compute.resourcePolicies().insert(
+        project=project_id, region=region, body=placement_policy).execute()
+    response = wait_for_compute_region_operation(project_id, region, operation,
+                                                 compute)
+    return response
+
+
+def _get_placement_policy(project_id: str, region: str, compute, name: str):
+    try:
+        placement_policy = (compute.resourcePolicies().get(
+            project=project_id,
+            region=region,
+            resourcePolicy=name,
+        ).execute())
+    except gcp.http_error_exception() as e:
+        if e.resp.status == 404:
+            return None
+        raise
+    return placement_policy
diff --git a/sky/provision/gcp/constants.py b/sky/provision/gcp/constants.py
index 11cd7ed86..4103a2001 100644
--- a/sky/provision/gcp/constants.py
+++ b/sky/provision/gcp/constants.py
@@ -41,6 +41,70 @@ HAS_TPU_PROVIDER_FIELD = '_has_tpus'
 # with ServiceAccounts.
 
 SKYPILOT_VPC_NAME = 'skypilot-vpc'
+SKYPILOT_GPU_DIRECT_VPC_NUM = 5
+SKYPILOT_GPU_DIRECT_VPC_CIDR_PREFIX = '10.129'
+GPU_DIRECT_TCPX_INSTANCE_TYPES = [
+    'a3-edgegpu-8g',
+    'a3-highgpu-8g',
+]
+# The prefix length of the cluster name.
+# To make sure the VPC and subnet names are within the GCP limits.
+CLUSTER_PREFIX_LENGTH = 10
+
+COMPACT_GROUP_PLACEMENT_POLICY = 'compact'
+COLLOCATED_COLLOCATION = 'COLLOCATED'
+GPU_DIRECT_TCPX_USER_DATA = """#!/bin/bash
+    set -e
+    set -x
+    # Install GPU Direct TCPX
+    cos-extensions install gpu -- --version=latest;
+    sudo mount --bind /var/lib/nvidia /var/lib/nvidia;
+    sudo mount -o remount,exec /var/lib/nvidia;
+    docker ps -a | grep -q receive-datapath-manager || \
+    docker run \
+    --detach \
+    --pull=always \
+    --name receive-datapath-manager \
+    --privileged \
+    --cap-add=NET_ADMIN --network=host \
+    --volume /var/lib/nvidia/lib64:/usr/local/nvidia/lib64 \
+    --device /dev/nvidia0:/dev/nvidia0 --device /dev/nvidia1:/dev/nvidia1 \
+    --device /dev/nvidia2:/dev/nvidia2 --device /dev/nvidia3:/dev/nvidia3 \
+    --device /dev/nvidia4:/dev/nvidia4 --device /dev/nvidia5:/dev/nvidia5 \
+    --device /dev/nvidia6:/dev/nvidia6 --device /dev/nvidia7:/dev/nvidia7 \
+    --device /dev/nvidia-uvm:/dev/nvidia-uvm --device /dev/nvidiactl:/dev/nvidiactl \
+    --env LD_LIBRARY_PATH=/usr/local/nvidia/lib64 \
+    --volume /run/tcpx:/run/tcpx \
+    --entrypoint /tcpgpudmarxd/build/app/tcpgpudmarxd \
+    us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/tcpgpudmarxd \
+    --gpu_nic_preset a3vm --gpu_shmem_type fd --uds_path "/run/tcpx" --setup_param "--verbose 128 2 0";
+    sudo iptables -I INPUT -p tcp -m tcp -j ACCEPT;
+    docker run --rm -v /var/lib:/var/lib us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx install --install-nccl;
+    sudo mount --bind /var/lib/tcpx /var/lib/tcpx;
+    sudo mount -o remount,exec /var/lib/tcpx;
+    echo "GPU Direct TCPX installed"
+    """
+
+GPU_DIRECT_TCPX_SPECIFIC_OPTIONS = [
+    '--cap-add=IPC_LOCK',
+    '--userns=host',
+    '--volume /run/tcpx:/run/tcpx',
+    '--volume /var/lib/nvidia/lib64:/usr/local/nvidia/lib64',
+    '--volume /var/lib/tcpx/lib64:/usr/local/tcpx/lib64',
+    '--volume /var/lib/nvidia/bin:/usr/local/nvidia/bin',
+    '--shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864',
+    '--device /dev/nvidia0:/dev/nvidia0',
+    '--device /dev/nvidia1:/dev/nvidia1',
+    '--device /dev/nvidia2:/dev/nvidia2',
+    '--device /dev/nvidia3:/dev/nvidia3',
+    '--device /dev/nvidia4:/dev/nvidia4',
+    '--device /dev/nvidia5:/dev/nvidia5',
+    '--device /dev/nvidia6:/dev/nvidia6',
+    '--device /dev/nvidia7:/dev/nvidia7',
+    '--device /dev/nvidia-uvm:/dev/nvidia-uvm',
+    '--device /dev/nvidiactl:/dev/nvidiactl',
+    '--env LD_LIBRARY_PATH=/usr/local/nvidia/lib64:/usr/local/tcpx/lib64',
+]
 
 # Below parameters are from the default VPC on GCP.
 # https://cloud.google.com/vpc/docs/firewalls#more_rules_default_vpc
diff --git a/sky/resources.py b/sky/resources.py
index a33fc3653..ced4c6c29 100644
--- a/sky/resources.py
+++ b/sky/resources.py
@@ -499,6 +499,10 @@ class Resources:
     def requires_fuse(self, value: Optional[bool]) -> None:
         self._requires_fuse = value
 
+    @property
+    def docker_login_config(self) -> Optional[docker_utils.DockerLoginConfig]:
+        return self._docker_login_config
+
     @property
     def docker_username_for_runpod(self) -> Optional[str]:
         return self._docker_username_for_runpod
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/constants.py b/sky/skylet/constants.py
index 9514c8625..2f95e59a8 100644
--- a/sky/skylet/constants.py
+++ b/sky/skylet/constants.py
@@ -370,6 +370,9 @@ OVERRIDEABLE_CONFIG_KEYS_IN_TASK: List[Tuple[str, ...]] = [
     ('kubernetes', 'pod_config'),
     ('kubernetes', 'provision_timeout'),
     ('gcp', 'managed_instance_group'),
+    ('gcp', 'enable_gvnic'),
+    ('gcp', 'enable_gpu_direct'),
+    ('gcp', 'placement_policy'),
 ]
 # When overriding the SkyPilot configs on the API server with the client one,
 # we skip the following keys because they are meant to be client-side configs.
diff --git a/sky/skylet/log_lib.py b/sky/skylet/log_lib.py
index 24df608cd..d2a45edd4 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
@@ -77,6 +80,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 = context.get()
+                if ctx is not None and ctx.is_canceled():
+                    return
                 line = out_io.readline()
                 if not line:
                     break
@@ -111,30 +117,29 @@ 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
 
 
+@context_utils.cancellation_guard
 def run_with_log(
     cmd: Union[List[str], str],
     log_path: str,
@@ -176,7 +181,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 = context.get()
+    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 +207,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 +235,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/templates/gcp-ray.yml.j2 b/sky/templates/gcp-ray.yml.j2
index b1cdb064c..014526e8a 100644
--- a/sky/templates/gcp-ray.yml.j2
+++ b/sky/templates/gcp-ray.yml.j2
@@ -69,6 +69,12 @@ provider:
 {%- if enable_gvnic %}
   enable_gvnic: {{ enable_gvnic }}
 {%- endif %}
+{%- if enable_gpu_direct %}
+  enable_gpu_direct: {{ enable_gpu_direct }}
+{%- endif %}
+{%- if placement_policy %}
+  placement_policy: {{ placement_policy }}
+{%- endif %}
 
 auth:
   ssh_user: gcpuser
@@ -148,6 +154,11 @@ available_node_types:
           - key: install-nvidia-driver
             value: "True"
   {%- endif %}
+  {%- if user_data is not none %}
+          - key: user-data
+            value: |-
+              {{ user_data | indent(10) }}
+  {%- endif %}
   {%- if use_spot or gpu is not none %}
       scheduling:
   {%- if use_spot %}
diff --git a/sky/templates/websocket_proxy.py b/sky/templates/websocket_proxy.py
index e470304d9..7847099ba 100644
--- a/sky/templates/websocket_proxy.py
+++ b/sky/templates/websocket_proxy.py
@@ -16,8 +16,11 @@ from typing import Dict
 from urllib.request import Request
 
 import websockets
+from websockets.asyncio.client import ClientConnection
 from websockets.asyncio.client import connect
 
+BUFFER_SIZE = 2**16  # 64KB
+
 
 def _get_cookie_header(url: str) -> Dict[str, str]:
     """Extract Cookie header value from a cookie jar for a specific URL"""
@@ -51,19 +54,36 @@ async def main(url: str) -> None:
             old_settings = None
 
         try:
-            await asyncio.gather(stdin_to_websocket(websocket),
-                                 websocket_to_stdout(websocket))
+            loop = asyncio.get_running_loop()
+            # Use asyncio.Stream primitives to wrap stdin and stdout, this is to
+            # avoid creating a new thread for each read/write operation
+            # excessively.
+            stdin_reader = asyncio.StreamReader()
+            protocol = asyncio.StreamReaderProtocol(stdin_reader)
+            await loop.connect_read_pipe(lambda: protocol, sys.stdin)
+            transport, protocol = await loop.connect_write_pipe(
+                asyncio.streams.FlowControlMixin, sys.stdout)  # type: ignore
+            stdout_writer = asyncio.StreamWriter(transport, protocol, None,
+                                                 loop)
+
+            await asyncio.gather(stdin_to_websocket(stdin_reader, websocket),
+                                 websocket_to_stdout(websocket, stdout_writer))
         finally:
             if old_settings:
                 termios.tcsetattr(sys.stdin.fileno(), termios.TCSADRAIN,
                                   old_settings)
 
 
-async def stdin_to_websocket(websocket):
+async def stdin_to_websocket(reader: asyncio.StreamReader,
+                             websocket: ClientConnection):
     try:
         while True:
-            data = await asyncio.get_event_loop().run_in_executor(
-                None, sys.stdin.buffer.read, 1)
+            # Read at most BUFFER_SIZE bytes, this not affect
+            # responsiveness since it will return as soon as
+            # there is at least one byte.
+            # The BUFFER_SIZE is chosen to be large enough to improve
+            # throughput.
+            data = await reader.read(BUFFER_SIZE)
             if not data:
                 break
             await websocket.send(data)
@@ -73,13 +93,13 @@ async def stdin_to_websocket(websocket):
         await websocket.close()
 
 
-async def websocket_to_stdout(websocket):
+async def websocket_to_stdout(websocket: ClientConnection,
+                              writer: asyncio.StreamWriter):
     try:
         while True:
             message = await websocket.recv()
-            sys.stdout.buffer.write(message)
-            await asyncio.get_event_loop().run_in_executor(
-                None, sys.stdout.buffer.flush)
+            writer.write(message)
+            await writer.drain()
     except websockets.exceptions.ConnectionClosed:
         print('WebSocket connection closed', file=sys.stderr)
     except Exception as e:  # pylint: disable=broad-except
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/schemas.py b/sky/utils/schemas.py
index 8f35806bb..4ed551e08 100644
--- a/sky/utils/schemas.py
+++ b/sky/utils/schemas.py
@@ -837,6 +837,12 @@ def get_config_schema():
                 'enable_gvnic': {
                     'type': 'boolean'
                 },
+                'enable_gpu_direct': {
+                    'type': 'boolean'
+                },
+                'placement_policy': {
+                    'type': 'string',
+                },
                 'vpc_name': {
                     'oneOf': [
                         {
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