skypilot-nightly 1.0.0.dev20250709__py3-none-any.whl → 1.0.0.dev20250711__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

@@ -15,6 +15,7 @@ from sky import sky_logging
 from sky.skylet import constants
 from sky.users import rbac
 from sky.utils import common_utils
+from sky.utils import db_utils
 
 logging.getLogger('casbin.policy').setLevel(sky_logging.ERROR)
 logging.getLogger('casbin.role').setLevel(sky_logging.ERROR)
@@ -33,11 +34,18 @@ 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:
                 _enforcer_instance = self
                 engine = global_user_state.initialize_and_get_db()
+                db_utils.add_tables_to_db_sqlalchemy(
+                    sqlalchemy_adapter.Base.metadata, engine)
                 adapter = sqlalchemy_adapter.Adapter(engine)
                 model_path = os.path.join(os.path.dirname(__file__),
                                           'model.conf')
@@ -70,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()
 
@@ -149,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)
 
@@ -168,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()
@@ -181,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()
@@ -213,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)
 
@@ -225,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):
@@ -233,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()
 
@@ -248,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.
@@ -304,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}, '
@@ -321,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
@@ -334,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/common_utils.py

@@ -11,6 +11,7 @@ import platform
 import random
 import re
 import socket
+import subprocess
 import sys
 import time
 import typing
@@ -87,6 +88,18 @@ def generate_user_hash() -> str:
     return user_hash
 
 
+def get_git_commit(path: Optional[str] = None) -> Optional[str]:
+    try:
+        result = subprocess.run(['git', 'rev-parse', 'HEAD'],
+                                capture_output=True,
+                                text=True,
+                                cwd=path,
+                                check=True)
+        return result.stdout.strip()
+    except subprocess.CalledProcessError:
+        return None
+
+
 def get_user_hash() -> str:
     """Returns a unique user-machine specific hash as a user id.
 

sky/utils/db_utils.py

@@ -84,6 +84,22 @@ def add_column_to_table(
     conn.commit()
 
 
+def add_tables_to_db_sqlalchemy(
+    metadata: sqlalchemy.MetaData,
+    engine: sqlalchemy.Engine,
+):
+    """Add tables to the database."""
+    for table in metadata.tables.values():
+        try:
+            table.create(bind=engine, checkfirst=True)
+        except (sqlalchemy_exc.OperationalError,
+                sqlalchemy_exc.ProgrammingError) as e:
+            if 'already exists' in str(e):
+                pass
+            else:
+                raise
+
+
 def add_column_to_table_sqlalchemy(
     session: 'Session',
     table_name: str,

sky/utils/schemas.py

@@ -870,6 +870,9 @@ def get_task_schema():
                 'type': 'array',
                 'items': get_volume_mount_schema(),
             },
+            '_metadata': {
+                'type': 'object',
+            },
             **_experimental_task_schema(),
         }
     }
@@ -1103,6 +1106,9 @@ _CONTEXT_CONFIG_SCHEMA_KUBERNETES = {
             },
         },
     },
+    'remote_identity': {
+        'type': 'string',
+    }
 }
 
 

sky/utils/ux_utils.py

@@ -253,9 +253,7 @@ def command_hint_messages(hint_type: CommandHintType,
                 f'{BOLD}sky jobs logs {job_id}{RESET_BOLD}'
                 f'\n{INDENT_SYMBOL}To stream controller logs:\t\t'
                 f'{BOLD}sky jobs logs --controller {job_id}{RESET_BOLD}'
-                f'\n{INDENT_SYMBOL}To view all managed jobs:\t\t'
-                f'{BOLD}sky jobs queue{RESET_BOLD}'
-                f'\n{INDENT_LAST_SYMBOL}To view managed job dashboard:\t\t'
-                f'{BOLD}sky jobs dashboard{RESET_BOLD}')
+                f'\n{INDENT_LAST_SYMBOL}To view all managed jobs:\t\t'
+                f'{BOLD}sky jobs queue{RESET_BOLD}')
     else:
         raise ValueError(f'Invalid hint type: {hint_type}')

{skypilot_nightly-1.0.0.dev20250709.dist-info → skypilot_nightly-1.0.0.dev20250711.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: skypilot-nightly
-Version: 1.0.0.dev20250709
+Version: 1.0.0.dev20250711
 Summary: SkyPilot: Run AI on Any Infra — Unified, Faster, Cheaper.
 Author: SkyPilot Team
 License: Apache 2.0