| 1 | #!/bin/bash |
| 2 | #SBATCH --job-name=vllm-disagg-pd |
| 3 | #SBATCH --nodes=2 |
| 4 | #SBATCH --ntasks-per-node=1 |
| 5 | #SBATCH --gres=gpu:8 |
| 6 | #SBATCH --exclusive |
| 7 | #SBATCH --chdir=/tmp |
| 8 | #SBATCH --output=/tmp/spur-%j.out |
| 9 | #SBATCH --error=/tmp/spur-%j.err |
| 10 | |
| 11 | set -euo pipefail |
| 12 | |
| 13 | |
| 14 | # ---- inputs (env > default; forwarded by run-slurm-disagg-test.sh via --export) |
| 15 | IMAGE="${IMAGE:-vllm/vllm-openai-rocm:nightly}" |
| 16 | WIDE_EP_MODE="${WIDE_EP_MODE:-0}" |
| 17 | xP="${xP:-1}" |
| 18 | yD="${yD:-1}" |
| 19 | GPUS_PER_NODE="${GPUS_PER_NODE:-8}" |
| 20 | |
| 21 | SHARED_MOUNT="${SHARED_MOUNT:-/data}" |
| 22 | LOG_ROOT="${LOG_ROOT:-/data/${USER:-$(id -un)}/disagg_logs}" |
| 23 | RUN_AFTER_HEALTH="${RUN_AFTER_HEALTH:-accuracy}" # bench | accuracy |
| 24 | DRY_RUN="${DRY_RUN:-0}" |
| 25 | MORIIO_READ_MODE="${MORIIO_READ_MODE:-0}" # 0 -> prefill writes (default); 1 -> decode reads KV (read_mode) |
| 26 | |
| 27 | # ---- router / gateway (host-side; also forwarded into the container) -------- |
| 28 | # vllm-router -> external `vllm/vllm-router` container started on the rank-0 node (default) |
| 29 | # proxy -> in-container MoRIIO proxy (no extra container) |
| 30 | ROUTER_TYPE="${ROUTER_TYPE:-vllm-router}" |
| 31 | ROUTER_PORT="${ROUTER_PORT:-30000}" |
| 32 | ROUTER_POLICY="${ROUTER_POLICY:-round_robin}" |
| 33 | VLLM_ROUTER_IMAGE="${VLLM_ROUTER_IMAGE:-vllm/vllm-router:nightly}" |
| 34 | PROXY_PING_PORT="${PROXY_PING_PORT:-36367}" # must equal vllm-router --vllm-discovery-address port |
| 35 | PROXY_PORT="${PROXY_PORT:-10001}" |
| 36 | # Intra-node data-parallel size the router uses to route the MoRIIO KV-notify to |
| 37 | # the right DP rank. Required for wideEP DP (else decode hangs: "remote blocks |
| 38 | # never arrived"). EP -> GPUS_PER_NODE (matches --data-parallel-size-local); TP -> 1. |
| 39 | if [[ "${WIDE_EP_MODE}" == "1" ]]; then |
| 40 | ROUTER_DP_LOCAL="${ROUTER_DP_LOCAL:-${GPUS_PER_NODE}}" |
| 41 | else |
| 42 | ROUTER_DP_LOCAL="${ROUTER_DP_LOCAL:-1}" |
| 43 | fi |
| 44 | if [[ "${ROUTER_TYPE}" == "vllm-router" ]]; then |
| 45 | GATEWAY_PORT="${GATEWAY_PORT:-${ROUTER_PORT}}" |
| 46 | else |
| 47 | GATEWAY_PORT="${GATEWAY_PORT:-${PROXY_PORT}}" |
| 48 | fi |
| 49 | |
| 50 | # This line is a recommended template from infra side. |
| 51 | exec > "$LOG_ROOT/${SLURM_JOB_NAME}-${SLURM_JOB_ID}.log" 2>&1 |
| 52 | |
| 53 | # Where the disagg scripts live ON THE HOST (bind-mounted into the container). |
| 54 | # Spur runs the job body from /tmp, so $(pwd)/SLURM_SUBMIT_DIR aren't reliable; |
| 55 | # default to the known host scripts dir. Override via env if you relocate it. |
| 56 | DISAGG_SCRIPTS_DIR="${DISAGG_SCRIPTS_DIR:-${SLURM_SUBMIT_DIR:-$(pwd)}}" |
| 57 | echo "DISAGG SCRIPTS DIR is ----" |
| 58 | echo "$DISAGG_SCRIPTS_DIR" |
| 59 | |
| 60 | # In-container path the scripts are mounted to (and sourced from). |
| 61 | CONTAINER_SCRIPTS="${CONTAINER_SCRIPTS:-/vllm-workspace/.buildkite/amd-disagg}" |
| 62 | CLUSTER_ENV_IN="${CLUSTER_ENV_IN:-${CONTAINER_SCRIPTS}/cluster.sh}" |
| 63 | |
| 64 | |
| 65 | # Health endpoints (TP-mode ports from cluster.sh defaults). |
| 66 | PREFILL_PORT="${PREFILL_PORT:-8100}" |
| 67 | DECODE_PORT="${DECODE_PORT:-8200}" |
| 68 | HEALTH_TIMEOUT_S="${HEALTH_TIMEOUT_S:-3600}" |
| 69 | |
| 70 | # EP-mode ports, mirrored from cluster.sh purely so the preflight below can test |
| 71 | # them host-side; the in-container values still come from cluster.sh. |
| 72 | SERVE_PORT="${SERVE_PORT:-20005}" |
| 73 | RPC_PORT="${RPC_PORT:-13345}" |
| 74 | KV_PORT="${KV_PORT:-9711}" |
| 75 | LOCAL_PING_PORT="${LOCAL_PING_PORT:-61555}" |
| 76 | # vllm-router binds this for Prometheus in addition to ROUTER_PORT, and panics |
| 77 | # at startup if it is taken. |
| 78 | ROUTER_METRICS_PORT="${ROUTER_METRICS_PORT:-29000}" |
| 79 | |
| 80 | # Node preflight (see preflight_node): 0 disables, VRAM% above which a GPU counts |
| 81 | # as busy, and how long to let a reaped container's VRAM drain before judging. |
| 82 | PREFLIGHT="${PREFLIGHT:-1}" |
| 83 | PREFLIGHT_MAX_VRAM_PCT="${PREFLIGHT_MAX_VRAM_PCT:-10}" |
| 84 | PREFLIGHT_GPU_SETTLE_S="${PREFLIGHT_GPU_SETTLE_S:-30}" |
| 85 | |
| 86 | |
| 87 | NUM_NODES=$((xP + yD)) |
| 88 | echo "Calculated NUM_NODES: $NUM_NODES (xP=$xP + yD=$yD) WIDE_EP_MODE=$WIDE_EP_MODE" |
| 89 | |
| 90 | # ------------------------ |
| 91 | # Extract NUM_NODES from SLURM allocation and select nodes |
| 92 | # ------------------------ |
| 93 | echo "Original SLURM allocation:" |
| 94 | echo "SLURM_JOB_NODELIST: $SLURM_JOB_NODELIST" |
| 95 | echo "SLURM_NNODES: $SLURM_NNODES" |
| 96 | echo "SLURM_NTASKS: $SLURM_NTASKS" |
| 97 | |
| 98 | # scontrol is available on classic Slurm but not in the Spur job body. Use it to |
| 99 | # expand compressed ranges when present, else split the (already comma-separated, |
| 100 | # full-hostname) SLURM_JOB_NODELIST on commas. |
| 101 | if command -v scontrol >/dev/null 2>&1; then |
| 102 | FULL_NODELIST=$(scontrol show hostnames "$SLURM_JOB_NODELIST") |
| 103 | else |
| 104 | FULL_NODELIST=$(echo "$SLURM_JOB_NODELIST" | tr ',' '\n') |
| 105 | fi |
| 106 | AVAILABLE_COUNT=$(echo "$FULL_NODELIST" | wc -l) |
| 107 | |
| 108 | if (( AVAILABLE_COUNT < NUM_NODES )); then |
| 109 | echo "ERROR: Need ${NUM_NODES} nodes (xP=${xP} + yD=${yD}) but only ${AVAILABLE_COUNT} allocated." >&2 |
| 110 | exit 1 |
| 111 | fi |
| 112 | |
| 113 | # Pick the first NUM_NODES nodes and order them to match how IPs get assigned: |
| 114 | # * classic Slurm: sort alphabetically -- srun assigns PROCID=0 to the |
| 115 | # alphabetically-first hostname, so IPS[] (built per node below) lines up. |
| 116 | # * Spur: preserve RANK order (SPUR_NODELIST), so the hostnames pair 1:1 with |
| 117 | # the rank-ordered IPs in SPUR_PEER_NODES; an alphabetical sort here would |
| 118 | # desync MASTER_NODE / MASTER_ADDR / IPADDRS from the actual ranks. |
| 119 | # IPADDRS[0..xP-1] = prefill, IPADDRS[xP..] = decode, so this order must match. |
| 120 | if command -v srun >/dev/null 2>&1; then |
| 121 | SELECTED_NODES=$(echo "$FULL_NODELIST" | head -n "$NUM_NODES" | sort) |
| 122 | elif [[ -n "${SPUR_NODELIST:-}" ]]; then |
| 123 | SELECTED_NODES=$(echo "${SPUR_NODELIST}" | tr ',' '\n' | sed '/^$/d' | head -n "$NUM_NODES") |
| 124 | else |
| 125 | SELECTED_NODES=$(echo "$FULL_NODELIST" | head -n "$NUM_NODES") |
| 126 | fi |
| 127 | SELECTED_NODELIST_STR=$(echo "$SELECTED_NODES" | tr '\n' ',' | sed 's/,$//') |
| 128 | |
| 129 | # Create new nodelist in SLURM format |
| 130 | # This is a simplified approach - for complex ranges, you might need more sophisticated parsing |
| 131 | NEW_SLURM_NODELIST=$(echo "$SELECTED_NODES" | paste -sd, | sed 's/,/,/g') |
| 132 | |
| 133 | # Update SLURM environment variables |
| 134 | export SLURM_NNODES=$NUM_NODES |
| 135 | export SLURM_NTASKS=$NUM_NODES |
| 136 | export SLURM_JOB_NUM_NODES=$NUM_NODES |
| 137 | export SLURM_NPROCS=$NUM_NODES |
| 138 | export SLURM_JOB_NODELIST="$NEW_SLURM_NODELIST" |
| 139 | export SLURM_NODELIST="$NEW_SLURM_NODELIST" |
| 140 | |
| 141 | # Keep other SLURM variables as they were or set defaults |
| 142 | export SLURM_TASKS_PER_NODE="1(x$NUM_NODES)" |
| 143 | export SLURM_SUBMIT_DIR="${SLURM_SUBMIT_DIR:-${HOME}}" |
| 144 | export SLURM_CLUSTER_NAME="${SLURM_CLUSTER_NAME:-vllm-di-ci}" |
| 145 | export SLURM_JOB_CPUS_PER_NODE="${SLURM_JOB_CPUS_PER_NODE:-1}" |
| 146 | export SLURM_JOB_PARTITION="${SLURM_JOB_PARTITION:-amd-rccl}" |
| 147 | export SLURM_JOBID="${SLURM_JOBID:-$SLURM_JOB_ID}" |
| 148 | export SLURM_JOB_QOS="${SLURM_JOB_QOS:-normal}" |
| 149 | export SLURM_JOB_ACCOUNT="${SLURM_JOB_ACCOUNT:-amd-rccl}" |
| 150 | export SLURM_NTASKS_PER_NODE=1 |
| 151 | export SLURM_SUBMIT_HOST="${SLURM_SUBMIT_HOST:-$(hostname)}" |
| 152 | export SLURM_JOB_ID="${SLURM_JOB_ID:-}" |
| 153 | export SLURM_CONF="${SLURM_CONF:-/etc/slurm/slurm.conf}" |
| 154 | export SLURM_JOB_NAME="${SLURM_JOB_NAME:-${xP}p${yD}d_bench-serving}" |
| 155 | |
| 156 | echo "" |
| 157 | echo "Updated SLURM Environment Variables:" |
| 158 | echo "SLURM_JOB_ID: $SLURM_JOB_ID" |
| 159 | echo "SLURM_JOB_NODELIST: $SLURM_JOB_NODELIST" |
| 160 | echo "SLURM_NNODES: $SLURM_NNODES" |
| 161 | echo "SLURM_NTASKS: $SLURM_NTASKS" |
| 162 | echo "SLURM_TASKS_PER_NODE: $SLURM_TASKS_PER_NODE" |
| 163 | echo "SLURM_JOB_CPUS_PER_NODE: $SLURM_JOB_CPUS_PER_NODE" |
| 164 | echo "SLURM_JOB_PARTITION: $SLURM_JOB_PARTITION" |
| 165 | echo "SLURM_JOB_NUM_NODES: $SLURM_JOB_NUM_NODES" |
| 166 | echo "SLURM_JOBID: $SLURM_JOBID" |
| 167 | echo "SLURM_JOB_QOS: $SLURM_JOB_QOS" |
| 168 | echo "SLURM_NODELIST: $SLURM_NODELIST" |
| 169 | echo "SLURM_JOB_ACCOUNT: $SLURM_JOB_ACCOUNT" |
| 170 | echo "SLURM_NPROCS: $SLURM_NPROCS" |
| 171 | echo "SLURM_SUBMIT_HOST: $SLURM_SUBMIT_HOST" |
| 172 | echo "SLURM_CONF: $SLURM_CONF" |
| 173 | echo "SLURM_JOB_NAME: $SLURM_JOB_NAME" |
| 174 | echo "SLURM_NTASKS_PER_NODE: $SLURM_NTASKS_PER_NODE" |
| 175 | echo "SLURM_SUBMIT_DIR: $SLURM_SUBMIT_DIR" |
| 176 | echo "SLURM_CLUSTER_NAME: $SLURM_CLUSTER_NAME" |
| 177 | echo "ulimit: $(ulimit -a)" |
| 178 | echo "" |
| 179 | echo "Selected nodes for execution:" |
| 180 | echo "$SELECTED_NODES" |
| 181 | echo "" |
| 182 | |
| 183 | # Node information |
| 184 | MASTER_NODE=$(echo "$SELECTED_NODES" | head -n 1) |
| 185 | MASTER_PORT=39566 # Choose an open port |
| 186 | |
| 187 | IPS=() |
| 188 | if command -v srun >/dev/null 2>&1; then |
| 189 | # Classic Slurm: resolve each node's IP via srun (original behavior). |
| 190 | for NODE in $SELECTED_NODES; do |
| 191 | IP=$(srun --nodes=1 --ntasks=1 --time=00:20:00 --nodelist="$NODE" bash -c 'hostname -I' | awk 'NR==1 {print $1}') |
| 192 | IPS+=("$IP") |
| 193 | done |
| 194 | else |
| 195 | # Spur job body (no srun): use the rank-ordered IPs Spur provides in |
| 196 | # SPUR_PEER_NODES ("ip:port,ip:port,..."). These are the cross-node-routable |
| 197 | # bootstrap addresses; hostnames resolve to public IPs and the local node to |
| 198 | # 127.0.1.1, so DNS is not usable here. |
| 199 | mapfile -t IPS < <(echo "${SPUR_PEER_NODES:-}" | tr ',' '\n' | cut -d: -f1 | sed '/^$/d') |
| 200 | if (( ${#IPS[@]} == 0 )); then |
| 201 | echo "ERROR: no srun and SPUR_PEER_NODES is empty; cannot determine node IPs." >&2 |
| 202 | exit 1 |
| 203 | fi |
| 204 | fi |
| 205 | MASTER_ADDR="${IPS[0]}" |
| 206 | |
| 207 | IPADDRS="${IPS[*]}" |
| 208 | IPADDRS="${IPADDRS// /,}" |
| 209 | echo "Selected node IPs: ${IPADDRS}" |
| 210 | |
| 211 | NNODES=$NUM_NODES |
| 212 | |
| 213 | echo "MASTER_NODE is ${MASTER_NODE}" |
| 214 | echo "MASTER_ADDR is ${MASTER_ADDR}" |
| 215 | echo "MASTER_PORT is ${MASTER_PORT}" |
| 216 | echo "NNODES is ${NNODES}" |
| 217 | |
| 218 | # Per-job log folder (MAD-style): all per-role server/proxy/bench logs |
| 219 | export LOG_PATH="${LOG_ROOT}/${SLURM_JOB_ID}" |
| 220 | mkdir -p "$LOG_PATH" |
| 221 | echo "[disagg-pd] LOG_PATH=$LOG_PATH" |
| 222 | |
| 223 | # ---------------------------------------------------------------------------- |
| 224 | # Launch one container per node via a single srun and hand off to |
| 225 | # vllm_disagg.sh node (rank-based self-select). |
| 226 | # ---------------------------------------------------------------------------- |
| 227 | # IPADDRS is the ordered node IP list (prefill IPs first, then decode). IPS[] |
| 228 | # was built in SELECTED_NODES (alphabetical) order — the same order srun assigns |
| 229 | # SLURM_PROCID — so the rank->role mapping lines up: |
| 230 | # rank 0 -> prefill master + MoRIIO proxy + orchestrator |
| 231 | # rank 1 .. xP-1 -> prefill |
| 232 | # rank xP .. end -> decode |
| 233 | export IPADDRS |
| 234 | |
| 235 | export WIDE_EP_MODE xP yD GPUS_PER_NODE RUN_AFTER_HEALTH HEALTH_TIMEOUT_S DRY_RUN MORIIO_READ_MODE |
| 236 | export IMAGE SHARED_MOUNT CONTAINER_SCRIPTS DISAGG_SCRIPTS_DIR |
| 237 | export CLUSTER_ENV="${CLUSTER_ENV_IN}" |
| 238 | export DOCKER_CONT_NAME="vllm_disagg_${MODEL_NAME:-model}_${SLURM_JOB_ID}" |
| 239 | |
| 240 | # Router env: forwarded into the main container (cluster.sh honors these) and used |
| 241 | # host-side to launch the external router container on rank 0. |
| 242 | export ROUTER_TYPE ROUTER_PORT ROUTER_POLICY VLLM_ROUTER_IMAGE PROXY_PING_PORT GATEWAY_PORT ROUTER_DP_LOCAL |
| 243 | export ROUTER_CONT_NAME="vllm_router_${SLURM_JOB_ID}" |
| 244 | |
| 245 | # Consumed by preflight_node on each node (srun forwards the environment). |
| 246 | export PREFILL_PORT DECODE_PORT PROXY_PORT ROUTER_METRICS_PORT |
| 247 | export SERVE_PORT RPC_PORT KV_PORT LOCAL_PING_PORT |
| 248 | export PREFLIGHT PREFLIGHT_MAX_VRAM_PCT PREFLIGHT_GPU_SETTLE_S |
| 249 | |
| 250 | echo "[disagg-pd] IPADDRS=${IPADDRS}" |
| 251 | echo "[disagg-pd] launching ${NUM_NODES} containers (1/node) -> vllm_disagg.sh node" |
| 252 | |
| 253 | # Refuse a node that is not actually free. $1 = this node's global rank. |
| 254 | # |
| 255 | # `#SBATCH --exclusive` only keeps out other SLURM jobs. A plain `docker run` on |
| 256 | # a node is invisible to the scheduler, so we get allocated nodes whose GPUs and |
| 257 | # ports are already in use. Without this the job pulls the image, loads the model |
| 258 | # and dies minutes later inside vLLM ("Free memory on device cuda:N ... is less |
| 259 | # than desired GPU memory utilization"); worse, on rank 0 a foreign router |
| 260 | # already listening on ROUTER_PORT can satisfy the health gate and take the |
| 261 | # accuracy traffic. Called after the stale-container reap, with a settle window |
| 262 | # so VRAM the reap is still releasing is not counted against us. |
| 263 | preflight_node() { |
| 264 | local rank="$1" |
| 265 | if [ "${PREFLIGHT:-1}" != "1" ]; then |
| 266 | echo "Rank $rank preflight: skipped (PREFLIGHT=0)" |
| 267 | return 0 |
| 268 | fi |
| 269 | |
| 270 | local host max_pct settle problems="" |
| 271 | local -a issues=() # one short tag per fault, for the one-line summary |
| 272 | host="$(hostname)" |
| 273 | max_pct="${PREFLIGHT_MAX_VRAM_PCT:-10}" |
| 274 | settle="${PREFLIGHT_GPU_SETTLE_S:-30}" |
| 275 | |
| 276 | # GPUs: vLLM asks for a fraction of TOTAL, so memory already resident is |
| 277 | # subtracted from what it can get and pushes it under its own threshold. |
| 278 | if command -v rocm-smi >/dev/null 2>&1; then |
| 279 | local csv busy nbusy ncards deadline |
| 280 | deadline=$(( $(date +%s) + settle )) |
| 281 | while :; do |
| 282 | csv="$(rocm-smi --showmeminfo vram --csv 2>/dev/null || true)" |
| 283 | busy="$(printf '%s\n' "$csv" | |
| 284 | awk -F, -v m="$max_pct" '/^card/ && $2+0 > 0 { |
| 285 | pct = $3 * 100 / $2 |
| 286 | if (pct > m) printf "%s=%d%%(%.0fGiB) ", $1, pct, $3/1073741824 |
| 287 | }' || true)" |
| 288 | if [ -z "$busy" ]; then break; fi |
| 289 | if [ "$(date +%s)" -ge "$deadline" ]; then break; fi |
| 290 | sleep 5 |
| 291 | done |
| 292 | if [ -n "$busy" ]; then |
| 293 | nbusy="$(printf '%s' "$busy" | tr ' ' '\n' | grep -c 'card' || true)" |
| 294 | ncards="$(printf '%s\n' "$csv" | grep -c '^card' || true)" |
| 295 | issues+=("${nbusy}/${ncards} GPUs hold VRAM above ${max_pct}%") |
| 296 | problems="${problems} GPU VRAM in use above ${max_pct}%: ${busy}"$'\n' |
| 297 | fi |
| 298 | else |
| 299 | echo "Rank $rank preflight: no rocm-smi on $host, skipping GPU check" |
| 300 | fi |
| 301 | |
| 302 | # Ports this rank will bind. Every container here is --network host, so a |
| 303 | # foreign listener either collides or silently serves our traffic. |
| 304 | local -a want=() |
| 305 | if [ "${WIDE_EP_MODE:-0}" = "1" ]; then |
| 306 | want+=("$SERVE_PORT" "$RPC_PORT" "$KV_PORT" "$LOCAL_PING_PORT") |
| 307 | elif [ "$rank" -lt "${xP:-1}" ]; then |
| 308 | want+=("$PREFILL_PORT") |
| 309 | else |
| 310 | want+=("$DECODE_PORT") |
| 311 | fi |
| 312 | if [ "$rank" = "0" ]; then |
| 313 | want+=("$PROXY_PING_PORT") |
| 314 | if [ "$ROUTER_TYPE" = "vllm-router" ]; then |
| 315 | want+=("$ROUTER_PORT" "$ROUTER_METRICS_PORT") |
| 316 | else |
| 317 | want+=("$PROXY_PORT") |
| 318 | fi |
| 319 | fi |
| 320 | |
| 321 | if command -v ss >/dev/null 2>&1; then |
| 322 | local bound taken="" p |
| 323 | bound="$(ss -lnt 2>/dev/null | awk '{ n = split($4, a, ":"); print a[n] }' | sort -u || true)" |
| 324 | for p in "${want[@]}"; do |
| 325 | if printf '%s\n' "$bound" | grep -qx -- "$p"; then |
| 326 | taken="${taken}${p} " |
| 327 | fi |
| 328 | done |
| 329 | if [ -n "$taken" ]; then |
| 330 | issues+=("ports already bound: ${taken% }") |
| 331 | problems="${problems} ports already bound: ${taken}"$'\n' |
| 332 | fi |
| 333 | else |
| 334 | echo "Rank $rank preflight: no ss on $host, skipping port check" |
| 335 | fi |
| 336 | |
| 337 | if [ -n "$problems" ]; then |
| 338 | mkdir -p "$LOG_PATH" 2>/dev/null || true |
| 339 | local summary |
| 340 | summary="$(printf '%s; ' "${issues[@]}")" |
| 341 | summary="${summary%; }" |
| 342 | # Every node writes the shared job log concurrently, so this diagnostic can |
| 343 | # be overwritten there. Keep a single-writer copy per node as well, or the |
| 344 | # job fails in seconds with no surviving explanation. The submitter reads |
| 345 | # the PREFLIGHT-REJECTED line back out of it to use as the failure reason, |
| 346 | # so keep that prefix in sync with run-slurm-disagg-test.sh. It must not |
| 347 | # contain "FAIL:" — that is the accuracy gate's pattern. |
| 348 | { |
| 349 | echo "PREFLIGHT-REJECTED: rank ${rank} on ${host}: ${summary}" |
| 350 | printf '%s' "$problems" |
| 351 | echo " containers here: $(docker ps --format '{{.Names}} ({{.Image}})' 2>/dev/null | paste -sd'; ' - || true)" |
| 352 | echo " --exclusive does not exclude non-SLURM workloads; clear the node or exclude it from the allocation." |
| 353 | echo " Set PREFLIGHT=0 to submit anyway." |
| 354 | } | tee -a "$LOG_PATH/preflight_NODE${rank}.log" 2>/dev/null || true |
| 355 | # Unblock the other ranks: they park on this sentinel, and the submitter |
| 356 | # polls it. Without it a peer waits out HEALTH_TIMEOUT_S for an endpoint |
| 357 | # that is never coming. |
| 358 | echo 1 > "$LOG_PATH/.disagg_done" 2>/dev/null || true |
| 359 | return 1 |
| 360 | fi |
| 361 | |
| 362 | echo "Rank $rank preflight ok on $host: GPUs idle, ports free (${want[*]})" |
| 363 | return 0 |
| 364 | } |
| 365 | |
| 366 | # Per-node container launch, auto-dispatched by srun availability: |
| 367 | # Launch this node's disagg container and hand off to vllm_disagg.sh. $1 = this |
| 368 | # node's global rank. Relies on the env exported above (forwarded by srun on |
| 369 | # classic Slurm; already in scope when called directly on Spur). |
| 370 | run_node_container() { |
| 371 | local rank="$1" |
| 372 | set -o pipefail # propagate docker run exit code through the tee pipe |
| 373 | echo "Rank $rank on $(hostname)" |
| 374 | local -a stale=() |
| 375 | mapfile -t stale < <(docker ps -aq --filter name=vllm_disagg_ --filter name=vllm_router_ 2>/dev/null || true) |
| 376 | if (( ${#stale[@]} > 0 )); then |
| 377 | echo "Reaping stale containers on $(hostname): ${stale[*]}" |
| 378 | docker stop -t 30 "${stale[@]}" >/dev/null 2>&1 || true |
| 379 | docker rm "${stale[@]}" >/dev/null 2>&1 || true |
| 380 | fi |
| 381 | |
| 382 | preflight_node "$rank" || return 1 |
| 383 | |
| 384 | # Bind-mount the host AMD Pensando (ionic) RoCE userspace provider into the |
| 385 | # container IF present: the image ships an ABI-incompatible copy and the NICs |
| 386 | # are invisible without the host libs. Each entry is added only when the host |
| 387 | # file exists, so this is a no-op on non-ionic (e.g. mlx5) clusters. |
| 388 | local IONIC_MOUNTS=() |
| 389 | local host_ionic |
| 390 | host_ionic=$(readlink -f /usr/lib/x86_64-linux-gnu/libionic.so.1 2>/dev/null || true) |
| 391 | [ -n "$host_ionic" ] && IONIC_MOUNTS+=(-v "$host_ionic:/usr/lib/x86_64-linux-gnu/libionic.so.1:ro") |
| 392 | [ -f /usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so ] && \ |
| 393 | IONIC_MOUNTS+=(-v "/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so:/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so:ro") |
| 394 | [ -f /etc/libibverbs.d/ionic.driver ] && \ |
| 395 | IONIC_MOUNTS+=(-v "/etc/libibverbs.d/ionic.driver:/etc/libibverbs.d/ionic.driver:ro") |
| 396 | echo "Rank $rank ionic mounts: ${IONIC_MOUNTS[*]:-<none>}" |
| 397 | |
| 398 | # Pull latest container. |
| 399 | if [ "$DRY_RUN" != "1" ]; then |
| 400 | echo "Rank $rank pulling image on $(hostname): $IMAGE" |
| 401 | docker pull "$IMAGE" || echo "WARN: docker pull failed on $(hostname) for $IMAGE (using local copy if present)" |
| 402 | if [ "$ROUTER_TYPE" = "vllm-router" ] && [ "$rank" = "0" ]; then |
| 403 | echo "Rank $rank pulling router image on $(hostname): $VLLM_ROUTER_IMAGE" |
| 404 | docker pull "$VLLM_ROUTER_IMAGE" || echo "WARN: docker pull failed on $(hostname) for $VLLM_ROUTER_IMAGE (using local copy if present)" |
| 405 | fi |
| 406 | fi |
| 407 | |
| 408 | # Start the external vLLM router as a SEPARATE container on the rank-0 node. |
| 409 | if [ "$ROUTER_TYPE" = "vllm-router" ] && [ "$rank" = "0" ] && [ "$DRY_RUN" = "1" ]; then |
| 410 | echo "DRY_RUN: would start vllm-router on $(hostname): $VLLM_ROUTER_IMAGE (discovery :$PROXY_PING_PORT, port :$ROUTER_PORT, policy $ROUTER_POLICY, dp_local=$ROUTER_DP_LOCAL)" |
| 411 | elif [ "$ROUTER_TYPE" = "vllm-router" ] && [ "$rank" = "0" ]; then |
| 412 | mkdir -p "$LOG_PATH" 2>/dev/null || true |
| 413 | { |
| 414 | echo "Starting vllm-router on $(hostname): $VLLM_ROUTER_IMAGE (discovery :$PROXY_PING_PORT, port :$ROUTER_PORT, policy $ROUTER_POLICY, dp_local=$ROUTER_DP_LOCAL)" |
| 415 | docker run -d \ |
| 416 | --name "$ROUTER_CONT_NAME" \ |
| 417 | --network host \ |
| 418 | --user "$(id -u):$(id -g)" \ |
| 419 | -v "$LOG_PATH":/run_logs \ |
| 420 | -v "$SHARED_MOUNT":"$SHARED_MOUNT" \ |
| 421 | -v "/data/$(id -un):/workspace" \ |
| 422 | -e HOME=/workspace \ |
| 423 | -e USER="$(id -un)" \ |
| 424 | -e LOGNAME="$(id -un)" \ |
| 425 | "$VLLM_ROUTER_IMAGE" \ |
| 426 | bash -lc "exec vllm-router --vllm-pd-disaggregation --kv-connector moriio --vllm-discovery-address 0.0.0.0:$PROXY_PING_PORT --port $ROUTER_PORT --host 0.0.0.0 --intra-node-data-parallel-size $ROUTER_DP_LOCAL --policy $ROUTER_POLICY --prefill-policy $ROUTER_POLICY --decode-policy $ROUTER_POLICY --log-level info >> $LOG_PATH/vllm_router.log 2>&1" \ |
| 427 | && echo "vllm-router container started ok" \ |
| 428 | || echo "WARN: failed to start vllm-router container on $(hostname)" |
| 429 | } 2>&1 | tee -a "$LOG_PATH/vllm_router_start.log" |
| 430 | fi |
| 431 | |
| 432 | docker run --rm \ |
| 433 | --user "$(id -u):$(id -g)" \ |
| 434 | --device /dev/dri \ |
| 435 | --device /dev/kfd \ |
| 436 | --device /dev/infiniband \ |
| 437 | --network host \ |
| 438 | --ipc host \ |
| 439 | --group-add "$(getent group video | cut -d: -f3)" \ |
| 440 | --group-add "$(getent group render | cut -d: -f3)" \ |
| 441 | --cap-add SYS_PTRACE \ |
| 442 | --cap-add IPC_LOCK \ |
| 443 | --security-opt seccomp=unconfined \ |
| 444 | --shm-size 64G \ |
| 445 | --ulimit nofile=1048576:1048576 \ |
| 446 | -v "$SHARED_MOUNT":"$SHARED_MOUNT" \ |
| 447 | -v "$DISAGG_SCRIPTS_DIR":"$CONTAINER_SCRIPTS" \ |
| 448 | -v "/data/$(id -un):/workspace" \ |
| 449 | "${IONIC_MOUNTS[@]}" \ |
| 450 | --ulimit memlock=-1:-1 \ |
| 451 | -e HF_TOKEN \ |
| 452 | -e HOME=/workspace \ |
| 453 | -e USER="$(id -un)" \ |
| 454 | -e LOGNAME="$(id -un)" \ |
| 455 | -e TRITON_CACHE_DIR=/tmp/triton_cache \ |
| 456 | -e TORCHINDUCTOR_CACHE_DIR=/tmp/inductor_cache \ |
| 457 | -e VLLM_CACHE_ROOT=/tmp/vllm_cache \ |
| 458 | -e AITER_JIT_DIR=/tmp/aiter_jit \ |
| 459 | -e FLYDSL_RUNTIME_CACHE_DIR=/tmp/flydsl_cache \ |
| 460 | -e SLURM_JOB_ID="$SLURM_JOB_ID" \ |
| 461 | -e SLURM_JOB_NODELIST="$SLURM_JOB_NODELIST" \ |
| 462 | -e NODE_RANK="$rank" \ |
| 463 | -e IPADDRS="$IPADDRS" \ |
| 464 | -e WIDE_EP_MODE="$WIDE_EP_MODE" \ |
| 465 | -e xP="$xP" -e yD="$yD" \ |
| 466 | -e GPUS_PER_NODE="$GPUS_PER_NODE" \ |
| 467 | -e CLUSTER_ENV="$CLUSTER_ENV" \ |
| 468 | -e MODEL_NAME="${MODEL_NAME:-}" \ |
| 469 | -e MODEL_DIR="${MODEL_DIR:-}" \ |
| 470 | -e LOG_PATH="$LOG_PATH" \ |
| 471 | -e RUN_AFTER_HEALTH="$RUN_AFTER_HEALTH" \ |
| 472 | -e HEALTH_TIMEOUT_S="$HEALTH_TIMEOUT_S" \ |
| 473 | -e DRY_RUN="$DRY_RUN" \ |
| 474 | -e MORIIO_READ_MODE="$MORIIO_READ_MODE" \ |
| 475 | -e ROUTER_TYPE="$ROUTER_TYPE" \ |
| 476 | -e ROUTER_PORT="$ROUTER_PORT" \ |
| 477 | -e ROUTER_POLICY="$ROUTER_POLICY" \ |
| 478 | -e GATEWAY_PORT="$GATEWAY_PORT" \ |
| 479 | --name "$DOCKER_CONT_NAME" \ |
| 480 | --entrypoint /bin/bash \ |
| 481 | "$IMAGE" -c "bash $CONTAINER_SCRIPTS/vllm_disagg.sh node" \ |
| 482 | 2>&1 | tee "$LOG_PATH/main_NODE${rank}.log" |
| 483 | local node_rc=${PIPESTATUS[0]} # docker run exit code (not tee), for rank-0 pass/fail |
| 484 | # Stop the external router (rank-0 node) once the main container returns. |
| 485 | if [ "$ROUTER_TYPE" = "vllm-router" ] && [ "$rank" = "0" ]; then |
| 486 | docker stop "$ROUTER_CONT_NAME" 2>/dev/null || true; docker rm "$ROUTER_CONT_NAME" 2>/dev/null || true; |
| 487 | fi |
| 488 | return "$node_rc" |
| 489 | } |
| 490 | |
| 491 | # Teardown of the main + external router containers on every node ($1 = docker |
| 492 | # stop timeout). `docker run` containers outlive the client that started them, |
| 493 | # so on scancel the srun tasks die but the containers keep the GPUs until the |
| 494 | # walltime expires unless they are stopped explicitly. |
| 495 | teardown_containers() { |
| 496 | export TEARDOWN_STOP_T="${1:-30}" |
| 497 | if command -v srun >/dev/null 2>&1; then |
| 498 | # shellcheck disable=SC2016 # must expand on the remote node, not here |
| 499 | srun --nodelist="$SELECTED_NODELIST_STR" \ |
| 500 | --nodes="$NUM_NODES" --ntasks="$NUM_NODES" --ntasks-per-node=1 \ |
| 501 | bash -c 'docker stop -t "${TEARDOWN_STOP_T:-30}" "$DOCKER_CONT_NAME" >/dev/null 2>&1 || true; docker rm "$DOCKER_CONT_NAME" >/dev/null 2>&1 || true |
| 502 | [ "$ROUTER_TYPE" = "vllm-router" ] && { docker stop -t "${TEARDOWN_STOP_T:-30}" "$ROUTER_CONT_NAME" >/dev/null 2>&1 || true; docker rm "$ROUTER_CONT_NAME" >/dev/null 2>&1 || true; }' || true |
| 503 | else |
| 504 | docker stop -t "$TEARDOWN_STOP_T" "$DOCKER_CONT_NAME" >/dev/null 2>&1 || true |
| 505 | docker rm "$DOCKER_CONT_NAME" >/dev/null 2>&1 || true |
| 506 | if [ "$ROUTER_TYPE" = "vllm-router" ]; then |
| 507 | docker stop -t "$TEARDOWN_STOP_T" "$ROUTER_CONT_NAME" >/dev/null 2>&1 || true |
| 508 | docker rm "$ROUTER_CONT_NAME" >/dev/null 2>&1 || true |
| 509 | fi |
| 510 | fi |
| 511 | return 0 |
| 512 | } |
| 513 | |
| 514 | # SLURM sends SIGTERM then SIGKILL after KillWait (30s by default), so the trap |
| 515 | # uses a short stop timeout to fit inside that window. |
| 516 | trap 'echo "[disagg-pd] signal received, tearing down containers"; teardown_containers 10; exit 143' TERM INT |
| 517 | |
| 518 | RC=0 |
| 519 | if command -v srun >/dev/null 2>&1; then |
| 520 | # classic Slurm: fan out one task/node; ship the fn and call it with the |
| 521 | # srun-supplied per-task SLURM_PROCID (escaped so the remote shell expands it). |
| 522 | srun --label \ |
| 523 | --nodelist="$SELECTED_NODELIST_STR" \ |
| 524 | --nodes="$NUM_NODES" --ntasks="$NUM_NODES" --ntasks-per-node=1 \ |
| 525 | bash -c "$(declare -f preflight_node run_node_container); run_node_container \"\$SLURM_PROCID\"" || RC=$? |
| 526 | else |
| 527 | # Spur: the body already runs once per node; this node's rank is |
| 528 | # SPUR_TASK_OFFSET (fall back to NODE_RANK, then 0 for a single-node run). |
| 529 | run_node_container "${SPUR_TASK_OFFSET:-${NODE_RANK:-0}}" || RC=$? |
| 530 | fi |
| 531 | |
| 532 | teardown_containers 30 |
| 533 | |
| 534 | echo "[disagg-pd] launch finished rc=${RC}; logs in ${LOG_PATH}" |
| 535 | exit "${RC}" |