skypilot-nightly 1.0.0.dev20251203__py3-none-any.whl → 1.0.0.dev20260112__py3-none-any.whl

sky/ssh_node_pools/deploy/utils.py

@@ -0,0 +1,173 @@
+"""Utilities for SSH Node Pools Deployment"""
+import os
+import subprocess
+from typing import List, Optional
+
+import colorama
+
+from sky import sky_logging
+from sky.utils import ux_utils
+
+logger = sky_logging.init_logger(__name__)
+
+
+def check_ssh_cluster_dependencies(
+        raise_error: bool = True) -> Optional[List[str]]:
+    """Checks if the dependencies for ssh cluster are installed.
+
+    Args:
+        raise_error: set to true when the dependency needs to be present.
+            set to false for `sky check`, where reason strings are compiled
+            at the end.
+
+    Returns: the reasons list if there are missing dependencies.
+    """
+    # error message
+    jq_message = ('`jq` is required to setup ssh cluster.')
+
+    # save
+    reasons = []
+    required_binaries = []
+
+    # Ensure jq is installed
+    try:
+        subprocess.run(['jq', '--version'],
+                       stdout=subprocess.DEVNULL,
+                       stderr=subprocess.DEVNULL,
+                       check=True)
+    except (FileNotFoundError, subprocess.CalledProcessError):
+        required_binaries.append('jq')
+        reasons.append(jq_message)
+
+    if required_binaries:
+        reasons.extend([
+            'On Debian/Ubuntu, install the missing dependenc(ies) with:',
+            f'  $ sudo apt install {" ".join(required_binaries)}',
+            'On MacOS, install with: ',
+            f'  $ brew install {" ".join(required_binaries)}',
+        ])
+        if raise_error:
+            with ux_utils.print_exception_no_traceback():
+                raise RuntimeError('\n'.join(reasons))
+        return reasons
+    return None
+
+
+def run_command(cmd, shell=False, silent=False):
+    """Run a local command and return the output."""
+    process = subprocess.run(cmd,
+                             shell=shell,
+                             capture_output=True,
+                             text=True,
+                             check=False)
+    if process.returncode != 0:
+        if not silent:
+            logger.error(f'{colorama.Fore.RED}Error executing command: {cmd}\n'
+                         f'{colorama.Style.RESET_ALL}STDOUT: {process.stdout}\n'
+                         f'STDERR: {process.stderr}')
+        return None
+    return process.stdout.strip()
+
+
+def get_effective_host_ip(hostname: str) -> str:
+    """Get the effective IP for a hostname from SSH config."""
+    try:
+        result = subprocess.run(['ssh', '-G', hostname],
+                                capture_output=True,
+                                text=True,
+                                check=False)
+        if result.returncode == 0:
+            for line in result.stdout.splitlines():
+                if line.startswith('hostname '):
+                    return line.split(' ', 1)[1].strip()
+    except Exception:  # pylint: disable=broad-except
+        pass
+    return hostname  # Return the original hostname if lookup fails
+
+
+def run_remote(node,
+               cmd,
+               user='',
+               ssh_key='',
+               connect_timeout=30,
+               use_ssh_config=False,
+               print_output=False,
+               use_shell=False,
+               silent=False):
+    """Run a command on a remote machine via SSH."""
+    ssh_cmd: List[str]
+    if use_ssh_config:
+        # Use SSH config for connection parameters
+        ssh_cmd = ['ssh', node, cmd]
+    else:
+        # Use explicit parameters
+        ssh_cmd = [
+            'ssh', '-o', 'StrictHostKeyChecking=no', '-o', 'IdentitiesOnly=yes',
+            '-o', f'ConnectTimeout={connect_timeout}', '-o',
+            'ServerAliveInterval=10', '-o', 'ServerAliveCountMax=3'
+        ]
+
+        if ssh_key:
+            if not os.path.isfile(ssh_key):
+                raise ValueError(f'SSH key not found: {ssh_key}')
+            ssh_cmd.extend(['-i', ssh_key])
+
+        ssh_cmd.append(f'{user}@{node}' if user else node)
+        ssh_cmd.append(cmd)
+
+    subprocess_cmd = ' '.join(ssh_cmd) if use_shell else ssh_cmd
+    process = subprocess.run(subprocess_cmd,
+                             capture_output=True,
+                             text=True,
+                             check=False,
+                             shell=use_shell)
+    if process.returncode != 0:
+        if not silent:
+            logger.error(f'{colorama.Fore.RED}Error executing command {cmd} on '
+                         f'{node}:{colorama.Style.RESET_ALL} {process.stderr}')
+        return None
+    if print_output:
+        logger.info(process.stdout)
+    return process.stdout.strip()
+
+
+def ensure_directory_exists(path):
+    """Ensure the directory for the specified file path exists."""
+    directory = os.path.dirname(path)
+    if directory and not os.path.exists(directory):
+        os.makedirs(directory, exist_ok=True)
+
+
+def check_gpu(node, user, ssh_key, use_ssh_config=False, is_head=False):
+    """Check if a node has a GPU."""
+    cmd = ('command -v nvidia-smi &> /dev/null && '
+           'nvidia-smi --query-gpu=gpu_name --format=csv,noheader')
+    result = run_remote(node,
+                        cmd,
+                        user,
+                        ssh_key,
+                        use_ssh_config=use_ssh_config,
+                        silent=True)
+    if result is not None:
+        # Check that all GPUs have the same type.
+        # Currently, SkyPilot does not support heterogeneous GPU node
+        # (i.e. more than one GPU type on the same node).
+        gpu_names = {
+            line.strip() for line in result.splitlines() if line.strip()
+        }
+        if not gpu_names:
+            # This can happen if nvidia-smi returns only whitespace.
+            # Set result to None to ensure this function returns False.
+            result = None
+        elif len(gpu_names) > 1:
+            # Sort for a deterministic error message.
+            sorted_gpu_names = sorted(list(gpu_names))
+            raise RuntimeError(
+                f'Node {node} has more than one GPU types '
+                f'({", ".join(sorted_gpu_names)}). '
+                'SkyPilot does not support a node with multiple GPU types.')
+        else:
+            logger.info(f'{colorama.Fore.YELLOW}➜ GPU {list(gpu_names)[0]} '
+                        f'detected on {"head" if is_head else "worker"} '
+                        f'node ({node}).{colorama.Style.RESET_ALL}')
+    return result is not None

sky/ssh_node_pools/server.py

@@ -4,12 +4,11 @@ from typing import Any, Dict, List
 
 import fastapi
 
-from sky import core as sky_core
 from sky.server.requests import executor
 from sky.server.requests import payloads
 from sky.server.requests import request_names
 from sky.server.requests import requests as requests_lib
-from sky.ssh_node_pools import core as ssh_node_pools_core
+from sky.ssh_node_pools import core
 from sky.utils import common_utils
 
 router = fastapi.APIRouter()
@@ -19,7 +18,7 @@ router = fastapi.APIRouter()
 def get_ssh_node_pools() -> Dict[str, Any]:
     """Get all SSH Node Pool configurations."""
     try:
-        return ssh_node_pools_core.get_all_pools()
+        return core.get_all_pools()
     except Exception as e:
         raise fastapi.HTTPException(
             status_code=500,
@@ -31,7 +30,7 @@ def get_ssh_node_pools() -> Dict[str, Any]:
 def update_ssh_node_pools(pools_config: Dict[str, Any]) -> Dict[str, str]:
     """Update SSH Node Pool configurations."""
     try:
-        ssh_node_pools_core.update_pools(pools_config)
+        core.update_pools(pools_config)
         return {'status': 'success'}
     except Exception as e:
         raise fastapi.HTTPException(status_code=400,
@@ -43,7 +42,7 @@ def update_ssh_node_pools(pools_config: Dict[str, Any]) -> Dict[str, str]:
 def delete_ssh_node_pool(pool_name: str) -> Dict[str, str]:
     """Delete a SSH Node Pool configuration."""
     try:
-        if ssh_node_pools_core.delete_pool(pool_name):
+        if core.delete_pool(pool_name):
             return {'status': 'success'}
         else:
             raise fastapi.HTTPException(
@@ -70,8 +69,7 @@ async def upload_ssh_key(request: fastapi.Request) -> Dict[str, str]:
                                         detail='Missing key_name or key_file')
 
         key_content = await key_file.read()
-        key_path = ssh_node_pools_core.upload_ssh_key(key_name,
-                                                      key_content.decode())
+        key_path = core.upload_ssh_key(key_name, key_content.decode())
 
         return {'status': 'success', 'key_path': key_path}
     except fastapi.HTTPException:
@@ -87,7 +85,7 @@ async def upload_ssh_key(request: fastapi.Request) -> Dict[str, str]:
 def list_ssh_keys() -> List[str]:
     """List available SSH keys."""
     try:
-        return ssh_node_pools_core.list_ssh_keys()
+        return core.list_ssh_keys()
     except Exception as e:
         exception_msg = common_utils.format_exception(e)
         raise fastapi.HTTPException(
@@ -104,7 +102,7 @@ async def deploy_ssh_node_pool(request: fastapi.Request,
             request_id=request.state.request_id,
             request_name=request_names.RequestName.SSH_NODE_POOLS_UP,
             request_body=ssh_up_body,
-            func=sky_core.ssh_up,
+            func=core.ssh_up,
             schedule_type=requests_lib.ScheduleType.LONG,
         )
 
@@ -129,7 +127,7 @@ async def deploy_ssh_node_pool_general(
             request_id=request.state.request_id,
             request_name=request_names.RequestName.SSH_NODE_POOLS_UP,
             request_body=ssh_up_body,
-            func=sky_core.ssh_up,
+            func=core.ssh_up,
             schedule_type=requests_lib.ScheduleType.LONG,
         )
 
@@ -155,7 +153,7 @@ async def down_ssh_node_pool(request: fastapi.Request,
             request_id=request.state.request_id,
             request_name=request_names.RequestName.SSH_NODE_POOLS_DOWN,
             request_body=ssh_up_body,
-            func=sky_core.ssh_up,  # Reuse ssh_up function with cleanup=True
+            func=core.ssh_up,  # Reuse ssh_up function with cleanup=True
             schedule_type=requests_lib.ScheduleType.LONG,
         )
 
@@ -183,7 +181,7 @@ async def down_ssh_node_pool_general(
             request_id=request.state.request_id,
             request_name=request_names.RequestName.SSH_NODE_POOLS_DOWN,
             request_body=ssh_up_body,
-            func=sky_core.ssh_up,  # Reuse ssh_up function with cleanup=True
+            func=core.ssh_up,  # Reuse ssh_up function with cleanup=True
             schedule_type=requests_lib.ScheduleType.LONG,
         )
 
@@ -206,7 +204,7 @@ def get_ssh_node_pool_status(pool_name: str) -> Dict[str, str]:
     try:
         # Call ssh_status to check the context
         context_name = f'ssh-{pool_name}'
-        is_ready, reason = sky_core.ssh_status(context_name)
+        is_ready, reason = core.ssh_status(context_name)
 
         # Strip ANSI escape codes from the reason text
         def strip_ansi_codes(text):

sky/{utils/kubernetes/ssh_utils.py → ssh_node_pools/utils.py}

@@ -5,13 +5,14 @@ import subprocess
 from typing import Any, Callable, Dict, List, Optional
 import uuid
 
+import colorama
 import yaml
 
+from sky import sky_logging
+from sky.ssh_node_pools import constants
 from sky.utils import ux_utils
 
-DEFAULT_SSH_NODE_POOLS_PATH = os.path.expanduser('~/.sky/ssh_node_pools.yaml')
-RED = '\033[0;31m'
-NC = '\033[0m'  # No color
+logger = sky_logging.init_logger(__name__)
 
 
 def check_host_in_ssh_config(hostname: str) -> bool:
@@ -92,7 +93,8 @@ def load_ssh_targets(file_path: str) -> Dict[str, Any]:
 def get_cluster_config(
         targets: Dict[str, Any],
         cluster_name: Optional[str] = None,
-        file_path: str = DEFAULT_SSH_NODE_POOLS_PATH) -> Dict[str, Any]:
+        file_path: str = constants.DEFAULT_SSH_NODE_POOLS_PATH
+) -> Dict[str, Any]:
     """Get configuration for specific clusters or all clusters."""
     if not targets:
         with ux_utils.print_exception_no_traceback():
@@ -186,8 +188,9 @@ def prepare_hosts_info(
         else:
             # It's a dict with potential overrides
             if 'ip' not in host:
-                print(f'{RED}Warning: Host missing \'ip\' field, '
-                      f'skipping: {host}{NC}')
+                logger.warning(f'{colorama.Fore.RED}Warning: Host missing'
+                               f'\'ip\' field, skipping: {host}'
+                               f'{colorama.Style.RESET_ALL}')
                 continue
 
             # Check if this is an SSH config hostname

sky/templates/kubernetes-ray.yml.j2

@@ -523,6 +523,14 @@ available_node_types:
                 resourceFieldRef:
                   containerName: ray-node
                   resource: requests.memory
+            # Disable Ray memory monitor to prevent Ray's memory manager 
+            # from interfering with kubernetes resource manager.
+            # If ray memory monitor is enabled, the ray memory monitor kills
+            # the running job is the job uses more than 95% of allocated memory,
+            # even if the job is not misbehaving or using its full allocated memory.
+            # This behavior does not give a chance for k8s scheduler to evict the pod.
+            - name: RAY_memory_monitor_refresh_ms
+              value: "0"
             {% for key, value in k8s_env_vars.items() if k8s_env_vars is not none %}
             - name: {{ key }}
               value: {{ value }}
@@ -912,19 +920,17 @@ available_node_types:
                 {{ ray_installation_commands }}
 
                 # set UV_SYSTEM_PYTHON to false in case the user provided docker image set it to true.
-                # unset PYTHONPATH in case the user provided docker image set it.
-                VIRTUAL_ENV=~/skypilot-runtime UV_SYSTEM_PYTHON=false env -u PYTHONPATH ~/.local/bin/uv pip install skypilot[kubernetes,remote]
+                # unset PYTHONPATH and set CWD to $HOME to avoid user image interfering with SkyPilot runtime.
+                VIRTUAL_ENV=~/skypilot-runtime UV_SYSTEM_PYTHON=false {{sky_unset_pythonpath_and_set_cwd}} ~/.local/bin/uv pip install skypilot[kubernetes,remote]
                 # Wait for `patch` package to be installed before applying ray patches
                 until dpkg -l | grep -q "^ii  patch "; do
                   sleep 0.1
                   echo "Waiting for patch package to be installed..."
                 done
                 # Apply Ray patches for progress bar fix
-                # set UV_SYSTEM_PYTHON to false in case the user provided docker image set it to true.
-                # unset PYTHONPATH in case the user provided docker image set it.
                 # ~/.sky/python_path is seeded by conda_installation_commands
-                VIRTUAL_ENV=~/skypilot-runtime UV_SYSTEM_PYTHON=false env -u PYTHONPATH ~/.local/bin/uv pip list | grep "ray " | grep 2.9.3 2>&1 > /dev/null && {
-                  env -u PYTHONPATH $(cat ~/.sky/python_path) -c "from sky.skylet.ray_patches import patch; patch()" || exit 1;
+                VIRTUAL_ENV=~/skypilot-runtime UV_SYSTEM_PYTHON=false {{sky_unset_pythonpath_and_set_cwd}} ~/.local/bin/uv pip list | grep "ray " | grep 2.9.3 2>&1 > /dev/null && {
+                  {{sky_unset_pythonpath_and_set_cwd}} $(cat ~/.sky/python_path) -c "from sky.skylet.ray_patches import patch; patch()" || exit 1;
                 }
                 touch /tmp/ray_skypilot_installation_complete
                 echo "=== Ray and skypilot installation completed ==="

sky/templates/slurm-ray.yml.j2

@@ -0,0 +1,115 @@
+cluster_name: {{cluster_name_on_cloud}}
+
+# The maximum number of workers nodes to launch in addition to the head node.
+max_workers: {{num_nodes - 1}}
+upscaling_speed: {{num_nodes - 1}}
+idle_timeout_minutes: 60
+
+provider:
+  type: external
+  module: sky.provision.slurm
+
+  cluster: {{slurm_cluster}}
+  partition: {{slurm_partition}}
+
+  ssh:
+    hostname: {{ssh_hostname}}
+    port: {{ssh_port}}
+    user: {{ssh_user}}
+    private_key: {{slurm_private_key}}
+{% if slurm_proxy_command is not none %}
+    proxycommand: {{slurm_proxy_command | tojson }}
+{% endif %}
+{% if slurm_proxy_jump is not none %}
+    proxyjump: {{slurm_proxy_jump | tojson }}
+{% endif %}
+
+auth:
+  ssh_user: {{ssh_user}}
+  # TODO(jwj,kevin): Modify this tmp workaround.
+  # Right now there's a chicken-and-egg problem:
+  # 1. ssh_credential_from_yaml reads from the auth.ssh_private_key: ~/.sky/clients/.../ssh/sky-key
+  # 2. This is SkyPilot's generated key, not the Slurm cluster's key
+  # 3. The internal_file_mounts stage tries to rsync using sky-key, but its public key isn't on the remote yet
+  # 4. The public key only gets added by setup_commands, which runs AFTER file_mounts
+  # ssh_private_key: {{ssh_private_key}}
+  ssh_private_key: {{slurm_private_key}}
+  ssh_proxy_command: {{slurm_proxy_command | tojson }}
+
+available_node_types:
+  ray_head_default:
+    resources: {}
+    node_config:
+      # From clouds/slurm.py::Slurm.make_deploy_resources_variables.
+      instance_type: {{instance_type}}
+      disk_size: {{disk_size}}
+      cpus: {{cpus}}
+      memory: {{memory}}
+      accelerator_type: {{accelerator_type}}
+      accelerator_count: {{accelerator_count}}
+
+      # TODO: more configs that is required by the provisioner to create new
+      # instances on the FluffyCloud:
+      # sky/provision/fluffycloud/instance.py::run_instances
+
+head_node_type: ray_head_default
+
+# Format: `REMOTE_PATH : LOCAL_PATH`
+file_mounts: {
+  "{{sky_ray_yaml_remote_path}}": "{{sky_ray_yaml_local_path}}",
+  "{{sky_remote_path}}/{{sky_wheel_hash}}": "{{sky_local_path}}",
+{%- for remote_path, local_path in credentials.items() %}
+  "{{remote_path}}": "{{local_path}}",
+{%- endfor %}
+}
+
+rsync_exclude: []
+
+initialization_commands: []
+
+# List of shell commands to run to set up nodes.
+# NOTE: these are very performance-sensitive. Each new item opens/closes an SSH
+# connection, which is expensive. Try your best to co-locate commands into fewer
+# items!
+#
+# Increment the following for catching performance bugs easier:
+#   current num items (num SSH connections): 1
+setup_commands:
+  - |
+    {%- for initial_setup_command in initial_setup_commands %}
+    {{ initial_setup_command }}
+    {%- endfor %}
+    # Generate host key for sshd -i if not exists
+    mkdir -p ~{{ssh_user}}/.ssh && chmod 700 ~{{ssh_user}}/.ssh
+    [ -f ~{{ssh_user}}/.ssh/{{slurm_sshd_host_key_filename}} ] || ssh-keygen -t ed25519 -f ~{{ssh_user}}/.ssh/{{slurm_sshd_host_key_filename}} -N "" -q
+    # Add public key to user's authorized_keys if not already present
+    grep -qF 'skypilot:ssh_public_key_content' ~{{ssh_user}}/.ssh/authorized_keys 2>/dev/null || cat >> ~{{ssh_user}}/.ssh/authorized_keys <<'SKYPILOT_SSH_KEY_EOF'
+    skypilot:ssh_public_key_content
+    SKYPILOT_SSH_KEY_EOF
+    chmod 600 ~{{ssh_user}}/.ssh/authorized_keys
+
+    mkdir -p ~{{ssh_user}}/.sky
+    cat > ~{{ssh_user}}/.sky_ssh_rc <<'SKYPILOT_SSH_RC'
+    # Added by SkyPilot: override HOME for Slurm interactive sessions
+    if [ -n "${{slurm_cluster_name_env_var}}" ]; then
+        CLUSTER_DIR=~/.sky_clusters/${{slurm_cluster_name_env_var}}
+        if [ -d "$CLUSTER_DIR" ]; then
+            cd "$CLUSTER_DIR"
+            export HOME=$(pwd)
+        fi
+    fi
+    SKYPILOT_SSH_RC
+    grep -q "source ~/.sky_ssh_rc" ~{{ssh_user}}/.bashrc 2>/dev/null || (echo "" >> ~{{ssh_user}}/.bashrc && echo "source ~/.sky_ssh_rc" >> ~{{ssh_user}}/.bashrc)
+    {{ setup_sky_dirs_commands }}
+    {{ conda_installation_commands }}
+    {{ skypilot_wheel_installation_commands }}
+    {{ copy_skypilot_templates_commands }}
+
+head_node: {}
+worker_nodes: {}
+
+# These fields are required for external cloud providers.
+head_setup_commands: []
+worker_setup_commands: []
+cluster_synced_files: []
+file_mounts_sync_continuously: False

sky/templates/vast-ray.yml.j2

@@ -10,6 +10,7 @@ provider:
   module: sky.provision.vast
   region: "{{region}}"
   disable_launch_config_check: true
+  secure_only: {{secure_only}}
 
 auth:
   ssh_user: root

sky/templates/websocket_proxy.py

@@ -9,13 +9,11 @@
 This script is useful for users who do not have local Kubernetes credentials.
 """
 import asyncio
-from http.cookiejar import MozillaCookieJar
 import os
 import struct
 import sys
 import time
 from typing import Dict, Optional
-from urllib.request import Request
 
 import requests
 import websockets
@@ -24,46 +22,19 @@ from websockets.asyncio.client import connect
 
 from sky import exceptions
 from sky.client import service_account_auth
+from sky.server import common as server_common
 from sky.server import constants
-from sky.server.server import KubernetesSSHMessageType
+from sky.server.server import SSHMessageType
 from sky.skylet import constants as skylet_constants
 
 BUFFER_SIZE = 2**16  # 64KB
 HEARTBEAT_INTERVAL_SECONDS = 10
-
-# Environment variable for a file path to the API cookie file.
-# Keep in sync with server/constants.py
-API_COOKIE_FILE_ENV_VAR = 'SKYPILOT_API_COOKIE_FILE'
-# Default file if unset.
-# Keep in sync with server/constants.py
-API_COOKIE_FILE_DEFAULT_LOCATION = '~/.sky/cookies.txt'
-
 MAX_UNANSWERED_PINGS = 100
 
 
-def _get_cookie_header(url: str) -> Dict[str, str]:
-    """Extract Cookie header value from a cookie jar for a specific URL"""
-    cookie_path = os.environ.get(API_COOKIE_FILE_ENV_VAR)
-    if cookie_path is None:
-        cookie_path = API_COOKIE_FILE_DEFAULT_LOCATION
-    cookie_path = os.path.expanduser(cookie_path)
-    if not os.path.exists(cookie_path):
-        return {}
-
-    request = Request(url)
-    cookie_jar = MozillaCookieJar(os.path.expanduser(cookie_path))
-    cookie_jar.load(ignore_discard=True, ignore_expires=True)
-    cookie_jar.add_cookie_header(request)
-    cookie_header = request.get_header('Cookie')
-    # if cookie file is empty, return empty dict
-    if cookie_header is None:
-        return {}
-    return {'Cookie': cookie_header}
-
-
 async def main(url: str, timestamps_supported: bool, login_url: str) -> None:
     headers = {}
-    headers.update(_get_cookie_header(url))
+    headers.update(server_common.get_cookie_header_for_url(url))
     headers.update(service_account_auth.get_service_account_headers())
     try:
         async with connect(url, ping_interval=None,
@@ -142,8 +113,9 @@ async def latency_monitor(websocket: ClientConnection,
             ping_time = time.time()
             next_id += 1
             last_ping_time_dict[next_id] = ping_time
-            message_header_bytes = struct.pack(
-                '!BI', KubernetesSSHMessageType.PINGPONG.value, next_id)
+            message_header_bytes = struct.pack('!BI',
+                                               SSHMessageType.PINGPONG.value,
+                                               next_id)
             try:
                 async with websocket_lock:
                     await websocket.send(message_header_bytes)
@@ -176,7 +148,7 @@ async def stdin_to_websocket(reader: asyncio.StreamReader,
             if timestamps_supported:
                 # Send message with type 0 to indicate data.
                 message_type_bytes = struct.pack(
-                    '!B', KubernetesSSHMessageType.REGULAR_DATA.value)
+                    '!B', SSHMessageType.REGULAR_DATA.value)
                 data = message_type_bytes + data
             async with websocket_lock:
                 await websocket.send(data)
@@ -201,10 +173,10 @@ async def websocket_to_stdout(websocket: ClientConnection,
             if (timestamps_supported and len(message) > 0 and
                     last_ping_time_dict is not None):
                 message_type = struct.unpack('!B', message[:1])[0]
-                if message_type == KubernetesSSHMessageType.REGULAR_DATA.value:
+                if message_type == SSHMessageType.REGULAR_DATA.value:
                     # Regular data - strip type byte and write to stdout
                     message = message[1:]
-                elif message_type == KubernetesSSHMessageType.PINGPONG.value:
+                elif message_type == SSHMessageType.PINGPONG.value:
                     # PONG response - calculate latency and send measurement
                     if not len(message) == struct.calcsize('!BI'):
                         raise ValueError(
@@ -222,8 +194,7 @@ async def websocket_to_stdout(websocket: ClientConnection,
 
                     # Send latency measurement (type 2)
                     message_type_bytes = struct.pack(
-                        '!B',
-                        KubernetesSSHMessageType.LATENCY_MEASUREMENT.value)
+                        '!B', SSHMessageType.LATENCY_MEASUREMENT.value)
                     latency_bytes = struct.pack('!Q', latency_ms)
                     message = message_type_bytes + latency_bytes
                     # Send to server.
@@ -255,7 +226,7 @@ if __name__ == '__main__':
         # TODO(aylei): remove the separate /api/health call and use the header
         # during websocket handshake to determine the server version.
         health_url = f'{server_url}/api/health'
-        cookie_hdr = _get_cookie_header(health_url)
+        cookie_hdr = server_common.get_cookie_header_for_url(health_url)
         health_response = requests.get(health_url, headers=cookie_hdr)
         health_data = health_response.json()
         timestamps_are_supported = int(health_data.get('api_version', 0)) > 21
@@ -272,7 +243,13 @@ if __name__ == '__main__':
     client_version_str = (f'&client_version={constants.API_VERSION}'
                           if timestamps_are_supported else '')
 
-    websocket_url = (f'{server_url}/kubernetes-pod-ssh-proxy'
+    # For backwards compatibility, fallback to kubernetes-pod-ssh-proxy if
+    # no endpoint is provided.
+    endpoint = sys.argv[3] if len(sys.argv) > 3 else 'kubernetes-pod-ssh-proxy'
+    # Worker index for Slurm.
+    worker_idx = sys.argv[4] if len(sys.argv) > 4 else '0'
+    websocket_url = (f'{server_url}/{endpoint}'
                      f'?cluster_name={sys.argv[2]}'
+                     f'&worker={worker_idx}'
                      f'{client_version_str}')
     asyncio.run(main(websocket_url, timestamps_are_supported, _login_url))

sky/users/model.conf

@@ -12,4 +12,4 @@ g = _, _
 e = some(where (p.eft == allow))
 
 [matchers]
-m = (g(r.sub, p.sub)|| p.sub == '*') && r.obj == p.obj && r.act == p.act
+m = (g(r.sub, p.sub)|| p.sub == '*') && keyMatch2(r.obj, p.obj) && r.act == p.act