Repository: skypilot-org/skypilot
Pull request #5839 (2025-06-03 00:35:48)

Gke high speed network
<!-- Describe the changes in this PR -->

Deploying TCPX on GKE nodes

<!-- 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):

- [x] Code formatting: install pre-commit (auto-check on commit) or `bash format.sh`
- [ ] Any manual or new tests for this PR (please specify below)
- [ ] All smoke tests: `/smoke-test` (CI) or `pytest tests/test_smoke.py` (local)
- [ ] Relevant individual tests: `/smoke-test -k test_name` (CI) or `pytest tests/test_smoke.py::test_name` (local)
- [ ] 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 -->


Maknee: /smoke-test --gcp
Maknee: 
/smoke-test --gcp


Maknee: 
/smoke-test --gcp


Maknee: 
/smoke-test --gcp


Maknee: 
/smoke-test --gcp


romilbhardwaj: /smoke-test --kubernetes
Maknee: /smoke-test --kubernetes
Michaelvll: @Maknee did we test this PR with NCCL test? Seems there is no nccl test related file in this PR

Patch:
diff --git a/docs/source/examples/performance/index.rst b/docs/source/examples/performance/index.rst
index 7d6e81f45..1671b7035 100644
--- a/docs/source/examples/performance/index.rst
+++ b/docs/source/examples/performance/index.rst
@@ -5,5 +5,5 @@ AI Performance
    :maxdepth: 1
 
    AWS EFA <aws_efa>
-   GCP GPUDirect-TCPX <gcp_gpu_direct_tcpx>
+   GCP/GKE GPUDirect <gcp_gpu_direct_tcpx>
    Nebius with InfiniBand <nebius_infiniband>
diff --git a/docs/source/reference/yaml-spec.rst b/docs/source/reference/yaml-spec.rst
index dc19787b1..f61373b79 100644
--- a/docs/source/reference/yaml-spec.rst
+++ b/docs/source/reference/yaml-spec.rst
@@ -546,7 +546,8 @@ If ``'best'`` is specified, use the best network tier available on the specified
 
 - ``infra: gcp``: Enable GPUDirect-TCPX for high-performance node-to-node GPU communication
 - ``infra: nebius``: Enable Infiniband for high-performance GPU communication across Nebius VMs
-- ``infra: k8s/my-nebius-cluster``: Enable InfiniBand for high-performance GPU communication across pods on Nebius managed Kubernetes.
+- ``infra: k8s/my-nebius-cluster``: Enable InfiniBand for high-performance GPU communication across pods on Nebius managed Kubernetes
+- ``infra: k8s/my-gke-cluster``: Enable GPUDirect-TCPX/TCPXO/RDMA for high-performance GPU communication across pods on Google Kubernetes Engine (GKE).
 
 .. code-block:: yaml
 
diff --git a/examples/gcp_gpu_direct_tcpx/README.md b/examples/gcp_gpu_direct_tcpx/README.md
index 1c51ec171..0dcada58c 100644
--- a/examples/gcp_gpu_direct_tcpx/README.md
+++ b/examples/gcp_gpu_direct_tcpx/README.md
@@ -1,12 +1,170 @@
-# Using GCP GPUDirect-TCPX on A3 VM with SkyPilot
+# Using High-Performance GPU Networking on GCP and GKE 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.
+SkyPilot supports advanced GPU networking technologies on both GCP VMs and GKE clusters, enabling high-performance inter-GPU communication for distributed deep learning and HPC workloads. This includes support for:
 
-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.
+- **GPUDirect-TCPX**: a3-highgpu-8g, a3-edgegpu-8g (H100)
+- **GPUDirect-RDMA**: a3-ultragpu-8g (H200), a4-highgpu-8g (B200)
 
-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.
+## NCCL Test Example YAMLs
 
-## TL;DR: enable GPUDirect-TCPX with SkyPilot
+We offer example YAMLs for running NCCL tests to verify high-performance GPU networking on GCP, covering both VM-based and GKE-based deployments:
+
+| Configuration | Target Platform | GPU Networking Technology | VM Types |
+|---------------|----------------|---------------------------|----------|
+| [`nccl_tcpx_gcpvm_h100.yaml`](https://github.com/skypilot-org/skypilot/blob/master/examples/gcp_gpu_direct_tcpx/nccl_tcpx_gcpvm_h100.yaml) | GCP VM | GPUDirect-TCPX | a3-highgpu-8g, a3-edgegpu-8g |
+| [`nccl_tcpx_gke_h100.yaml`](https://github.com/skypilot-org/skypilot/blob/master/examples/gcp_gpu_direct_tcpx/nccl_tcpx_gke.yaml) | GKE | GPUDirect-TCPX | a3-highgpu-8g, a3-edgegpu-8g |
+| [`nccl_rdma_gke_h200.yaml`](https://github.com/skypilot-org/skypilot/blob/master/examples/gcp_gpu_direct_tcpx/nccl_rdma_gke_h200.yaml) | GKE | GPUDirect-RDMA | a3-ultragpu-8g |
+
+
+## GKE: Using High-Performance GPU Networking on GKE
+
+SkyPilot supports advanced GPU networking GKE clusters, including GPUDirect-TCPX and, GPUDirect-RDMA by simply setting `network_tier: best`:
+
+```
+resources:
+  ...
+  network_tier: best # Turn on GPUDirect if available 
+```
+
+To make sure your cluster is set up correctly, refer to the [GKE documentation](https://cloud.google.com/kubernetes-engine/docs/how-to/gpu-bandwidth-gpudirect-tcpx#create-vpcs-subnets) to setup the cluster with the appropriate networking configuration.
+
+In addition to creating a node pool with fixed node size to request the desired GPU instances, you can also use Dynamic Workload Scheduler (DWS) on GKE to provision the nodes, refer to [using DWS on GKE](https://docs.skypilot.co/en/latest/reservations/reservations.html#using-dws-on-gke) for more details.
+
+After setting up the GKE cluster, you can run the appropriate NCCL tests for your GPUs:
+* H100: [`sky launch -c nccl nccl_tcpx_gke_h100.yaml`](https://github.com/skypilot-org/skypilot/blob/master/examples/gcp_gpu_direct_tcpx/nccl_tcpx_gke.yaml)
+* H200: [`sky launch -c nccl nccl_rdma_gke_h200.yaml`](https://github.com/skypilot-org/skypilot/blob/master/examples/gcp_gpu_direct_tcpx/nccl_rdma_gke_h200.yaml)
+
+### GPUDirect-RDMA on A3-ultragpu-8g (H200)
+
+We validated GPUDirect-RDMA performance on `a3-ultragpu-8g` instances with H200 GPUs. Testing was conducted on a 2-node cluster with 16x H200 GPUs (8 per node) using the configuration in [`nccl_rdma_gke_h200.yaml`](https://github.com/skypilot-org/skypilot/blob/master/examples/gcp_gpu_direct_tcpx/nccl_rdma_gke_h200.yaml).
+
+The scaling curves and bandwidth measurements match exactly with Google's official benchmarks shown in their [documentation](https://cloud.google.com/ai-hypercomputer/docs/create/gke-ai-hypercompute-custom#flex-start).
+
+### Running [`nccl_rdma_gke_h200.yaml`](https://github.com/skypilot-org/skypilot/blob/master/examples/gcp_gpu_direct_tcpx/nccl_rdma_gke_h200.yaml) on SkyPilot
+
+```console
+$ sky launch -c nccl nccl_rdma_gke_h200.yaml
+...
+(head, rank=0, pid=3808) All nodes: 10.100.9.12:8,10.100.10.12:8
+(worker1, rank=1, pid=2769, ip=10.100.10.12) Worker nodes
+(head, rank=0, pid=3808) [1,0]<stdout>:# nThread 1 nGpus 1 minBytes 1024 maxBytes 8589934592 step: 2(factor) warmup iters: 5 iters: 100 agg iters: 1 validation: 1 graph: 0
+(head, rank=0, pid=3808) [1,0]<stdout>:#
+(head, rank=0, pid=3808) [1,0]<stdout>:# Using devices
+(head, rank=0, pid=3808) [1,0]<stdout>:#  Rank  0 Group  0 Pid   7774 on nc-d87e1263-head device  0 [0000:8f:00] NVIDIA H200
+(head, rank=0, pid=3808) [1,0]<stdout>:#  Rank  1 Group  0 Pid   7775 on nc-d87e1263-head device  1 [0000:90:00] NVIDIA H200
+(head, rank=0, pid=3808) [1,0]<stdout>:#  Rank  2 Group  0 Pid   7776 on nc-d87e1263-head device  2 [0000:96:00] NVIDIA H200
+(head, rank=0, pid=3808) [1,0]<stdout>:#  Rank  3 Group  0 Pid   7778 on nc-d87e1263-head device  3 [0000:97:00] NVIDIA H200
+(head, rank=0, pid=3808) [1,0]<stdout>:#  Rank  4 Group  0 Pid   7783 on nc-d87e1263-head device  4 [0000:c4:00] NVIDIA H200
+(head, rank=0, pid=3808) [1,0]<stdout>:#  Rank  5 Group  0 Pid   7786 on nc-d87e1263-head device  5 [0000:c5:00] NVIDIA H200
+(head, rank=0, pid=3808) [1,0]<stdout>:#  Rank  6 Group  0 Pid   7789 on nc-d87e1263-head device  6 [0000:cb:00] NVIDIA H200
+(head, rank=0, pid=3808) [1,0]<stdout>:#  Rank  7 Group  0 Pid   7792 on nc-d87e1263-head device  7 [0000:cc:00] NVIDIA H200
+(head, rank=0, pid=3808) [1,0]<stdout>:#  Rank  8 Group  0 Pid   4926 on nc-d87e1263-worker1 device  0 [0000:8f:00] NVIDIA H200
+(head, rank=0, pid=3808) [1,0]<stdout>:#  Rank  9 Group  0 Pid   4927 on nc-d87e1263-worker1 device  1 [0000:90:00] NVIDIA H200
+(head, rank=0, pid=3808) [1,0]<stdout>:#  Rank 10 Group  0 Pid   4928 on nc-d87e1263-worker1 device  2 [0000:96:00] NVIDIA H200
+(head, rank=0, pid=3808) [1,0]<stdout>:#  Rank 11 Group  0 Pid   4929 on nc-d87e1263-worker1 device  3 [0000:97:00] NVIDIA H200
+(head, rank=0, pid=3808) [1,0]<stdout>:#  Rank 12 Group  0 Pid   4931 on nc-d87e1263-worker1 device  4 [0000:c4:00] NVIDIA H200
+(head, rank=0, pid=3808) [1,0]<stdout>:#  Rank 13 Group  0 Pid   4934 on nc-d87e1263-worker1 device  5 [0000:c5:00] NVIDIA H200
+(head, rank=0, pid=3808) [1,0]<stdout>:#  Rank 14 Group  0 Pid   4937 on nc-d87e1263-worker1 device  6 [0000:cb:00] NVIDIA H200
+(head, rank=0, pid=3808) [1,0]<stdout>:#  Rank 15 Group  0 Pid   4940 on nc-d87e1263-worker1 device  7 [0000:cc:00] NVIDIA H200
+(head, rank=0, pid=3808) [1,0]<stdout>:#
+(head, rank=0, pid=3808) [1,0]<stdout>:#                                                              out-of-place                       in-place
+(head, rank=0, pid=3808) [1,0]<stdout>:#       size         count      type   redop    root     time   algbw   busbw #wrong     time   algbw   busbw #wrong
+(head, rank=0, pid=3808) [1,0]<stdout>:#        (B)    (elements)                               (us)  (GB/s)  (GB/s)            (us)  (GB/s)  (GB/s)
+(head, rank=0, pid=3808) [1,0]<stdout>:        1024            16     float    none      -1    28.39    0.04    0.03      0    28.07    0.04    0.03      0
+(head, rank=0, pid=3808) [1,0]<stdout>:        2048            32     float    none      -1    28.29    0.07    0.07      0    28.24    0.07    0.07      0
+(head, rank=0, pid=3808) [1,0]<stdout>:        4096            64     float    none      -1    28.85    0.14    0.13      0    28.64    0.14    0.13      0
+(head, rank=0, pid=3808) [1,0]<stdout>:        8192           128     float    none      -1    34.12    0.24    0.23      0    32.47    0.25    0.24      0
+(head, rank=0, pid=3808) [1,0]<stdout>:       16384           256     float    none      -1    33.28    0.49    0.46      0    33.29    0.49    0.46      0
+(head, rank=0, pid=3808) [1,0]<stdout>:       32768           512     float    none      -1    34.80    0.94    0.88      0    34.86    0.94    0.88      0
+(head, rank=0, pid=3808) [1,0]<stdout>:       65536          1024     float    none      -1    40.62    1.61    1.51      0    41.22    1.59    1.49      0
+(head, rank=0, pid=3808) [1,0]<stdout>:      131072          2048     float    none      -1    36.17    3.62    3.40      0    40.87    3.21    3.01      0
+(head, rank=0, pid=3808) [1,0]<stdout>:      262144          4096     float    none      -1    41.81    6.27    5.88      0    37.99    6.90    6.47      0
+(head, rank=0, pid=3808) [1,0]<stdout>:      524288          8192     float    none      -1    43.98   11.92   11.18      0    45.96   11.41   10.69      0
+(head, rank=0, pid=3808) [1,0]<stdout>:     1048576         16384     float    none      -1    58.46   17.94   16.82      0    54.81   19.13   17.93      0
+(head, rank=0, pid=3808) [1,0]<stdout>:     2097152         32768     float    none      -1    68.40   30.66   28.74      0    70.87   29.59   27.74      0
+(head, rank=0, pid=3808) [1,0]<stdout>:     4194304         65536     float    none      -1    76.56   54.78   51.36      0    76.13   55.09   51.65      0
+(head, rank=0, pid=3808) [1,0]<stdout>:     8388608        131072     float    none      -1    86.92   96.51   90.47      0    85.77   97.81   91.70      0
+(head, rank=0, pid=3808) [1,0]<stdout>:    16777216        262144     float    none      -1    116.2  144.43  135.41      0    114.9  146.00  136.87      0
+(head, rank=0, pid=3808) [1,0]<stdout>:    33554432        524288     float    none      -1    174.4  192.45  180.42      0    172.4  194.66  182.49      0
+(head, rank=0, pid=3808) [1,0]<stdout>:    67108864       1048576     float    none      -1    278.1  241.27  226.19      0    270.5  248.10  232.59      0
+(head, rank=0, pid=3808) [1,0]<stdout>:   134217728       2097152     float    none      -1    499.6  268.67  251.88      0    483.8  277.44  260.10      0
+(head, rank=0, pid=3808) [1,0]<stdout>:   268435456       4194304     float    none      -1    885.5  303.16  284.21      0    870.7  308.30  289.03      0
+(head, rank=0, pid=3808) [1,0]<stdout>:   536870912       8388608     float    none      -1   1575.6  340.75  319.45      0   1568.7  342.24  320.85      0
+(head, rank=0, pid=3808) [1,0]<stdout>:  1073741824      16777216     float    none      -1   3123.9  343.72  322.23      0   3079.7  348.65  326.86      0
+(head, rank=0, pid=3808) [1,0]<stdout>:  2147483648      33554432     float    none      -1   6229.6  344.72  323.18      0   6107.6  351.61  329.63      0
+(head, rank=0, pid=3808) [1,0]<stdout>:  4294967296      67108864     float    none      -1    12416  345.92  324.30      0    12133  354.00  331.87      0
+(head, rank=0, pid=3808) [1,0]<stdout>:  8589934592     134217728     float    none      -1    24724  347.44  325.72      0    24214  354.75  332.58      0
+(head, rank=0, pid=3808) [1,0]<stdout>:# Out of bounds values : 0 OK
+(head, rank=0, pid=3808) [1,0]<stdout>:# Avg bus bandwidth    : 122.073
+(head, rank=0, pid=3808) [1,0]<stdout>:#
+```
+
+### Comparing with raw NCCL test pods from GCP documentation
+
+```console
+$ kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/container-engine-accelerators/refs/heads/master/gpudirect-rdma/nccl-test-a4.yaml
+...
+$ kubectl exec nccl-test-host-1 -it -- /usr/local/gib/scripts/run_nccl_tests.sh -t all_gather -b 1K -e 8G nccl-host-1 nccl-host-2
+...
+# nThread 1 nGpus 1 minBytes 1024 maxBytes 8589934592 step: 2(factor) warmup iters: 50 iters: 100 agg iters: 1 validation: 1 graph: 0
+#
+# Using devices
+#  Rank  0 Group  0 Pid  13581 on nccl-test-host-1 device  0 [0000:8f:00] NVIDIA H200
+#  Rank  1 Group  0 Pid  13779 on nccl-test-host-1 device  1 [0000:90:00] NVIDIA H200
+#  Rank  2 Group  0 Pid  13629 on nccl-test-host-1 device  2 [0000:96:00] NVIDIA H200
+#  Rank  3 Group  0 Pid  13795 on nccl-test-host-1 device  3 [0000:97:00] NVIDIA H200
+#  Rank  4 Group  0 Pid  13790 on nccl-test-host-1 device  4 [0000:c4:00] NVIDIA H200
+#  Rank  5 Group  0 Pid  13750 on nccl-test-host-1 device  5 [0000:c5:00] NVIDIA H200
+#  Rank  6 Group  0 Pid  13751 on nccl-test-host-1 device  6 [0000:cb:00] NVIDIA H200
+#  Rank  7 Group  0 Pid  13754 on nccl-test-host-1 device  7 [0000:cc:00] NVIDIA H200
+#  Rank  8 Group  0 Pid  13708 on nccl-test-host-2 device  0 [0000:8f:00] NVIDIA H200
+#  Rank  9 Group  0 Pid  13749 on nccl-test-host-2 device  1 [0000:90:00] NVIDIA H200
+#  Rank 10 Group  0 Pid  13728 on nccl-test-host-2 device  2 [0000:96:00] NVIDIA H200
+#  Rank 11 Group  0 Pid  13735 on nccl-test-host-2 device  3 [0000:97:00] NVIDIA H200
+#  Rank 12 Group  0 Pid  13648 on nccl-test-host-2 device  4 [0000:c4:00] NVIDIA H200
+#  Rank 13 Group  0 Pid  13685 on nccl-test-host-2 device  5 [0000:c5:00] NVIDIA H200
+#  Rank 14 Group  0 Pid  13653 on nccl-test-host-2 device  6 [0000:cb:00] NVIDIA H200
+#  Rank 15 Group  0 Pid  13751 on nccl-test-host-2 device  7 [0000:cc:00] NVIDIA H200
+NCCL version 2.26.6+cuda12.8
+#
+#                                                              out-of-place                       in-place
+#       size         count      type   redop    root     time   algbw   busbw #wrong     time   algbw   busbw #wrong
+#        (B)    (elements)                               (us)  (GB/s)  (GB/s)            (us)  (GB/s)  (GB/s)
+        1024            16     float    none      -1    28.70    0.04    0.03      0    27.82    0.04    0.03      0
+        2048            32     float    none      -1    28.05    0.07    0.07      0    28.03    0.07    0.07      0
+        4096            64     float    none      -1    28.39    0.14    0.14      0    28.35    0.14    0.14      0
+        8192           128     float    none      -1    31.44    0.26    0.24      0    31.23    0.26    0.25      0
+       16384           256     float    none      -1    31.61    0.52    0.49      0    31.70    0.52    0.48      0
+       32768           512     float    none      -1    32.81    1.00    0.94      0    32.76    1.00    0.94      0
+       65536          1024     float    none      -1    34.77    1.88    1.77      0    34.58    1.90    1.78      0
+      131072          2048     float    none      -1    34.65    3.78    3.55      0    37.37    3.51    3.29      0
+      262144          4096     float    none      -1    39.47    6.64    6.23      0    37.89    6.92    6.49      0
+      524288          8192     float    none      -1    42.75   12.26   11.50      0    40.59   12.92   12.11      0
+     1048576         16384     float    none      -1    54.70   19.17   17.97      0    53.14   19.73   18.50      0
+     2097152         32768     float    none      -1    69.52   30.16   28.28      0    66.52   31.52   29.55      0
+     4194304         65536     float    none      -1    77.41   54.18   50.79      0    71.70   58.50   54.84      0
+     8388608        131072     float    none      -1    87.98   95.35   89.39      0    86.18   97.34   91.25      0
+    16777216        262144     float    none      -1    117.6  142.68  133.76      0    123.6  135.69  127.21      0
+    33554432        524288     float    none      -1    177.2  189.36  177.53      0    176.9  189.66  177.81      0
+    67108864       1048576     float    none      -1    277.8  241.56  226.47      0    271.8  246.94  231.50      0
+   134217728       2097152     float    none      -1    493.7  271.86  254.87      0    486.3  276.01  258.76      0
+   268435456       4194304     float    none      -1    876.1  306.38  287.23      0    870.5  308.38  289.11      0
+   536870912       8388608     float    none      -1   1580.2  339.74  318.51      0   1568.3  342.32  320.93      0
+  1073741824      16777216     float    none      -1   3126.8  343.40  321.93      0   3084.1  348.16  326.40      0
+  2147483648      33554432     float    none      -1   6218.0  345.36  323.78      0   6097.1  352.22  330.20      0
+  4294967296      67108864     float    none      -1    12400  346.38  324.73      0    12135  353.94  331.82      0
+  8589934592     134217728     float    none      -1    24739  347.22  325.52      0    24244  354.31  332.16      0
+# Out of bounds values : 0 OK
+# Avg bus bandwidth    : 121.902
+#
+```
+
+## GCP VMs: GPUDirect-TCPX on A3 VMs (H100)
+
+> **Note:** The following instructions apply only to GCP VMs (a3-highgpu-8g and a3-edgegpu-8g). For GKE clusters, see the [GKE: Using High-Performance GPU Networking on GKE](#gke-using-high-performance-gpu-networking-on-gke) above.
+
+
+This example demonstrates how to run [NCCL tests](https://github.com/NVIDIA/nccl-tests) on a GCP cluster with A3 VMs, comparing performance with and without GPUDirect-TCPX enabled.
 
 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:
 
@@ -30,16 +188,16 @@ With `enable_gpu_direct: true`, SkyPilot automatically:
    - 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
+### Running NCCL Tests with GPUDirect-TCPX
 
-The complete configuration is available in [`gpu_direct_tcpx.yaml`](https://github.com/skypilot-org/skypilot/blob/master/examples/gcp_gpu_direct_tcpx/gpu_direct_tcpx.yaml). The configuration includes:
+The complete configuration is available in [`nccl_tcpx_gcpvm_h100.yaml`](https://github.com/skypilot-org/skypilot/blob/master/examples/gcp_gpu_direct_tcpx/nccl_tcpx_gcpvm_h100.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
+sky launch -c tcpx nccl_tcpx_gcpvm_h100.yaml
 ```
 
 SkyPilot will:
@@ -56,7 +214,7 @@ 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
+> sky launch -c tcpx --env USE_GPU_DIRECT=false nccl_tcpx_gcpvm_h100.yaml
 > ```
 
 ### Performance Benchmark Results
diff --git a/examples/gcp_gpu_direct_tcpx/nccl_rdma_gke_h200.yaml b/examples/gcp_gpu_direct_tcpx/nccl_rdma_gke_h200.yaml
new file mode 100644
index 000000000..e12bbf561
--- /dev/null
+++ b/examples/gcp_gpu_direct_tcpx/nccl_rdma_gke_h200.yaml
@@ -0,0 +1,68 @@
+# This example is used to test the NCCL performance with
+# GPUDirect RDMA on GKE for H200 instances (a3-ultragpu-8g).
+#
+# Usage:
+# sky launch -c nccl examples/gcp_gpu_direct_tcpx/nccl_rdma_gke_h200.yaml
+name: nccl-gpudirect-rdma-h200
+
+resources:
+  infra: k8s
+  accelerators: H200:8
+  image_id: us-docker.pkg.dev/gce-ai-infra/gpudirect-gib/nccl-plugin-gib-diagnostic:v1.0.6
+  network_tier: best
+
+num_nodes: 2
+
+envs:
+  USE_GPU_DIRECT: "true"
+
+run: |
+  /scripts/container_entry.sh shell
+  source /usr/local/gib/scripts/set_nccl_env.sh
+
+  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}"
+
+    # World Level = 0x0, Rail Aligned = 0x7
+    export NCCL_TESTS_SPLIT_MASK="0x0";
+
+    # Force use of libnccl-gib
+    export NCCL_NET=gIB
+
+    # Set all the correct libnccl-gib environment variables
+    source /usr/local/gib/scripts/set_nccl_env.sh
+
+    # Get all relevant NCCL / env vars to pass to all workers
+    ENV_VARS=$(echo ${!NCCL*} ${!OMPI*} LD_LIBRARY_PATH PATH | sed 's/ / -x /g')
+
+    mpirun \
+      --allow-run-as-root \
+      --tag-output \
+      -H $nodes \
+      -np $NP \
+      -N $SKYPILOT_NUM_GPUS_PER_NODE \
+      -x PATH \
+      -x LD_LIBRARY_PATH \
+      -x $ENV_VARS  \
+      -mca plm_rsh_no_tree_spawn 1 \
+      --mca orte_keep_fqdn_hostnames 1 \
+      --mca btl self,tcp \
+      --mca btl_tcp_if_include eth0 \
+      --bind-to none \
+      --mca plm_rsh_agent "ssh -q -o LogLevel=ERROR -o StrictHostKeyChecking=no -p 222" \
+      /third_party/nccl-tests/build/all_gather_perf -b 1K -e 8G -f 2 -g 1 -w 5 --iters 100 -c 1
+
+  else
+    echo "Worker nodes"
+  fi
diff --git a/examples/gcp_gpu_direct_tcpx/gpu_direct_tcpx.yaml b/examples/gcp_gpu_direct_tcpx/nccl_tcpx_gcpvm_h100.yaml
similarity index 100%
rename from examples/gcp_gpu_direct_tcpx/gpu_direct_tcpx.yaml
rename to examples/gcp_gpu_direct_tcpx/nccl_tcpx_gcpvm_h100.yaml
diff --git a/examples/gcp_gpu_direct_tcpx/nccl_tcpx_gke_h100.yaml b/examples/gcp_gpu_direct_tcpx/nccl_tcpx_gke_h100.yaml
new file mode 100644
index 000000000..ef74273dd
--- /dev/null
+++ b/examples/gcp_gpu_direct_tcpx/nccl_tcpx_gke_h100.yaml
@@ -0,0 +1,86 @@
+# This example is used to test the NCCL performance with
+# GPUDirect TCPX on GKE.#
+#
+# Usage:
+# sky launch -c nccl examples/gcp_gpu_direct_tcpx/nccl_tcpx_gke_h100.yaml
+name: nccl-gpudirect-tcpx
+
+resources:
+  infra: k8s
+  accelerators: H100:8
+  image_id: docker:us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/nccl-plugin-gpudirecttcpx
+  network_tier: best
+
+num_nodes: 2
+
+envs:
+  USE_GPU_DIRECT: "true"
+
+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}"
+
+    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 222" \
+      /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
diff --git a/examples/gcp_gpu_direct_tcpx/tcpx_sglang_serving_gke.yaml b/examples/gcp_gpu_direct_tcpx/tcpx_sglang_serving_gke.yaml
new file mode 100644
index 000000000..dc3687673
--- /dev/null
+++ b/examples/gcp_gpu_direct_tcpx/tcpx_sglang_serving_gke.yaml
@@ -0,0 +1,45 @@
+name: tcpx-sglang-serving-gke
+
+resources:
+  infra: k8s
+  gpus: H100:8
+  network_tier: best
+  disk_size: 1024
+  disk_tier: best
+  ports: 30000
+
+num_nodes: 2 # Specify number of nodes to launch; requirements may vary based on accelerators
+
+setup: |
+  sudo apt install -y nvidia-cuda-toolkit libnuma1
+  
+  source ~/sky-venv/bin/activate
+  if [ $? -eq 0 ]; then
+    echo 'venv exists'
+  else
+    uv venv ~/sky-venv --seed --python=3.10
+    source ~/sky-venv/bin/activate
+  fi
+  
+  # Install sglang with all dependencies using uv
+  uv pip install "sglang[all]>=0.4.7.post1"
+  uv pip install setuptools
+
+run: |
+  source ~/sky-venv/bin/activate
+  
+  # Launch the server with appropriate configuration
+  MASTER_ADDR=$(echo "$SKYPILOT_NODE_IPS" | head -n1)
+  # TP should be number of GPUs per node times number of nodes
+  TP=$(($SKYPILOT_NUM_GPUS_PER_NODE * $SKYPILOT_NUM_NODES))
+
+  python -m sglang.launch_server \
+    --model deepseek-ai/DeepSeek-R1-Distill-Llama-8B \
+    --tp $TP \
+    --dist-init-addr ${MASTER_ADDR}:5000 \
+    --nnodes ${SKYPILOT_NUM_NODES} \
+    --node-rank ${SKYPILOT_NODE_RANK} \
+    --trust-remote-code \
+    --host 0.0.0.0 \
+    --port 30000
+
diff --git a/llm/llama-4-finetuning/configs/llama4_full_sft.yaml b/llm/llama-4-finetuning/configs/llama4_full_sft.yaml
new file mode 100644
index 000000000..26a7714ea
--- /dev/null
+++ b/llm/llama-4-finetuning/configs/llama4_full_sft.yaml
@@ -0,0 +1,47 @@
+# pip install git+https://github.com/hiyouga/transformers.git@llama4_train
+
+### model
+model_name_or_path: meta-llama/Llama-4-Maverick-17B-128E-Instruct
+trust_remote_code: true
+
+### method
+stage: sft
+do_train: true
+finetuning_type: full
+deepspeed: examples/deepspeed/ds_z3_config.json  # choices: [ds_z0_config.json, ds_z2_config.json, ds_z3_config.json]
+
+### dataset
+dataset: mllm_demo,identity,alpaca_en_demo
+template: llama4
+cutoff_len: 1024
+max_samples: 5000
+overwrite_cache: false
+preprocessing_num_workers: 4
+dataloader_num_workers: 1
+
+### output
+output_dir: saves/llama4-8b/lora/sft
+logging_steps: 10
+save_steps: 500
+plot_loss: true
+overwrite_output_dir: true
+save_only_model: false
+report_to: none  # choices: [none, wandb, tensorboard, swanlab, mlflow]
+
+### train
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 2
+learning_rate: 1.0e-4
+num_train_epochs: 1.0
+lr_scheduler_type: cosine
+warmup_ratio: 0.1
+bf16: true
+ddp_timeout: 180000000
+resume_from_checkpoint: null
+
+### eval
+# eval_dataset: alpaca_en_demo
+# val_size: 0.1
+# per_device_eval_batch_size: 1
+# eval_strategy: steps
+# eval_steps: 500
diff --git a/llm/llama-4-finetuning/configs/llama4_lora_sft.yaml b/llm/llama-4-finetuning/configs/llama4_lora_sft.yaml
new file mode 100644
index 000000000..5d75aa49c
--- /dev/null
+++ b/llm/llama-4-finetuning/configs/llama4_lora_sft.yaml
@@ -0,0 +1,49 @@
+# pip install git+https://github.com/hiyouga/transformers.git@llama4_train
+
+### model
+model_name_or_path: meta-llama/Llama-4-Maverick-17B-128E-Instruct
+trust_remote_code: true
+
+### method
+stage: sft
+do_train: true
+finetuning_type: lora
+lora_rank: 8
+lora_target: all
+deepspeed: examples/deepspeed/ds_z3_offload_config.json  # choices: [ds_z0_config.json, ds_z2_config.json, ds_z3_config.json]
+
+### dataset
+dataset: mllm_demo,identity,alpaca_en_demo
+template: llama4
+cutoff_len: 2048
+max_samples: 1000
+overwrite_cache: true
+preprocessing_num_workers: 16
+dataloader_num_workers: 4
+
+### output
+output_dir: saves/llama4-8b/lora/sft
+logging_steps: 10
+save_steps: 500
+plot_loss: true
+overwrite_output_dir: true
+save_only_model: false
+report_to: none  # choices: [none, wandb, tensorboard, swanlab, mlflow]
+
+### train
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 1
+learning_rate: 1.0e-4
+num_train_epochs: 1.0
+lr_scheduler_type: cosine
+warmup_ratio: 0.1
+bf16: true
+ddp_timeout: 180000000
+resume_from_checkpoint: null
+
+### eval
+# eval_dataset: alpaca_en_demo
+# val_size: 0.1
+# per_device_eval_batch_size: 1
+# eval_strategy: steps
+# eval_steps: 500
diff --git a/llm/llama-4-finetuning/configs/llama4_scout_full_sft.yaml b/llm/llama-4-finetuning/configs/llama4_scout_full_sft.yaml
new file mode 100644
index 000000000..9cbeecb03
--- /dev/null
+++ b/llm/llama-4-finetuning/configs/llama4_scout_full_sft.yaml
@@ -0,0 +1,50 @@
+# pip install git+https://github.com/hiyouga/transformers.git@llama4_train
+
+### model
+model_name_or_path: meta-llama/Llama-4-Scout-17B-16E-Instruct
+trust_remote_code: true
+
+### method
+stage: sft
+do_train: true
+finetuning_type: full
+# deepspeed: examples/deepspeed/ds_z2_offload_config.json  # choices: [ds_z0_config.json, ds_z2_config.json, ds_z3_config.json]
+deepspeed: /configs/offload_scout.yaml
+flash_attn: fa2
+enable_liger_kernel: True
+
+### dataset
+dataset: alpaca_en_demo
+template: llama4
+cutoff_len: 128
+max_samples: 100
+overwrite_cache: true
+preprocessing_num_workers: 4
+dataloader_num_workers: 1
+
+### output
+output_dir: saves/llama4-8b/lora/sft
+logging_steps: 10
+save_steps: 500
+plot_loss: true
+overwrite_output_dir: true
+save_only_model: false
+report_to: none  # choices: [none, wandb, tensorboard, swanlab, mlflow]
+
+### train
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 1
+learning_rate: 1.0e-4
+num_train_epochs: 1.0
+lr_scheduler_type: cosine
+warmup_ratio: 0.0
+bf16: true
+ddp_timeout: 180000000
+resume_from_checkpoint: null
+
+### eval
+# eval_dataset: alpaca_en_demo
+# val_size: 0.1
+# per_device_eval_batch_size: 1
+# eval_strategy: steps
+# eval_steps: 500
diff --git a/llm/llama-4-finetuning/llama-4-maverick_gke_llama_factory_lora.yaml b/llm/llama-4-finetuning/llama-4-maverick_gke_llama_factory_lora.yaml
new file mode 100644
index 000000000..86002496e
--- /dev/null
+++ b/llm/llama-4-finetuning/llama-4-maverick_gke_llama_factory_lora.yaml
@@ -0,0 +1,59 @@
+# Full finetuning of Llama-4 Maverick 17B MoE model with 128 experts.
+#
+# Usage:
+#
+#  HF_TOKEN=xxx sky launch llama-4-maverick.yaml -c maverick --env HF_TOKEN
+#
+# This config requires at least 2 nodes with 8x H200 GPUs each.
+
+envs:
+  HF_TOKEN: 
+
+resources:
+  infra: k8s
+  cpus: 100+
+  memory: 1000+
+  accelerators: H100:8
+  disk_tier: best
+  network_tier: best
+
+num_nodes: 2
+
+# Optional: configure buckets for dataset and checkpoints. You can then use the /outputs directory to write checkpoints.
+# file_mounts:
+#  /dataset:
+#    source: s3://my-dataset-bucket
+#    mode: COPY  # COPY mode will prefetch the dataset to the node for faster access
+#  /checkpoints:
+#    source: s3://my-checkpoint-bucket
+#    mode: MOUNT_CACHED  # MOUNT_CACHED mode will intelligently cache the checkpoint for faster writes
+
+file_mounts:
+  /configs: ./configs
+
+setup: |
+  # Download the repository configuration package
+  wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
+
+  # Install the keyring package
+  sudo dpkg -i cuda-keyring_1.1-1_all.deb
+
+  # Update package list
+  sudo apt-get update
+
+  #sudo apt-get install cuda-minimal-build-12-6 -y
+  sudo apt-get install cuda-toolkit-12-6 -y
+
+  git clone --depth 1 https://github.com/hiyouga/LLaMA-Factory.git
+  cd LLaMA-Factory
+  pip install -e ".[torch,metrics,deepspeed]" --no-build-isolation
+  pip install "transformers>=4.51.1"
+
+
+run: |
+  MASTER_ADDR=$(echo "$SKYPILOT_NODE_IPS" | head -n1)
+  echo "Starting distributed finetuning, head node: $MASTER_ADDR"
+
+  cd LLaMA-Factory
+
+  HF_TOKEN=$HF_TOKEN FORCE_TORCHRUN=1 NNODES=2 NODE_RANK=$SKYPILOT_NODE_RANK MASTER_ADDR=$MASTER_ADDR MASTER_PORT=29500 llamafactory-cli train /configs/llama4_lora_sft.yaml
diff --git a/llm/llama-4-finetuning/llama-4-scout_gke_llama_factory_sft.yaml b/llm/llama-4-finetuning/llama-4-scout_gke_llama_factory_sft.yaml
new file mode 100644
index 000000000..13a97d20e
--- /dev/null
+++ b/llm/llama-4-finetuning/llama-4-scout_gke_llama_factory_sft.yaml
@@ -0,0 +1,60 @@
+# Full finetuning of Llama-4 Scout 17B MoE model with 128 experts.
+#
+# Usage:
+#
+#  HF_TOKEN=xxx sky launch llama-4-scout_gke_llama_factory_sft.yaml -c scout --env HF_TOKEN
+#
+# This config uses H100 GPUs with 80GB memory each
+
+envs:
+  HF_TOKEN: 
+
+resources:
+  infra: k8s
+  cpus: 100+
+  memory: 1000+
+  accelerators: H100:8
+  disk_tier: best
+  network_tier: best
+
+num_nodes: 2
+
+# Optional: configure buckets for dataset and checkpoints. You can then use the /outputs directory to write checkpoints.
+# file_mounts:
+#  /dataset:
+#    source: s3://my-dataset-bucket
+#    mode: COPY  # COPY mode will prefetch the dataset to the node for faster access
+#  /checkpoints:
+#    source: s3://my-checkpoint-bucket
+#    mode: MOUNT_CACHED  # MOUNT_CACHED mode will intelligently cache the checkpoint for faster writes
+
+file_mounts:
+  /configs: ./configs
+
+setup: |
+  # Download the repository configuration package
+  wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb
+
+  # Install the keyring package
+  sudo dpkg -i cuda-keyring_1.1-1_all.deb
+
+  # Update package list
+  sudo apt-get update
+
+  # Install CUDA and build tools
+  sudo apt-get install cuda-toolkit-12-6 -y
+  sudo apt-get install -y nvidia-cuda-dev libaio-dev libssl-dev build-essential cmake ninja-build
+
+  git clone --depth 1 https://github.com/hiyouga/LLaMA-Factory.git
+  cd LLaMA-Factory
+  pip install -e ".[torch,metrics,deepspeed,liger-kernel]" --no-build-isolation
+  pip install "transformers>=4.51.1"
+  pip install deepspeed==0.14.5 --upgrade --force-reinstall --no-cache-dir
+
+run: |
+  MASTER_ADDR=$(echo "$SKYPILOT_NODE_IPS" | head -n1)
+  echo "Starting distributed finetuning, head node: $MASTER_ADDR"
+
+  cd LLaMA-Factory
+
+  HF_TOKEN=$HF_TOKEN FORCE_TORCHRUN=1 NNODES=2 NODE_RANK=$SKYPILOT_NODE_RANK MASTER_ADDR=$MASTER_ADDR MASTER_PORT=29500 llamafactory-cli train /configs/llama4_scout_full_sft.yaml
diff --git a/sky/clouds/kubernetes.py b/sky/clouds/kubernetes.py
index 6bbf220bd..1b8475c24 100644
--- a/sky/clouds/kubernetes.py
+++ b/sky/clouds/kubernetes.py
@@ -16,9 +16,11 @@ from sky import skypilot_config
 from sky.adaptors import kubernetes
 from sky.clouds.utils import gcp_utils
 from sky.provision import instance_setup
+from sky.provision.gcp import constants as gcp_constants
 from sky.provision.kubernetes import network_utils
 from sky.provision.kubernetes import utils as kubernetes_utils
 from sky.provision.kubernetes.utils import is_tpu_on_gke
+from sky.provision.kubernetes.utils import KubernetesHighPerformanceNetworkType
 from sky.provision.kubernetes.utils import normalize_tpu_accelerator_name
 from sky.skylet import constants
 from sky.utils import annotations
@@ -127,7 +129,9 @@ class Kubernetes(clouds.Cloud):
                     clouds.CloudImplementationFeatures.SPOT_INSTANCE, None)
             # Allow custom network tier if supported by the cluster
             # (e.g., Nebius clusters with high performance networking)
-            if cls._cluster_supports_high_performance_networking(context):
+            network_type, _ = cls._detect_network_type(context,
+                                                       resources.network_tier)
+            if network_type.supports_high_performance_networking():
                 unsupported_features.pop(
                     clouds.CloudImplementationFeatures.CUSTOM_NETWORK_TIER,
                     None)
@@ -597,22 +601,25 @@ class Kubernetes(clouds.Cloud):
         if resources.use_spot:
             spot_label_key, spot_label_value = kubernetes_utils.get_spot_label()
 
+        network_type, machine_type = self._detect_network_type(
+            context, resources.network_tier)
+
         # Check if this cluster supports high performance networking and
-        # configure IPC_LOCK capability for clusters like Nebius that support it
-        k8s_ipc_lock_capability = False
+        # configure appropriate settings for different cluster types
         if (resources.network_tier is not None and
                 resources.network_tier == resources_utils.NetworkTier.BEST):
             # Only proceed if CUSTOM_NETWORK_TIER is supported by this cluster
             unsupported_features = self._unsupported_features_for_resources(
                 resources)
             if clouds.CloudImplementationFeatures.CUSTOM_NETWORK_TIER \
-                not in unsupported_features:
-                k8s_ipc_lock_capability = True
-
-        if k8s_ipc_lock_capability:
-            k8s_env_vars['NCCL_IB_HCA'] = 'mlx5'
-            k8s_env_vars['UCX_NET_DEVICES'] = \
-                'mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1,mlx5_4:1,mlx5_5:1,mlx5_6:1,mlx5_7:1'  # pylint: disable=line-too-long
+                    not in unsupported_features:
+                # Add high-performance networking environment variables for
+                # Nebius (GCP environment variables are handled directly in
+                # the template)
+                if (network_type == KubernetesHighPerformanceNetworkType.NEBIUS
+                   ):
+                    network_env_vars = network_type.get_network_env_vars()
+                    k8s_env_vars.update(network_env_vars)
 
         # We specify object-store-memory to be 500MB to avoid taking up too
         # much memory on the head node. 'num-cpus' should be set to limit
@@ -732,7 +739,6 @@ class Kubernetes(clouds.Cloud):
             'k8s_high_availability_storage_class_name':
                 (k8s_ha_storage_class_name),
             'avoid_label_keys': avoid_label_keys,
-            'k8s_ipc_lock_capability': k8s_ipc_lock_capability,
             'k8s_enable_flex_start': enable_flex_start,
             'k8s_max_run_duration_seconds': max_run_duration_seconds,
         }
@@ -745,6 +751,22 @@ class Kubernetes(clouds.Cloud):
         namespace = kubernetes_utils.get_kube_config_context_namespace(context)
         deploy_vars['k8s_namespace'] = namespace
 
+        # Add backward compatibility template variables for GPUDirect variants
+        deploy_vars['k8s_enable_gpudirect_tcpx'] = (
+            network_type == KubernetesHighPerformanceNetworkType.GCP_TCPX)
+        deploy_vars['k8s_enable_gpudirect_tcpxo'] = (
+            network_type == KubernetesHighPerformanceNetworkType.GCP_TCPXO)
+        rdma_enabled = (network_type ==
+                        KubernetesHighPerformanceNetworkType.GCP_GPUDIRECT_RDMA)
+        deploy_vars['k8s_enable_gpudirect_rdma'] = rdma_enabled
+        if rdma_enabled and machine_type.startswith('a4'):
+            deploy_vars['k8s_enable_gpudirect_rdma_a4'] = True
+        else:
+            deploy_vars['k8s_enable_gpudirect_rdma_a4'] = False
+
+        deploy_vars['k8s_ipc_lock_capability'] = (
+            network_type.requires_ipc_lock_capability())
+
         return deploy_vars
 
     def _get_feasible_launchable_resources(
@@ -1039,28 +1061,120 @@ class Kubernetes(clouds.Cloud):
         ]
 
     @classmethod
-    @annotations.lru_cache(scope='request', maxsize=10)
-    def _cluster_supports_high_performance_networking(cls,
-                                                      context: str) -> bool:
-        """Check if the cluster supports high performance networking.
-
-        Currently detects Nebius clusters by checking for nebius.com/ labels
-        on cluster nodes.
+    def _detect_network_type(
+        cls,
+        context: str,
+        network_tier: Optional['resources_utils.NetworkTier'] = None
+    ) -> Tuple[KubernetesHighPerformanceNetworkType, str]:
+        """Detect the type of Kubernetes network based on node labels.
 
         Args:
             context: The Kubernetes context to check.
+            network_tier: The network tier requested. If None or not BEST,
+                         returns NONE (no high-performance networking).
 
         Returns:
-            True if the cluster supports high performance networking.
+            A tuple of the detected network type and the instance type.
         """
+        # If network_tier is None or not BEST, return NONE
+        if (network_tier is None or
+                network_tier != resources_utils.NetworkTier.BEST):
+            return KubernetesHighPerformanceNetworkType.NONE, ''
+
         try:
             nodes = kubernetes_utils.get_kubernetes_nodes(context=context)
             for node in nodes:
                 if node.metadata.labels:
-                    for label_key in node.metadata.labels.keys():
+                    # Check for Nebius clusters
+                    for label_key, _ in node.metadata.labels.items():
                         if label_key.startswith('nebius.com/'):
-                            return True
+                            return (KubernetesHighPerformanceNetworkType.NEBIUS,
+                                    '')
+
+                    # Check for GKE clusters with specific GPUDirect variants
+                    machine_family = node.metadata.labels.get(
+                        'cloud.google.com/machine-family', '')
+                    instance_type = node.metadata.labels.get(
+                        'node.kubernetes.io/instance-type', '')
+                    gke_accelerator = node.metadata.labels.get(
+                        'cloud.google.com/gke-accelerator', '')
+
+                    # Check if this is a GKE cluster with A3/A4 machine family
+                    if machine_family in ['a3', 'a4']:
+                        # Check instance type to determine specific GPUDirect
+                        # variant
+                        if 'a3-highgpu-8g' in instance_type:
+                            return (
+                                KubernetesHighPerformanceNetworkType.GCP_TCPX,
+                                'a3-highgpu-8g')
+                        elif 'a3-edgegpu-8g' in instance_type:
+                            return (
+                                KubernetesHighPerformanceNetworkType.GCP_TCPX,
+                                'a3-edgegpu-8g')
+                        elif 'a3-megagpu-8g' in instance_type:
+                            return (
+                                KubernetesHighPerformanceNetworkType.GCP_TCPXO,
+                                'a3-megagpu-8g')
+                        elif 'a4-highgpu-8g' in instance_type:
+                            return (KubernetesHighPerformanceNetworkType.
+                                    GCP_GPUDIRECT_RDMA, 'a4-highgpu-8g')
+                        elif 'a3-ultragpu-8g' in instance_type:
+                            return (KubernetesHighPerformanceNetworkType.
+                                    GCP_GPUDIRECT_RDMA, 'a3-ultragpu-8g')
+                        # Generic A3/A4 detection as fallback
+                        elif machine_family == 'a4':
+                            return (KubernetesHighPerformanceNetworkType.
+                                    GCP_GPUDIRECT_RDMA, 'a4')
+
+                    # Fallback: Check for GPU Direct TCPX capable instance
+                    # types with high-perf GPUs
+                    is_gpu_direct_tcpx_instance = (
+                        instance_type
+                        in gcp_constants.GPU_DIRECT_TCPX_INSTANCE_TYPES)
+                    has_high_perf_gpu = ('nvidia-h100' in gke_accelerator or
+                                         'nvidia-h200' in gke_accelerator or
+                                         'nvidia-b200' in gke_accelerator)
+
+                    if is_gpu_direct_tcpx_instance and has_high_perf_gpu:
+                        # Default to TCPX if we can't determine the specific
+                        # variant
+                        return (KubernetesHighPerformanceNetworkType.GCP_TCPX,
+                                instance_type)
+
         except exceptions.KubeAPIUnreachableError:
             # If we can't reach the cluster, assume no high perf networking
-            return False
-        return False
+            pass
+
+        # If we cannot determine the network type based on nodes
+        # Check if the cluster has any node pools with autoscaling enabled
+        # with machine types that support high perf networking for GKE.
+        autoscaler_type = skypilot_config.get_effective_region_config(
+            cloud='kubernetes',
+            region=context,
+            keys=('autoscaler',),
+            default_value=None)
+        if (autoscaler_type !=
+                kubernetes_enums.KubernetesAutoscalerType.GKE.value):
+            return KubernetesHighPerformanceNetworkType.NONE, ''
+        autoscaler = kubernetes_utils.get_autoscaler(
+            kubernetes_enums.KubernetesAutoscalerType(autoscaler_type))
+        logger.debug(f'{context} has autoscaler of type: {autoscaler_type}')
+        machine_types = autoscaler.get_available_machine_types(context)
+        # Check if any machine type supports high perf networking for GKE.
+        if 'a3-highgpu-8g' in machine_types:
+            return (KubernetesHighPerformanceNetworkType.GCP_TCPX,
+                    'a3-highgpu-8g')
+        elif 'a3-edgegpu-8g' in machine_types:
+            return (KubernetesHighPerformanceNetworkType.GCP_TCPX,
+                    'a3-edgegpu-8g')
+        elif 'a3-megagpu-8g' in machine_types:
+            return (KubernetesHighPerformanceNetworkType.GCP_TCPXO,
+                    'a3-megagpu-8g')
+        elif 'a4-highgpu-8g' in machine_types:
+            return (KubernetesHighPerformanceNetworkType.GCP_GPUDIRECT_RDMA,
+                    'a4-highgpu-8g')
+        elif 'a3-ultragpu-8g' in machine_types:
+            return (KubernetesHighPerformanceNetworkType.GCP_GPUDIRECT_RDMA,
+                    'a3-ultragpu-8g')
+
+        return KubernetesHighPerformanceNetworkType.NONE, ''
diff --git a/sky/core.py b/sky/core.py
index d34939bea..af59ac2f0 100644
--- a/sky/core.py
+++ b/sky/core.py
@@ -878,8 +878,10 @@ def cancel(
         f'handle for cluster {cluster_name!r} should not be None')
 
     backend = backend_utils.get_backend_from_handle(handle)
+    user_hash: Optional[str] = common_utils.get_current_user().id
 
     if all_users:
+        user_hash = None
         sky_logging.print(
             f'{colorama.Fore.YELLOW}'
             f'Cancelling all users\' jobs on cluster {cluster_name!r}...'
@@ -904,7 +906,7 @@ def cancel(
     backend.cancel_jobs(handle,
                         job_ids,
                         cancel_all=all or all_users,
-                        user_hash=common_utils.get_current_user().id)
+                        user_hash=user_hash)
 
 
 @usage_lib.entrypoint
diff --git a/sky/dashboard/src/components/clusters.jsx b/sky/dashboard/src/components/clusters.jsx
index be276c4bb..6559b4cbe 100755
--- a/sky/dashboard/src/components/clusters.jsx
+++ b/sky/dashboard/src/components/clusters.jsx
@@ -35,7 +35,7 @@ import { getClusters, getClusterHistory } from '@/data/connectors/clusters';
 import { getWorkspaces } from '@/data/connectors/workspaces';
 import { getUsers } from '@/data/connectors/users';
 import { sortData } from '@/data/utils';
-import { SquareCode, Terminal, RotateCwIcon } from 'lucide-react';
+import { SquareCode, Terminal, RotateCwIcon, Brackets } from 'lucide-react';
 import { relativeTime } from '@/components/utils';
 import { Layout } from '@/components/elements/layout';
 import {
@@ -66,25 +66,29 @@ import { UserDisplay } from '@/components/elements/UserDisplay';
 //   return cost.toFixed(2);
 // };
 
-const ALL_WORKSPACES_VALUE = '__ALL_WORKSPACES__'; // Define constant for "All Workspaces"
-
-// Define constant for "All Users" similar to workspaces
-const ALL_USERS_VALUE = '__ALL_USERS__';
-
-// Helper function to filter clusters by name
-export function filterClustersByName(clusters, nameFilter) {
-  // If no name filter, return all clusters
-  if (!nameFilter || nameFilter.trim() === '') {
-    return clusters;
-  }
-
-  // Filter clusters by the name filter (case-insensitive partial match)
-  const filterLower = nameFilter.toLowerCase().trim();
-  return clusters.filter((cluster) => {
-    const clusterName = cluster.cluster || '';
-    return clusterName.toLowerCase().includes(filterLower);
-  });
-}
+// Define filter options for the filter dropdown
+const PROPERTY_OPTIONS = [
+  {
+    label: 'Status',
+    value: 'status',
+  },
+  {
+    label: 'Cluster',
+    value: 'cluster',
+  },
+  {
+    label: 'User',
+    value: 'user',
+  },
+  {
+    label: 'Infra',
+    value: 'infra',
+  },
+  {
+    label: 'Workspace',
+    value: 'workspace',
+  },
+];
 
 // Helper function to format autostop information, similar to _get_autostop in CLI utils
 const formatAutostop = (autostop, toDown) => {
@@ -169,96 +173,24 @@ export function Clusters() {
   const [isSSHModalOpen, setIsSSHModalOpen] = useState(false);
   const [isVSCodeModalOpen, setIsVSCodeModalOpen] = useState(false);
   const [selectedCluster, setSelectedCluster] = useState(null);
-  const [workspaceFilter, setWorkspaceFilter] = useState(ALL_WORKSPACES_VALUE);
-  const [userFilter, setUserFilter] = useState(ALL_USERS_VALUE);
-  const [nameFilter, setNameFilter] = useState('');
-  const [workspaces, setWorkspaces] = useState([]);
-  const [users, setUsers] = useState([]);
   const [showHistory, setShowHistory] = useState(false); // 'active' or 'history'
   const isMobile = useMobile();
 
+  const [filters, setFilters] = useState([]);
+  const [optionValues, setOptionValues] = useState({
+    status: [],
+    cluster: [],
+    user: [],
+    workspace: [],
+    infra: [],
+  }); /// Option values for properties
+
   // Handle URL query parameters for workspace and user filtering
   useEffect(() => {
     if (router.isReady) {
-      if (router.query.workspace) {
-        const workspaceParam = Array.isArray(router.query.workspace)
-          ? router.query.workspace[0]
-          : router.query.workspace;
-        setWorkspaceFilter(workspaceParam);
-      }
-      if (router.query.user) {
-        const userParam = Array.isArray(router.query.user)
-          ? router.query.user[0]
-          : router.query.user;
-        setUserFilter(userParam);
-      }
-      if (router.query.name) {
-        const nameParam = Array.isArray(router.query.name)
-          ? router.query.name[0]
-          : router.query.name;
-        setNameFilter(nameParam);
-      }
+      updateFiltersByURLParams();
     }
-  }, [
-    router.isReady,
-    router.query.workspace,
-    router.query.user,
-    router.query.name,
-  ]);
-
-  // Helper function to update URL query parameters
-  const updateURLParams = (newWorkspace, newUser, newName) => {
-    const query = { ...router.query };
-
-    // Update workspace parameter
-    if (newWorkspace && newWorkspace !== ALL_WORKSPACES_VALUE) {
-      query.workspace = newWorkspace;
-    } else {
-      delete query.workspace;
-    }
-
-    // Update user parameter
-    if (newUser && newUser !== ALL_USERS_VALUE) {
-      query.user = newUser;
-    } else {
-      delete query.user;
-    }
-
-    // Update name parameter
-    if (newName && newName.trim() !== '') {
-      query.name = newName.trim();
-    } else {
-      delete query.name;
-    }
-
-    // Use replace to avoid adding to browser history for filter changes
-    router.replace(
-      {
-        pathname: router.pathname,
-        query,
-      },
-      undefined,
-      { shallow: true }
-    );
-  };
-
-  // Handle workspace filter change
-  const handleWorkspaceFilterChange = (newWorkspace) => {
-    setWorkspaceFilter(newWorkspace);
-    updateURLParams(newWorkspace, userFilter, nameFilter);
-  };
-
-  // Handle user filter change
-  const handleUserFilterChange = (newUser) => {
-    setUserFilter(newUser);
-    updateURLParams(workspaceFilter, newUser, nameFilter);
-  };
-
-  // Handle name filter change
-  const handleNameFilterChange = (newName) => {
-    setNameFilter(newName);
-    updateURLParams(workspaceFilter, userFilter, newName);
-  };
+  }, [router.isReady]);
 
   useEffect(() => {
     const fetchFilterData = async () => {
@@ -294,8 +226,6 @@ export function Clusters() {
           finalWorkspaces.add(wsName)
         );
 
-        setWorkspaces(Array.from(finalWorkspaces).sort());
-
         // Fetch users for the filter dropdown
         const fetchedUsers = await dashboardCache.get(getUsers);
         const uniqueClusterUsers = [
@@ -331,21 +261,85 @@ export function Clusters() {
             });
           }
         });
-
-        setUsers(
-          Array.from(finalUsers.values()).sort((a, b) =>
-            a.display.localeCompare(b.display)
-          )
-        );
       } catch (error) {
         console.error('Error fetching data for filters:', error);
-        setWorkspaces(['default']); // Fallback or error state
-        setUsers([]); // Fallback or error state
       }
     };
+
     fetchFilterData();
   }, []);
 
+  // Helper function to update URL query parameters
+  const updateURLParams = (filters) => {
+    const query = { ...router.query };
+
+    let properties = [];
+    let operators = [];
+    let values = [];
+
+    filters.map((filter, _index) => {
+      properties.push(filter.property.toLowerCase() ?? '');
+      operators.push(filter.operator);
+      values.push(filter.value);
+    });
+
+    query.property = properties;
+    query.operator = operators;
+    query.value = values;
+
+    // Use replace to avoid adding to browser history for filter changes
+    router.replace(
+      {
+        pathname: router.pathname,
+        query,
+      },
+      undefined,
+      { shallow: true }
+    );
+  };
+
+  const updateFiltersByURLParams = () => {
+    const query = { ...router.query };
+
+    const properties = query.property;
+    const operators = query.operator;
+    const values = query.value;
+
+    if (properties === undefined) {
+      return;
+    }
+
+    let filters = [];
+
+    const length = Array.isArray(properties) ? properties.length : 1;
+
+    const propertyMap = new Map();
+    propertyMap.set('', '');
+    propertyMap.set('status', 'Status');
+    propertyMap.set('cluster', 'Cluster');
+    propertyMap.set('user', 'User');
+    propertyMap.set('workspace', 'Workspace');
+    propertyMap.set('infra', 'Infra');
+
+    if (length === 1) {
+      filters.push({
+        property: propertyMap.get(properties),
+        operator: operators,
+        value: values,
+      });
+    } else {
+      for (let i = 0; i < length; i++) {
+        filters.push({
+          property: propertyMap.get(properties[i]),
+          operator: operators[i],
+          value: values[i],
+        });
+      }
+    }
+
+    setFilters(filters);
+  };
+
   const handleRefresh = () => {
     // Invalidate cache to ensure fresh data is fetched
     dashboardCache.invalidate(getClusters);
@@ -368,6 +362,7 @@ export function Clusters() {
           >
             Sky Clusters
           </Link>
+
           <div className="flex items-center ml-6 space-x-3">
             <label className="flex items-center cursor-pointer">
               <input
@@ -392,76 +387,14 @@ export function Clusters() {
               </span>
             </label>
           </div>
-          <div className="relative ml-4 mr-2">
-            <input
-              type="text"
-              placeholder="Filter by cluster name"
-              value={nameFilter}
-              onChange={(e) => handleNameFilterChange(e.target.value)}
-              className="h-8 w-32 sm:w-48 px-3 pr-8 text-sm border border-gray-300 rounded-md focus:ring-1 focus:ring-sky-500 focus:border-sky-500 outline-none"
-            />
-            {nameFilter && (
-              <button
-                onClick={() => handleNameFilterChange('')}
-                className="absolute right-2 top-1/2 transform -translate-y-1/2 text-gray-400 hover:text-gray-600"
-                title="Clear filter"
-              >
-                <svg
-                  className="h-4 w-4"
-                  fill="none"
-                  stroke="currentColor"
-                  viewBox="0 0 24 24"
-                >
-                  <path
-                    strokeLinecap="round"
-                    strokeLinejoin="round"
-                    strokeWidth={2}
-                    d="M6 18L18 6M6 6l12 12"
-                  />
-                </svg>
-              </button>
-            )}
-          </div>
-          <Select
-            value={workspaceFilter}
-            onValueChange={handleWorkspaceFilterChange}
-          >
-            <SelectTrigger className="h-8 w-48 ml-2 mr-2 text-sm border-none focus:ring-0 focus:outline-none">
-              <SelectValue placeholder="Filter by workspace...">
-                {workspaceFilter === ALL_WORKSPACES_VALUE
-                  ? 'All Workspaces'
-                  : workspaceFilter}
-              </SelectValue>
-            </SelectTrigger>
-            <SelectContent>
-              <SelectItem value={ALL_WORKSPACES_VALUE}>
-                All Workspaces
-              </SelectItem>
-              {workspaces.map((ws) => (
-                <SelectItem key={ws} value={ws}>
-                  {ws}
-                </SelectItem>
-              ))}
-            </SelectContent>
-          </Select>
-          <Select value={userFilter} onValueChange={handleUserFilterChange}>
-            <SelectTrigger className="h-8 w-48 ml-2 mr-2 text-sm border-none focus:ring-0 focus:outline-none">
-              <SelectValue placeholder="Filter by user...">
-                {userFilter === ALL_USERS_VALUE
-                  ? 'All Users'
-                  : users.find((u) => u.userId === userFilter)?.display ||
-                    userFilter}
-              </SelectValue>
-            </SelectTrigger>
-            <SelectContent>
-              <SelectItem value={ALL_USERS_VALUE}>All Users</SelectItem>
-              {users.map((user) => (
-                <SelectItem key={user.userId} value={user.userId}>
-                  {user.display}
-                </SelectItem>
-              ))}
-            </SelectContent>
-          </Select>
+
+          <FilterDropdown
+            propertyList={PROPERTY_OPTIONS}
+            valueList={optionValues}
+            setFilters={setFilters}
+            updateURLParams={updateURLParams}
+            placeholder="Filter clusters"
+          />
         </div>
         <div className="flex items-center">
           {loading && (
@@ -480,13 +413,18 @@ export function Clusters() {
           </button>
         </div>
       </div>
+
+      <Filters
+        filters={filters}
+        setFilters={setFilters}
+        updateURLParams={updateURLParams}
+      />
+
       <ClusterTable
         refreshInterval={REFRESH_INTERVAL}
         setLoading={setLoading}
         refreshDataRef={refreshDataRef}
-        workspaceFilter={workspaceFilter}
-        userFilter={userFilter}
-        nameFilter={nameFilter}
+        filters={filters}
         showHistory={showHistory}
         onOpenSSHModal={(cluster) => {
           setSelectedCluster(cluster);
@@ -496,6 +434,7 @@ export function Clusters() {
           setSelectedCluster(cluster);
           setIsVSCodeModalOpen(true);
         }}
+        setOptionValues={setOptionValues}
       />
 
       {/* SSH Instructions Modal */}
@@ -518,12 +457,11 @@ export function ClusterTable({
   refreshInterval,
   setLoading,
   refreshDataRef,
-  workspaceFilter,
-  userFilter,
-  nameFilter,
+  filters,
   showHistory,
   onOpenSSHModal,
   onOpenVSCodeModal,
+  setOptionValues,
 }) {
   const [data, setData] = useState([]);
   const [sortConfig, setSortConfig] = useState({
@@ -535,6 +473,32 @@ export function ClusterTable({
   const [currentPage, setCurrentPage] = useState(1);
   const [pageSize, setPageSize] = useState(10);
 
+  const fetchOptionValuesFromClusters = (clusters) => {
+    let optionValues = {
+      status: [],
+      cluster: [],
+      user: [],
+      workspace: [],
+      infra: [],
+    };
+
+    const pushWithoutDuplication = (array, item) => {
+      if (array.includes(item)) return;
+
+      array.push(item);
+    };
+
+    clusters.map((cluster) => {
+      pushWithoutDuplication(optionValues.status, cluster.status);
+      pushWithoutDuplication(optionValues.cluster, cluster.cluster);
+      pushWithoutDuplication(optionValues.user, cluster.user);
+      pushWithoutDuplication(optionValues.workspace, cluster.workspace);
+      pushWithoutDuplication(optionValues.infra, cluster.infra);
+    });
+
+    return optionValues;
+  };
+
   const fetchData = React.useCallback(async () => {
     setLoading(true);
     setLocalLoading(true);
@@ -565,6 +529,9 @@ export function ClusterTable({
             combinedData.push(histCluster);
           }
         });
+
+        setOptionValues(fetchOptionValuesFromClusters(combinedData));
+
         setData(combinedData);
       } else {
         // Mark active clusters for consistency
@@ -572,10 +539,14 @@ export function ClusterTable({
           ...cluster,
           isHistorical: false,
         }));
+
+        setOptionValues(fetchOptionValuesFromClusters(markedActiveClusters));
+
         setData(markedActiveClusters);
       }
     } catch (error) {
       console.error('Error fetching cluster data:', error);
+      setOptionValues(fetchOptionValuesFromClusters([]));
       setData([]);
     }
 
@@ -584,29 +555,63 @@ export function ClusterTable({
     setIsInitialLoad(false);
   }, [setLoading, showHistory]);
 
+  // Utility: checks a condition based on operator
+  const evaluateCondition = (item, filter) => {
+    const { property, operator, value } = filter;
+
+    if (!value) return true; // skip empty filters
+
+    // Global search: check all values
+    if (!property) {
+      const strValue = value.toLowerCase();
+      return Object.values(item).some((val) =>
+        val?.toString().toLowerCase().includes(strValue)
+      );
+    }
+
+    const itemValue = item[property.toLowerCase()]?.toString().toLowerCase();
+    const filterValue = value.toString().toLowerCase();
+
+    switch (operator) {
+      case '=':
+        return itemValue === filterValue;
+      case ':':
+        return itemValue?.includes(filterValue);
+      default:
+        return true;
+    }
+  };
+
   // Use useMemo to compute sorted data
   const sortedData = React.useMemo(() => {
-    let filteredData = data;
-    // Filter by workspace if workspaceFilter is set and not 'ALL_WORKSPACES_VALUE'
-    if (workspaceFilter && workspaceFilter !== ALL_WORKSPACES_VALUE) {
-      filteredData = filteredData.filter((item) => {
-        const itemWorkspace = item.workspace || 'default'; // Treat missing/empty workspace as 'default'
-        return itemWorkspace.toLowerCase() === workspaceFilter.toLowerCase();
-      });
-    }
-    // Filter by user if userFilter is set and not 'ALL_USERS_VALUE'
-    if (userFilter && userFilter !== ALL_USERS_VALUE) {
-      filteredData = filteredData.filter((item) => {
-        const itemUserId = item.user_hash || item.user;
-        return itemUserId === userFilter;
+    // Main filter function
+    const filterData = (data, filters) => {
+      if (filters.length === 0) {
+        return data;
+      }
+
+      return data.filter((item) => {
+        let result = null;
+
+        for (let i = 0; i < filters.length; i++) {
+          const filter = filters[i];
+          const current = evaluateCondition(item, filter);
+
+          if (result === null) {
+            result = current;
+          } else {
+            result = result && current;
+          }
+        }
+
+        return result;
       });
-    }
-    // Filter by name if nameFilter is set
-    if (nameFilter) {
-      filteredData = filterClustersByName(filteredData, nameFilter);
-    }
+    };
+
+    const filteredData = filterData(data, filters);
+
     return sortData(filteredData, sortConfig.key, sortConfig.direction);
-  }, [data, sortConfig, workspaceFilter, userFilter, nameFilter]);
+  }, [data, sortConfig, filters]);
 
   // Expose fetchData to parent component
   React.useEffect(() => {
@@ -893,8 +898,7 @@ export function ClusterTable({
               </div>
             </div>
             <div>
-              {startIndex + 1} – {Math.min(endIndex, data.length)} of{' '}
-              {data.length}
+              {`${startIndex + 1} - ${Math.min(endIndex, sortedData.length)} of ${sortedData.length}`}
             </div>
             <div className="flex items-center space-x-2">
               <Button
@@ -1059,3 +1063,256 @@ export function Status2Actions({
     </>
   );
 }
+
+const FilterDropdown = ({
+  propertyList = [],
+  valueList,
+  setFilters,
+  updateURLParams,
+  placeholder = 'Filter clusters',
+}) => {
+  const inputRef = useRef(null);
+
+  const [isOpen, setIsOpen] = useState(false);
+  const [value, setValue] = useState('');
+  const [propertyValue, setPropertValue] = useState('cluster');
+  const [valueOptions, setValueOptions] = useState([]);
+
+  let filteredOptions = [];
+
+  useEffect(() => {
+    let updatedValueOptions = [];
+
+    switch (propertyValue) {
+      case 'status':
+        updatedValueOptions = valueList.status;
+        break;
+
+      case 'user':
+        updatedValueOptions = valueList.user;
+        break;
+
+      case 'cluster':
+        updatedValueOptions = valueList.cluster;
+        break;
+
+      case 'workspace':
+        updatedValueOptions = valueList.workspace;
+        break;
+
+      case 'infra':
+        updatedValueOptions = valueList.infra;
+        break;
+
+      default:
+        break;
+    }
+
+    updatedValueOptions = updatedValueOptions.filter((item) =>
+      item.includes(value.toLowerCase())
+    );
+
+    setValueOptions(updatedValueOptions);
+  }, [propertyValue, valueList, value]);
+
+  const handleValueChange = (e) => {
+    setValue(e.target.value);
+  };
+
+  return (
+    <>
+      <div className="flex flex-row ml-4 mr-2 border border-gray-300 rounded-md">
+        <div className="border-r">
+          <Select
+            onValueChange={(value) => setPropertValue(value)}
+            value={propertyValue}
+          >
+            <SelectTrigger
+              aria-label="Node"
+              className="focus:ring-0 focus:ring-offset-0 border-none w-36"
+            >
+              <SelectValue placeholder="Select Property" />
+            </SelectTrigger>
+            <SelectContent>
+              {propertyList.map((item, index) => (
+                <SelectItem key={`property-item-${index}`} value={item.value}>
+                  {item.label}
+                </SelectItem>
+              ))}
+            </SelectContent>
+          </Select>
+        </div>
+        <div className="relative">
+          <input
+            type="text"
+            ref={inputRef}
+            placeholder={placeholder}
+            value={value}
+            onChange={(e) => handleValueChange(e)}
+            onFocus={() => setIsOpen(true)}
+            onBlur={() => setIsOpen(false)}
+            onKeyDown={(e) => {
+              if (e.key === 'Enter' && value !== '') {
+                setFilters((prevFilters) => {
+                  const updatedFilters = [
+                    ...prevFilters,
+                    {
+                      property: propertyValue,
+                      operator: ':',
+                      value: value,
+                    },
+                  ];
+
+                  updateURLParams(updatedFilters);
+
+                  return updatedFilters;
+                });
+                setValue('');
+                setIsOpen(false);
+                inputRef.current.blur();
+              }
+            }}
+            className="h-10 w-32 sm:w-96 px-3 pr-8 text-sm rounded-md outline-none my-auto"
+            autoComplete="off"
+          />
+          {value && (
+            <button
+              onClick={() => {
+                setValue('');
+              }}
+              className="absolute right-2 top-1/2 transform -translate-y-1/2 text-gray-400 hover:text-gray-600"
+              title="Clear filter"
+              tabIndex={-1}
+            >
+              <svg
+                className="h-4 w-4"
+                fill="none"
+                stroke="currentColor"
+                viewBox="0 0 24 24"
+              >
+                <path
+                  strokeLinecap="round"
+                  strokeLinejoin="round"
+                  strokeWidth={2}
+                  d="M6 18L18 6M6 6l12 12"
+                />
+              </svg>
+            </button>
+          )}
+          {isOpen && valueOptions.length > 0 && (
+            <div className="flex flex-col absolute z-10 mt-1 w-full bg-white border border-gray-200 rounded-md shadow-lg">
+              {valueOptions.map((option, index) => (
+                <div
+                  key={option}
+                  className={`flex flex-col pl-7 py-2 cursor-pointer hover:bg-sky-50 text-sm ${index != filteredOptions.length - 1} && border-b`}
+                  onMouseDown={(e) => {
+                    e.preventDefault();
+                    setFilters((prevFilters) => {
+                      const updatedFilters = [
+                        ...prevFilters,
+                        {
+                          property: propertyValue,
+                          operator: '=',
+                          value: option,
+                        },
+                      ];
+
+                      updateURLParams(updatedFilters);
+
+                      return updatedFilters;
+                    });
+                    setIsOpen(false);
+                    setValue('');
+                    inputRef.current.blur();
+                  }}
+                >
+                  <span className="text-sm text-gray-600">{option}</span>
+                </div>
+              ))}
+            </div>
+          )}
+        </div>
+      </div>
+    </>
+  );
+};
+
+const Filters = ({ filters = [], setFilters, updateURLParams }) => {
+  const onRemove = (index) => {
+    setFilters((prevFilters) => {
+      const updatedFilters = prevFilters.filter(
+        (_, _index) => _index !== index
+      );
+
+      updateURLParams(updatedFilters);
+
+      return updatedFilters;
+    });
+  };
+
+  const clearFilters = () => {
+    updateURLParams([]);
+    setFilters([]);
+  };
+
+  return (
+    <>
+      <div className="flex items-center gap-4 p-2">
+        <div className="flex flex-wrap items-content gap-2">
+          {filters.map((filter, _index) => (
+            <FilterItem
+              key={`filteritem-${_index}`}
+              filter={filter}
+              onRemove={() => onRemove(_index)}
+            />
+          ))}
+
+          {filters.length > 0 && (
+            <>
+              <div className="border border-gray-400 h-8"></div>
+              <button
+                onClick={clearFilters}
+                className="rounded-full px-4 py-1 text-blue-600 bg-blue-100 hover:bg-blue-200"
+              >
+                Clear filters
+              </button>
+            </>
+          )}
+        </div>
+      </div>
+    </>
+  );
+};
+
+const FilterItem = ({ filter, onRemove }) => {
+  return (
+    <>
+      <div className="flex items-center gap-2 text-blue-600 bg-blue-100 px-2 rounded-full">
+        <div className="flex items-center gap-1 px-2 py-1 rounded-l-lg">
+          <span>{`${filter.property} `}</span>
+          <span className="font-bold">{`${filter.operator} ${filter.value}`}</span>
+        </div>
+
+        <button
+          onClick={() => onRemove()}
+          className="p-1 transform text-gray-400 hover:text-gray-600 bg-blue-500 hover:bg-blue-600 rounded-full flex flex-col items-center"
+          title="Clear filter"
+        >
+          <svg
+            className="h-3 w-3"
+            fill="none"
+            stroke="white"
+            viewBox="0 0 24 24"
+          >
+            <path
+              strokeLinecap="round"
+              strokeLinejoin="round"
+              strokeWidth={5}
+              d="M6 18L18 6M6 6l12 12"
+            />
+          </svg>
+        </button>
+      </div>
+    </>
+  );
+};
diff --git a/sky/global_user_state.py b/sky/global_user_state.py
index 6559b6a0c..14543ef75 100644
--- a/sky/global_user_state.py
+++ b/sky/global_user_state.py
@@ -220,17 +220,16 @@ def _glob_to_similar(glob_pattern):
     return like_pattern
 
 
-def create_table():
+def create_table(engine: sqlalchemy.engine.Engine):
     # Enable WAL mode to avoid locking issues.
     # See: issue #1441 and PR #1509
     # https://github.com/microsoft/WSL/issues/2395
     # TODO(romilb): We do not enable WAL for WSL because of known issue in WSL.
     #  This may cause the database locked problem from WSL issue #1441.
-    if (_SQLALCHEMY_ENGINE.dialect.name
-            == db_utils.SQLAlchemyDialect.SQLITE.value and
+    if (engine.dialect.name == db_utils.SQLAlchemyDialect.SQLITE.value and
             not common_utils.is_wsl()):
         try:
-            with orm.Session(_SQLALCHEMY_ENGINE) as session:
+            with orm.Session(engine) as session:
                 session.execute(sqlalchemy.text('PRAGMA journal_mode=WAL'))
                 session.commit()
         except sqlalchemy_exc.OperationalError as e:
@@ -240,12 +239,12 @@ def create_table():
             # is not critical and is likely to be enabled by other processes.
 
     # Create tables if they don't exist
-    db_utils.add_tables_to_db_sqlalchemy(Base.metadata, _SQLALCHEMY_ENGINE)
+    db_utils.add_tables_to_db_sqlalchemy(Base.metadata, engine)
 
     # For backward compatibility.
     # TODO(zhwu): Remove this function after all users have migrated to
     # the latest version of SkyPilot.
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(engine) as session:
         # Add autostop column to clusters table
         db_utils.add_column_to_table_sqlalchemy(session,
                                                 'clusters',
@@ -391,15 +390,15 @@ def initialize_and_get_db() -> sqlalchemy.engine.Engine:
                 conn_string = skypilot_config.get_nested(('db',), None)
             if conn_string:
                 logger.debug(f'using db URI from {conn_string}')
-                _SQLALCHEMY_ENGINE = sqlalchemy.create_engine(
-                    conn_string, poolclass=sqlalchemy.NullPool)
+                engine = sqlalchemy.create_engine(conn_string,
+                                                  poolclass=sqlalchemy.NullPool)
             else:
                 db_path = os.path.expanduser('~/.sky/state.db')
                 pathlib.Path(db_path).parents[0].mkdir(parents=True,
                                                        exist_ok=True)
-                _SQLALCHEMY_ENGINE = sqlalchemy.create_engine('sqlite:///' +
-                                                              db_path)
-            create_table()
+                engine = sqlalchemy.create_engine('sqlite:///' + db_path)
+            create_table(engine)
+            _SQLALCHEMY_ENGINE = engine
     return _SQLALCHEMY_ENGINE
 
 
diff --git a/sky/jobs/state.py b/sky/jobs/state.py
index aac6ca82f..5e437686f 100644
--- a/sky/jobs/state.py
+++ b/sky/jobs/state.py
@@ -112,17 +112,16 @@ ha_recovery_script_table = sqlalchemy.Table(
 )
 
 
-def create_table():
+def create_table(engine: sqlalchemy.engine.Engine):
     # Enable WAL mode to avoid locking issues.
     # See: issue #3863, #1441 and PR #1509
     # https://github.com/microsoft/WSL/issues/2395
     # TODO(romilb): We do not enable WAL for WSL because of known issue in WSL.
     #  This may cause the database locked problem from WSL issue #1441.
-    if (_SQLALCHEMY_ENGINE.dialect.name
-            == db_utils.SQLAlchemyDialect.SQLITE.value and
+    if (engine.dialect.name == db_utils.SQLAlchemyDialect.SQLITE.value and
             not common_utils.is_wsl()):
         try:
-            with orm.Session(_SQLALCHEMY_ENGINE) as session:
+            with orm.Session(engine) as session:
                 session.execute(sqlalchemy.text('PRAGMA journal_mode=WAL'))
                 session.commit()
         except sqlalchemy_exc.OperationalError as e:
@@ -132,10 +131,10 @@ def create_table():
             # is not critical and is likely to be enabled by other processes.
 
     # Create tables if they don't exist
-    db_utils.add_tables_to_db_sqlalchemy(Base.metadata, _SQLALCHEMY_ENGINE)
+    db_utils.add_tables_to_db_sqlalchemy(Base.metadata, engine)
 
     # Backward compatibility: add columns that not exist in older databases
-    with orm.Session(_SQLALCHEMY_ENGINE) as session:
+    with orm.Session(engine) as session:
         db_utils.add_column_to_table_sqlalchemy(session, 'spot',
                                                 'failure_reason',
                                                 sqlalchemy.Text())
@@ -228,15 +227,15 @@ def initialize_and_get_db() -> sqlalchemy.engine.Engine:
                 conn_string = skypilot_config.get_nested(('db',), None)
             if conn_string:
                 logger.debug(f'using db URI from {conn_string}')
-                _SQLALCHEMY_ENGINE = sqlalchemy.create_engine(
-                    conn_string, poolclass=sqlalchemy.NullPool)
+                engine = sqlalchemy.create_engine(conn_string,
+                                                  poolclass=sqlalchemy.NullPool)
             else:
                 db_path = os.path.expanduser('~/.sky/spot_jobs.db')
                 pathlib.Path(db_path).parents[0].mkdir(parents=True,
                                                        exist_ok=True)
-                _SQLALCHEMY_ENGINE = sqlalchemy.create_engine('sqlite:///' +
-                                                              db_path)
-            create_table()
+                engine = sqlalchemy.create_engine('sqlite:///' + db_path)
+            create_table(engine)
+            _SQLALCHEMY_ENGINE = engine
     return _SQLALCHEMY_ENGINE
 
 
diff --git a/sky/jobs/utils.py b/sky/jobs/utils.py
index 76c8adca3..b0c9044d2 100644
--- a/sky/jobs/utils.py
+++ b/sky/jobs/utils.py
@@ -30,6 +30,7 @@ from sky.backends import backend_utils
 from sky.jobs import constants as managed_job_constants
 from sky.jobs import scheduler
 from sky.jobs import state as managed_job_state
+from sky.server import common as server_common
 from sky.skylet import constants
 from sky.skylet import job_lib
 from sky.skylet import log_lib
@@ -38,6 +39,7 @@ from sky.utils import annotations
 from sky.utils import command_runner
 from sky.utils import common_utils
 from sky.utils import controller_utils
+from sky.utils import env_options
 from sky.utils import infra_utils
 from sky.utils import log_utils
 from sky.utils import message_utils
@@ -128,9 +130,15 @@ def terminate_cluster(cluster_name: str, max_retry: int = 6) -> None:
             time.sleep(backoff.current_backoff())
 
 
-def _check_consolidation_mode_consistency(
+def _validate_consolidation_mode_config(
         current_is_consolidation_mode: bool) -> None:
-    """Check the consistency of the consolidation mode."""
+    """Validate the consolidation mode config."""
+    if (current_is_consolidation_mode and
+            not env_options.Options.IS_DEVELOPER.get() and
+            server_common.is_api_server_local()):
+        with ux_utils.print_exception_no_traceback():
+            raise exceptions.NotSupportedError(
+                'Consolidation mode is not supported when running locally.')
     # Check whether the consolidation mode config is changed.
     if current_is_consolidation_mode:
         controller_cn = (
@@ -176,7 +184,7 @@ def _check_consolidation_mode_consistency(
 def is_consolidation_mode() -> bool:
     consolidation_mode = skypilot_config.get_nested(
         ('jobs', 'controller', 'consolidation_mode'), default_value=False)
-    _check_consolidation_mode_consistency(consolidation_mode)
+    _validate_consolidation_mode_config(consolidation_mode)
     return consolidation_mode
 
 
diff --git a/sky/provision/kubernetes/utils.py b/sky/provision/kubernetes/utils.py
index c4ba5315b..a328659ec 100644
--- a/sky/provision/kubernetes/utils.py
+++ b/sky/provision/kubernetes/utils.py
@@ -1,5 +1,6 @@
 """Kubernetes utilities for SkyPilot."""
 import dataclasses
+import enum
 import functools
 import hashlib
 import json
@@ -57,6 +58,69 @@ HIGH_AVAILABILITY_DEPLOYMENT_VOLUME_MOUNT_NAME = 'sky-data'
 # and store all data that needs to be persisted in future.
 HIGH_AVAILABILITY_DEPLOYMENT_VOLUME_MOUNT_PATH = '/home/sky'
 
+
+class KubernetesHighPerformanceNetworkType(enum.Enum):
+    """Enum for different Kubernetes cluster types with high performance
+    network configurations.
+
+    This enum defines cluster types that support optimized networking for
+    distributed ML workloads:
+    - GCP_TCPX: GKE clusters with GPUDirect-TCPX support
+      (A3 High instances: a3-highgpu-8g)
+    - GCP_TCPXO: GKE clusters with GPUDirect-TCPXO support
+      (A3 Mega instances: a3-megagpu-8g)
+    - GCP_GPUDIRECT_RDMA: GKE clusters with GPUDirect-RDMA support
+      (A4/A3 Ultra instances)
+    - NEBIUS: Nebius clusters with InfiniBand support for high-throughput,
+      low-latency networking
+    - NONE: Standard clusters without specialized networking optimizations
+
+    The network configurations align with corresponding VM-based
+    implementations:
+    - GCP settings match
+      sky.provision.gcp.constants.GPU_DIRECT_TCPX_SPECIFIC_OPTIONS
+    - Nebius settings match the InfiniBand configuration used in Nebius VMs
+    """
+
+    GCP_TCPX = 'gcp_tcpx'
+    GCP_TCPXO = 'gcp_tcpxo'
+    GCP_GPUDIRECT_RDMA = 'gcp_gpudirect_rdma'
+    NEBIUS = 'nebius'
+    NONE = 'none'
+
+    def get_network_env_vars(self) -> Dict[str, str]:
+        """Get network environment variables for this cluster type."""
+        if self == KubernetesHighPerformanceNetworkType.NEBIUS:
+            # Nebius cluster with InfiniBand - use InfiniBand optimizations
+            return {
+                'NCCL_IB_HCA': 'mlx5',
+                'UCX_NET_DEVICES': ('mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1,'
+                                    'mlx5_4:1,mlx5_5:1,mlx5_6:1,mlx5_7:1')
+            }
+        else:
+            # GCP clusters and generic clusters - environment variables are
+            # handled directly in the template
+            return {}
+
+    def supports_high_performance_networking(self) -> bool:
+        """Check if this cluster type supports high performance networking."""
+        return self is not KubernetesHighPerformanceNetworkType.NONE
+
+    def supports_gpu_direct(self) -> bool:
+        """Check if this cluster type supports GPUDirect networking."""
+        return self in (KubernetesHighPerformanceNetworkType.GCP_TCPX,
+                        KubernetesHighPerformanceNetworkType.GCP_TCPXO,
+                        KubernetesHighPerformanceNetworkType.GCP_GPUDIRECT_RDMA)
+
+    def requires_ipc_lock_capability(self) -> bool:
+        """Check if this cluster type requires IPC_LOCK capability."""
+        return self.supports_high_performance_networking()
+
+    def requires_tcpxo_daemon(self) -> bool:
+        """Check if this cluster type requires TCPXO daemon."""
+        return self == KubernetesHighPerformanceNetworkType.GCP_TCPXO
+
+
 # TODO(romilb): Move constants to constants.py
 DEFAULT_NAMESPACE = 'default'
 
@@ -758,6 +822,74 @@ class GKEAutoscaler(Autoscaler):
                     return True
         return False
 
+    @classmethod
+    @annotations.lru_cache(scope='request', maxsize=10)
+    def get_available_machine_types(cls, context: str) -> List[str]:
+        """Returns the list of machine types that are available in the cluster.
+        """
+        # Assume context naming convention of
+        # gke_PROJECT-ID_LOCATION_CLUSTER-NAME
+        valid, project_id, location, cluster_name = cls._validate_context_name(
+            context)
+        if not valid:
+            # Context name is not in the format of
+            # gke_PROJECT-ID_LOCATION_CLUSTER-NAME.
+            # Cannot determine if the context can autoscale.
+            # Return empty list.
+            logger.debug(f'Context {context} is not in the format of '
+                         f'gke_PROJECT-ID_LOCATION_CLUSTER-NAME. '
+                         'Returning empty machine type list.')
+            return []
+        try:
+            logger.debug(
+                f'Attempting to get information about cluster {cluster_name}')
+            container_service = gcp.build('container',
+                                          'v1',
+                                          credentials=None,
+                                          cache_discovery=False)
+            cluster = container_service.projects().locations().clusters().get(
+                name=f'projects/{project_id}'
+                f'/locations/{location}'
+                f'/clusters/{cluster_name}').execute()
+        except ImportError:
+            # If the gcp module is not installed, return empty list.
+            # Remind the user once per day to install the gcp module for better
+            # pod scheduling with GKE autoscaler.
+            if time.time() - cls._pip_install_gcp_hint_last_sent > 60 * 60 * 24:
+                logger.info(
+                    'Could not fetch autoscaler information from GKE. '
+                    'Run pip install "skypilot[gcp]" for more intelligent pod '
+                    'scheduling with GKE autoscaler.')
+                cls._pip_install_gcp_hint_last_sent = time.time()
+            return []
+        except gcp.http_error_exception() as e:
+            # Cluster information is not available.
+            # Return empty list.
+            logger.debug(f'{e.message}', exc_info=True)
+            return []
+
+        machine_types = []
+        # Get the list of machine types that are available in the cluster.
+        node_pools = cluster.get('nodePools', [])
+        for node_pool in node_pools:
+            name = node_pool.get('name', '')
+            logger.debug(f'Checking if node pool {name} '
+                         'has autoscaling enabled.')
+            autoscaling_enabled = (node_pool.get('autoscaling',
+                                                 {}).get('enabled', False))
+            if autoscaling_enabled:
+                logger.debug(f'Node pool {name} has autoscaling enabled.')
+                try:
+                    machine_type = node_pool.get('config',
+                                                 {}).get('machineType', '')
+                    if machine_type:
+                        machine_types.append(machine_type)
+                except KeyError:
+                    logger.debug(f'Encountered KeyError while checking machine '
+                                 f'type of node pool {name}.')
+                    continue
+        return machine_types
+
     @classmethod
     def _validate_context_name(cls, context: str) -> Tuple[bool, str, str, str]:
         """Validates the context name is in the format of
diff --git a/sky/skypilot_config.py b/sky/skypilot_config.py
index bb9907864..04c7f739b 100644
--- a/sky/skypilot_config.py
+++ b/sky/skypilot_config.py
@@ -52,6 +52,7 @@ import contextlib
 import copy
 import json
 import os
+import pathlib
 import tempfile
 import threading
 import typing
@@ -848,7 +849,9 @@ def update_api_server_config_no_lock(config: config_utils.Config) -> None:
 
     global_config_path = _resolve_server_config_path()
     if global_config_path is None:
-        global_config_path = get_user_config_path()
+        # Fallback to ~/.sky/config.yaml, and make sure it exists.
+        global_config_path = os.path.expanduser(get_user_config_path())
+        pathlib.Path(global_config_path).touch(exist_ok=True)
 
     db_updated = False
     if os.environ.get(constants.ENV_VAR_IS_SKYPILOT_SERVER) is not None:
diff --git a/sky/templates/kubernetes-ray.yml.j2 b/sky/templates/kubernetes-ray.yml.j2
index 292701a95..2e67d37e0 100644
--- a/sky/templates/kubernetes-ray.yml.j2
+++ b/sky/templates/kubernetes-ray.yml.j2
@@ -293,13 +293,82 @@ available_node_types:
             kueue.x-k8s.io/queue-name: {{k8s_kueue_local_queue_name}}
             kueue.x-k8s.io/pod-group-name: {{cluster_name_on_cloud}}
             {% endif %}
-        {% if k8s_kueue_local_queue_name %}
+        {% if k8s_kueue_local_queue_name or k8s_enable_gpudirect_tcpx or k8s_enable_gpudirect_tcpxo or k8s_enable_gpudirect_rdma %}
         annotations:
-            kueue.x-k8s.io/retriable-in-group: "false"
-            kueue.x-k8s.io/pod-group-total-count: "{{ num_nodes|string }}"
-            {% if k8s_max_run_duration_seconds %}
-            provreq.kueue.x-k8s.io/maxRunDurationSeconds: "{{k8s_max_run_duration_seconds|string}}"
-            {% endif %}
+          {% if k8s_kueue_local_queue_name %}
+          kueue.x-k8s.io/retriable-in-group: "false"
+          kueue.x-k8s.io/pod-group-total-count: "{{ num_nodes|string }}"
+          {% if k8s_max_run_duration_seconds %}
+          provreq.kueue.x-k8s.io/maxRunDurationSeconds: "{{k8s_max_run_duration_seconds|string}}"
+          {% endif %}
+          {% endif %}
+        # https://cloud.google.com/kubernetes-engine/docs/how-to/gpu-bandwidth-gpudirect-tcpx 
+        # Values from google cloud guide
+          {% if k8s_enable_gpudirect_tcpx %}
+          devices.gke.io/container.tcpx-daemon: |+
+            - path: /dev/nvidia0
+            - path: /dev/nvidia1
+            - path: /dev/nvidia2
+            - path: /dev/nvidia3
+            - path: /dev/nvidia4
+            - path: /dev/nvidia5
+            - path: /dev/nvidia6
+            - path: /dev/nvidia7
+            - path: /dev/nvidiactl
+            - path: /dev/nvidia-uvm
+          networking.gke.io/default-interface: 'eth0'
+          networking.gke.io/interfaces: |
+            [
+              {"interfaceName":"eth0","network":"default"},
+              {"interfaceName":"eth1","network":"vpc1"},
+              {"interfaceName":"eth2","network":"vpc2"},
+              {"interfaceName":"eth3","network":"vpc3"},
+              {"interfaceName":"eth4","network":"vpc4"}
+            ]
+          {% endif %}
+          {% if k8s_enable_gpudirect_tcpxo %}
+          devices.gke.io/container.tcpxo-daemon: |+
+            - path: /dev/nvidia0
+            - path: /dev/nvidia1
+            - path: /dev/nvidia2
+            - path: /dev/nvidia3
+            - path: /dev/nvidia4
+            - path: /dev/nvidia5
+            - path: /dev/nvidia6
+            - path: /dev/nvidia7
+            - path: /dev/nvidiactl
+            - path: /dev/nvidia-uvm
+            - path: /dev/dmabuf_import_helper
+          networking.gke.io/default-interface: 'eth0'
+          networking.gke.io/interfaces: |
+            [
+              {"interfaceName":"eth0","network":"default"},
+              {"interfaceName":"eth1","network":"vpc1"},
+              {"interfaceName":"eth2","network":"vpc2"},
+              {"interfaceName":"eth3","network":"vpc3"},
+              {"interfaceName":"eth4","network":"vpc4"},
+              {"interfaceName":"eth5","network":"vpc5"},
+              {"interfaceName":"eth6","network":"vpc6"},
+              {"interfaceName":"eth7","network":"vpc7"},
+              {"interfaceName":"eth8","network":"vpc8"}
+            ]
+          {% endif %}
+          {% if k8s_enable_gpudirect_rdma %}
+          networking.gke.io/default-interface: 'eth0'
+          networking.gke.io/interfaces: |
+            [
+              {"interfaceName":"eth0","network":"default"},
+              {"interfaceName":"eth1","network":"gvnic-1"},
+              {"interfaceName":"eth2","network":"rdma-0"},
+              {"interfaceName":"eth3","network":"rdma-1"},
+              {"interfaceName":"eth4","network":"rdma-2"},
+              {"interfaceName":"eth5","network":"rdma-3"},
+              {"interfaceName":"eth6","network":"rdma-4"},
+              {"interfaceName":"eth7","network":"rdma-5"},
+              {"interfaceName":"eth8","network":"rdma-6"},
+              {"interfaceName":"eth9","network":"rdma-7"}
+            ]
+          {% endif %}
         {% endif %}
       spec:
         # serviceAccountName: skypilot-service-account
@@ -396,6 +465,41 @@ available_node_types:
           persistentVolumeClaim:
             claimName: {{volume_mount.volume_name_on_cloud}}
         {% endfor %}
+        {% if k8s_enable_gpudirect_tcpx %}
+        - name: libraries
+          hostPath:
+            path: /home/kubernetes/bin/nvidia/lib64
+        - name: tcpx-socket
+          emptyDir: {}
+        - name: sys
+          hostPath:
+            path: /sys
+        - name: proc-sys
+          hostPath:
+            path: /proc/sys
+        {% endif %}
+        {% if k8s_enable_gpudirect_tcpxo %}
+        - name: libraries
+          hostPath:
+            path: /home/kubernetes/bin/nvidia
+        - name: sys
+          hostPath:
+            path: /sys
+        - name: proc-sys
+          hostPath:
+            path: /proc/sys
+        - name: aperture-devices
+          hostPath:
+            path: /dev/aperture_devices
+        {% endif %}
+        {% if k8s_enable_gpudirect_rdma %}
+        - name: library-dir-host
+          hostPath:
+            path: /home/kubernetes/bin/nvidia
+        - name: gib
+          hostPath:
+            path: /home/kubernetes/bin/gib
+        {% endif %}
         containers:
         - name: ray-node
           imagePullPolicy: Always
@@ -409,6 +513,113 @@ available_node_types:
             - name: {{ key }}
               value: {{ value }}
             {% endfor %}
+            # https://cloud.google.com/kubernetes-engine/docs/how-to/gpu-bandwidth-gpudirect-tcpx#environment-variables-nccl
+            # Page recommends setting NCCL values for GPUDirect TCPX for best performance.
+            {% if k8s_enable_gpudirect_tcpx %}
+            - name: LD_LIBRARY_PATH
+              value: /usr/local/nvidia/lib64:/usr/local/tcpx/lib64
+            - name: NCCL_GPUDIRECTTCPX_SOCKET_IFNAME
+              value: eth1,eth2,eth3,eth4
+            - name: NCCL_GPUDIRECTTCPX_CTRL_DEV
+              value: eth0
+            - name: NCCL_GPUDIRECTTCPX_TX_BINDINGS
+              value: "eth1:8-21,112-125;eth2:8-21,112-125;eth3:60-73,164-177;eth4:60-73,164-177"
+            - name: NCCL_GPUDIRECTTCPX_RX_BINDINGS
+              value: "eth1:22-35,126-139;eth2:22-35,126-139;eth3:74-87,178-191;eth4:74-87,178-191"
+            - name: NCCL_GPUDIRECTTCPX_PROGRAM_FLOW_STEERING_WAIT_MICROS
+              value: "500000"
+            - name: NCCL_GPUDIRECTTCPX_UNIX_CLIENT_PREFIX
+              value: "/tmp"
+            - name: NCCL_GPUDIRECTTCPX_FORCE_ACK
+              value: "0"
+            - name: NCCL_SOCKET_IFNAME
+              value: eth0
+            - name: NCCL_CROSS_NIC
+              value: "0"
+            - name: NCCL_ALGO
+              value: Ring
+            - name: NCCL_PROTO
+              value: Simple
+            - name: NCCL_NSOCKS_PERTHREAD
+              value: "4"
+            - name: NCCL_SOCKET_NTHREADS
+              value: "1"
+            - name: NCCL_NET_GDR_LEVEL
+              value: PIX
+            - name: NCCL_DYNAMIC_CHUNK_SIZE
+              value: "524288"
+            - name: NCCL_P2P_PXN_LEVEL
+              value: "0"
+            - name: NCCL_P2P_NET_CHUNKSIZE
+              value: "524288"
+            - name: NCCL_P2P_PCI_CHUNKSIZE
+              value: "524288"
+            - name: NCCL_P2P_NVL_CHUNKSIZE
+              value: "1048576"
+            - name: NCCL_BUFFSIZE
+              value: "4194304"
+            - name: NCCL_MAX_NCHANNELS
+              value: "8"
+            - name: NCCL_MIN_NCHANNELS
+              value: "8"
+            - name: CUDA_VISIBLE_DEVICES
+              value: "0,1,2,3,4,5,6,7"
+            {% endif %}
+            {% if k8s_enable_gpudirect_tcpxo %}
+            - name: LD_LIBRARY_PATH
+              value: /usr/local/nvidia/lib64
+            - name: NCCL_FASTRAK_LLCM_DEVICE_DIRECTORY
+              value: /dev/aperture_devices
+            - name: NCCL_FASTRAK_CTRL_DEV
+              value: eth0
+            - name: NCCL_FASTRAK_IFNAME
+              value: eth1,eth2,eth3,eth4,eth5,eth6,eth7,eth8
+            - name: NCCL_SOCKET_IFNAME
+              value: eth0
+            - name: NCCL_CROSS_NIC
+              value: "0"
+            - name: NCCL_ALGO
+              value: Ring,Tree
+            - name: NCCL_PROTO
+              value: Simple,LL128
+            - name: NCCL_MIN_NCHANNELS
+              value: "4"
+            - name: NCCL_TUNER_PLUGIN
+              value: libnccl-tuner.so
+            - name: NCCL_TUNER_CONFIG_PATH
+              value: /usr/local/nvidia/lib64/a3plus_tuner_config.textproto
+            - name: CUDA_VISIBLE_DEVICES
+              value: "0,1,2,3,4,5,6,7"
+            {% endif %}
+            {% if k8s_enable_gpudirect_rdma %}
+            - name: LD_LIBRARY_PATH
+              value: /usr/local/nvidia/lib64
+            - name: NCCL_NET
+              value: gIB
+            - name: NCCL_CROSS_NIC
+              value: "0"
+            - name: NCCL_NET_GDR_LEVEL
+              value: PIX
+            - name: NCCL_P2P_NET_CHUNKSIZE
+              value: "131072"
+            - name: NCCL_NVLS_CHUNKSIZE
+              value: "524288"
+            - name: NCCL_IB_ADAPTIVE_ROUTING
+              value: "1"
+            - name: NCCL_IB_QPS_PER_CONNECTION
+              value: "4"
+            - name: NCCL_IB_TC
+              value: "52"
+            - name: NCCL_IB_FIFO_TC
+              value: "84"
+            {% if k8s_enable_gpudirect_rdma_a4 %}
+            - name: NCCL_TUNER_CONFIG_PATH
+              value: /usr/local/gib/configs/tuner_config_a4.txtpb
+            {% else %}
+            - name: NCCL_TUNER_CONFIG_PATH
+              value: /usr/local/gib/configs/tuner_config_a3u.txtpb
+            {% endif %}
+            {% endif %}
             {% if k8s_fuse_device_required %}
             - name: FUSERMOUNT_SHARED_DIR
               value: {{k8s_fusermount_shared_dir}}
@@ -752,11 +963,27 @@ available_node_types:
           - name: secret-volume
             readOnly: true
             mountPath: "/etc/secret-volume"
-          # This volume allocates shared memory for Ray to use for its plasma
-          # object store. If you do not provide this, Ray will fall back to
-          # /tmp which cause slowdowns if is not a shared memory volume.
           - mountPath: /dev/shm
             name: dshm
+          {% if k8s_enable_gpudirect_tcpx %}
+          - name: tcpx-socket
+            mountPath: /tmp
+          - name: libraries
+            mountPath: /usr/local/nvidia/lib64
+            readOnly: true
+          {% endif %}
+          {% if k8s_enable_gpudirect_tcpxo %}
+          - name: libraries
+            mountPath: /usr/local/nvidia
+          - name: aperture-devices
+            mountPath: /dev/aperture_devices
+          {% endif %}
+          {% if k8s_enable_gpudirect_rdma %}
+          - name: library-dir-host
+            mountPath: /usr/local/nvidia
+          - name: gib
+            mountPath: /usr/local/gib
+          {% endif %}
           {% if high_availability %}
           - name: {{k8s_high_availability_deployment_volume_mount_name}}
             mountPath: {{k8s_high_availability_deployment_volume_mount_path}}
@@ -794,7 +1021,68 @@ available_node_types:
               add:
                 - IPC_LOCK
           {% endif %}
-
+        {% if k8s_enable_gpudirect_tcpx %}
+        # GPUDirect TCPX daemon sidecar container
+        - name: tcpx-daemon
+          image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpx/tcpgpudmarxd-dev:v2.0.11
+          imagePullPolicy: Always
+          command:
+            - /tcpgpudmarxd/build/app/tcpgpudmarxd
+            - --gpu_nic_preset
+            - a3vm
+            - --gpu_shmem_type
+            - fd
+            - --uds_path
+            - /run/tcpx
+            - --setup_param
+            - --verbose
+            - "128"
+            - "2"
+            - "0"
+          securityContext:
+            capabilities:
+              add:
+                - NET_ADMIN
+          volumeMounts:
+            - name: libraries
+              mountPath: /usr/local/nvidia/lib64
+              readOnly: true
+            - name: tcpx-socket
+              mountPath: /run/tcpx
+            - name: sys
+              mountPath: /hostsysfs
+            - name: proc-sys
+              mountPath: /hostprocsysfs
+          env:
+            - name: LD_LIBRARY_PATH
+              value: /usr/local/nvidia/lib64
+        {% endif %}
+        {% if k8s_enable_gpudirect_tcpxo %}
+        - name: tcpxo-daemon
+          image: us-docker.pkg.dev/gce-ai-infra/gpudirect-tcpxo/tcpgpudmarxd-dev:v1.0.17
+          imagePullPolicy: Always
+          command: ["/bin/sh", "-c"]
+          args:
+            - |
+              set -ex
+              chmod 755 /fts/entrypoint_rxdm_container.sh
+              /fts/entrypoint_rxdm_container.sh --num_hops=2 --num_nics=8 --uid= --alsologtostderr
+          securityContext:
+            capabilities:
+              add:
+                - NET_ADMIN
+                - NET_BIND_SERVICE
+          volumeMounts:
+            - name: libraries
+              mountPath: /usr/local/nvidia
+            - name: sys
+              mountPath: /hostsysfs
+            - name: proc-sys
+              mountPath: /hostprocsysfs
+          env:
+            - name: LD_LIBRARY_PATH
+              value: /usr/local/nvidia/lib64
+        {% endif %}
 
       {% if high_availability %}
       pvc_spec:
diff --git a/sky/users/permission.py b/sky/users/permission.py
index 9c7b2fdc6..87949f234 100644
--- a/sky/users/permission.py
+++ b/sky/users/permission.py
@@ -34,6 +34,11 @@ class PermissionService:
     """Permission service for SkyPilot API Server."""
 
     def __init__(self):
+        self.enforcer = None
+
+    def _lazy_initialize(self):
+        if self.enforcer is not None:
+            return
         with _policy_lock():
             global _enforcer_instance
             if _enforcer_instance is None:
@@ -73,7 +78,6 @@ class PermissionService:
 
     def _maybe_initialize_policies(self) -> None:
         """Initialize policies if they don't already exist."""
-        # TODO(zhwu): we should avoid running this on client side.
         logger.debug(f'Initializing policies in process: {os.getpid()}')
         self._load_policy_no_lock()
 
@@ -152,6 +156,7 @@ class PermissionService:
 
     def add_user_if_not_exists(self, user_id: str) -> None:
         """Add user role relationship."""
+        self._lazy_initialize()
         with _policy_lock():
             self._add_user_if_not_exists_no_lock(user_id)
 
@@ -171,6 +176,7 @@ class PermissionService:
 
     def delete_user(self, user_id: str) -> None:
         """Delete user role relationship."""
+        self._lazy_initialize()
         with _policy_lock():
             # Get current roles
             self._load_policy_no_lock()
@@ -184,6 +190,7 @@ class PermissionService:
 
     def update_role(self, user_id: str, new_role: str) -> None:
         """Update user role relationship."""
+        self._lazy_initialize()
         with _policy_lock():
             # Get current roles
             self._load_policy_no_lock()
@@ -216,6 +223,7 @@ class PermissionService:
         Returns:
             A list of role names that the user has.
         """
+        self._lazy_initialize()
         self._load_policy_no_lock()
         return self.enforcer.get_roles_for_user(user_id)
 
@@ -228,6 +236,7 @@ class PermissionService:
         # it is a hot path in every request. It is ok to have a stale policy,
         # as long as it is eventually consistent.
         # self._load_policy_no_lock()
+        self._lazy_initialize()
         return self.enforcer.enforce(user_id, path, method)
 
     def _load_policy_no_lock(self):
@@ -236,6 +245,7 @@ class PermissionService:
 
     def load_policy(self):
         """Load policy from storage with lock."""
+        self._lazy_initialize()
         with _policy_lock():
             self._load_policy_no_lock()
 
@@ -251,6 +261,7 @@ class PermissionService:
         For public workspaces, the permission is granted via a wildcard policy
         ('*').
         """
+        self._lazy_initialize()
         if os.getenv(constants.ENV_VAR_IS_SKYPILOT_SERVER) is None:
             # When it is not on API server, we allow all users to access all
             # workspaces, as the workspace check has been done on API server.
@@ -307,6 +318,7 @@ class PermissionService:
                    For public workspaces, this should be ['*'].
                    For private workspaces, this should be specific user IDs.
         """
+        self._lazy_initialize()
         with _policy_lock():
             for user in users:
                 logger.debug(f'Adding workspace policy: user={user}, '
@@ -324,6 +336,7 @@ class PermissionService:
                    For public workspaces, this should be ['*'].
                    For private workspaces, this should be specific user IDs.
         """
+        self._lazy_initialize()
         with _policy_lock():
             self._load_policy_no_lock()
             # Remove all existing policies for this workspace
@@ -337,6 +350,7 @@ class PermissionService:
 
     def remove_workspace_policy(self, workspace_name: str) -> None:
         """Remove workspace policy."""
+        self._lazy_initialize()
         with _policy_lock():
             self.enforcer.remove_filtered_policy(1, workspace_name)
             self.enforcer.save_policy()
diff --git a/sky/users/token_service.py b/sky/users/token_service.py
index 5935c94f1..12582038f 100644
--- a/sky/users/token_service.py
+++ b/sky/users/token_service.py
@@ -5,6 +5,7 @@ import datetime
 import hashlib
 import os
 import secrets
+import threading
 from typing import Any, Dict, Generator, Optional
 
 import filelock
@@ -44,12 +45,21 @@ class TokenService:
     """Service for managing JWT-based service account tokens."""
 
     def __init__(self):
-        self.secret_key = self._get_or_generate_secret()
+        self.secret_key = None
+        self.init_lock = threading.Lock()
+
+    def _lazy_initialize(self):
+        if self.secret_key is not None:
+            return
+        with self.init_lock:
+            if self.secret_key is not None:
+                return
+            self.secret_key = self._get_or_generate_secret()
 
     def _get_or_generate_secret(self) -> str:
         """Get JWT secret from database or generate a new one."""
-        with _jwt_secret_lock():
-            # Try to get from database (persistent across deployments)
+
+        def _get_secret_from_db():
             try:
                 db_secret = global_user_state.get_system_config(
                     JWT_SECRET_DB_KEY)
@@ -58,7 +68,17 @@ class TokenService:
                     return db_secret
             except Exception as e:  # pylint: disable=broad-except
                 logger.debug(f'Failed to get JWT secret from database: {e}')
+            return None
+
+        # Try to get from database (persistent across deployments)
+        token_from_db = _get_secret_from_db()
+        if token_from_db:
+            return token_from_db
 
+        with _jwt_secret_lock():
+            token_from_db = _get_secret_from_db()
+            if token_from_db:
+                return token_from_db
             # Generate a new secret and store in database
             new_secret = secrets.token_urlsafe(64)
             try:
@@ -91,6 +111,7 @@ class TokenService:
         Returns:
             Dict containing token info including the JWT token
         """
+        self._lazy_initialize()
         now = datetime.datetime.now(datetime.timezone.utc)
         token_id = secrets.token_urlsafe(12)  # Shorter ID for JWT
 
@@ -144,6 +165,7 @@ class TokenService:
         Returns:
             Decoded token payload or None if invalid
         """
+        self._lazy_initialize()
         if not token.startswith('sky_'):
             return None
 
diff --git a/sky/utils/schemas.py b/sky/utils/schemas.py
index a78cba4ef..fe2ecacdf 100644
--- a/sky/utils/schemas.py
+++ b/sky/utils/schemas.py
@@ -1106,6 +1106,9 @@ _CONTEXT_CONFIG_SCHEMA_KUBERNETES = {
             },
         },
     },
+    'remote_identity': {
+        'type': 'string',
+    }
 }
 
 


Test patch:
diff --git a/tests/smoke_tests/backward_compat/test_backward_compat.py b/tests/smoke_tests/backward_compat/test_backward_compat.py
index 1ec2fd3df..8d4316af6 100644
--- a/tests/smoke_tests/backward_compat/test_backward_compat.py
+++ b/tests/smoke_tests/backward_compat/test_backward_compat.py
@@ -66,20 +66,29 @@ class TestBackwardCompatibility:
             job_name=job_name, job_status=status, timeout=300)
 
     def _get_base_skylet_version(self) -> str:
-        """Get SKYLET_VERSION from the base environment directory."""
-        base_constants_file = self.BASE_SKY_DIR / 'sky' / 'skylet' / 'constants.py'
-        try:
-            with open(base_constants_file, 'r') as f:
-                content = f.read()
-                # Find the SKYLET_VERSION line
-                for line in content.split('\n'):
-                    if line.strip().startswith('SKYLET_VERSION = '):
-                        # Extract version string between quotes
-                        version = line.split('=')[1].strip().strip('\'"')
-                        return version
-            raise ValueError('SKYLET_VERSION not found in base constants.py')
-        except Exception as e:
-            raise ValueError(f'Failed to read base SKYLET_VERSION: {e}') from e
+        """Get SKYLET_VERSION from the base environment by running Python code."""
+        base_skylet_version = subprocess.run(
+            f'{self.ACTIVATE_BASE} && python -c "'
+            'from sky.skylet import constants; '
+            'print(constants.SKYLET_VERSION);'
+            '"',
+            shell=True,
+            check=False,
+            executable='/bin/bash',
+            text=True,
+            capture_output=True)
+
+        if base_skylet_version.returncode != 0:
+            raise ValueError(f'Failed to get base SKYLET_VERSION. '
+                             f'Return code: {base_skylet_version.returncode}, '
+                             f'stderr: {base_skylet_version.stderr}, '
+                             f'stdout: {base_skylet_version.stdout}')
+
+        version = base_skylet_version.stdout.strip().split('\n')[-1]
+        if not version:
+            raise ValueError('SKYLET_VERSION is empty')
+
+        return version
 
     @pytest.fixture(scope="session", autouse=True)
     def session_setup(self, request):
diff --git a/tests/test_failover.py b/tests/test_failover.py
index 404088873..1f8a8e09c 100644
--- a/tests/test_failover.py
+++ b/tests/test_failover.py
@@ -28,7 +28,7 @@ def _mock_db_conn(tmp_path, monkeypatch):
     monkeypatch.setattr(global_user_state, '_SQLALCHEMY_ENGINE',
                         sqlalchemy_engine)
 
-    global_user_state.create_table()
+    global_user_state.create_table(sqlalchemy_engine)
 
 
 @pytest.mark.parametrize('enable_all_clouds', [[sky.AWS()]], indirect=True)
diff --git a/tests/test_jobs.py b/tests/test_jobs.py
index f2e337297..1136d9878 100644
--- a/tests/test_jobs.py
+++ b/tests/test_jobs.py
@@ -19,7 +19,7 @@ def _mock_db_conn(tmp_path, monkeypatch):
     monkeypatch.setattr(global_user_state, '_SQLALCHEMY_ENGINE',
                         sqlalchemy_engine)
 
-    global_user_state.create_table()
+    global_user_state.create_table(sqlalchemy_engine)
 
 
 @pytest.fixture
diff --git a/tests/test_jobs_and_serve.py b/tests/test_jobs_and_serve.py
index d13431ef2..e9e8f2dfa 100644
--- a/tests/test_jobs_and_serve.py
+++ b/tests/test_jobs_and_serve.py
@@ -47,7 +47,7 @@ def _mock_db_conn(tmp_path, monkeypatch):
     monkeypatch.setattr(global_user_state, '_SQLALCHEMY_ENGINE',
                         sqlalchemy_engine)
 
-    global_user_state.create_table()
+    global_user_state.create_table(sqlalchemy_engine)
 
 
 def _generate_tmp_yaml(tmp_path, filename: str) -> str:
diff --git a/tests/unit_tests/test_core.py b/tests/unit_tests/test_core.py
new file mode 100644
index 000000000..65eb075e2
--- /dev/null
+++ b/tests/unit_tests/test_core.py
@@ -0,0 +1,35 @@
+from unittest import mock
+
+from sky.backends.cloud_vm_ray_backend import CloudVmRayResourceHandle
+from sky.core import cancel
+from sky.utils import common_utils
+
+
+@mock.patch('sky.backends.backend_utils.check_cluster_available')
+@mock.patch('sky.backends.cloud_vm_ray_backend.CloudVmRayBackend.cancel_jobs')
+def test_cancel_jobs_for_current_user(mock_cancel_jobs,
+                                      mock_check_cluster_available) -> None:
+    mock_handle = mock.create_autospec(CloudVmRayResourceHandle, instance=True)
+    mock_check_cluster_available.return_value = mock_handle
+    cancel('test-cluster', all=True)
+    mock_cancel_jobs.assert_called_once_with(
+        mock_handle,
+        None,
+        cancel_all=True,
+        user_hash=common_utils.get_current_user().id,
+    )
+
+
+@mock.patch('sky.backends.backend_utils.check_cluster_available')
+@mock.patch('sky.backends.cloud_vm_ray_backend.CloudVmRayBackend.cancel_jobs')
+def test_cancel_jobs_for_all_users(mock_cancel_jobs,
+                                   mock_check_cluster_available) -> None:
+    mock_handle = mock.create_autospec(CloudVmRayResourceHandle, instance=True)
+    mock_check_cluster_available.return_value = mock_handle
+    cancel('test-cluster', all_users=True)
+    mock_cancel_jobs.assert_called_once_with(
+        mock_handle,
+        None,
+        cancel_all=True,
+        user_hash=None,
+    )
diff --git a/tests/unit_tests/test_sky/clouds/test_kubernetes.py b/tests/unit_tests/test_sky/clouds/test_kubernetes.py
index 2e9653e91..c48dab584 100644
--- a/tests/unit_tests/test_sky/clouds/test_kubernetes.py
+++ b/tests/unit_tests/test_sky/clouds/test_kubernetes.py
@@ -221,8 +221,8 @@ class TestKubernetesExistingAllowedContexts(unittest.TestCase):
     @patch('sky.provision.kubernetes.utils.get_all_kube_context_names')
     @patch('sky.skypilot_config.get_workspace_cloud')
     @patch('sky.skypilot_config.get_nested')
-    @patch('sky.provision.kubernetes.utils.'
-           'get_current_kube_config_context_name')
+    @patch('sky.provision.kubernetes.utils.get_current_kube_config_context_name'
+          )
     @patch('sky.provision.kubernetes.utils.is_incluster_config_available')
     def test_empty_result_when_no_contexts_found(
             self, mock_is_incluster_available, mock_get_current_context,
@@ -296,9 +296,6 @@ class TestKubernetesSecurityContextMerging(unittest.TestCase):
         self.region = mock.MagicMock()
         self.region.name = "test-context"
 
-    @patch(
-        'sky.clouds.kubernetes.Kubernetes._cluster_supports_high_performance_networking'
-    )
     @patch('sky.provision.kubernetes.utils.get_kubernetes_nodes')
     @patch('sky.provision.kubernetes.utils.get_current_kube_config_context_name'
           )
@@ -310,16 +307,23 @@ class TestKubernetesSecurityContextMerging(unittest.TestCase):
     @patch('sky.provision.kubernetes.network_utils.get_port_mode')
     @patch('sky.provision.kubernetes.network_utils.get_networking_mode')
     @patch('sky.catalog.get_image_id_from_tag')
+    @patch('sky.clouds.kubernetes.Kubernetes._detect_network_type')
     def test_ipc_lock_capability_enabled_with_user_security_context(
-            self, mock_get_image, mock_get_networking_mode, mock_get_port_mode,
+            self, mock_detect_network_type, mock_get_image,
+            mock_get_networking_mode, mock_get_port_mode,
             mock_get_workspace_cloud, mock_get_cloud_config_value,
             mock_is_exec_auth, mock_get_accelerator_label_keys,
-            mock_get_namespace, mock_get_current_context, mock_get_k8s_nodes,
-            mock_cluster_supports_high_perf):
+            mock_get_namespace, mock_get_current_context, mock_get_k8s_nodes):
         """Test that IPC_LOCK capability is enabled when network tier is BEST and cluster supports it."""
 
         # Setup mocks
-        mock_cluster_supports_high_perf.return_value = True
+        from sky.provision.kubernetes.utils import (
+            KubernetesHighPerformanceNetworkType)
+        mock_cluster_type = mock.MagicMock()
+        mock_cluster_type.supports_high_performance_networking.return_value = True
+        mock_cluster_type.requires_ipc_lock_capability.return_value = True
+        mock_detect_network_type.return_value = (mock_cluster_type, '')
+
         mock_get_current_context.return_value = "test-context"
         mock_get_namespace.return_value = "default"
         mock_get_accelerator_label_keys.return_value = []
@@ -360,15 +364,13 @@ class TestKubernetesSecurityContextMerging(unittest.TestCase):
         # Verify that k8s_ipc_lock_capability is True
         self.assertTrue(deploy_vars['k8s_ipc_lock_capability'])
 
-        # Verify that NCCL environment variables are set
+        # For GCP clusters, NCCL environment variables are set in the template, not Python code
+        # We should not expect them in k8s_env_vars for GCP cluster types
         k8s_env_vars = deploy_vars['k8s_env_vars']
-        self.assertIn('NCCL_IB_HCA', k8s_env_vars)
-        self.assertEqual(k8s_env_vars['NCCL_IB_HCA'], 'mlx5')
-        self.assertIn('UCX_NET_DEVICES', k8s_env_vars)
+        self.assertIn('SKYPILOT_IN_CLUSTER_CONTEXT_NAME', k8s_env_vars)
+        self.assertEqual(k8s_env_vars['SKYPILOT_IN_CLUSTER_CONTEXT_NAME'],
+                         'test-context')
 
-    @patch(
-        'sky.clouds.kubernetes.Kubernetes._cluster_supports_high_performance_networking'
-    )
     @patch('sky.provision.kubernetes.utils.get_kubernetes_nodes')
     @patch('sky.provision.kubernetes.utils.get_current_kube_config_context_name'
           )
@@ -380,16 +382,21 @@ class TestKubernetesSecurityContextMerging(unittest.TestCase):
     @patch('sky.provision.kubernetes.network_utils.get_port_mode')
     @patch('sky.provision.kubernetes.network_utils.get_networking_mode')
     @patch('sky.catalog.get_image_id_from_tag')
+    @patch('sky.clouds.kubernetes.Kubernetes._detect_network_type')
     def test_ipc_lock_capability_disabled_when_no_high_perf_networking(
-            self, mock_get_image, mock_get_networking_mode, mock_get_port_mode,
+            self, mock_detect_network_type, mock_get_image,
+            mock_get_networking_mode, mock_get_port_mode,
             mock_get_workspace_cloud, mock_get_cloud_config_value,
             mock_is_exec_auth, mock_get_accelerator_label_keys,
-            mock_get_namespace, mock_get_current_context, mock_get_k8s_nodes,
-            mock_cluster_supports_high_perf):
+            mock_get_namespace, mock_get_current_context, mock_get_k8s_nodes):
         """Test that IPC_LOCK capability is disabled when cluster doesn't support high performance networking."""
 
         # Setup mocks - cluster does NOT support high performance networking
-        mock_cluster_supports_high_perf.return_value = False
+        from sky.provision.kubernetes.utils import (
+            KubernetesHighPerformanceNetworkType)
+        mock_detect_network_type.return_value = (
+            KubernetesHighPerformanceNetworkType.NONE, '')
+
         mock_get_current_context.return_value = "test-context"
         mock_get_namespace.return_value = "default"
         mock_get_accelerator_label_keys.return_value = []
@@ -430,14 +437,11 @@ class TestKubernetesSecurityContextMerging(unittest.TestCase):
         # Verify that k8s_ipc_lock_capability is False
         self.assertFalse(deploy_vars['k8s_ipc_lock_capability'])
 
-        # Verify that NCCL environment variables are not set
+        # Verify that high performance networking environment variables are not set
         k8s_env_vars = deploy_vars['k8s_env_vars']
         self.assertNotIn('NCCL_IB_HCA', k8s_env_vars)
         self.assertNotIn('UCX_NET_DEVICES', k8s_env_vars)
 
-    @patch(
-        'sky.clouds.kubernetes.Kubernetes._cluster_supports_high_performance_networking'
-    )
     @patch('sky.provision.kubernetes.utils.get_kubernetes_nodes')
     @patch('sky.provision.kubernetes.utils.get_current_kube_config_context_name'
           )
@@ -449,20 +453,25 @@ class TestKubernetesSecurityContextMerging(unittest.TestCase):
     @patch('sky.provision.kubernetes.network_utils.get_port_mode')
     @patch('sky.provision.kubernetes.network_utils.get_networking_mode')
     @patch('sky.catalog.get_image_id_from_tag')
+    @patch('sky.clouds.kubernetes.Kubernetes._detect_network_type')
     def test_ipc_lock_capability_disabled_when_network_tier_not_best(
-            self, mock_get_image, mock_get_networking_mode, mock_get_port_mode,
+            self, mock_detect_network_type, mock_get_image,
+            mock_get_networking_mode, mock_get_port_mode,
             mock_get_workspace_cloud, mock_get_cloud_config_value,
             mock_is_exec_auth, mock_get_accelerator_label_keys,
-            mock_get_namespace, mock_get_current_context, mock_get_k8s_nodes,
-            mock_cluster_supports_high_perf):
+            mock_get_namespace, mock_get_current_context, mock_get_k8s_nodes):
         """Test that IPC_LOCK capability is disabled when network tier is not BEST."""
 
         # Modify resources to not use BEST network tier
         from sky.utils import resources_utils
         self.resources.network_tier = resources_utils.NetworkTier.STANDARD
 
-        # Setup mocks
-        mock_cluster_supports_high_perf.return_value = True  # Cluster supports it but network tier is not BEST
+        # Setup mocks - when network tier is not BEST, _detect_network_type returns NONE
+        from sky.provision.kubernetes.utils import (
+            KubernetesHighPerformanceNetworkType)
+        mock_detect_network_type.return_value = (
+            KubernetesHighPerformanceNetworkType.NONE, '')
+
         mock_get_current_context.return_value = "test-context"
         mock_get_namespace.return_value = "default"
         mock_get_accelerator_label_keys.return_value = []
@@ -503,9 +512,6 @@ class TestKubernetesSecurityContextMerging(unittest.TestCase):
         # Verify that k8s_ipc_lock_capability is False
         self.assertFalse(deploy_vars['k8s_ipc_lock_capability'])
 
-    @patch(
-        'sky.clouds.kubernetes.Kubernetes._cluster_supports_high_performance_networking'
-    )
     @patch('sky.provision.kubernetes.utils.get_kubernetes_nodes')
     @patch('sky.provision.kubernetes.utils.get_current_kube_config_context_name'
           )
@@ -519,14 +525,15 @@ class TestKubernetesSecurityContextMerging(unittest.TestCase):
     @patch('sky.provision.kubernetes.network_utils.get_port_mode')
     @patch('sky.provision.kubernetes.network_utils.get_networking_mode')
     @patch('sky.catalog.get_image_id_from_tag')
+    @patch('sky.clouds.kubernetes.Kubernetes._detect_network_type')
     def test_nebius_network_tier_with_gpu_environment_variables(
-            self, mock_get_image, mock_get_networking_mode, mock_get_port_mode,
+            self, mock_detect_network_type, mock_get_image,
+            mock_get_networking_mode, mock_get_port_mode,
             mock_get_workspace_cloud, mock_get_cloud_config_value,
             mock_is_exec_auth, mock_get_gpu_resource_key,
             mock_get_accelerator_label_key_values,
             mock_get_accelerator_label_keys, mock_get_namespace,
-            mock_get_current_context, mock_get_k8s_nodes,
-            mock_cluster_supports_high_perf):
+            mock_get_current_context, mock_get_k8s_nodes):
         """Test Nebius network tier with GPU functionality and environment variables."""
 
         # Create resources with GPU and BEST network tier (mimicking H100:8 --network-tier best)
@@ -547,7 +554,11 @@ class TestKubernetesSecurityContextMerging(unittest.TestCase):
         gpu_resources.network_tier = resources_utils.NetworkTier.BEST
 
         # Setup mocks - cluster supports high performance networking (Nebius)
-        mock_cluster_supports_high_perf.return_value = True
+        from sky.provision.kubernetes.utils import (
+            KubernetesHighPerformanceNetworkType)
+        mock_detect_network_type.return_value = (
+            KubernetesHighPerformanceNetworkType.NEBIUS, '')
+
         mock_get_current_context.return_value = "nebius-context"
         mock_get_namespace.return_value = "default"
         mock_get_accelerator_label_keys.return_value = []
diff --git a/tests/unit_tests/test_sky/users/test_permission.py b/tests/unit_tests/test_sky/users/test_permission.py
index 953abaad4..4ef0273b2 100644
--- a/tests/unit_tests/test_sky/users/test_permission.py
+++ b/tests/unit_tests/test_sky/users/test_permission.py
@@ -65,6 +65,7 @@ class TestPermissionService:
         with mock.patch.object(permission.PermissionService,
                                '_maybe_initialize_policies'):
             service = permission.PermissionService()
+            service._lazy_initialize()
 
             # Verify SQLAlchemy adapter was created with the correct engine
             mock_adapter_class.assert_called_once_with(mock_engine)
@@ -122,30 +123,26 @@ class TestPermissionService:
         # Mock users
         mock_get_users.return_value = mock_users
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
-            service._maybe_initialize_policies()
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
+        service._maybe_initialize_policies()
 
-            # Verify policies were added for user role blocklist
-            mock_enforcer.add_policy.assert_any_call('user',
-                                                     '/workspaces/config',
-                                                     'POST')
+        # Verify policies were added for user role blocklist
+        mock_enforcer.add_policy.assert_any_call('user', '/workspaces/config',
+                                                 'POST')
 
-            # Verify workspace policies were added
-            mock_enforcer.add_policy.assert_any_call('*', 'default', '*')
-            mock_enforcer.add_policy.assert_any_call('user1', 'private-ws', '*')
-            mock_enforcer.add_policy.assert_any_call('user2', 'private-ws', '*')
+        # Verify workspace policies were added
+        mock_enforcer.add_policy.assert_any_call('*', 'default', '*')
+        mock_enforcer.add_policy.assert_any_call('user1', 'private-ws', '*')
+        mock_enforcer.add_policy.assert_any_call('user2', 'private-ws', '*')
 
-            # Verify users were assigned default roles
-            for user in mock_users:
-                mock_enforcer.add_grouping_policy.assert_any_call(
-                    user.id, rbac.get_default_role())
+        # Verify users were assigned default roles
+        for user in mock_users:
+            mock_enforcer.add_grouping_policy.assert_any_call(
+                user.id, rbac.get_default_role())
 
-            # Verify policy was saved
-            mock_enforcer.save_policy.assert_called()
+        # Verify policy was saved
+        mock_enforcer.save_policy.assert_called()
 
     @mock.patch('sky.users.permission._policy_lock')
     @mock.patch('sky.users.rbac.get_workspace_policy_permissions')
@@ -198,17 +195,14 @@ class TestPermissionService:
         # Mock users
         mock_get_users.return_value = mock_users
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
-            service._maybe_initialize_policies()
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
+        service._maybe_initialize_policies()
 
-            # Verify no new policies were added (since they already exist)
-            mock_enforcer.remove_filtered_policy.assert_not_called()
-            # save_policy should not be called if no updates were made
-            # (users already have roles, policies already exist)
+        # Verify no new policies were added (since they already exist)
+        mock_enforcer.remove_filtered_policy.assert_not_called()
+        # save_policy should not be called if no updates were made
+        # (users already have roles, policies already exist)
 
     def test_add_user_if_not_exists_new_user(self):
         """Test adding a new user that doesn't exist."""
@@ -216,17 +210,14 @@ class TestPermissionService:
         mock_enforcer.get_roles_for_user.return_value = []  # No existing roles
         mock_enforcer.add_grouping_policy.return_value = True
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
 
-            result = service._add_user_if_not_exists_no_lock('new_user')
+        result = service._add_user_if_not_exists_no_lock('new_user')
 
-            assert result is True
-            mock_enforcer.add_grouping_policy.assert_called_once_with(
-                'new_user', rbac.get_default_role())
+        assert result is True
+        mock_enforcer.add_grouping_policy.assert_called_once_with(
+            'new_user', rbac.get_default_role())
 
     def test_add_user_if_not_exists_existing_user(self):
         """Test adding a user that already exists."""
@@ -235,16 +226,13 @@ class TestPermissionService:
             'user'
         ]  # User already has roles
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
 
-            result = service._add_user_if_not_exists_no_lock('existing_user')
+        result = service._add_user_if_not_exists_no_lock('existing_user')
 
-            assert result is False
-            mock_enforcer.add_grouping_policy.assert_not_called()
+        assert result is False
+        mock_enforcer.add_grouping_policy.assert_not_called()
 
     @mock.patch('sky.users.permission._policy_lock')
     def test_update_role_new_role(self, mock_policy_lock):
@@ -257,21 +245,18 @@ class TestPermissionService:
         mock_enforcer.remove_grouping_policy.return_value = True
         mock_enforcer.add_grouping_policy.return_value = True
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
-            service._load_policy_no_lock = mock.Mock()
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
+        service._load_policy_no_lock = mock.Mock()
 
-            service.update_role('user1', 'admin')
+        service.update_role('user1', 'admin')
 
-            # Verify old role was removed and new role was added
-            mock_enforcer.remove_grouping_policy.assert_called_once_with(
-                'user1', 'user')
-            mock_enforcer.add_grouping_policy.assert_called_once_with(
-                'user1', 'admin')
-            mock_enforcer.save_policy.assert_called_once()
+        # Verify old role was removed and new role was added
+        mock_enforcer.remove_grouping_policy.assert_called_once_with(
+            'user1', 'user')
+        mock_enforcer.add_grouping_policy.assert_called_once_with(
+            'user1', 'admin')
+        mock_enforcer.save_policy.assert_called_once()
 
     @mock.patch('sky.users.permission._policy_lock')
     def test_update_role_same_role(self, mock_policy_lock):
@@ -283,54 +268,44 @@ class TestPermissionService:
         mock_enforcer.get_roles_for_user.return_value = ['admin'
                                                         ]  # Current role
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
-            service._load_policy_no_lock = mock.Mock()
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
+        service._load_policy_no_lock = mock.Mock()
 
-            service.update_role('user1', 'admin')
+        service.update_role('user1', 'admin')
 
-            # Verify no changes were made
-            mock_enforcer.remove_grouping_policy.assert_not_called()
-            mock_enforcer.add_grouping_policy.assert_not_called()
-            mock_enforcer.save_policy.assert_not_called()
+        # Verify no changes were made
+        mock_enforcer.remove_grouping_policy.assert_not_called()
+        mock_enforcer.add_grouping_policy.assert_not_called()
+        mock_enforcer.save_policy.assert_not_called()
 
     def test_get_user_roles(self):
         """Test getting user roles."""
         mock_enforcer = mock.Mock()
         mock_enforcer.get_roles_for_user.return_value = ['admin', 'user']
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
-            service._load_policy_no_lock = mock.Mock()
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
+        service._load_policy_no_lock = mock.Mock()
 
-            roles = service.get_user_roles('user1')
+        roles = service.get_user_roles('user1')
 
-            assert roles == ['admin', 'user']
-            mock_enforcer.get_roles_for_user.assert_called_once_with('user1')
+        assert roles == ['admin', 'user']
+        mock_enforcer.get_roles_for_user.assert_called_once_with('user1')
 
     def test_check_endpoint_permission(self):
         """Test checking endpoint permissions."""
         mock_enforcer = mock.Mock()
         mock_enforcer.enforce.return_value = True
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
 
-            result = service.check_endpoint_permission('user1', '/api/test',
-                                                       'GET')
+        result = service.check_endpoint_permission('user1', '/api/test', 'GET')
 
-            assert result is True
-            mock_enforcer.enforce.assert_called_once_with(
-                'user1', '/api/test', 'GET')
+        assert result is True
+        mock_enforcer.enforce.assert_called_once_with('user1', '/api/test',
+                                                      'GET')
 
     def test_check_workspace_permission_non_server(self):
         """Test workspace permission check when not on API server."""
@@ -338,16 +313,13 @@ class TestPermissionService:
         if constants.ENV_VAR_IS_SKYPILOT_SERVER in os.environ:
             del os.environ[constants.ENV_VAR_IS_SKYPILOT_SERVER]
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
+        service = permission.PermissionService()
+        service._lazy_initialize()
 
-            result = service.check_workspace_permission('user1',
-                                                        'test-workspace')
+        result = service.check_workspace_permission('user1', 'test-workspace')
 
-            # Should always return True when not on API server
-            assert result is True
+        # Should always return True when not on API server
+        assert result is True
 
     def test_check_workspace_permission_admin_user(self):
         """Test workspace permission check for admin user."""
@@ -355,18 +327,15 @@ class TestPermissionService:
 
         mock_enforcer = mock.Mock()
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
-            service.get_user_roles = mock.Mock(return_value=['admin'])
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
+        service.get_user_roles = mock.Mock(return_value=['admin'])
 
-            result = service.check_workspace_permission('admin_user',
-                                                        'test-workspace')
+        result = service.check_workspace_permission('admin_user',
+                                                    'test-workspace')
 
-            # Admin should always have access
-            assert result is True
+        # Admin should always have access
+        assert result is True
 
     def test_check_workspace_permission_regular_user(self):
         """Test workspace permission check for regular user."""
@@ -375,19 +344,15 @@ class TestPermissionService:
         mock_enforcer = mock.Mock()
         mock_enforcer.enforce.return_value = True
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
-            service.get_user_roles = mock.Mock(return_value=['user'])
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
+        service.get_user_roles = mock.Mock(return_value=['user'])
 
-            result = service.check_workspace_permission('user1',
-                                                        'test-workspace')
+        result = service.check_workspace_permission('user1', 'test-workspace')
 
-            assert result is True
-            mock_enforcer.enforce.assert_called_once_with(
-                'user1', 'test-workspace', '*')
+        assert result is True
+        mock_enforcer.enforce.assert_called_once_with('user1', 'test-workspace',
+                                                      '*')
 
     @mock.patch('sky.users.permission._policy_lock')
     def test_add_workspace_policy(self, mock_policy_lock):
@@ -398,20 +363,15 @@ class TestPermissionService:
         mock_enforcer = mock.Mock()
         mock_enforcer.add_policy.return_value = True
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
 
-            service.add_workspace_policy('test-workspace', ['user1', 'user2'])
+        service.add_workspace_policy('test-workspace', ['user1', 'user2'])
 
-            # Verify policies were added for each user
-            mock_enforcer.add_policy.assert_any_call('user1', 'test-workspace',
-                                                     '*')
-            mock_enforcer.add_policy.assert_any_call('user2', 'test-workspace',
-                                                     '*')
-            mock_enforcer.save_policy.assert_called_once()
+        # Verify policies were added for each user
+        mock_enforcer.add_policy.assert_any_call('user1', 'test-workspace', '*')
+        mock_enforcer.add_policy.assert_any_call('user2', 'test-workspace', '*')
+        mock_enforcer.save_policy.assert_called_once()
 
     @mock.patch('sky.users.permission.filelock.FileLock')
     def test_policy_lock_context_manager(self, mock_filelock):
@@ -510,9 +470,11 @@ class TestPermissionServiceMultiProcess:
                                '_maybe_initialize_policies'):
             # Create first instance
             service1 = permission.PermissionService()
+            service1._lazy_initialize()
 
             # Create second instance
             service2 = permission.PermissionService()
+            service2._lazy_initialize()
 
             # Both should share the same enforcer
             assert service1.enforcer is service2.enforcer
@@ -579,13 +541,10 @@ class TestPermissionServiceMultiProcess:
 
         def create_service():
             try:
-                with mock.patch.object(permission.PermissionService,
-                                       '__init__',
-                                       return_value=None):
-                    service = permission.PermissionService()
-                    service.enforcer = mock_enforcer
-                    service._maybe_initialize_policies()
-                    services.append(service)
+                service = permission.PermissionService()
+                service.enforcer = mock_enforcer
+                service._maybe_initialize_policies()
+                services.append(service)
             except Exception as e:
                 errors.append(e)
 
@@ -674,26 +633,23 @@ class TestPermissionServiceMultiProcess:
         mock_get_workspace_perms.return_value = {}
         mock_get_users.return_value = mock_users
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
 
-            # Initialize policies multiple times
-            service._maybe_initialize_policies()
-            service._maybe_initialize_policies()
-            service._maybe_initialize_policies()
+        # Initialize policies multiple times
+        service._maybe_initialize_policies()
+        service._maybe_initialize_policies()
+        service._maybe_initialize_policies()
 
-            # Each user should only be added once (3 users total)
-            assert len(grouping_policy_calls) == len(mock_users)
+        # Each user should only be added once (3 users total)
+        assert len(grouping_policy_calls) == len(mock_users)
 
-            # Verify each user was added exactly once
-            expected_calls = {
-                (user.id, rbac.get_default_role()) for user in mock_users
-            }
-            actual_calls = set(grouping_policy_calls)
-            assert actual_calls == expected_calls
+        # Verify each user was added exactly once
+        expected_calls = {
+            (user.id, rbac.get_default_role()) for user in mock_users
+        }
+        actual_calls = set(grouping_policy_calls)
+        assert actual_calls == expected_calls
 
     @mock.patch('sky.users.permission._policy_lock')
     def test_concurrent_policy_updates_use_lock(self, mock_policy_lock):
@@ -718,20 +674,17 @@ class TestPermissionServiceMultiProcess:
         mock_enforcer.add_policy.return_value = True
         mock_enforcer.save_policy.return_value = True
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
 
-            # Test a single workspace policy update to verify lock is used
-            service.add_workspace_policy('test-workspace', ['user1'])
+        # Test a single workspace policy update to verify lock is used
+        service.add_workspace_policy('test-workspace', ['user1'])
 
-            # Verify that lock was acquired and released
-            assert 'enter' in lock_calls
-            assert 'exit' in lock_calls
-            assert len([call for call in lock_calls if call == 'enter']) == 1
-            assert len([call for call in lock_calls if call == 'exit']) == 1
+        # Verify that lock was acquired and released
+        assert 'enter' in lock_calls
+        assert 'exit' in lock_calls
+        assert len([call for call in lock_calls if call == 'enter']) == 1
+        assert len([call for call in lock_calls if call == 'exit']) == 1
 
     @mock.patch('sky.users.permission._policy_lock')
     @mock.patch('sky.users.rbac.get_workspace_policy_permissions')
@@ -775,18 +728,15 @@ class TestPermissionServiceMultiProcess:
 
         mock_get_users.return_value = mock_users
 
-        with mock.patch.object(permission.PermissionService,
-                               '__init__',
-                               return_value=None):
-            service = permission.PermissionService()
-            service.enforcer = mock_enforcer
+        service = permission.PermissionService()
+        service.enforcer = mock_enforcer
 
-            # Initialize policies when they already exist
-            service._maybe_initialize_policies()
+        # Initialize policies when they already exist
+        service._maybe_initialize_policies()
 
-            # Since policies already exist and users already have roles,
-            # no new policies or grouping policies should be added
-            mock_enforcer.remove_filtered_policy.assert_not_called()
+        # Since policies already exist and users already have roles,
+        # no new policies or grouping policies should be added
+        mock_enforcer.remove_filtered_policy.assert_not_called()
 
-            # The save_policy might not be called if no changes were made
-            # We mainly want to ensure no duplicate additions occurred
+        # The save_policy might not be called if no changes were made
+        # We mainly want to ensure no duplicate additions occurred
diff --git a/tests/unit_tests/test_sky/users/test_token_service.py b/tests/unit_tests/test_sky/users/test_token_service.py
index ca93ea7dc..6d83c3512 100644
--- a/tests/unit_tests/test_sky/users/test_token_service.py
+++ b/tests/unit_tests/test_sky/users/test_token_service.py
@@ -21,6 +21,9 @@ class TestTokenService:
 
             service = token_service.TokenService()
 
+            assert service.secret_key is None
+
+            service._lazy_initialize()
             assert service.secret_key is not None
             assert len(service.secret_key) >= 32  # Ensure sufficient entropy
             mock_global_state.set_system_config.assert_called_once()
@@ -33,7 +36,7 @@ class TestTokenService:
             mock_global_state.get_system_config.return_value = existing_secret
 
             service = token_service.TokenService()
-
+            service._lazy_initialize()
             assert service.secret_key == existing_secret
 
     def test_create_token_basic(self):
diff --git a/tests/unit_tests/test_sky/utils/test_schemas.py b/tests/unit_tests/test_sky/utils/test_schemas.py
index 08351d9bd..e97d6a6d6 100644
--- a/tests/unit_tests/test_sky/utils/test_schemas.py
+++ b/tests/unit_tests/test_sky/utils/test_schemas.py
@@ -423,5 +423,24 @@ class TestWorkspaceSchema(unittest.TestCase):
                                     schema=self.workspaces_schema)
 
 
+class TestKubernetesSchema(unittest.TestCase):
+    """Tests for the kubernetes schema in schemas.py."""
+
+    def setUp(self):
+        self.config_schema = schemas.get_config_schema()
+        self.k8s_schema = self.config_schema['properties']['kubernetes']
+
+    def test_context_configs_allows_remote_identity(self):
+        """Test that context_configs allows remote_identity."""
+        valid_config = {
+            'context_configs': {
+                'my-context': {
+                    'remote_identity': 'my-service-account'
+                }
+            }
+        }
+        jsonschema.validate(instance=valid_config, schema=self.k8s_schema)
+
+
 if __name__ == "__main__":
     unittest.main()