skypilot-nightly 1.0.0.dev20250710__py3-none-any.whl → 1.0.0.dev20250712__py3-none-any.whl

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:

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()

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
 

sky/utils/schemas.py

@@ -1106,6 +1106,9 @@ _CONTEXT_CONFIG_SCHEMA_KUBERNETES = {
             },
         },
     },
+    'remote_identity': {
+        'type': 'string',
+    }
 }
 
 

{skypilot_nightly-1.0.0.dev20250710.dist-info → skypilot_nightly-1.0.0.dev20250712.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: skypilot-nightly
-Version: 1.0.0.dev20250710
+Version: 1.0.0.dev20250712
 Summary: SkyPilot: Run AI on Any Infra — Unified, Faster, Cheaper.
 Author: SkyPilot Team
 License: Apache 2.0
@@ -67,6 +67,7 @@ Requires-Dist: azure-mgmt-network>=27.0.0; extra == "azure"
 Requires-Dist: azure-mgmt-compute>=33.0.0; extra == "azure"
 Requires-Dist: azure-storage-blob>=12.23.1; extra == "azure"
 Requires-Dist: msgraph-sdk; extra == "azure"
+Requires-Dist: msrestazure; extra == "azure"
 Requires-Dist: ray[default]!=2.6.0,>=2.2.0; extra == "azure"
 Provides-Extra: gcp
 Requires-Dist: google-api-python-client>=2.69.0; extra == "gcp"
@@ -140,6 +141,7 @@ Requires-Dist: azure-mgmt-network>=27.0.0; extra == "all"
 Requires-Dist: azure-mgmt-compute>=33.0.0; extra == "all"
 Requires-Dist: azure-storage-blob>=12.23.1; extra == "all"
 Requires-Dist: msgraph-sdk; extra == "all"
+Requires-Dist: msrestazure; extra == "all"
 Requires-Dist: ray[default]!=2.6.0,>=2.2.0; extra == "all"
 Requires-Dist: google-api-python-client>=2.69.0; extra == "all"
 Requires-Dist: google-cloud-storage; extra == "all"