skypilot-nightly 1.0.0.dev20250509__py3-none-any.whl → 1.0.0.dev20251107__py3-none-any.whl

sky/templates/kubernetes-ray.yml.j2

@@ -33,14 +33,11 @@ provider:
     networking_mode: {{k8s_networking_mode}}
 
     # We use internal IPs since we set up a port-forward between the kubernetes
-    # cluster and the local machine, or directly use NodePort to reach the
-    # head node.
+    # cluster and the local machine.
     use_internal_ips: true
 
     timeout: {{timeout}}
 
-    ssh_jump_image: {{k8s_ssh_jump_image}}
-
     # Namespace used to host SkyPilot system components, such as fuse device
     # manager.
     skypilot_system_namespace: {{k8s_skypilot_system_namespace}}
@@ -212,7 +209,9 @@ provider:
         metadata:
           labels:
             parent: skypilot
+            # TODO (kyuds): remove this label for v0.12.0, as skypilot-cluster label is deprecated in favor of skypilot-cluster-name.
             skypilot-cluster: {{cluster_name_on_cloud}}
+            skypilot-cluster-name: {{cluster_name_on_cloud}}
             skypilot-user: {{ user }}
           name: {{cluster_name_on_cloud}}-head-ssh
         spec:
@@ -230,7 +229,9 @@ provider:
         metadata:
             labels:
               parent: skypilot
+              # TODO (kyuds): remove this label for v0.12.0, as skypilot-cluster label is deprecated in favor of skypilot-cluster-name.
               skypilot-cluster: {{cluster_name_on_cloud}}
+              skypilot-cluster-name: {{cluster_name_on_cloud}}
               skypilot-user: {{ user }}
             # NOTE: If you're running multiple Ray clusters with services
             # on one Kubernetes cluster, they must have unique service
@@ -243,6 +244,24 @@ provider:
             # This selector must match the head node pod's selector below.
             selector:
                 component: {{cluster_name_on_cloud}}-head
+      # Headless service mapping hostnames to rest of the worker nodes
+      {% for worker_id in range(1, num_nodes) %}
+      - apiVersion: v1
+        kind: Service
+        metadata:
+          labels:
+            parent: skypilot
+            # TODO (kyuds): remove this label for v0.12.0, as skypilot-cluster label is deprecated in favor of skypilot-cluster-name.
+            skypilot-cluster: {{cluster_name_on_cloud}}
+            skypilot-cluster-name: {{cluster_name_on_cloud}}
+            skypilot-user: {{ user }}
+          name: {{cluster_name_on_cloud}}-worker{{ worker_id }}
+        spec:
+          selector:
+              component: {{cluster_name_on_cloud}}-worker{{ worker_id }}
+          clusterIP: None
+      {% endfor %}
+
 
 # Specify the pod type for the ray head node (as configured below).
 head_node_type: ray_head_default
@@ -255,13 +274,12 @@ available_node_types:
       metadata:
         # name will be filled in the provisioner
         # head node name will be {{cluster_name_on_cloud}}-head, which will match the head node service selector above if a head node
-        # service is required.
+        # service is required. Worker nodes are named {{cluster_name_on_cloud}}-worker{{ node_id }}
         labels:
             parent: skypilot
             # component will be set for the head node pod to be the same as the head node service selector above if a
+            # TODO (kyuds): remove this label for v0.11.0, as skypilot-cluster label is deprecated in favor of skypilot-cluster-name.
             skypilot-cluster: {{cluster_name_on_cloud}}
-            # Identifies the SSH jump pod used by this pod. Used in life cycle management of the ssh jump pod.
-            skypilot-ssh-jump: {{k8s_ssh_jump_name}}
             skypilot-user: {{ user }}
             # Custom tags for the pods
             {%- for label_key, label_value in labels.items() %}
@@ -273,14 +291,100 @@ available_node_types:
             {% if (k8s_acc_label_key is not none and k8s_acc_label_values is not none) %}
             skypilot-binpack: "gpu"
             {% endif %}
+            {% if k8s_kueue_local_queue_name %}
+            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 or k8s_enable_gpudirect_tcpx or k8s_enable_gpudirect_tcpxo or k8s_enable_gpudirect_rdma %}
+        annotations:
+          {% 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
         serviceAccountName: {{k8s_service_account_name}}
         automountServiceAccountToken: {{k8s_automount_sa_token}}
         restartPolicy: {{ "Always" if high_availability else "Never" }}
+        {% if volume_mounts %}
+        securityContext:
+          fsGroup: 1000
+          fsGroupChangePolicy: OnRootMismatch
+        {% endif %}
 
         # Add node selector if GPU/TPUs are requested:
-        {% if (k8s_topology_label_key is not none and k8s_topology_label_value is not none) or (k8s_spot_label_key is not none) %}
+        {% if (k8s_topology_label_key is not none and k8s_topology_label_value is not none) or (k8s_spot_label_key is not none) or (k8s_enable_flex_start) %}
         nodeSelector:
             {% if k8s_topology_label_key is not none and k8s_topology_label_value is not none %}
             {{k8s_topology_label_key}}: {{k8s_topology_label_value}}
@@ -288,6 +392,9 @@ available_node_types:
             {% if k8s_spot_label_key is not none %}
             {{k8s_spot_label_key}}: {{k8s_spot_label_value|tojson}}
             {% endif %}
+            {% if k8s_enable_flex_start %}
+            cloud.google.com/gke-flex-start: "true"
+            {% endif %}
         {% endif %}
         {% if (k8s_acc_label_key is not none and k8s_acc_label_values is not none) or (avoid_label_keys is not none) %}
         affinity:
@@ -339,9 +446,6 @@ available_node_types:
         # object store. If you do not provide this, Ray will fall back to
         # /tmp which cause slowdowns if is not a shared memory volume.
         volumes:
-        - name: secret-volume
-          secret:
-            secretName: {{k8s_ssh_key_secret_name}}
         - name: dshm
           emptyDir:
             medium: Memory
@@ -356,19 +460,176 @@ available_node_types:
           persistentVolumeClaim:
             claimName: {{cluster_name_on_cloud}}-{{k8s_high_availability_deployment_volume_mount_name}}
         {% endif %}
+        {% for volume_mount in volume_mounts %}
+        - name: {{volume_mount.name}}
+          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: IfNotPresent
+          imagePullPolicy: Always
           image: {{image_id}}
           env:
             - name: SKYPILOT_POD_NODE_TYPE
               valueFrom:
                 fieldRef:
                   fieldPath: metadata.labels['ray-node-type']
+            - name: SKYPILOT_POD_CPU_CORE_LIMIT
+              valueFrom:
+                resourceFieldRef:
+                  containerName: ray-node
+                  resource: requests.cpu
+            - name: SKYPILOT_POD_MEMORY_BYTES_LIMIT
+              valueFrom:
+                resourceFieldRef:
+                  containerName: ray-node
+                  resource: requests.memory
             {% for key, value in k8s_env_vars.items() if k8s_env_vars is not none %}
             - 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}}
@@ -378,13 +639,9 @@ available_node_types:
           command: ["/bin/bash", "-c", "--"]
           args:
             - |
-              # For backwards compatibility, we put a marker file in the pod
-              # to indicate that the pod is running with the changes introduced
-              # in project nimbus: https://github.com/skypilot-org/skypilot/pull/4393
-              # TODO: Remove this marker file and it's usage in setup_commands
-              # after v0.10.0 release.
-              touch /tmp/skypilot_is_nimbus
-
+              # Set -x to print the commands and their arguments as they are executed.
+              # Useful for debugging.
+              set -x
               # Helper function to conditionally use sudo
               # TODO(zhwu): consolidate the two prefix_cmd and sudo replacements
               prefix_cmd() { if [ $(id -u) -ne 0 ]; then echo "sudo"; else echo ""; fi; }
@@ -395,14 +652,138 @@ available_node_types:
               # STEP 1: Run apt update, install missing packages, and set up ssh.
               (
                 (
-                DEBIAN_FRONTEND=noninteractive $(prefix_cmd) apt-get update > /tmp/apt-update.log 2>&1 || \
-                echo "Warning: apt-get update failed. Continuing anyway..." >> /tmp/apt-update.log
+                # For backwards compatibility, we put a marker file in the pod
+                # to indicate that the apt ssh setup step will write a completion
+                # marker file (/tmp/apt_ssh_setup_complete) to the pod.
+                # TODO: Remove this marker file and its usage in setup_commands
+                # after v0.11.0 release.
+                touch /tmp/apt_ssh_setup_started
+
+                # Helper: run apt-get update with retries
+                apt_update_with_retries() {
+                  # do not fail the whole shell; we handle return codes
+                  set +e
+                  local log=/tmp/apt-update.log
+                  local tries=3
+                  local delay=1
+                  local i
+                  for i in $(seq 1 $tries); do
+                    DEBIAN_FRONTEND=noninteractive $(prefix_cmd) apt-get update >> "$log" 2>&1 && { set -e; return 0; }
+                    echo "apt-get update attempt $i/$tries failed; retrying in ${delay}s" >> "$log"
+                    sleep $delay
+                    delay=$((delay * 2))
+                  done
+                  set -e
+                  return 1
+                }
+                apt_install_with_retries() {
+                  local packages="$@"
+                  [ -z "$packages" ] && return 0
+                  set +e
+                  local log=/tmp/apt-update.log
+                  local tries=3
+                  local delay=1
+                  local i
+                  for i in $(seq 1 $tries); do
+                    DEBIAN_FRONTEND=noninteractive $(prefix_cmd) apt-get install -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" $packages && { set -e; return 0; }
+                    echo "apt-get install failed for: $packages (attempt $i/$tries). Running -f install and retrying..." >> "$log"
+                    DEBIAN_FRONTEND=noninteractive $(prefix_cmd) apt-get -f install -y >> "$log" 2>&1 || true
+                    DEBIAN_FRONTEND=noninteractive $(prefix_cmd) apt-get clean >> "$log" 2>&1 || true
+                    sleep $delay
+                    delay=$((delay * 2))
+                  done
+                  set -e
+                  return 1
+                }
+                apt_update_install_with_retries() {
+                  apt_update_with_retries
+                  apt_install_with_retries "$@"
+                }
+                backup_dir=/etc/apt/sources.list.backup_skypilot
+                backup_source() {
+                  $(prefix_cmd) mkdir -p "$backup_dir"
+                  if [ -f /etc/apt/sources.list ] && [ ! -f "$backup_dir/sources.list" ]; then
+                    $(prefix_cmd) cp -a /etc/apt/sources.list "$backup_dir/sources.list" || true
+                  fi
+                }
+                restore_source() {
+                  if [ -f "$backup_dir/sources.list" ]; then
+                    $(prefix_cmd) cp -a "$backup_dir/sources.list" /etc/apt/sources.list || true
+                  fi
+                }
+                update_apt_sources() {
+                  local host=$1
+                  local apt_file=$2
+                  $(prefix_cmd) sed -i -E "s|https?://[a-zA-Z0-9.-]+\.ubuntu\.com/ubuntu|http://$host/ubuntu|g" $apt_file
+                }
+                # Helper: install packages across mirrors with retries
+                apt_install_with_mirrors() {
+                  local required=$1; shift
+                  local packages="$@"
+                  [ -z "$packages" ] && return 0
+                  set +e
+                  # Install packages with default sources first
+                  local log=/tmp/apt-update.log
+                  echo "$(date +%Y-%m-%d\ %H:%M:%S) Installing packages: $packages" >> "$log"
+                  restore_source
+                  apt_update_install_with_retries $packages >> "$log" 2>&1 && { set -e; return 0; }
+                  echo "Install failed with default sources: $packages" >> "$log"
+                  # Detect distro (ubuntu/debian)
+                  local APT_OS="unknown"
+                  if [ -f /etc/os-release ]; then
+                    . /etc/os-release
+                    case "$ID" in
+                      debian) APT_OS="debian" ;;
+                      ubuntu) APT_OS="ubuntu" ;;
+                      *)
+                        if [ -n "$ID_LIKE" ]; then
+                          case " $ID $ID_LIKE " in
+                            *ubuntu*) APT_OS="ubuntu" ;;
+                            *debian*) APT_OS="debian" ;;
+                          esac
+                        fi
+                        ;;
+                    esac
+                  fi
+                  # Build mirror candidates
+                  # deb.debian.org is a CDN endpoint, if one backend goes down,
+                  # the CDN automatically fails over to another mirror,
+                  # so we only retry for ubuntu here.
+                  if [ "$APT_OS" = "ubuntu" ]; then
+                    # Backup current sources once
+                    backup_source
+                    # Selected from https://launchpad.net/ubuntu/+archivemirrors
+                    # and results from apt-select
+                    local MIRROR_CANDIDATES="mirrors.wikimedia.org mirror.umd.edu"
+                    for host in $MIRROR_CANDIDATES; do
+                      echo "Trying APT mirror ($APT_OS): $host" >> "$log"
+                      if [ -f /etc/apt/sources.list ]; then
+                        update_apt_sources $host /etc/apt/sources.list
+                      else
+                        echo "Error: /etc/apt/sources.list not found" >> "$log"
+                        break
+                      fi
+                      apt_update_install_with_retries $packages >> "$log" 2>&1 && { set -e; return 0; }
+                      echo "Install failed with mirror ($APT_OS): $host" >> "$log"
+                      # Restore to default sources
+                      restore_source
+                    done
+                  fi
+                  set -e
+                  if [ "$required" = "1" ]; then
+                    echo "Error: required package install failed across all mirrors: $packages" >> "$log"
+                    return 1
+                  else
+                    echo "Optional package install failed across all mirrors: $packages; skipping." >> "$log"
+                    return 0
+                  fi
+                }
                 # Install both fuse2 and fuse3 for compatibility for all possible fuse adapters in advance,
                 # so that both fusemount and fusermount3 can be masked before enabling SSH access.
                 PACKAGES="rsync curl wget netcat gcc patch pciutils fuse fuse3 openssh-server";
 
                 # Separate packages into two groups: packages that are installed first
-                # so that curl, rsync and wget are available sooner to unblock the following
+                # so that curl, rsync, ssh and wget are available sooner to unblock the following
                 # conda installation and rsync.
                 # Also, we install fuse first to avoid confliction with fuse3.
                 set -e
@@ -423,7 +804,7 @@ available_node_types:
                 done;
                 if [ ! -z "$INSTALL_FIRST" ]; then
                   echo "Installing core packages: $INSTALL_FIRST";
-                  DEBIAN_FRONTEND=noninteractive $(prefix_cmd) apt-get install -y -o Dpkg::Options::="--force-confdef" -o Dpkg::Options::="--force-confold" $INSTALL_FIRST;
+                  apt_install_with_mirrors 1 $INSTALL_FIRST || { echo "Error: core package installation failed." >> /tmp/apt-update.log; exit 1; }
                 fi;
                 # SSH and other packages are not necessary, so we disable set -e
                 set +e
@@ -447,7 +828,8 @@ available_node_types:
                 fi
                 $(prefix_cmd) cp -p "$FUSERMOUNT_PATH" "${FUSERMOUNT_PATH}-original"
                 $(prefix_cmd) ln -sf {{k8s_fusermount_shared_dir}}/fusermount-shim "$FUSERMOUNT_PATH"
-                FUSERMOUNT3_PATH=$(which fusermount3)
+                # "|| true" because fusermount3 is not always available
+                FUSERMOUNT3_PATH=$(which fusermount3) || true
                 if [ -z "$FUSERMOUNT3_PATH" ]; then
                     FUSERMOUNT3_PATH="${FUSERMOUNT_PATH}3"
                 fi
@@ -489,16 +871,23 @@ available_node_types:
                 $(prefix_cmd) mkdir -p ~/.ssh;
                 $(prefix_cmd) chown -R $(whoami) ~/.ssh;
                 $(prefix_cmd) chmod 700 ~/.ssh;
-                $(prefix_cmd) cat /etc/secret-volume/ssh-publickey* > ~/.ssh/authorized_keys;
+                $(prefix_cmd) cat > ~/.ssh/authorized_keys <<'SKYPILOT_SSH_KEY_EOF'
+              skypilot:ssh_public_key_content
+              SKYPILOT_SSH_KEY_EOF
                 $(prefix_cmd) chmod 644 ~/.ssh/authorized_keys;
                 $(prefix_cmd) service ssh restart;
                 $(prefix_cmd) sed -i "s/mesg n/tty -s \&\& mesg n/" ~/.profile;
 
-                ) > /tmp/${STEPS[0]}.log 2>&1 || {
-                  echo "Error: ${STEPS[0]} failed. Continuing anyway..." > /tmp/${STEPS[0]}.failed
+                touch /tmp/apt_ssh_setup_complete
+                echo "=== SSH setup completed ==="
+                ) > /tmp/${STEPS[0]}.log 2>&1
+                if [ "$?" -ne "0" ]; then
+                  {
+                  echo "Error: ${STEPS[0]} failed. Continuing anyway..." > /tmp/${STEPS[0]}.failed 2>&1
                   cat /tmp/${STEPS[0]}.log
                   exit 1
-                }
+                  }
+                fi
               ) &
 
               # STEP 2: Install conda, ray and skypilot (for dependencies); start
@@ -516,7 +905,21 @@ available_node_types:
                 {{ conda_installation_commands }}
                 {{ ray_installation_commands }}
 
-                VIRTUAL_ENV=~/skypilot-runtime ~/.local/bin/uv pip install skypilot[kubernetes,remote]
+                # 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]
+                # 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;
+                }
                 touch /tmp/ray_skypilot_installation_complete
                 echo "=== Ray and skypilot installation completed ==="
 
@@ -544,11 +947,14 @@ available_node_types:
                   set +e
                   {{ ray_worker_start_command }}
                 fi
-                ) > /tmp/${STEPS[1]}.log 2>&1 || {
-                  echo "Error: ${STEPS[1]} failed. Continuing anyway..." > /tmp/${STEPS[1]}.failed
+                ) > /tmp/${STEPS[1]}.log 2>&1
+                if [ "$?" -ne "0" ]; then
+                  {
+                  echo "Error: ${STEPS[1]} failed. Continuing anyway..." > /tmp/${STEPS[1]}.failed 2>&1
                   cat /tmp/${STEPS[1]}.log
                   exit 1
-                }
+                  }
+                fi
               ) &
 
 
@@ -566,11 +972,14 @@ available_node_types:
                   fi;
                 fi;
                 export -p > ~/container_env_var.sh && $(prefix_cmd) mv ~/container_env_var.sh /etc/profile.d/container_env_var.sh
-                ) > /tmp/${STEPS[2]}.log 2>&1 || {
-                  echo "Error: ${STEPS[2]} failed. Continuing anyway..." > /tmp/${STEPS[2]}.failed
+                ) > /tmp/${STEPS[2]}.log 2>&1
+                if [ "$?" -ne "0" ]; then
+                  {
+                  echo "Error: ${STEPS[2]} failed. Continuing anyway..." > /tmp/${STEPS[2]}.failed 2>&1
                   cat /tmp/${STEPS[2]}.log
                   exit 1
-                }
+                  }
+                fi
               ) &
 
               function mylsof { p=$(for pid in /proc/{0..9}*; do i=$(basename "$pid"); for file in "$pid"/fd/*; do link=$(readlink -e "$file"); if [ "$link" = "$1" ]; then echo "$i"; fi; done; done); echo "$p"; };
@@ -623,23 +1032,72 @@ available_node_types:
               {% if high_availability %}
               mkdir -p {{k8s_high_availability_deployment_run_script_dir}}
               if [ -f {{k8s_high_availability_deployment_volume_mount_path}}/k8s_container_ready ]; then
+                SKYPILOT_HA_RECOVERY_LOG="{{ha_recovery_log_path}}"
+                echo "Starting HA recovery at $(date)" >> $SKYPILOT_HA_RECOVERY_LOG
+                start_time=$SECONDS
+                retry_count=0
+
+                # Wait for Ray to be ready, as the following commands is depending on Ray.
+                GET_RAY_STATUS_CMD=$({{sky_python_cmd}} -c 'from sky.provision import instance_setup; print(instance_setup.RAY_STATUS_WITH_SKY_RAY_PORT_COMMAND)')
+                while true; do
+                  retry_count=$((retry_count + 1))
+                  current_duration=$(( SECONDS - start_time ))
+                  echo "Attempt $retry_count to get Ray status after $current_duration seconds..." >> $SKYPILOT_HA_RECOVERY_LOG
+
+                  bash --login -c "$GET_RAY_STATUS_CMD"
+                  if [ $? -eq 0 ]; then
+                    wait_duration=$(( SECONDS - start_time ))
+                    echo "Ray ready after waiting $wait_duration seconds (took $retry_count attempts)" >> $SKYPILOT_HA_RECOVERY_LOG
+                    break
+                  fi
+                  echo "Waiting for Ray to be ready..." >> $SKYPILOT_HA_RECOVERY_LOG
+                  sleep 2
+                done
+
                 # ! Keep this aligned with `CloudVmRayBackend._setup()`
-                # Suppose all `task.setup` are the same for skyserve controller task.
+                # Suppose all `task.setup` are the same for sky serve / managed jobs controller task.
                 # So be careful for compatibility issue once you change it.
                 chmod +x {{k8s_high_availability_deployment_setup_script_path}}
                 /bin/bash --login -c "true && export OMP_NUM_THREADS=1 PYTHONWARNINGS='ignore' && {{k8s_high_availability_deployment_setup_script_path}} > /tmp/controller_recovery_setup_commands.log 2>&1"
-                echo "=== Controller setup commands completed for recovery ==="
+                echo "=== Controller setup commands completed for recovery at $(date) ===" >> $SKYPILOT_HA_RECOVERY_LOG
 
+                touch {{k8s_high_availability_restarting_signal_file}}
+                # Get all in-progress jobs from managed jobs controller. We skip any jobs that are already done.
+                # Also, skip the jobs that are waiting to be scheduled as those does not have a controller process running.
+                # For SkyServe, this will be None and every service will be recovered. This is because SkyServe
+                # will delete the service from the database after it is terminated so everything in the database is running.
+                ALL_IN_PROGRESS_JOBS=$({{sky_python_cmd}} -c "from sky.jobs import state; jobs, _ = state.get_managed_jobs_with_filters(fields=['job_id', 'schedule_state']); print(' '.join({str(job['job_id']) for job in jobs if job['schedule_state'] not in [state.ManagedJobScheduleState.DONE, state.ManagedJobScheduleState.WAITING]}) if jobs else None)")
+                if [ "$ALL_IN_PROGRESS_JOBS" != "None" ]; then
+                  read -ra ALL_IN_PROGRESS_JOBS_SEQ <<< "$ALL_IN_PROGRESS_JOBS"
+                fi
                 for file in {{k8s_high_availability_deployment_run_script_dir}}/*; do
+                  # This is the cluster job id on managed jobs controller, but it is guaranteed to be the same as the managed job id,
+                  # so we directly use it here. See `CloudVmRayBackend._exec_code_on_head::_dump_code_to_file` for more details.
+                  JOB_ID=$(basename $file | sed 's/sky_job_//')
+                  # If the list of in-progress jobs is not None (meaning this is a managed job HA controller) and job is not in-progress, skip.
+                  if [ "$ALL_IN_PROGRESS_JOBS" != "None" ]; then
+                    if [[ ! " ${ALL_IN_PROGRESS_JOBS_SEQ[@]} " =~ " ${JOB_ID} " ]]; then
+                      continue
+                    fi
+                  fi
                   # ! Keep this aligned with `CloudVmRayBackend._execute()`
                   chmod +x $file
+                  # TODO(tian): This logic may run a lot of things if the jobs controller previously had many jobs.
+                  # We should do more tests and make sure it will scale well.
                   /bin/bash --login -c "true && export OMP_NUM_THREADS=1 PYTHONWARNINGS='ignore' && $file > /tmp/task_run_$(basename $file).log 2>&1"
-                  echo "=== Controller task run for service (file: $file) completed for recovery ==="
+                  echo "=== Controller task run for service / job (file: $file) completed for recovery at $(date) ===" >> $SKYPILOT_HA_RECOVERY_LOG
                 done
+                rm {{k8s_high_availability_restarting_signal_file}}
+
+                duration=$(( SECONDS - start_time ))
+                echo "HA recovery completed at $(date)" >> $SKYPILOT_HA_RECOVERY_LOG
+                echo "Total recovery time: $duration seconds" >> $SKYPILOT_HA_RECOVERY_LOG
               fi
 
               touch {{k8s_high_availability_deployment_volume_mount_path}}/k8s_container_ready
               {% endif %}
+              # Set +x to stop printing the commands and their arguments as they are executed.
+              set +x
 
               trap : TERM INT; log_tail || sleep infinity & wait
 
@@ -653,14 +1111,27 @@ available_node_types:
           # object store. If you do not provide this, Ray will fall back to
           # /tmp which cause slowdowns if is not a shared memory volume.
           volumeMounts:
-          - 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}}
@@ -669,6 +1140,10 @@ available_node_types:
           - name: fusermount-shared-dir
             mountPath: {{k8s_fusermount_shared_dir}}
           {% endif %}
+          {% for volume_mount in volume_mounts %}
+          - name: {{volume_mount.name}}
+            mountPath: {{volume_mount.path}}
+          {% endfor %}
           resources:
             requests:
               cpu: {{cpus}}
@@ -681,13 +1156,87 @@ available_node_types:
               # https://cloud.google.com/kubernetes-engine/docs/concepts/tpus#how_tpus_work
               {{k8s_resource_key}}: {{accelerator_count}}
               {% endif %}
+              {% if k8s_network_type == 'coreweave' %}
+              rdma/ib: 1
+              {% endif %}
             {% if k8s_resource_key is not none %}
             limits:
               # Limits need to be defined for GPU/TPU requests
               {% if k8s_resource_key is not none %}
               {{k8s_resource_key}}: {{accelerator_count}}
               {% endif %}
+              {% if k8s_network_type == 'coreweave' %}
+              rdma/ib: 1
+              {% endif %}
             {% endif %}
+          {% if k8s_ipc_lock_capability %}
+          securityContext:
+            capabilities:
+              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:
@@ -724,7 +1273,7 @@ available_node_types:
             spec:
               securityContext:
                 fsGroup: 1000
-              # To prevent the home dir provided by the docker image from being overriden by pvc mounting,
+              # To prevent the home dir provided by the docker image from being overridden by pvc mounting,
               # we use initContainers to copy it first to /mnt/home, which will later be mounted to home dir.
               initContainers:
                 - name: init-copy-home
@@ -791,13 +1340,20 @@ setup_commands:
     {%- endfor %}
     STEPS=("apt-ssh-setup" "runtime-setup" "env-setup")
     start_epoch=$(date +%s);
-    echo "=== Logs for asynchronous ray and skypilot installation ===";
-    if [ -f /tmp/skypilot_is_nimbus ]; then
-      echo "=== Logs for asynchronous ray and skypilot installation ===";
-      [ -f /tmp/ray_skypilot_installation_complete ] && cat /tmp/${STEPS[1]}.log ||
-      { tail -f -n +1 /tmp/${STEPS[1]}.log & TAIL_PID=$!; echo "Tail PID: $TAIL_PID"; until [ -f /tmp/ray_skypilot_installation_complete ]; do sleep 0.5; done; kill $TAIL_PID || true; };
-      [ -f /tmp/${STEPS[1]}.failed ] && { echo "Error: ${STEPS[1]} failed. Exiting."; exit 1; } || true;
+
+    # Wait for SSH setup to complete before proceeding
+    if [ -f /tmp/apt_ssh_setup_started ]; then
+      echo "=== Logs for asynchronous SSH setup ===";
+      ([ -f /tmp/apt_ssh_setup_complete ]|| [ -f /tmp/${STEPS[0]}.failed ]) && cat /tmp/${STEPS[0]}.log ||
+      { tail -f -n +1 /tmp/${STEPS[0]}.log & TAIL_PID=$!; echo "Tail PID: $TAIL_PID"; sleep 0.5; until [ -f /tmp/apt_ssh_setup_complete ] || [ -f /tmp/${STEPS[0]}.failed ]; do sleep 0.5; done; kill $TAIL_PID || true; };
+      [ -f /tmp/${STEPS[0]}.failed ] && { echo "Error: ${STEPS[0]} failed. Exiting."; exit 1; } || true;
     fi
+
+    echo "=== Logs for asynchronous ray and skypilot installation ===";
+    ([ -f /tmp/ray_skypilot_installation_complete ]|| [ -f /tmp/${STEPS[1]}.failed ]) && cat /tmp/${STEPS[1]}.log ||
+    { tail -f -n +1 /tmp/${STEPS[1]}.log & TAIL_PID=$!; echo "Tail PID: $TAIL_PID"; sleep 0.5; until [ -f /tmp/ray_skypilot_installation_complete ] || [ -f /tmp/${STEPS[1]}.failed ]; do sleep 0.5; done; kill $TAIL_PID || true; };
+    [ -f /tmp/${STEPS[1]}.failed ] && { echo "Error: ${STEPS[1]} failed. Exiting."; exit 1; } || true;
+
     end_epoch=$(date +%s);
     echo "=== Ray and skypilot dependencies installation completed in $(($end_epoch - $start_epoch)) secs ===";
     start_epoch=$(date +%s);