skypilot-nightly 1.0.0.dev20250905__py3-none-any.whl → 1.0.0.dev20251210__py3-none-any.whl

sky/templates/seeweb-ray.yml.j2

@@ -0,0 +1,171 @@
+cluster_name: {{ cluster_name_on_cloud }}
+
+max_workers: {{ num_nodes - 1 }}
+upscaling_speed: {{ num_nodes - 1 }}
+idle_timeout_minutes: 5
+
+{%- if docker_image is not none %}
+docker:
+  image: {{docker_image}}
+  container_name: {{docker_container_name}}
+  run_options:
+    - --ulimit nofile=1048576:1048576
+    {%- for run_option in docker_run_options %}
+    - {{run_option}}
+    {%- endfor %}
+  {%- if docker_login_config is not none %}
+  docker_login_config:
+    username: |-
+      {{docker_login_config.username}}
+    password: |-
+      {{docker_login_config.password | indent(6) }}
+    server: |-
+      {{docker_login_config.server}}
+  {%- endif %}
+{%- endif %}
+
+provider:
+  type: external
+  module: sky.provision.seeweb
+  region: "{{ region }}"
+
+auth:
+  ssh_user: ecuser
+  ssh_private_key: {{ ssh_private_key }}
+
+available_node_types:
+  ray_head_default:
+    resources: {}
+    node_config:
+      plan: {{ instance_type }}
+      image: {{ image_id }}
+      location: {{ region }}
+      {% if seeweb_gpu_config is not none %}
+      gpu: {{ seeweb_gpu_config.gpu }}
+      gpu_label: "{{ seeweb_gpu_config.gpu_label }}"
+      {% endif %}
+      disk: {{ disk_size }}
+      {% if docker_image is not none %}
+      user_customize: |
+        #!/bin/bash
+        # Auto-generated Docker installation script for Seeweb
+        LOG_FILE=/var/log/user_customize.log
+        sudo mkdir -p "$(dirname "$LOG_FILE")"
+        {
+          echo "[$(date -Is)] Cloud script: start"
+          sudo apt-get update
+          sudo apt-get install -y \
+            apt-transport-https \
+            ca-certificates \
+            curl \
+            gnupg-agent \
+            lsb-release \
+            software-properties-common
+          sudo mkdir -p /usr/share/keyrings
+          curl -fsSL https://download.docker.com/linux/ubuntu/gpg | \
+            sudo gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg
+          UBU_CODENAME="$(. /etc/os-release && echo "$VERSION_CODENAME")"
+          echo "deb [arch=amd64 signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu ${UBU_CODENAME} stable" | \
+            sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
+          sudo apt-get update
+          sudo apt-get install -y docker-ce docker-ce-cli containerd.io
+          echo "[$(date -Is)] Cloud script: docker installed"
+          sudo usermod -aG docker ecuser || true
+          sudo systemctl enable docker || true
+          sudo systemctl start docker || true
+          command -v docker && docker --version || echo "[$(date -Is)] docker still missing"
+          echo "[$(date -Is)] Cloud script: complete"
+        } | sudo tee -a "$LOG_FILE"
+        sudo touch /var/log/docker_install_done
+      {% endif %}
+
+head_node_type: ray_head_default
+
+# Format: `REMOTE_PATH : LOCAL_PATH`
+file_mounts: {
+  "~/.seeweb_cloud/seeweb_keys": "~/.seeweb_cloud/seeweb_keys",
+  "{{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 %}
+  "~/.ssh/sky-cluster-key": "{{ssh_private_key}}",
+}
+
+rsync_exclude: []
+
+setup_commands:
+  - |
+    {%- for initial_setup_command in initial_setup_commands %}
+    {{ initial_setup_command }}
+    {%- endfor %}
+    touch ~/.bashrc;
+    echo "127.0.0.1 $(hostname)" | sudo tee -a /etc/hosts || true;
+    echo "127.0.0.1 localhost" | sudo tee -a /etc/hosts || true;
+    sudo systemctl stop unattended-upgrades || true;
+    sudo systemctl disable unattended-upgrades || true;
+    sudo apt update && sudo apt install -y patch || sudo yum install -y patch || true;
+
+    {%- if docker_image is not none %}
+    # Docker installed via cloud-init; ensure service will be started by cloud-init
+    {%- endif %}
+
+    {{ conda_installation_commands }}
+    {{ ray_skypilot_installation_commands }}
+    {{ copy_skypilot_templates_commands }}
+
+head_start_ray_commands:
+  - |
+    retry_ray() {
+      local n=0; local max=30
+      until [ $n -ge $max ]; do
+        export SKYPILOT_NUM_GPUS=0
+        command -v nvidia-smi >/dev/null 2>&1 && \
+          SKYPILOT_NUM_GPUS=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)
+
+        ray stop || true
+        RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 \
+        ray start --disable-usage-stats --head \
+          --port={{ ray_port }} --dashboard-port={{ ray_dashboard_port }} \
+          --object-manager-port=8076 \
+          --autoscaling-config=~/ray_bootstrap_config.yaml \
+          --num-gpus=$SKYPILOT_NUM_GPUS --temp-dir {{ ray_temp_dir }} && break
+
+        echo "[head] Ray failed to start ($((++n))/$max), retrying in 5s..."
+        sleep 5
+      done
+      [ $n -eq $max ] && { echo "Ray head failed"; exit 1; }
+    }
+    retry_ray
+
+worker_start_ray_commands:
+  - |
+    retry_ray() {
+      local n=0; local max=30
+      until [ $n -ge $max ]; do
+        SKYPILOT_NUM_GPUS=0
+        command -v nvidia-smi >/dev/null 2>&1 && \
+          SKYPILOT_NUM_GPUS=$(nvidia-smi --query-gpu=index --format=csv,noheader | wc -l)
+
+        ray stop || true
+        RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 \
+        ray start --disable-usage-stats \
+          --address=$RAY_HEAD_IP:{{ ray_port }} \
+          --object-manager-port=8076 \
+          --num-gpus=$SKYPILOT_NUM_GPUS --temp-dir {{ ray_temp_dir }} && break
+
+        echo "[worker] Ray failed to start ($((++n))/$max), retrying in 5s..."
+        sleep 5
+      done
+      [ $n -eq $max ] && { echo "Ray worker failed"; exit 1; }
+    }
+    retry_ray
+
+head_node: {}
+worker_nodes: {}
+
+head_setup_commands: []
+worker_setup_commands: []
+
+cluster_synced_files: []
+file_mounts_sync_continuously: False

sky/templates/shadeform-ray.yml.j2

@@ -0,0 +1,73 @@
+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.shadeform
+  region: "{{region}}"
+  disable_launch_config_check: true
+
+auth:
+  ssh_user: shadeform
+  ssh_private_key: {{ssh_private_key}}
+  ssh_key_id: {{ssh_key_id}}
+
+available_node_types:
+  ray_head_default:
+{%- if custom_resources %}
+    resources: {{custom_resources}}
+{%- else %}
+    resources: {}
+{%- endif %}
+    node_config:
+      InstanceType: {{instance_type}}
+      PublicKey: |-
+        skypilot:ssh_public_key_content
+
+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:
+  # Create ~/.ssh/config file in case the file does not exist in the image.
+  # Line 'rm ..': there is another installation of pip.
+  # Line 'sudo bash ..': set the ulimit as suggested by ray docs for performance. https://docs.ray.io/en/latest/cluster/vms/user-guides/large-cluster-best-practices.html#system-configuration
+  # Line 'sudo grep ..': set the number of threads per process to unlimited to avoid ray job submit stucking issue when the number of running ray jobs increase.
+  # Line 'mkdir -p ..': disable host key check
+  # Line 'python3 -c ..': patch the buggy ray files and enable `-o allow_other` option for `goofys`
+  - {%- for initial_setup_command in initial_setup_commands %}
+    {{ initial_setup_command }}
+    {%- endfor %}
+    mkdir -p ~/.ssh; touch ~/.ssh/config; which patch > /dev/null || sudo apt install -y patch;
+    {{ conda_installation_commands }}
+    {{ ray_skypilot_installation_commands }}
+    {{ copy_skypilot_templates_commands }}
+    sudo bash -c 'rm -rf /etc/security/limits.d; echo "* soft nofile 1048576" >> /etc/security/limits.conf; echo "* hard nofile 1048576" >> /etc/security/limits.conf';
+    sudo grep -e '^DefaultTasksMax' /etc/systemd/system.conf || (sudo bash -c 'echo "DefaultTasksMax=infinity" >> /etc/systemd/system.conf'); sudo systemctl set-property user-$(id -u $(whoami)).slice TasksMax=infinity; sudo systemctl daemon-reload;
+    (grep -Pzo -q "Host \*\n  StrictHostKeyChecking no" ~/.ssh/config) || printf "Host *\n  StrictHostKeyChecking no\n" >> ~/.ssh/config;
+    {{ ssh_max_sessions_config }}
+
+# Command to start ray clusters are now placed in `sky.provision.instance_setup`.
+# We do not need to list it here anymore.

sky/templates/slurm-ray.yml.j2

@@ -0,0 +1,85 @@
+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 %}
+
+auth:
+  ssh_user: {{ssh_user}}
+  # TODO(jwj): Modify this tmp workaround.
+  # 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 %}
+    {{ 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
@@ -61,6 +62,7 @@ setup_commands:
     mkdir -p ~/.ssh; touch ~/.ssh/config; which patch > /dev/null || sudo apt install -y patch;
     {{ conda_installation_commands }}
     {{ ray_skypilot_installation_commands }}
+    {{ copy_skypilot_templates_commands }}
     sudo bash -c 'rm -rf /etc/security/limits.d; echo "* soft nofile 1048576" >> /etc/security/limits.conf; echo "* hard nofile 1048576" >> /etc/security/limits.conf';
     sudo grep -e '^DefaultTasksMax' /etc/systemd/system.conf || (sudo bash -c 'echo "DefaultTasksMax=infinity" >> /etc/systemd/system.conf'); sudo systemctl set-property user-$(id -u $(whoami)).slice TasksMax=infinity; sudo systemctl daemon-reload;
     (grep -Pzo -q "Host \*\n  StrictHostKeyChecking no" ~/.ssh/config) || printf "Host *\n  StrictHostKeyChecking no\n" >> ~/.ssh/config;

sky/templates/vsphere-ray.yml.j2

@@ -67,6 +67,7 @@ setup_commands:
     pip3 --version > /dev/null 2>&1 || (curl -sSL https://bootstrap.pypa.io/get-pip.py -o get-pip.py && python3 get-pip.py && echo "PATH=$HOME/.local/bin:$PATH" >> ~/.bashrc);
     {{ conda_installation_commands }}
     {{ ray_skypilot_installation_commands }}
+    {{ copy_skypilot_templates_commands }}
     sudo bash -c 'rm -rf /etc/security/limits.d; echo "* soft nofile 1048576" >> /etc/security/limits.conf; echo "* hard nofile 1048576" >> /etc/security/limits.conf';
     sudo grep -e '^DefaultTasksMax' /etc/systemd/system.conf || (sudo bash -c 'echo "DefaultTasksMax=infinity" >> /etc/systemd/system.conf'); sudo systemctl set-property user-$(id -u $(whoami)).slice TasksMax=infinity; sudo systemctl daemon-reload;
     mkdir -p ~/.ssh; (grep -Pzo -q "Host \*\n  StrictHostKeyChecking no\n  IdentityFile ~/.ssh/sky-cluster-key\n  IdentityFile ~/.ssh/id_rsa" ~/.ssh/config) || printf "Host *\n  StrictHostKeyChecking no\n  IdentityFile ~/.ssh/sky-cluster-key\n  IdentityFile ~/.ssh/id_rsa\n" >> ~/.ssh/config;

sky/templates/websocket_proxy.py

@@ -11,15 +11,25 @@ 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
-from typing import Dict
+import time
+from typing import Dict, Optional
 from urllib.request import Request
 
+import requests
 import websockets
 from websockets.asyncio.client import ClientConnection
 from websockets.asyncio.client import connect
 
+from sky import exceptions
+from sky.client import service_account_auth
+from sky.server import constants
+from sky.server.server import KubernetesSSHMessageType
+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
@@ -28,6 +38,8 @@ API_COOKIE_FILE_ENV_VAR = 'SKYPILOT_API_COOKIE_FILE'
 # 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"""
@@ -49,85 +61,218 @@ def _get_cookie_header(url: str) -> Dict[str, str]:
     return {'Cookie': cookie_header}
 
 
-async def main(url: str) -> None:
-    cookie_header = _get_cookie_header(url)
-    async with connect(url,
-                       ping_interval=None,
-                       additional_headers=cookie_header) as websocket:
-        if os.isatty(sys.stdin.fileno()):
-            # pylint: disable=import-outside-toplevel
-            import termios
-            import tty
-            old_settings = termios.tcgetattr(sys.stdin.fileno())
-            tty.setraw(sys.stdin.fileno())
+async def main(url: str, timestamps_supported: bool, login_url: str) -> None:
+    headers = {}
+    headers.update(_get_cookie_header(url))
+    headers.update(service_account_auth.get_service_account_headers())
+    try:
+        async with connect(url, ping_interval=None,
+                           additional_headers=headers) as websocket:
+            await run_websocket_proxy(websocket, timestamps_supported)
+    except websockets.exceptions.InvalidStatus as e:
+        if e.response.status_code == 403:
+            print(str(exceptions.ApiServerAuthenticationError(login_url)),
+                  file=sys.stderr)
         else:
-            old_settings = None
+            print(f'Error ssh into cluster: {e}', file=sys.stderr)
+        sys.exit(1)
+
+
+async def run_websocket_proxy(websocket: ClientConnection,
+                              timestamps_supported: bool) -> None:
+    if os.isatty(sys.stdin.fileno()):
+        # pylint: disable=import-outside-toplevel
+        import termios
+        import tty
+        old_settings = termios.tcgetattr(sys.stdin.fileno())
+        tty.setraw(sys.stdin.fileno())
+    else:
+        old_settings = None
+
+    try:
+        loop = asyncio.get_running_loop()
+        # Use asyncio.Stream primitives to wrap stdin and stdout, this is to
+        # avoid creating a new thread for each read/write operation
+        # excessively.
+        stdin_reader = asyncio.StreamReader()
+        protocol = asyncio.StreamReaderProtocol(stdin_reader)
+        await loop.connect_read_pipe(lambda: protocol, sys.stdin)
+        transport, protocol = await loop.connect_write_pipe(
+            asyncio.streams.FlowControlMixin, sys.stdout)  # type: ignore
+        stdout_writer = asyncio.StreamWriter(transport, protocol, None, loop)
+        # Dictionary to store last ping time for latency measurement
+        last_ping_time_dict: Optional[Dict[int, float]] = None
+        if timestamps_supported:
+            last_ping_time_dict = {}
+
+        # Use an Event to signal when websocket is closed
+        websocket_closed_event = asyncio.Event()
+        websocket_lock = asyncio.Lock()
+
+        await asyncio.gather(
+            stdin_to_websocket(stdin_reader, websocket, timestamps_supported,
+                               websocket_closed_event, websocket_lock),
+            websocket_to_stdout(websocket, stdout_writer, timestamps_supported,
+                                last_ping_time_dict, websocket_closed_event,
+                                websocket_lock),
+            latency_monitor(websocket, last_ping_time_dict,
+                            websocket_closed_event, websocket_lock),
+            return_exceptions=True)
+    finally:
+        if old_settings:
+            termios.tcsetattr(sys.stdin.fileno(), termios.TCSADRAIN,
+                              old_settings)
+
 
+async def latency_monitor(websocket: ClientConnection,
+                          last_ping_time_dict: Optional[dict],
+                          websocket_closed_event: asyncio.Event,
+                          websocket_lock: asyncio.Lock):
+    """Periodically send PING messages (type 1) to measure latency."""
+    if last_ping_time_dict is None:
+        return
+    next_id = 0
+    while not websocket_closed_event.is_set():
         try:
-            loop = asyncio.get_running_loop()
-            # Use asyncio.Stream primitives to wrap stdin and stdout, this is to
-            # avoid creating a new thread for each read/write operation
-            # excessively.
-            stdin_reader = asyncio.StreamReader()
-            protocol = asyncio.StreamReaderProtocol(stdin_reader)
-            await loop.connect_read_pipe(lambda: protocol, sys.stdin)
-            transport, protocol = await loop.connect_write_pipe(
-                asyncio.streams.FlowControlMixin, sys.stdout)  # type: ignore
-            stdout_writer = asyncio.StreamWriter(transport, protocol, None,
-                                                 loop)
-
-            await asyncio.gather(stdin_to_websocket(stdin_reader, websocket),
-                                 websocket_to_stdout(websocket, stdout_writer))
-        finally:
-            if old_settings:
-                termios.tcsetattr(sys.stdin.fileno(), termios.TCSADRAIN,
-                                  old_settings)
+            await asyncio.sleep(HEARTBEAT_INTERVAL_SECONDS)
+            if len(last_ping_time_dict) >= MAX_UNANSWERED_PINGS:
+                # We are not getting responses, clear the dictionary so
+                # as not to grow unbounded.
+                last_ping_time_dict.clear()
+            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)
+            try:
+                async with websocket_lock:
+                    await websocket.send(message_header_bytes)
+            except websockets.exceptions.ConnectionClosed as e:
+                # Websocket is already closed.
+                print(f'Failed to send PING message: {e}', file=sys.stderr)
+                break
+        except Exception as e:
+            print(f'Error in latency_monitor: {e}', file=sys.stderr)
+            websocket_closed_event.set()
+            raise e
 
 
 async def stdin_to_websocket(reader: asyncio.StreamReader,
-                             websocket: ClientConnection):
+                             websocket: ClientConnection,
+                             timestamps_supported: bool,
+                             websocket_closed_event: asyncio.Event,
+                             websocket_lock: asyncio.Lock):
     try:
-        while True:
+        while not websocket_closed_event.is_set():
             # Read at most BUFFER_SIZE bytes, this not affect
             # responsiveness since it will return as soon as
             # there is at least one byte.
             # The BUFFER_SIZE is chosen to be large enough to improve
             # throughput.
             data = await reader.read(BUFFER_SIZE)
+
             if not data:
                 break
-            await websocket.send(data)
+            if timestamps_supported:
+                # Send message with type 0 to indicate data.
+                message_type_bytes = struct.pack(
+                    '!B', KubernetesSSHMessageType.REGULAR_DATA.value)
+                data = message_type_bytes + data
+            async with websocket_lock:
+                await websocket.send(data)
+
     except Exception as e:  # pylint: disable=broad-except
         print(f'Error in stdin_to_websocket: {e}', file=sys.stderr)
     finally:
-        await websocket.close()
+        async with websocket_lock:
+            await websocket.close()
+        websocket_closed_event.set()
 
 
 async def websocket_to_stdout(websocket: ClientConnection,
-                              writer: asyncio.StreamWriter):
+                              writer: asyncio.StreamWriter,
+                              timestamps_supported: bool,
+                              last_ping_time_dict: Optional[dict],
+                              websocket_closed_event: asyncio.Event,
+                              websocket_lock: asyncio.Lock):
     try:
-        while True:
+        while not websocket_closed_event.is_set():
             message = await websocket.recv()
+            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:
+                    # Regular data - strip type byte and write to stdout
+                    message = message[1:]
+                elif message_type == KubernetesSSHMessageType.PINGPONG.value:
+                    # PONG response - calculate latency and send measurement
+                    if not len(message) == struct.calcsize('!BI'):
+                        raise ValueError(
+                            f'Invalid PONG message length: {len(message)}')
+                    pong_id = struct.unpack('!I', message[1:5])[0]
+                    pong_time = time.time()
+
+                    ping_time = last_ping_time_dict.pop(pong_id, None)
+
+                    if ping_time is None:
+                        continue
+
+                    latency_seconds = pong_time - ping_time
+                    latency_ms = int(latency_seconds * 1000)
+
+                    # Send latency measurement (type 2)
+                    message_type_bytes = struct.pack(
+                        '!B',
+                        KubernetesSSHMessageType.LATENCY_MEASUREMENT.value)
+                    latency_bytes = struct.pack('!Q', latency_ms)
+                    message = message_type_bytes + latency_bytes
+                    # Send to server.
+                    async with websocket_lock:
+                        await websocket.send(message)
+                    continue
+            # No timestamps support, write directly
             writer.write(message)
             await writer.drain()
     except websockets.exceptions.ConnectionClosed:
         print('WebSocket connection closed', file=sys.stderr)
     except Exception as e:  # pylint: disable=broad-except
         print(f'Error in websocket_to_stdout: {e}', file=sys.stderr)
+        raise e
+    finally:
+        async with websocket_lock:
+            await websocket.close()
+        websocket_closed_event.set()
 
 
 if __name__ == '__main__':
     server_url = sys.argv[1].strip('/')
-    if '://' not in server_url:
-        # Keep backward compatibility for legacy server URLs without protocol
-        # TODO(aylei): Remove this after 0.10.0
-        server_url = f'http://{server_url}'
 
+    disable_latency_measurement = os.environ.get(
+        skylet_constants.SSH_DISABLE_LATENCY_MEASUREMENT_ENV_VAR, '0') == '1'
+    if disable_latency_measurement:
+        timestamps_are_supported = False
+    else:
+        # 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)
+        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
+
+    # Capture the original API server URL for login hint if authentication
+    # is required.
+    _login_url = server_url
     server_proto, server_fqdn = server_url.split('://')
     websocket_proto = 'ws'
     if server_proto == 'https':
         websocket_proto = 'wss'
     server_url = f'{websocket_proto}://{server_fqdn}'
+
+    client_version_str = (f'&client_version={constants.API_VERSION}'
+                          if timestamps_are_supported else '')
+
     websocket_url = (f'{server_url}/kubernetes-pod-ssh-proxy'
-                     f'?cluster_name={sys.argv[2]}')
-    asyncio.run(main(websocket_url))
+                     f'?cluster_name={sys.argv[2]}'
+                     f'{client_version_str}')
+    asyncio.run(main(websocket_url, timestamps_are_supported, _login_url))