@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1482d2e015

[ROCm][DistInf] Enable vLLM DI CI with buildkite/slurm (#47030)

Signed-off-by: lcskrishna <lollachaitanya@gmail.com> Co-authored-by: Sheral Kumar <shekumar@amd.com> Co-authored-by: avininjamay8 <Avinash.Paul@amd.com> Co-authored-by: tej <37236721+itej89@users.noreply.github.com> Co-authored-by: Cursor <cursoragent@cursor.com>

Chaitanya Sri Krishna Lolla committed Aug 11, 2026 at 00:09 UTC 1482d2e015cc71e6ce0d3d92249cb78f4738f2e7
6 files changed +1974
.buildkite/amd-disagg/cluster.sh new
+166
@@ -0,0 +1,166 @@
1 +#!/bin/bash
2 +# =============================================================================
3 +# cluster.sh — single, generic config for the vLLM disaggregated (P/D) launcher.
4 +# -----------------------------------------------------------------------------
5 +# This script is *sourced* by vllm_disagg.sh.
6 +#
7 +# Every value is `${VAR:-default}`, so the environment always wins:
8 +# environment variable > built-in default below
9 +#
10 +# So you can override anything inline:
11 +# PREFILL_IP=10.0.0.1 DECODE_IP=10.0.0.2 ./vllm_disagg.sh prefill
12 +#
13 +# Site-specific values (model dir, IPs, NIC list, partition) are the defaults
14 +# in the "site defaults" sections below — edit those for a new cluster.
15 +# Model-SPECIFIC perf flags live in models.yaml, NOT here.
16 +# =============================================================================
17 +
18 +_CLUSTER_SH_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
19 +
20 +# ----------------------------------------------------------------- model / mode
21 +# MODEL_NAME indexes into models.yaml; MODEL_DIR is the parent dir holding it.
22 +# MODEL_PATH (resolved by the launcher) = ${MODEL_DIR}/${MODEL_NAME}. [site]
23 +export MODEL_NAME="${MODEL_NAME:-DeepSeek-V3}"
24 +export MODEL_DIR="${MODEL_DIR:-/data/models2}"
25 +
26 +# Shared NFS root (5 TB, visible on every node): model weights + per-run logs. [site]
27 +export SHARED_MOUNT="${SHARED_MOUNT:-/data}"
28 +
29 +# Parallelism mode (launcher derives PARALLEL_MODE tp|ep from this):
30 +# WIDE_EP_MODE=0 tp : each node is an independent TP server (TP8, 1P1D)
31 +# WIDE_EP_MODE=1 ep : data-parallel + expert-parallel across xP/yD nodes (wideep)
32 +export WIDE_EP_MODE="${WIDE_EP_MODE:-0}"
33 +
34 +# ----------------------------------------------------------------- topology
35 +# xP prefill nodes + yD decode nodes. IPADDRS is the ordered, comma-separated
36 +# node list (prefill IPs first, then decode IPs). NODE_RANK is this node's global
37 +# 0-based rank; under SLURM it defaults to $SLURM_PROCID. Leave IPADDRS empty to
38 +# use the PREFILL_IP/DECODE_IP fallback defaults below (1P1D only).
39 +export xP="${xP:-1}"
40 +export yD="${yD:-1}"
41 +export IPADDRS="${IPADDRS:-}"
42 +export NODE_RANK="${NODE_RANK:-${SLURM_PROCID:-}}"
43 +export PREFILL_IP="${PREFILL_IP:-10.0.0.1}"
44 +export DECODE_IP="${DECODE_IP:-10.0.0.2}"
45 +
46 +# Per-node GPU count and TP degree
47 +export GPUS_PER_NODE="${GPUS_PER_NODE:-8}"
48 +export TP_SIZE="${TP_SIZE:-${GPUS_PER_NODE}}"
49 +
50 +# ----------------------------------------------------------------- ports
51 +# TP-mode (WIDE_EP_MODE=0) server ports.
52 +export PREFILL_PORT="${PREFILL_PORT:-8100}"
53 +export DECODE_PORT="${DECODE_PORT:-8200}"
54 +
55 +# EP-mode (WIDE_EP_MODE=1) ports: API serve port, DP RPC port, KV transfer port, and
56 +# per-node MoRIIO local ping port (must differ from PROXY_PING_PORT).
57 +export SERVE_PORT="${SERVE_PORT:-20005}"
58 +export RPC_PORT="${RPC_PORT:-13345}"
59 +export KV_PORT="${KV_PORT:-9711}"
60 +export LOCAL_PING_PORT="${LOCAL_PING_PORT:-61555}"
61 +
62 +# MoRIIO proxy: HTTP port clients/benchmark hit, plus the connector control ports.
63 +# PROXY_PING_PORT MUST be 36367 — the proxy hardcodes its zmq service-discovery
64 +# socket on that port; prefill/decode register to PROXY_IP:PROXY_PING_PORT.
65 +export PROXY_IP="${PROXY_IP:-${PREFILL_IP}}"
66 +export PROXY_PORT="${PROXY_PORT:-10001}"
67 +export PROXY_PING_PORT="${PROXY_PING_PORT:-36367}"
68 +export HANDSHAKE_PORT="${HANDSHAKE_PORT:-6301}"
69 +export NOTIFY_PORT="${NOTIFY_PORT:-61005}"
70 +
71 +export PROXY_SCRIPT="${PROXY_SCRIPT:-/app/vllm/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py}"
72 +
73 +# MoRIIO KV transfer direction (injected into --kv-transfer-config by the launcher):
74 +# 0 -> omit read_mode (default; MoRIIO write mode: prefill pushes to decode)
75 +# 1 -> "read_mode": true (decode pulls KV from prefill; matches upstream disagg)
76 +export MORIIO_READ_MODE="${MORIIO_READ_MODE:-0}"
77 +
78 +# ----------------------------------------------------------------- router / gateway
79 +# Selection for client (bench/accuracy) traffic:
80 +# proxy -> the in-container MoRIIO proxy started by the launcher
81 +# vllm-router -> an external `vllm/vllm-router` container started by the SLURM job
82 +# on the rank-0 node (default)
83 +# Both use the SAME MoRIIO discovery mechanism (prefill/decode register to
84 +# PROXY_IP:PROXY_PING_PORT=36367); only the client HTTP front door differs.
85 +export ROUTER_TYPE="${ROUTER_TYPE:-vllm-router}"
86 +export ROUTER_PORT="${ROUTER_PORT:-30000}"
87 +export ROUTER_POLICY="${ROUTER_POLICY:-round_robin}"
88 +export VLLM_ROUTER_IMAGE="${VLLM_ROUTER_IMAGE:-vllm/vllm-router:nightly}"
89 +# Single client-facing port bench/accuracy target: the router port when routing,
90 +# else the proxy port. Env override always wins.
91 +if [[ "${ROUTER_TYPE}" == "vllm-router" ]]; then
92 + export GATEWAY_PORT="${GATEWAY_PORT:-${ROUTER_PORT}}"
93 +else
94 + export GATEWAY_PORT="${GATEWAY_PORT:-${PROXY_PORT}}"
95 +fi
96 +
97 +# Where per-run logs / benchmark results are written. A $SLURM_JOB_ID subdir is
98 +# appended so each CI run is self-scoped (falls back to 'local' off-SLURM). [site]
99 +_LOG_BASE="${LOG_BASE:-/data/${USER:-$(id -un)}/disagg_logs}"
100 +export LOG_PATH="${LOG_PATH:-${_LOG_BASE}/${SLURM_JOB_ID:-local}}"
101 +
102 +# ----------------------------------------------------------------- vLLM runtime
103 +# Engine/platform/transport-level env (NOT model-specific). Model-architecture
104 +# AITER kernel toggles live in models.yaml under each model's `env:` block.
105 +export VLLM_USE_V1="${VLLM_USE_V1:-1}"
106 +export HSA_NO_SCRATCH_RECLAIM="${HSA_NO_SCRATCH_RECLAIM:-1}"
107 +
108 +#export HF_HUB_OFFLINE="${HF_HUB_OFFLINE:-1}"
109 +#export TRANSFORMERS_OFFLINE="${TRANSFORMERS_OFFLINE:-1}"
110 +#
111 +#export HOME=/tmp
112 +export HF_HOME="${HF_HOME:-/tmp/hf_home}"
113 +export XDG_CACHE_HOME="${XDG_CACHE_HOME:-/tmp/.cache}"
114 +export VLLM_ENGINE_READY_TIMEOUT_S="${VLLM_ENGINE_READY_TIMEOUT_S:-3600}"
115 +
116 +# ----------------------------------------------------------------- scale
117 +# DP/EP group formation timeout across nodes (seconds).
118 +export DISTRIBUTED_TIMEOUT_SECONDS="${DISTRIBUTED_TIMEOUT_SECONDS:-7200}"
119 +
120 +# ----------------------------------------------------------------- RDMA / NCCL
121 +# AMD Pensando AINIC RoCE fabric: 8 NICs exposed as ionic_0..7 (netdevs eth2..9),
122 +# each rail on its own /24. GID index 1 + traffic class 104
123 +_IB_DEVICES="${IB_DEVICES:-ionic_0,ionic_1,ionic_2,ionic_3,ionic_4,ionic_5,ionic_6,ionic_7}"
124 +_IB_GID_INDEX="${NCCL_IB_GID_INDEX:-1}"
125 +export IB_DEVICES="${IB_DEVICES:-${_IB_DEVICES}}"
126 +export NCCL_IB_HCA="${NCCL_IB_HCA:-${_IB_DEVICES}}"
127 +export NCCL_IB_GID_INDEX="${NCCL_IB_GID_INDEX:-${_IB_GID_INDEX}}"
128 +export NCCL_IB_DISABLE="${NCCL_IB_DISABLE:-0}"
129 +# In-box RCCL net transport (no external plugin); pin bootstrap to the VPC iface.
130 +export NCCL_NET_PLUGIN="${NCCL_NET_PLUGIN:-none}"
131 +export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-eth0}"
132 +export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-${NCCL_SOCKET_IFNAME}}"
133 +export NCCL_CROSS_NIC="${NCCL_CROSS_NIC:-0}"
134 +export NCCL_PXN_DISABLE="${NCCL_PXN_DISABLE:-0}"
135 +export NCCL_NET_DISABLE_INTRA="${NCCL_NET_DISABLE_INTRA:-1}"
136 +export NCCL_IB_TC="${NCCL_IB_TC:-104}"
137 +export NCCL_IB_FIFO_TC="${NCCL_IB_FIFO_TC:-192}"
138 +export NCCL_IB_QPS_PER_CONNECTION="${NCCL_IB_QPS_PER_CONNECTION:-1}"
139 +export NCCL_IB_TIMEOUT="${NCCL_IB_TIMEOUT:-22}"
140 +export NCCL_IB_RETRY_CNT="${NCCL_IB_RETRY_CNT:-12}"
141 +
142 +# MoRI uses the same NIC set as NCCL.
143 +export MORI_RDMA_DEVICES="${MORI_RDMA_DEVICES:-${_IB_DEVICES}}"
144 +export MORI_IB_GID_INDEX="${MORI_IB_GID_INDEX:-${_IB_GID_INDEX}}"
145 +export MORI_SHMEM_HEAP_SIZE="${MORI_SHMEM_HEAP_SIZE:-16G}"
146 +
147 +# Pin to gfx950 to avoid jit compilation failures with other archs on this cluster.
148 +export MORI_GPU_ARCHS="gfx950"
149 +
150 +# ----------------------------------------------------------------- benchmark
151 +export BENCHMARK_COMBINATIONS="${BENCHMARK_COMBINATIONS:-1024/128 2048/128}"
152 +export BENCHMARK_CON="${BENCHMARK_CON:-32 64}"
153 +export NUM_PROMPTS_FACTOR="${NUM_PROMPTS_FACTOR:-2}"
154 +export BENCHMARK_MIN_PROMPTS="${BENCHMARK_MIN_PROMPTS:-32}"
155 +
156 +# ----------------------------------------------------------------- accuracy
157 +# `accuracy` role runs lm_eval (local-completions backend) against the proxy.
158 +export ACCURACY_TASKS="${ACCURACY_TASKS:-gsm8k}"
159 +export ACCURACY_NUM_CONCURRENT="${ACCURACY_NUM_CONCURRENT:-64}"
160 +export ACCURACY_MAX_RETRIES="${ACCURACY_MAX_RETRIES:-3}"
161 +export ACCURACY_METRIC="${ACCURACY_METRIC:-exact_match}"
162 +export ACCURACY_THRESHOLD="${ACCURACY_THRESHOLD:-0.90}"
163 +
164 +# ----------------------------------------------------------------- SLURM (submit)
165 +# Used by run-slurm-disagg-test.sh on the login node (harmless to export here). [site]
166 +export SLURM_PARTITION="${SLURM_PARTITION:-default}"
.buildkite/amd-disagg/models.yaml new
+97
@@ -0,0 +1,97 @@
1 +# =============================================================================
2 +# Model catalog for the vLLM disaggregated (P/D) inference CI.
3 +# -----------------------------------------------------------------------------
4 +# Loaded by vllm_disagg.sh for the selected MODEL_NAME. WIDE_EP_MODE selects the
5 +# flag set (0 -> tp keys, 1 -> ep keys):
6 +# tp -> base_flags + prefill.tp / decode.tp
7 +# ep -> base_flags + prefill.ep / decode.ep
8 +#
9 +# Keep ONLY model-specific perf flags here. The launcher owns everything
10 +# topological / mode-structural:
11 +# - tp: --host/--port, --tensor-parallel-size, --kv-transfer-config
12 +# - ep: --tp 1, --data-parallel-size[-local], --data-parallel-address/-rpc-port,
13 +# --data-parallel-start-rank/--headless (children), --enable-expert-parallel,
14 +# --all2all-backend mori, --no-enable-prefix-caching,
15 +# --api-server-count + --kv-transfer-config (masters only)
16 +# Do NOT put any of those in this file.
17 +#
18 +# Format: a top-level `models:` list. The launcher selects the entry whose
19 +# `model:` matches MODEL_NAME. Each entry:
20 +# model : name (must match MODEL_NAME; also the dir under MODEL_DIR)
21 +# env : model/arch-specific environment variables. The launcher
22 +# exports these BEFORE `vllm serve`, but only if they are
23 +# not already set, so precedence is:
24 +# caller/inline env > models.yaml env:
25 +# (Transport/fabric/engine env stays in cluster.sh,
26 +# e.g. MORI_RDMA_DEVICES, NCCL_IB_HCA, VLLM_USE_V1.)
27 +# base_flags : always applied (both roles, both modes)
28 +# prefill/decode.<tp|ep> : role + mode specific perf flags
29 +# experimental_flags : optional extra flags (omit if empty)
30 +# =============================================================================
31 +
32 +models:
33 + - model: DeepSeek-V3
34 + env:
35 + VLLM_ROCM_USE_AITER: "1"
36 + VLLM_ROCM_USE_AITER_MLA: "1"
37 + VLLM_ROCM_USE_AITER_MOE: "1"
38 + VLLM_ROCM_USE_AITER_RMSNORM: "1"
39 + VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: "0"
40 + base_flags: "--trust-remote-code --kv-cache-dtype fp8"
41 + prefill:
42 + tp: "--gpu-memory-utilization 0.85"
43 + ep: "--gpu-memory-utilization 0.85 --enforce-eager"
44 + decode:
45 + tp: "--gpu-memory-utilization 0.85"
46 + #ep: '--gpu-memory-utilization 0.75 --enforce-eager'
47 + ep: '--gpu-memory-utilization 0.75 --compilation-config {"cudagraph_mode":"PIECEWISE","custom_ops":["+quant_fp8"]}'
48 +
49 + - model: MiniMax-M3-MXFP8
50 + env:
51 + VLLM_ROCM_USE_AITER: "1"
52 + VLLM_ROCM_USE_AITER_MOE: "1"
53 + VLLM_ROCM_USE_AITER_RMSNORM: "1"
54 + VLLM_USE_BREAKABLE_CUDAGRAPH: "0"
55 + VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: "INT6"
56 + base_flags: "--trust-remote-code --attention-backend TRITON_ATTN --block-size 128 --language-model-only --kv-cache-dtype fp8"
57 + prefill:
58 + tp: "--gpu-memory-utilization 0.85 --enforce-eager"
59 + decode:
60 + tp: "--gpu-memory-utilization 0.85"
61 +
62 + - model: DeepSeek-R1-MXFP4
63 + env:
64 + VLLM_ROCM_USE_AITER: "1"
65 + VLLM_ROCM_USE_AITER_MLA: "1"
66 + VLLM_ROCM_USE_AITER_MOE: "1"
67 + VLLM_ROCM_USE_AITER_RMSNORM: "1"
68 + base_flags: "--trust-remote-code --kv-cache-dtype fp8"
69 + prefill:
70 + tp: "--gpu-memory-utilization 0.85"
71 + decode:
72 + tp: "--gpu-memory-utilization 0.85"
73 +
74 + - model: Kimi-K2.5-MXFP4
75 + env:
76 + VLLM_ROCM_USE_AITER: "1"
77 + VLLM_ROCM_USE_AITER_MOE: "1"
78 + VLLM_ROCM_USE_AITER_RMSNORM: "1"
79 + base_flags: "--trust-remote-code"
80 + prefill:
81 + tp: "--gpu-memory-utilization 0.85"
82 + decode:
83 + tp: "--gpu-memory-utilization 0.85"
84 +
85 + - model: Kimi-K2.6-MXFP4
86 + env:
87 + VLLM_ROCM_USE_AITER: "1"
88 + AMDGCN_USE_BUFFER_OPS: "1"
89 + VLLM_ROCM_USE_AITER_MLA: "1"
90 + VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: "INT4"
91 + VLLM_ROCM_USE_SKINNY_GEMM: "0"
92 + VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: "1"
93 + base_flags: "--trust-remote-code --kv-cache-dtype fp8 --mm-encoder-tp-mode data --block-size 1 --attention-backend ROCM_AITER_MLA"
94 + prefill:
95 + tp: "--gpu-memory-utilization 0.9"
96 + decode:
97 + tp: "--gpu-memory-utilization 0.9"
.buildkite/amd-disagg/pipeline-disagg.yaml new
+202
@@ -0,0 +1,202 @@
1 +# Native Buildkite pipeline — disaggregated PD SLURM tests (MoRIIO).
2 +#
3 +# ROUTER_TYPE must be set explicitly on every step: run-slurm-disagg-test.sh
4 +# defaults to `vllm-router`
5 +#
6 +
7 +steps:
8 +
9 +#-------------------------------------- mi350 · disagg 1P1D (TP8, 2 nodes) · SLURM --------------------------------------#
10 +
11 +- label: "DeepSeek-V3-PD-1P1D-TP8-MoRIIO-proxy"
12 + agents:
13 + queue: amd_mi350_ainic
14 + timeout_in_minutes: 120
15 + concurrency: 2
16 + concurrency_group: "amd-disagg/mi350/2node"
17 + commands:
18 + - MODEL_NAME=DeepSeek-V3 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
19 +
20 +- label: "DeepSeek-V3-PD-1P1D-TP8-MoRIIO-vllm-router"
21 + agents:
22 + queue: amd_mi350_ainic
23 + timeout_in_minutes: 120
24 + concurrency: 2
25 + concurrency_group: "amd-disagg/mi350/2node"
26 + commands:
27 + - MODEL_NAME=DeepSeek-V3 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
28 +
29 +- label: "DeepSeek-R1-MXFP4-PD-1P1D-TP8-MoRIIO-proxy"
30 + agents:
31 + queue: amd_mi350_ainic
32 + timeout_in_minutes: 120
33 + concurrency: 2
34 + concurrency_group: "amd-disagg/mi350/2node"
35 + commands:
36 + - MODEL_NAME=DeepSeek-R1-MXFP4 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
37 +
38 +- label: "DeepSeek-R1-MXFP4-PD-1P1D-TP8-MoRIIO-vllm-router"
39 + agents:
40 + queue: amd_mi350_ainic
41 + timeout_in_minutes: 120
42 + concurrency: 2
43 + concurrency_group: "amd-disagg/mi350/2node"
44 + commands:
45 + - MODEL_NAME=DeepSeek-R1-MXFP4 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
46 +
47 +- label: "Kimi-K2.5-MXFP4-PD-1P1D-TP8-MoRIIO-proxy"
48 + agents:
49 + queue: amd_mi350_ainic
50 + timeout_in_minutes: 120
51 + concurrency: 2
52 + concurrency_group: "amd-disagg/mi350/2node"
53 + commands:
54 + - MODEL_NAME=Kimi-K2.5-MXFP4 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
55 +
56 +- label: "Kimi-K2.5-MXFP4-PD-1P1D-TP8-MoRIIO-vllm-router"
57 + agents:
58 + queue: amd_mi350_ainic
59 + timeout_in_minutes: 120
60 + concurrency: 2
61 + concurrency_group: "amd-disagg/mi350/2node"
62 + commands:
63 + - MODEL_NAME=Kimi-K2.5-MXFP4 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
64 +
65 +- label: "Kimi-K2.6-MXFP4-PD-1P1D-TP8-MoRIIO-proxy"
66 + agents:
67 + queue: amd_mi350_ainic
68 + timeout_in_minutes: 120
69 + concurrency: 2
70 + concurrency_group: "amd-disagg/mi350/2node"
71 + commands:
72 + - MODEL_NAME=Kimi-K2.6-MXFP4 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
73 +
74 +- label: "Kimi-K2.6-MXFP4-PD-1P1D-TP8-MoRIIO-vllm-router"
75 + agents:
76 + queue: amd_mi350_ainic
77 + timeout_in_minutes: 120
78 + concurrency: 2
79 + concurrency_group: "amd-disagg/mi350/2node"
80 + commands:
81 + - MODEL_NAME=Kimi-K2.6-MXFP4 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
82 +
83 +- label: "MiniMax-M3-MXFP8-PD-1P1D-TP8-MoRIIO-proxy"
84 + agents:
85 + queue: amd_mi350_ainic
86 + timeout_in_minutes: 120
87 + concurrency: 2
88 + concurrency_group: "amd-disagg/mi350/2node"
89 + commands:
90 + - MODEL_NAME=MiniMax-M3-MXFP8 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
91 +
92 +- label: "MiniMax-M3-MXFP8-PD-1P1D-TP8-MoRIIO-vllm-router"
93 + agents:
94 + queue: amd_mi350_ainic
95 + timeout_in_minutes: 120
96 + concurrency: 2
97 + concurrency_group: "amd-disagg/mi350/2node"
98 + commands:
99 + - MODEL_NAME=MiniMax-M3-MXFP8 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
100 +
101 +#-------------------------------------- mi350 · disagg 2P2D (TP8, 4 nodes) · SLURM --------------------------------------#
102 +
103 +- label: "DeepSeek-V3-PD-2P2D-TP8-MoRIIO-proxy"
104 + agents:
105 + queue: amd_mi350_ainic
106 + timeout_in_minutes: 120
107 + concurrency: 2
108 + concurrency_group: "amd-disagg/mi350/4node"
109 + commands:
110 + - MODEL_NAME=DeepSeek-V3 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
111 +
112 +- label: "DeepSeek-V3-PD-2P2D-TP8-MoRIIO-vllm-router"
113 + agents:
114 + queue: amd_mi350_ainic
115 + timeout_in_minutes: 120
116 + concurrency: 2
117 + concurrency_group: "amd-disagg/mi350/4node"
118 + commands:
119 + - MODEL_NAME=DeepSeek-V3 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
120 +
121 +- label: "DeepSeek-R1-MXFP4-PD-2P2D-TP8-MoRIIO-proxy"
122 + agents:
123 + queue: amd_mi350_ainic
124 + timeout_in_minutes: 120
125 + concurrency: 2
126 + concurrency_group: "amd-disagg/mi350/4node"
127 + commands:
128 + - MODEL_NAME=DeepSeek-R1-MXFP4 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
129 +
130 +- label: "DeepSeek-R1-MXFP4-PD-2P2D-TP8-MoRIIO-vllm-router"
131 + agents:
132 + queue: amd_mi350_ainic
133 + timeout_in_minutes: 120
134 + concurrency: 2
135 + concurrency_group: "amd-disagg/mi350/4node"
136 + commands:
137 + - MODEL_NAME=DeepSeek-R1-MXFP4 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
138 +
139 +- label: "Kimi-K2.5-MXFP4-PD-2P2D-TP8-MoRIIO-proxy"
140 + agents:
141 + queue: amd_mi350_ainic
142 + timeout_in_minutes: 120
143 + concurrency: 2
144 + concurrency_group: "amd-disagg/mi350/4node"
145 + commands:
146 + - MODEL_NAME=Kimi-K2.5-MXFP4 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
147 +
148 +- label: "Kimi-K2.5-MXFP4-PD-2P2D-TP8-MoRIIO-vllm-router"
149 + agents:
150 + queue: amd_mi350_ainic
151 + timeout_in_minutes: 120
152 + concurrency: 2
153 + concurrency_group: "amd-disagg/mi350/4node"
154 + commands:
155 + - MODEL_NAME=Kimi-K2.5-MXFP4 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
156 +
157 +- label: "Kimi-K2.6-MXFP4-PD-2P2D-TP8-MoRIIO-proxy"
158 + agents:
159 + queue: amd_mi350_ainic
160 + timeout_in_minutes: 120
161 + concurrency: 2
162 + concurrency_group: "amd-disagg/mi350/4node"
163 + commands:
164 + - MODEL_NAME=Kimi-K2.6-MXFP4 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
165 +
166 +- label: "Kimi-K2.6-MXFP4-PD-2P2D-TP8-MoRIIO-vllm-router"
167 + agents:
168 + queue: amd_mi350_ainic
169 + timeout_in_minutes: 120
170 + concurrency: 2
171 + concurrency_group: "amd-disagg/mi350/4node"
172 + commands:
173 + - MODEL_NAME=Kimi-K2.6-MXFP4 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
174 +
175 +- label: "MiniMax-M3-MXFP8-PD-2P2D-TP8-MoRIIO-proxy"
176 + agents:
177 + queue: amd_mi350_ainic
178 + timeout_in_minutes: 120
179 + concurrency: 2
180 + concurrency_group: "amd-disagg/mi350/4node"
181 + commands:
182 + - MODEL_NAME=MiniMax-M3-MXFP8 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
183 +
184 +- label: "MiniMax-M3-MXFP8-PD-2P2D-TP8-MoRIIO-vllm-router"
185 + agents:
186 + queue: amd_mi350_ainic
187 + timeout_in_minutes: 120
188 + concurrency: 2
189 + concurrency_group: "amd-disagg/mi350/4node"
190 + commands:
191 + - MODEL_NAME=MiniMax-M3-MXFP8 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
192 +
193 +#-------------------------------------- mi350 · disagg wide-EP (disabled) --------------------------------------#
194 +
195 +# - label: "DeepSeek-V3-PD-1P1D-EP8/DP8-WideEP-MoRIIO-proxy"
196 +# agents:
197 +# queue: amd_mi350_ainic
198 +# timeout_in_minutes: 120
199 +# concurrency: 2
200 +# concurrency_group: "amd-disagg/mi350/2node"
201 +# commands:
202 +# - IMAGE=vllm/vllm-openai-rocm:nightly MODEL_NAME=DeepSeek-V3 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=1 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh
.buildkite/amd-disagg/run-slurm-disagg-test.sh new
+338
@@ -0,0 +1,338 @@
1 +#!/bin/bash
2 +# =============================================================================
3 +# run-slurm-disagg-test.sh — foreground (login-node) submitter for the disagg
4 +# P/D gate. The single command the Buildkite step runs.
5 +# -----------------------------------------------------------------------------
6 +# Thin glue: submits run_xPyD_disagg.slurm. By default (WAIT=0, Spur-safe) it
7 +# fire-and-forgets and returns 0 on a good submit; with WAIT=1 it polls the
8 +# scheduler and exits with the job's exit code so a CI step pass/fails.
9 +# stdout carries ONLY the SLURM job info
10 +# (image/nodes line, job id, log path); after sbatch nothing else is written to
11 +# stdout — the full job log persists at the --output path and pass/fail
12 +# diagnostics go to stderr. The actual work lives in run_xPyD_disagg.slurm: it
13 +# selects nodes,
14 +# fans out one container per node via a single srun, and hands off to
15 +# vllm_disagg.sh (rank-based prefill/decode self-select).
16 +#
17 +# Default target: 1P1D TP8 (NODES=2), nightly image, accuracy gate.
18 +#
19 +# Spur usage (fire-and-forget; the default here):
20 +# bash run-slurm-disagg-test.sh # 1P1D TP8
21 +# WIDE_EP_MODE=1 bash run-slurm-disagg-test.sh # 1P1D wide-EP
22 +# NODES=4 WIDE_EP_MODE=1 xP=2 yD=2 bash run-slurm-disagg-test.sh # 2P2D EP
23 +# Prints the job id + NFS log paths and returns immediately (does NOT poll
24 +# squeue/sacct, which hang on Spur). Track the run via the printed `tail -f`.
25 +#
26 +# Classic-Slurm / CI usage (block until done, exit with the gate's pass/fail):
27 +# WAIT=1 bash run-slurm-disagg-test.sh
28 +# =============================================================================
29 +set -euo pipefail
30 +
31 +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
32 +JOB_SCRIPT="${JOB_SCRIPT:-${SCRIPT_DIR}/run_xPyD_disagg.slurm}"
33 +
34 +# ---- knobs (override from the Buildkite step env) --------------------------
35 +# Defaults tuned for the Spur AMD MI350X cluster
36 +IMAGE="${IMAGE:-vllm/vllm-openai-rocm:nightly}"
37 +NODES="${NODES:-2}"
38 +GPUS_PER_NODE="${GPUS_PER_NODE:-8}"
39 +PARTITION="${SLURM_PARTITION:-}"
40 +TIME_LIMIT="${SLURM_TIME_LIMIT:-02:00:00}"
41 +WIDE_EP_MODE="${WIDE_EP_MODE:-0}" # 0 -> 1P1D TP8 (default); 1 -> wide-EP
42 +xP="${xP:-1}"
43 +yD="${yD:-1}"
44 +RUN_AFTER_HEALTH="${RUN_AFTER_HEALTH:-accuracy}"
45 +HEALTH_TIMEOUT_S="${HEALTH_TIMEOUT_S:-3600}" # P/D bring-up budget; +900s grace must fit the 2h wall
46 +SHARED_MOUNT="${SHARED_MOUNT:-/data}"
47 +LOG_ROOT="${LOG_ROOT:-${SHARED_MOUNT}/${USER:-$(whoami)}/disagg_logs}"
48 +DRY_RUN="${DRY_RUN:-0}"
49 +MORIIO_READ_MODE="${MORIIO_READ_MODE:-0}"
50 +
51 +WAIT="${WAIT:-0}"
52 +
53 +# ROUTER Type - defaults to vllm-router
54 +ROUTER_TYPE="${ROUTER_TYPE:-vllm-router}"
55 +ROUTER_PORT="${ROUTER_PORT:-30000}"
56 +VLLM_ROUTER_IMAGE="${VLLM_ROUTER_IMAGE:-vllm/vllm-router:nightly}"
57 +# Dry-run only validates wiring; cap its walltime low so it never holds the queue.
58 +[[ "${DRY_RUN}" == "1" ]] && TIME_LIMIT="${SLURM_TIME_LIMIT:-00:10:00}"
59 +
60 +mkdir -p "${LOG_ROOT}"
61 +
62 +# Spur - sbatch scheduler.
63 +DISAGG_SCRIPTS_STAGE="${DISAGG_SCRIPTS_STAGE:-/data/scratch/buildkite-agent}"
64 +STAGED_DIR="${DISAGG_SCRIPTS_STAGE}/${BUILDKITE_COMMIT:-local}"
65 +mkdir -p "${STAGED_DIR}"
66 +cp -rL --no-preserve=ownership,timestamps "${SCRIPT_DIR}/." "${STAGED_DIR}/"
67 +chmod -R u+rwX "${STAGED_DIR}" 2>/dev/null || true
68 +export DISAGG_SCRIPTS_DIR="${STAGED_DIR}"
69 +echo "[slurm-submit] staged scripts for compute nodes: ${DISAGG_SCRIPTS_DIR}" >&2
70 +export IMAGE MODEL_NAME WIDE_EP_MODE xP yD GPUS_PER_NODE RUN_AFTER_HEALTH HEALTH_TIMEOUT_S
71 +export SHARED_MOUNT LOG_ROOT DRY_RUN MORIIO_READ_MODE
72 +export ROUTER_TYPE ROUTER_PORT VLLM_ROUTER_IMAGE
73 +
74 +# Model selection.
75 +[[ -n "${MODEL_NAME:-}" ]] && export MODEL_NAME
76 +[[ -n "${MODEL_DIR:-}" ]] && export MODEL_DIR
77 +
78 +[[ -n "${PARTITION}" ]] && echo "[slurm-submit] NOTE: PARTITION='${PARTITION}' ignored (Spur sbatch has no --partition)" >&2
79 +
80 +SUBMIT_SCRIPT="${JOB_SCRIPT}"
81 +FILE_NODES="$(grep -oE '^#SBATCH[[:space:]]+--nodes=[0-9]+' "${JOB_SCRIPT}" | grep -oE '[0-9]+' | head -n1)"
82 +if [[ -n "${NODES}" && -n "${FILE_NODES}" && "${NODES}" != "${FILE_NODES}" ]]; then
83 + SUBMIT_SCRIPT="/tmp/$(basename "${JOB_SCRIPT%.slurm}")-n${NODES}-$$.slurm"
84 + sed -E "s/^#SBATCH([[:space:]]+)--nodes=[0-9]+/#SBATCH\\1--nodes=${NODES}/" "${JOB_SCRIPT}" > "${SUBMIT_SCRIPT}"
85 + echo "[slurm-submit] node count ${FILE_NODES}->${NODES}: submitting patched copy ${SUBMIT_SCRIPT}" >&2
86 +fi
87 +
88 +# Job name comes from the script's #SBATCH --job-name (SLURM sets SLURM_JOB_NAME
89 +JOB_NAME="$(grep -oE '^#SBATCH[[:space:]]+--job-name=[^[:space:]]+' "${JOB_SCRIPT}" | sed -E 's/.*--job-name=//' | head -n1)"
90 +JOB_NAME="${JOB_NAME:-vllm-disagg-pd}"
91 +
92 +echo "[slurm-submit] image=${IMAGE} nodes=${NODES} gpus/node=${GPUS_PER_NODE} mode=$([[ ${WIDE_EP_MODE} == 0 ]] && echo tp || echo ep) router=${ROUTER_TYPE} walltime=${TIME_LIMIT}"
93 +SUBMIT_OUT="$(sbatch --time="${TIME_LIMIT}" "${SUBMIT_SCRIPT}")"
94 +echo "${SUBMIT_OUT}"
95 +# "Submitted batch job 114" -> 114 (last integer on the line).
96 +JOB_ID="$(printf '%s\n' "${SUBMIT_OUT}" | grep -oE '[0-9]+' | tail -n1 || true)"
97 +if [[ -z "${JOB_ID}" ]]; then
98 + echo "[slurm-submit] ERROR: could not parse a job id from sbatch output above" >&2
99 + exit 1
100 +fi
101 +echo "[slurm-submit] submitted job ${JOB_ID}"
102 +
103 +# The real per-job log is written to NFS from inside the job body by
104 +# run_xPyD_disagg.slurm (`exec > "$LOG_ROOT/${SLURM_JOB_NAME}-${SLURM_JOB_ID}.log"`).
105 +# SLURM_JOB_NAME resolves to our --job-name, so we can compute the path here. The
106 +# per-role server/proxy/bench logs land under ${LOG_ROOT}/${JOB_ID}/.
107 +LOG_FILE="${LOG_ROOT}/${JOB_NAME}-${JOB_ID}.log"
108 +LOG_DIR="${LOG_ROOT}/${JOB_ID}"
109 +echo "[slurm-submit] job log: ${LOG_FILE}"
110 +echo "[slurm-submit] role logs: ${LOG_DIR}/"
111 +
112 +# --- Spur-safe by default: fire-and-forget --------------------------------------
113 +if [[ "${WAIT}" != "1" ]]; then
114 + echo "[slurm-submit] submitted (WAIT=0, not polling). Track with:" >&2
115 + echo " tail -f ${LOG_FILE}" >&2
116 + echo " grep -aE '(PASS|FAIL): |exact_match' ${LOG_FILE}" >&2
117 + exit 0
118 +fi
119 +
120 +# --- WAIT=1: phase-aware poll of the NFS job log (scontrol-based) -------------
121 +# Detection is phased so failures surface fast instead of waiting out the full
122 +# walltime:
123 +# 1) submitted -> running : scontrol catches infra/scheduler failures
124 +# (NODE_FAIL/BOOT_FAIL/CANCELLED/TIMEOUT/...) within
125 +# one poll, or a stuck-PENDING/never-started job.
126 +# 2) running -> healthy : advance once every endpoint reports healthy (+ the
127 +# vllm-router, when used); fail on bring-up errors or
128 +# the health budget.
129 +# 3) healthy -> verdict : PASS/FAIL from the accuracy gate, capped.
130 +# scontrol is the scheduler authority on this cluster
131 +echo "[slurm-submit] WAIT=1: phase-aware poll of ${LOG_FILE} (timeout ${TIME_LIMIT})" >&2
132 +
133 +_h=0 _m=0 _s=0
134 +IFS=: read -r _h _m _s <<< "${TIME_LIMIT}"
135 +WAIT_DEADLINE=$(( $(date +%s) + 10#${_h}*3600 + 10#${_m}*60 + 10#${_s:-0} ))
136 +unset -v _h _m _s
137 +
138 +# Per-phase budgets (all overridable from the Buildkite step env).
139 +POLL_INTERVAL="${POLL_INTERVAL:-20}"
140 +SUBMIT_GRACE_S="${SUBMIT_GRACE_S:-900}" # reach RUNNING within 15m
141 +PENDING_MAX_S="${PENDING_MAX_S:-1800}" # tolerate 30m queued
142 +HEALTH_PHASE_TIMEOUT_S="${HEALTH_PHASE_TIMEOUT_S:-$(( HEALTH_TIMEOUT_S + 900 ))}"
143 +WORKLOAD_TIMEOUT_S="${WORKLOAD_TIMEOUT_S:-1800}" # accuracy/bench cap
144 +COMPLETED_GRACE_S="${COMPLETED_GRACE_S:-60}" # verdict lag after job exit
145 +
146 +SENTINEL="${LOG_DIR}/.disagg_done"
147 +
148 +# scontrol field extractor (authoritative here; timeout-guarded so a momentary
149 +# scheduler stall can't wedge the poll). Returns the value or "".
150 +job_field() { # $1=jobid $2=field -> value | ""
151 + timeout 15 scontrol show job "$1" 2>/dev/null \
152 + | grep -oE "$2=[^ ]+" | head -n1 | cut -d= -f2- || true
153 +}
154 +have() { grep -aqE "$1" "${LOG_FILE}" 2>/dev/null; }
155 +
156 +# The job's own verdict sources, in authority order: the .disagg_done sentinel
157 +# (rank-0 rc, written by a single writer) and then the gate line in the log.
158 +# Sets STATE/RC/REASON and returns 0 when a verdict exists, 1 when it does not.
159 +read_verdict() {
160 + if [[ -f "${SENTINEL}" ]]; then
161 + RC="$(tr -dc '0-9' < "${SENTINEL}" 2>/dev/null || true)"; RC="${RC:-1}"
162 + if [[ "${RC}" == "0" ]]; then STATE="COMPLETED"; else STATE="FAILED"; fi
163 + REASON="sentinel"; return 0
164 + fi
165 + if have '(PASS|FAIL): '; then
166 + if have 'FAIL: '; then STATE="FAILED"; RC=1; else STATE="COMPLETED"; RC=0; fi
167 + REASON="gate"; return 0
168 + fi
169 + return 1
170 +}
171 +
172 +# Never let the job outlive this poller. Cancelling a Buildkite build kills the
173 +# agent's bootstrap, and without this the sbatch job keeps its whole allocation
174 +# until the walltime expires. Only armed on the WAIT=1 path — under WAIT=0,
175 +# leaving the job running is the point.
176 +CANCEL_GRACE_S="${CANCEL_GRACE_S:-120}"
177 +_CLEANED=0
178 +# shellcheck disable=SC2329 # invoked from cleanup_job, which the traps below call
179 +job_active() {
180 + case "$(job_field "${JOB_ID}" JobState)" in
181 + RUNNING|PENDING|COMPLETING|CONFIGURING|SUSPENDED|REQUEUED) return 0 ;;
182 + *) return 1 ;;
183 + esac
184 +}
185 +# shellcheck disable=SC2329 # invoked from the traps below
186 +cleanup_job() {
187 + [[ "${_CLEANED}" == "1" || -z "${JOB_ID:-}" ]] && return 0
188 + _CLEANED=1
189 + # A job that reported its own verdict is mid-teardown; give it a bounded
190 + # moment to finish. Timeouts and infra failures get cancelled immediately —
191 + # there is nothing to wait for and the nodes should come back now.
192 + if [[ "${REASON:-}" == "sentinel" || "${REASON:-}" == "gate" || "${STATE:-}" == "COMPLETED" ]]; then
193 + local deadline=$(( $(date +%s) + CANCEL_GRACE_S ))
194 + while (( $(date +%s) < deadline )) && job_active; do sleep 5; done
195 + fi
196 + if job_active; then
197 + echo "[slurm-submit] cleanup: cancelling job ${JOB_ID}" >&2
198 + scancel "${JOB_ID}" 2>/dev/null \
199 + || echo "[slurm-submit] WARN: scancel ${JOB_ID} failed; job may still hold nodes" >&2
200 + fi
201 + return 0
202 +}
203 +trap cleanup_job EXIT
204 +trap 'cleanup_job; exit 130' INT
205 +trap 'cleanup_job; exit 143' TERM
206 +trap 'cleanup_job; exit 129' HUP
207 +
208 +STATE=""
209 +RC=1
210 +REASON=""
211 +PHASE="submitted"
212 +T_PHASE=$(date +%s)
213 +
214 +while [[ $(date +%s) -lt ${WAIT_DEADLINE} ]]; do
215 + NOW=$(date +%s)
216 +
217 + # (1) Ultimate authority: terminal sentinel (holds rank-0 rc), then the
218 + # explicit accuracy gate line. Honored regardless of phase.
219 + read_verdict && break
220 +
221 + # (2) Scheduler state via scontrol: drives phase transitions and catches
222 + # infra/scheduler failures fast.
223 + ST="$(job_field "${JOB_ID}" JobState)"
224 + case "${ST}" in
225 + RUNNING|COMPLETING)
226 + if [[ "${PHASE}" == "submitted" ]]; then PHASE="bringup"; T_PHASE="${NOW}"; fi
227 + ;;
228 + NODE_FAIL|BOOT_FAIL|CANCELLED|TIMEOUT|OUT_OF_MEMORY|DEADLINE|PREEMPTED)
229 + STATE="infra-${ST}"; RC=1
230 + REASON="scontrol Reason=$(job_field "${JOB_ID}" Reason)"; break
231 + ;;
232 + FAILED)
233 + # Ambiguous (infra crash vs a legit accuracy exit=1). A gate/sentinel
234 + # line would have won above, so classify by the phase we reached.
235 + case "${PHASE}" in
236 + submitted) STATE="infra-FAILED" ;;
237 + bringup) STATE="server-failed" ;;
238 + *) STATE="workload-failed" ;;
239 + esac
240 + RC=1; REASON="scontrol JobState=FAILED phase=${PHASE}"; break
241 + ;;
242 + COMPLETED)
243 + _vd=$(( NOW + COMPLETED_GRACE_S ))
244 + _found=0
245 + while :; do
246 + if read_verdict; then _found=1; break; fi
247 + (( $(date +%s) >= _vd )) && break
248 + sleep 5
249 + done
250 + if (( _found == 1 )); then
251 + echo "[slurm-submit] verdict arrived after JobState=COMPLETED (${REASON})" >&2
252 + else
253 + STATE="completed-no-verdict"; RC=1
254 + REASON="COMPLETED with no sentinel/gate within ${COMPLETED_GRACE_S}s (a walltime kill reports COMPLETED here)"
255 + fi
256 + break
257 + ;;
258 + PENDING|CONFIGURING|RESV_DEL_HOLD|REQUEUED)
259 + if (( NOW - T_PHASE > PENDING_MAX_S )); then
260 + STATE="infra-stuck-${ST}"; RC=1; REASON="queued > ${PENDING_MAX_S}s"; break
261 + fi
262 + ;;
263 + "")
264 + # scontrol lost the job (purged past MinJobAge) with no log verdict:
265 + # rely on the terminal markers above + the per-phase deadlines below.
266 + :
267 + ;;
268 + esac
269 +
270 + # (3) Log-driven phase progress + per-phase deadlines. Fast-fails hangs where
271 + # the job is still RUNNING (so scontrol won't help) but bring-up/workload
272 + # is stuck.
273 + case "${PHASE}" in
274 + submitted)
275 + # Job body appearing in the log is a second "it started" signal for
276 + # when scontrol is briefly empty.
277 + if have 'Selected node IPs|health-gate:|\[disagg-pd\]'; then
278 + PHASE="bringup"; T_PHASE="${NOW}"
279 + elif [[ -z "${ST}" ]] && (( NOW - T_PHASE > SUBMIT_GRACE_S )); then
280 + # scontrol can't confirm the job exists (empty state) AND no log
281 + # output within the grace window -> treat as a lost/failed launch.
282 + # NB: a genuinely-queued job reports PENDING via scontrol above and
283 + # is governed by PENDING_MAX_S (default 30m), not this grace window.
284 + STATE="infra-nostart"; RC=1; REASON="no scheduler state or log within ${SUBMIT_GRACE_S}s"; break
285 + fi
286 + ;;
287 + bringup)
288 + if have 'FAIL:|TIMEOUT waiting for|exited while waiting|Traceback \(most recent'; then
289 + STATE="server-failed"; RC=1; REASON="bring-up failure in log"; break
290 + fi
291 + NEED="$(grep -aoE 'waiting on [0-9]+' "${LOG_FILE}" 2>/dev/null | grep -oE '[0-9]+' | tail -n1 || true)"
292 + GOT="$(grep -acE '\] healthy: ' "${LOG_FILE}" 2>/dev/null || true)"; GOT="${GOT:-0}"
293 + ROUTER_OK=1
294 + if [[ "${ROUTER_TYPE}" == "vllm-router" ]]; then
295 + if ! have 'vllm-router healthy on'; then ROUTER_OK=0; fi
296 + fi
297 + if [[ -n "${NEED}" ]] && (( GOT >= NEED )) && (( ROUTER_OK == 1 )); then
298 + PHASE="workload"; T_PHASE="${NOW}"
299 + elif (( NOW - T_PHASE > HEALTH_PHASE_TIMEOUT_S )); then
300 + STATE="bringup-timeout"; RC=1; REASON="no healthy within ${HEALTH_PHASE_TIMEOUT_S}s"; break
301 + fi
302 + ;;
303 + workload)
304 + # PASS/FAIL handled at the top; only enforce the cap here.
305 + if (( NOW - T_PHASE > WORKLOAD_TIMEOUT_S )); then
306 + STATE="workload-timeout"; RC=1; REASON="no verdict within ${WORKLOAD_TIMEOUT_S}s"; break
307 + fi
308 + ;;
309 + esac
310 +
311 + sleep "${POLL_INTERVAL}"
312 +done
313 +
314 +if [[ -z "${STATE}" ]]; then
315 + echo "[slurm-submit] WARN: no verdict before ${TIME_LIMIT}; failing" >&2
316 + STATE="deadline"; RC=1; REASON="poll deadline"
317 +fi
318 +
319 +# Pre-flight logic.
320 +PF_LINES=()
321 +if compgen -G "${LOG_DIR}/preflight_NODE*.log" >/dev/null 2>&1; then
322 + cat -- "${LOG_DIR}"/preflight_NODE*.log >&2 || true
323 + mapfile -t PF_LINES < <(grep -h '^PREFLIGHT-REJECTED: ' "${LOG_DIR}"/preflight_NODE*.log 2>/dev/null || true)
324 +fi
325 +if (( ${#PF_LINES[@]} > 0 )) && (( RC != 0 )); then
326 + STATE="preflight-rejected"
327 + REASON="${PF_LINES[0]#PREFLIGHT-REJECTED: }"
328 + if (( ${#PF_LINES[@]} > 1 )); then
329 + REASON="${REASON} [+$(( ${#PF_LINES[@]} - 1 )) more node(s), see ${LOG_DIR}/]"
330 + fi
331 +fi
332 +
333 +# Surface the accuracy gate verdict (if any) from the job log — to stderr.
334 +GATE_LINE=$(grep -aE '(PASS|FAIL): ' "${LOG_FILE}" 2>/dev/null | tail -n1 || true)
335 +[[ -n "${GATE_LINE}" ]] && echo "[slurm-submit] gate: ${GATE_LINE}" >&2
336 +
337 +echo "[slurm-submit] job ${JOB_ID} finished: state=${STATE:-unknown} phase=${PHASE} exit=${RC} reason=${REASON:-}" >&2
338 +exit "${RC}"
.buildkite/amd-disagg/run_xPyD_disagg.slurm new
+536
@@ -0,0 +1,536 @@
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 --time=02:00:00
8 +#SBATCH --chdir=/tmp
9 +#SBATCH --output=/tmp/spur-%j.out
10 +#SBATCH --error=/tmp/spur-%j.err
11 +
12 +set -euo pipefail
13 +
14 +
15 +# ---- inputs (env > default; forwarded by run-slurm-disagg-test.sh via --export)
16 +IMAGE="${IMAGE:-vllm/vllm-openai-rocm:nightly}"
17 +WIDE_EP_MODE="${WIDE_EP_MODE:-0}"
18 +xP="${xP:-1}"
19 +yD="${yD:-1}"
20 +GPUS_PER_NODE="${GPUS_PER_NODE:-8}"
21 +
22 +SHARED_MOUNT="${SHARED_MOUNT:-/data}"
23 +LOG_ROOT="${LOG_ROOT:-/data/${USER:-$(id -un)}/disagg_logs}"
24 +RUN_AFTER_HEALTH="${RUN_AFTER_HEALTH:-accuracy}" # bench | accuracy
25 +DRY_RUN="${DRY_RUN:-0}"
26 +MORIIO_READ_MODE="${MORIIO_READ_MODE:-0}" # 0 -> prefill writes (default); 1 -> decode reads KV (read_mode)
27 +
28 +# ---- router / gateway (host-side; also forwarded into the container) --------
29 +# vllm-router -> external `vllm/vllm-router` container started on the rank-0 node (default)
30 +# proxy -> in-container MoRIIO proxy (no extra container)
31 +ROUTER_TYPE="${ROUTER_TYPE:-vllm-router}"
32 +ROUTER_PORT="${ROUTER_PORT:-30000}"
33 +ROUTER_POLICY="${ROUTER_POLICY:-round_robin}"
34 +VLLM_ROUTER_IMAGE="${VLLM_ROUTER_IMAGE:-vllm/vllm-router:nightly}"
35 +PROXY_PING_PORT="${PROXY_PING_PORT:-36367}" # must equal vllm-router --vllm-discovery-address port
36 +PROXY_PORT="${PROXY_PORT:-10001}"
37 +# Intra-node data-parallel size the router uses to route the MoRIIO KV-notify to
38 +# the right DP rank. Required for wideEP DP (else decode hangs: "remote blocks
39 +# never arrived"). EP -> GPUS_PER_NODE (matches --data-parallel-size-local); TP -> 1.
40 +if [[ "${WIDE_EP_MODE}" == "1" ]]; then
41 + ROUTER_DP_LOCAL="${ROUTER_DP_LOCAL:-${GPUS_PER_NODE}}"
42 +else
43 + ROUTER_DP_LOCAL="${ROUTER_DP_LOCAL:-1}"
44 +fi
45 +if [[ "${ROUTER_TYPE}" == "vllm-router" ]]; then
46 + GATEWAY_PORT="${GATEWAY_PORT:-${ROUTER_PORT}}"
47 +else
48 + GATEWAY_PORT="${GATEWAY_PORT:-${PROXY_PORT}}"
49 +fi
50 +
51 +# This line is a recommended template from infra side.
52 +exec > "$LOG_ROOT/${SLURM_JOB_NAME}-${SLURM_JOB_ID}.log" 2>&1
53 +
54 +# Where the disagg scripts live ON THE HOST (bind-mounted into the container).
55 +# Spur runs the job body from /tmp, so $(pwd)/SLURM_SUBMIT_DIR aren't reliable;
56 +# default to the known host scripts dir. Override via env if you relocate it.
57 +DISAGG_SCRIPTS_DIR="${DISAGG_SCRIPTS_DIR:-${SLURM_SUBMIT_DIR:-$(pwd)}}"
58 +echo "DISAGG SCRIPTS DIR is ----"
59 +echo "$DISAGG_SCRIPTS_DIR"
60 +
61 +# In-container path the scripts are mounted to (and sourced from).
62 +CONTAINER_SCRIPTS="${CONTAINER_SCRIPTS:-/vllm-workspace/.buildkite/amd-disagg}"
63 +CLUSTER_ENV_IN="${CLUSTER_ENV_IN:-${CONTAINER_SCRIPTS}/cluster.sh}"
64 +
65 +
66 +# Health endpoints (TP-mode ports from cluster.sh defaults).
67 +PREFILL_PORT="${PREFILL_PORT:-8100}"
68 +DECODE_PORT="${DECODE_PORT:-8200}"
69 +HEALTH_TIMEOUT_S="${HEALTH_TIMEOUT_S:-3600}"
70 +
71 +# EP-mode ports, mirrored from cluster.sh purely so the preflight below can test
72 +# them host-side; the in-container values still come from cluster.sh.
73 +SERVE_PORT="${SERVE_PORT:-20005}"
74 +RPC_PORT="${RPC_PORT:-13345}"
75 +KV_PORT="${KV_PORT:-9711}"
76 +LOCAL_PING_PORT="${LOCAL_PING_PORT:-61555}"
77 +# vllm-router binds this for Prometheus in addition to ROUTER_PORT, and panics
78 +# at startup if it is taken.
79 +ROUTER_METRICS_PORT="${ROUTER_METRICS_PORT:-29000}"
80 +
81 +# Node preflight (see preflight_node): 0 disables, VRAM% above which a GPU counts
82 +# as busy, and how long to let a reaped container's VRAM drain before judging.
83 +PREFLIGHT="${PREFLIGHT:-1}"
84 +PREFLIGHT_MAX_VRAM_PCT="${PREFLIGHT_MAX_VRAM_PCT:-10}"
85 +PREFLIGHT_GPU_SETTLE_S="${PREFLIGHT_GPU_SETTLE_S:-30}"
86 +
87 +
88 +NUM_NODES=$((xP + yD))
89 +echo "Calculated NUM_NODES: $NUM_NODES (xP=$xP + yD=$yD) WIDE_EP_MODE=$WIDE_EP_MODE"
90 +
91 +# ------------------------
92 +# Extract NUM_NODES from SLURM allocation and select nodes
93 +# ------------------------
94 +echo "Original SLURM allocation:"
95 +echo "SLURM_JOB_NODELIST: $SLURM_JOB_NODELIST"
96 +echo "SLURM_NNODES: $SLURM_NNODES"
97 +echo "SLURM_NTASKS: $SLURM_NTASKS"
98 +
99 +# scontrol is available on classic Slurm but not in the Spur job body. Use it to
100 +# expand compressed ranges when present, else split the (already comma-separated,
101 +# full-hostname) SLURM_JOB_NODELIST on commas.
102 +if command -v scontrol >/dev/null 2>&1; then
103 + FULL_NODELIST=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
104 +else
105 + FULL_NODELIST=$(echo "$SLURM_JOB_NODELIST" | tr ',' '\n')
106 +fi
107 +AVAILABLE_COUNT=$(echo "$FULL_NODELIST" | wc -l)
108 +
109 +if (( AVAILABLE_COUNT < NUM_NODES )); then
110 + echo "ERROR: Need ${NUM_NODES} nodes (xP=${xP} + yD=${yD}) but only ${AVAILABLE_COUNT} allocated." >&2
111 + exit 1
112 +fi
113 +
114 +# Pick the first NUM_NODES nodes and order them to match how IPs get assigned:
115 +# * classic Slurm: sort alphabetically -- srun assigns PROCID=0 to the
116 +# alphabetically-first hostname, so IPS[] (built per node below) lines up.
117 +# * Spur: preserve RANK order (SPUR_NODELIST), so the hostnames pair 1:1 with
118 +# the rank-ordered IPs in SPUR_PEER_NODES; an alphabetical sort here would
119 +# desync MASTER_NODE / MASTER_ADDR / IPADDRS from the actual ranks.
120 +# IPADDRS[0..xP-1] = prefill, IPADDRS[xP..] = decode, so this order must match.
121 +if command -v srun >/dev/null 2>&1; then
122 + SELECTED_NODES=$(echo "$FULL_NODELIST" | head -n "$NUM_NODES" | sort)
123 +elif [[ -n "${SPUR_NODELIST:-}" ]]; then
124 + SELECTED_NODES=$(echo "${SPUR_NODELIST}" | tr ',' '\n' | sed '/^$/d' | head -n "$NUM_NODES")
125 +else
126 + SELECTED_NODES=$(echo "$FULL_NODELIST" | head -n "$NUM_NODES")
127 +fi
128 +SELECTED_NODELIST_STR=$(echo "$SELECTED_NODES" | tr '\n' ',' | sed 's/,$//')
129 +
130 +# Create new nodelist in SLURM format
131 +# This is a simplified approach - for complex ranges, you might need more sophisticated parsing
132 +NEW_SLURM_NODELIST=$(echo "$SELECTED_NODES" | paste -sd, | sed 's/,/,/g')
133 +
134 +# Update SLURM environment variables
135 +export SLURM_NNODES=$NUM_NODES
136 +export SLURM_NTASKS=$NUM_NODES
137 +export SLURM_JOB_NUM_NODES=$NUM_NODES
138 +export SLURM_NPROCS=$NUM_NODES
139 +export SLURM_JOB_NODELIST="$NEW_SLURM_NODELIST"
140 +export SLURM_NODELIST="$NEW_SLURM_NODELIST"
141 +
142 +# Keep other SLURM variables as they were or set defaults
143 +export SLURM_TASKS_PER_NODE="1(x$NUM_NODES)"
144 +export SLURM_SUBMIT_DIR="${SLURM_SUBMIT_DIR:-${HOME}}"
145 +export SLURM_CLUSTER_NAME="${SLURM_CLUSTER_NAME:-vllm-di-ci}"
146 +export SLURM_JOB_CPUS_PER_NODE="${SLURM_JOB_CPUS_PER_NODE:-1}"
147 +export SLURM_JOB_PARTITION="${SLURM_JOB_PARTITION:-amd-rccl}"
148 +export SLURM_JOBID="${SLURM_JOBID:-$SLURM_JOB_ID}"
149 +export SLURM_JOB_QOS="${SLURM_JOB_QOS:-normal}"
150 +export SLURM_JOB_ACCOUNT="${SLURM_JOB_ACCOUNT:-amd-rccl}"
151 +export SLURM_NTASKS_PER_NODE=1
152 +export SLURM_SUBMIT_HOST="${SLURM_SUBMIT_HOST:-$(hostname)}"
153 +export SLURM_JOB_ID="${SLURM_JOB_ID:-}"
154 +export SLURM_CONF="${SLURM_CONF:-/etc/slurm/slurm.conf}"
155 +export SLURM_JOB_NAME="${SLURM_JOB_NAME:-${xP}p${yD}d_bench-serving}"
156 +
157 +echo ""
158 +echo "Updated SLURM Environment Variables:"
159 +echo "SLURM_JOB_ID: $SLURM_JOB_ID"
160 +echo "SLURM_JOB_NODELIST: $SLURM_JOB_NODELIST"
161 +echo "SLURM_NNODES: $SLURM_NNODES"
162 +echo "SLURM_NTASKS: $SLURM_NTASKS"
163 +echo "SLURM_TASKS_PER_NODE: $SLURM_TASKS_PER_NODE"
164 +echo "SLURM_JOB_CPUS_PER_NODE: $SLURM_JOB_CPUS_PER_NODE"
165 +echo "SLURM_JOB_PARTITION: $SLURM_JOB_PARTITION"
166 +echo "SLURM_JOB_NUM_NODES: $SLURM_JOB_NUM_NODES"
167 +echo "SLURM_JOBID: $SLURM_JOBID"
168 +echo "SLURM_JOB_QOS: $SLURM_JOB_QOS"
169 +echo "SLURM_NODELIST: $SLURM_NODELIST"
170 +echo "SLURM_JOB_ACCOUNT: $SLURM_JOB_ACCOUNT"
171 +echo "SLURM_NPROCS: $SLURM_NPROCS"
172 +echo "SLURM_SUBMIT_HOST: $SLURM_SUBMIT_HOST"
173 +echo "SLURM_CONF: $SLURM_CONF"
174 +echo "SLURM_JOB_NAME: $SLURM_JOB_NAME"
175 +echo "SLURM_NTASKS_PER_NODE: $SLURM_NTASKS_PER_NODE"
176 +echo "SLURM_SUBMIT_DIR: $SLURM_SUBMIT_DIR"
177 +echo "SLURM_CLUSTER_NAME: $SLURM_CLUSTER_NAME"
178 +echo "ulimit: $(ulimit -a)"
179 +echo ""
180 +echo "Selected nodes for execution:"
181 +echo "$SELECTED_NODES"
182 +echo ""
183 +
184 +# Node information
185 +MASTER_NODE=$(echo "$SELECTED_NODES" | head -n 1)
186 +MASTER_PORT=39566 # Choose an open port
187 +
188 +IPS=()
189 +if command -v srun >/dev/null 2>&1; then
190 + # Classic Slurm: resolve each node's IP via srun (original behavior).
191 + for NODE in $SELECTED_NODES; do
192 + IP=$(srun --nodes=1 --ntasks=1 --time=00:20:00 --nodelist="$NODE" bash -c 'hostname -I' | awk 'NR==1 {print $1}')
193 + IPS+=("$IP")
194 + done
195 +else
196 + # Spur job body (no srun): use the rank-ordered IPs Spur provides in
197 + # SPUR_PEER_NODES ("ip:port,ip:port,..."). These are the cross-node-routable
198 + # bootstrap addresses; hostnames resolve to public IPs and the local node to
199 + # 127.0.1.1, so DNS is not usable here.
200 + mapfile -t IPS < <(echo "${SPUR_PEER_NODES:-}" | tr ',' '\n' | cut -d: -f1 | sed '/^$/d')
201 + if (( ${#IPS[@]} == 0 )); then
202 + echo "ERROR: no srun and SPUR_PEER_NODES is empty; cannot determine node IPs." >&2
203 + exit 1
204 + fi
205 +fi
206 +MASTER_ADDR="${IPS[0]}"
207 +
208 +IPADDRS="${IPS[*]}"
209 +IPADDRS="${IPADDRS// /,}"
210 +echo "Selected node IPs: ${IPADDRS}"
211 +
212 +NNODES=$NUM_NODES
213 +
214 +echo "MASTER_NODE is ${MASTER_NODE}"
215 +echo "MASTER_ADDR is ${MASTER_ADDR}"
216 +echo "MASTER_PORT is ${MASTER_PORT}"
217 +echo "NNODES is ${NNODES}"
218 +
219 +# Per-job log folder (MAD-style): all per-role server/proxy/bench logs
220 +export LOG_PATH="${LOG_ROOT}/${SLURM_JOB_ID}"
221 +mkdir -p "$LOG_PATH"
222 +echo "[disagg-pd] LOG_PATH=$LOG_PATH"
223 +
224 +# ----------------------------------------------------------------------------
225 +# Launch one container per node via a single srun and hand off to
226 +# vllm_disagg.sh node (rank-based self-select).
227 +# ----------------------------------------------------------------------------
228 +# IPADDRS is the ordered node IP list (prefill IPs first, then decode). IPS[]
229 +# was built in SELECTED_NODES (alphabetical) order — the same order srun assigns
230 +# SLURM_PROCID — so the rank->role mapping lines up:
231 +# rank 0 -> prefill master + MoRIIO proxy + orchestrator
232 +# rank 1 .. xP-1 -> prefill
233 +# rank xP .. end -> decode
234 +export IPADDRS
235 +
236 +export WIDE_EP_MODE xP yD GPUS_PER_NODE RUN_AFTER_HEALTH HEALTH_TIMEOUT_S DRY_RUN MORIIO_READ_MODE
237 +export IMAGE SHARED_MOUNT CONTAINER_SCRIPTS DISAGG_SCRIPTS_DIR
238 +export CLUSTER_ENV="${CLUSTER_ENV_IN}"
239 +export DOCKER_CONT_NAME="vllm_disagg_${MODEL_NAME:-model}_${SLURM_JOB_ID}"
240 +
241 +# Router env: forwarded into the main container (cluster.sh honors these) and used
242 +# host-side to launch the external router container on rank 0.
243 +export ROUTER_TYPE ROUTER_PORT ROUTER_POLICY VLLM_ROUTER_IMAGE PROXY_PING_PORT GATEWAY_PORT ROUTER_DP_LOCAL
244 +export ROUTER_CONT_NAME="vllm_router_${SLURM_JOB_ID}"
245 +
246 +# Consumed by preflight_node on each node (srun forwards the environment).
247 +export PREFILL_PORT DECODE_PORT PROXY_PORT ROUTER_METRICS_PORT
248 +export SERVE_PORT RPC_PORT KV_PORT LOCAL_PING_PORT
249 +export PREFLIGHT PREFLIGHT_MAX_VRAM_PCT PREFLIGHT_GPU_SETTLE_S
250 +
251 +echo "[disagg-pd] IPADDRS=${IPADDRS}"
252 +echo "[disagg-pd] launching ${NUM_NODES} containers (1/node) -> vllm_disagg.sh node"
253 +
254 +# Refuse a node that is not actually free. $1 = this node's global rank.
255 +#
256 +# `#SBATCH --exclusive` only keeps out other SLURM jobs. A plain `docker run` on
257 +# a node is invisible to the scheduler, so we get allocated nodes whose GPUs and
258 +# ports are already in use. Without this the job pulls the image, loads the model
259 +# and dies minutes later inside vLLM ("Free memory on device cuda:N ... is less
260 +# than desired GPU memory utilization"); worse, on rank 0 a foreign router
261 +# already listening on ROUTER_PORT can satisfy the health gate and take the
262 +# accuracy traffic. Called after the stale-container reap, with a settle window
263 +# so VRAM the reap is still releasing is not counted against us.
264 +preflight_node() {
265 + local rank="$1"
266 + if [ "${PREFLIGHT:-1}" != "1" ]; then
267 + echo "Rank $rank preflight: skipped (PREFLIGHT=0)"
268 + return 0
269 + fi
270 +
271 + local host max_pct settle problems=""
272 + local -a issues=() # one short tag per fault, for the one-line summary
273 + host="$(hostname)"
274 + max_pct="${PREFLIGHT_MAX_VRAM_PCT:-10}"
275 + settle="${PREFLIGHT_GPU_SETTLE_S:-30}"
276 +
277 + # GPUs: vLLM asks for a fraction of TOTAL, so memory already resident is
278 + # subtracted from what it can get and pushes it under its own threshold.
279 + if command -v rocm-smi >/dev/null 2>&1; then
280 + local csv busy nbusy ncards deadline
281 + deadline=$(( $(date +%s) + settle ))
282 + while :; do
283 + csv="$(rocm-smi --showmeminfo vram --csv 2>/dev/null || true)"
284 + busy="$(printf '%s\n' "$csv" |
285 + awk -F, -v m="$max_pct" '/^card/ && $2+0 > 0 {
286 + pct = $3 * 100 / $2
287 + if (pct > m) printf "%s=%d%%(%.0fGiB) ", $1, pct, $3/1073741824
288 + }' || true)"
289 + if [ -z "$busy" ]; then break; fi
290 + if [ "$(date +%s)" -ge "$deadline" ]; then break; fi
291 + sleep 5
292 + done
293 + if [ -n "$busy" ]; then
294 + nbusy="$(printf '%s' "$busy" | tr ' ' '\n' | grep -c 'card' || true)"
295 + ncards="$(printf '%s\n' "$csv" | grep -c '^card' || true)"
296 + issues+=("${nbusy}/${ncards} GPUs hold VRAM above ${max_pct}%")
297 + problems="${problems} GPU VRAM in use above ${max_pct}%: ${busy}"$'\n'
298 + fi
299 + else
300 + echo "Rank $rank preflight: no rocm-smi on $host, skipping GPU check"
301 + fi
302 +
303 + # Ports this rank will bind. Every container here is --network host, so a
304 + # foreign listener either collides or silently serves our traffic.
305 + local -a want=()
306 + if [ "${WIDE_EP_MODE:-0}" = "1" ]; then
307 + want+=("$SERVE_PORT" "$RPC_PORT" "$KV_PORT" "$LOCAL_PING_PORT")
308 + elif [ "$rank" -lt "${xP:-1}" ]; then
309 + want+=("$PREFILL_PORT")
310 + else
311 + want+=("$DECODE_PORT")
312 + fi
313 + if [ "$rank" = "0" ]; then
314 + want+=("$PROXY_PING_PORT")
315 + if [ "$ROUTER_TYPE" = "vllm-router" ]; then
316 + want+=("$ROUTER_PORT" "$ROUTER_METRICS_PORT")
317 + else
318 + want+=("$PROXY_PORT")
319 + fi
320 + fi
321 +
322 + if command -v ss >/dev/null 2>&1; then
323 + local bound taken="" p
324 + bound="$(ss -lnt 2>/dev/null | awk '{ n = split($4, a, ":"); print a[n] }' | sort -u || true)"
325 + for p in "${want[@]}"; do
326 + if printf '%s\n' "$bound" | grep -qx -- "$p"; then
327 + taken="${taken}${p} "
328 + fi
329 + done
330 + if [ -n "$taken" ]; then
331 + issues+=("ports already bound: ${taken% }")
332 + problems="${problems} ports already bound: ${taken}"$'\n'
333 + fi
334 + else
335 + echo "Rank $rank preflight: no ss on $host, skipping port check"
336 + fi
337 +
338 + if [ -n "$problems" ]; then
339 + mkdir -p "$LOG_PATH" 2>/dev/null || true
340 + local summary
341 + summary="$(printf '%s; ' "${issues[@]}")"
342 + summary="${summary%; }"
343 + # Every node writes the shared job log concurrently, so this diagnostic can
344 + # be overwritten there. Keep a single-writer copy per node as well, or the
345 + # job fails in seconds with no surviving explanation. The submitter reads
346 + # the PREFLIGHT-REJECTED line back out of it to use as the failure reason,
347 + # so keep that prefix in sync with run-slurm-disagg-test.sh. It must not
348 + # contain "FAIL:" — that is the accuracy gate's pattern.
349 + {
350 + echo "PREFLIGHT-REJECTED: rank ${rank} on ${host}: ${summary}"
351 + printf '%s' "$problems"
352 + echo " containers here: $(docker ps --format '{{.Names}} ({{.Image}})' 2>/dev/null | paste -sd'; ' - || true)"
353 + echo " --exclusive does not exclude non-SLURM workloads; clear the node or exclude it from the allocation."
354 + echo " Set PREFLIGHT=0 to submit anyway."
355 + } | tee -a "$LOG_PATH/preflight_NODE${rank}.log" 2>/dev/null || true
356 + # Unblock the other ranks: they park on this sentinel, and the submitter
357 + # polls it. Without it a peer waits out HEALTH_TIMEOUT_S for an endpoint
358 + # that is never coming.
359 + echo 1 > "$LOG_PATH/.disagg_done" 2>/dev/null || true
360 + return 1
361 + fi
362 +
363 + echo "Rank $rank preflight ok on $host: GPUs idle, ports free (${want[*]})"
364 + return 0
365 +}
366 +
367 +# Per-node container launch, auto-dispatched by srun availability:
368 +# Launch this node's disagg container and hand off to vllm_disagg.sh. $1 = this
369 +# node's global rank. Relies on the env exported above (forwarded by srun on
370 +# classic Slurm; already in scope when called directly on Spur).
371 +run_node_container() {
372 + local rank="$1"
373 + set -o pipefail # propagate docker run exit code through the tee pipe
374 + echo "Rank $rank on $(hostname)"
375 + local -a stale=()
376 + mapfile -t stale < <(docker ps -aq --filter name=vllm_disagg_ --filter name=vllm_router_ 2>/dev/null || true)
377 + if (( ${#stale[@]} > 0 )); then
378 + echo "Reaping stale containers on $(hostname): ${stale[*]}"
379 + docker stop -t 30 "${stale[@]}" >/dev/null 2>&1 || true
380 + docker rm "${stale[@]}" >/dev/null 2>&1 || true
381 + fi
382 +
383 + preflight_node "$rank" || return 1
384 +
385 + # Bind-mount the host AMD Pensando (ionic) RoCE userspace provider into the
386 + # container IF present: the image ships an ABI-incompatible copy and the NICs
387 + # are invisible without the host libs. Each entry is added only when the host
388 + # file exists, so this is a no-op on non-ionic (e.g. mlx5) clusters.
389 + local IONIC_MOUNTS=()
390 + local host_ionic
391 + host_ionic=$(readlink -f /usr/lib/x86_64-linux-gnu/libionic.so.1 2>/dev/null || true)
392 + [ -n "$host_ionic" ] && IONIC_MOUNTS+=(-v "$host_ionic:/usr/lib/x86_64-linux-gnu/libionic.so.1:ro")
393 + [ -f /usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so ] && \
394 + IONIC_MOUNTS+=(-v "/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so:/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so:ro")
395 + [ -f /etc/libibverbs.d/ionic.driver ] && \
396 + IONIC_MOUNTS+=(-v "/etc/libibverbs.d/ionic.driver:/etc/libibverbs.d/ionic.driver:ro")
397 + echo "Rank $rank ionic mounts: ${IONIC_MOUNTS[*]:-<none>}"
398 +
399 + # Pull latest container.
400 + if [ "$DRY_RUN" != "1" ]; then
401 + echo "Rank $rank pulling image on $(hostname): $IMAGE"
402 + docker pull "$IMAGE" || echo "WARN: docker pull failed on $(hostname) for $IMAGE (using local copy if present)"
403 + if [ "$ROUTER_TYPE" = "vllm-router" ] && [ "$rank" = "0" ]; then
404 + echo "Rank $rank pulling router image on $(hostname): $VLLM_ROUTER_IMAGE"
405 + docker pull "$VLLM_ROUTER_IMAGE" || echo "WARN: docker pull failed on $(hostname) for $VLLM_ROUTER_IMAGE (using local copy if present)"
406 + fi
407 + fi
408 +
409 + # Start the external vLLM router as a SEPARATE container on the rank-0 node.
410 + if [ "$ROUTER_TYPE" = "vllm-router" ] && [ "$rank" = "0" ] && [ "$DRY_RUN" = "1" ]; then
411 + 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)"
412 + elif [ "$ROUTER_TYPE" = "vllm-router" ] && [ "$rank" = "0" ]; then
413 + mkdir -p "$LOG_PATH" 2>/dev/null || true
414 + {
415 + echo "Starting vllm-router on $(hostname): $VLLM_ROUTER_IMAGE (discovery :$PROXY_PING_PORT, port :$ROUTER_PORT, policy $ROUTER_POLICY, dp_local=$ROUTER_DP_LOCAL)"
416 + docker run -d \
417 + --name "$ROUTER_CONT_NAME" \
418 + --network host \
419 + --user "$(id -u):$(id -g)" \
420 + -v "$LOG_PATH":/run_logs \
421 + -v "$SHARED_MOUNT":"$SHARED_MOUNT" \
422 + -v "/data/$(id -un):/workspace" \
423 + -e HOME=/workspace \
424 + -e USER="$(id -un)" \
425 + -e LOGNAME="$(id -un)" \
426 + "$VLLM_ROUTER_IMAGE" \
427 + 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" \
428 + && echo "vllm-router container started ok" \
429 + || echo "WARN: failed to start vllm-router container on $(hostname)"
430 + } 2>&1 | tee -a "$LOG_PATH/vllm_router_start.log"
431 + fi
432 +
433 + docker run --rm \
434 + --user "$(id -u):$(id -g)" \
435 + --device /dev/dri \
436 + --device /dev/kfd \
437 + --device /dev/infiniband \
438 + --network host \
439 + --ipc host \
440 + --group-add "$(getent group video | cut -d: -f3)" \
441 + --group-add "$(getent group render | cut -d: -f3)" \
442 + --cap-add SYS_PTRACE \
443 + --cap-add IPC_LOCK \
444 + --security-opt seccomp=unconfined \
445 + --shm-size 64G \
446 + --ulimit nofile=1048576:1048576 \
447 + -v "$SHARED_MOUNT":"$SHARED_MOUNT" \
448 + -v "$DISAGG_SCRIPTS_DIR":"$CONTAINER_SCRIPTS" \
449 + -v "/data/$(id -un):/workspace" \
450 + "${IONIC_MOUNTS[@]}" \
451 + --ulimit memlock=-1:-1 \
452 + -e HF_TOKEN \
453 + -e HOME=/workspace \
454 + -e USER="$(id -un)" \
455 + -e LOGNAME="$(id -un)" \
456 + -e TRITON_CACHE_DIR=/tmp/triton_cache \
457 + -e TORCHINDUCTOR_CACHE_DIR=/tmp/inductor_cache \
458 + -e VLLM_CACHE_ROOT=/tmp/vllm_cache \
459 + -e AITER_JIT_DIR=/tmp/aiter_jit \
460 + -e FLYDSL_RUNTIME_CACHE_DIR=/tmp/flydsl_cache \
461 + -e SLURM_JOB_ID="$SLURM_JOB_ID" \
462 + -e SLURM_JOB_NODELIST="$SLURM_JOB_NODELIST" \
463 + -e NODE_RANK="$rank" \
464 + -e IPADDRS="$IPADDRS" \
465 + -e WIDE_EP_MODE="$WIDE_EP_MODE" \
466 + -e xP="$xP" -e yD="$yD" \
467 + -e GPUS_PER_NODE="$GPUS_PER_NODE" \
468 + -e CLUSTER_ENV="$CLUSTER_ENV" \
469 + -e MODEL_NAME="${MODEL_NAME:-}" \
470 + -e MODEL_DIR="${MODEL_DIR:-}" \
471 + -e LOG_PATH="$LOG_PATH" \
472 + -e RUN_AFTER_HEALTH="$RUN_AFTER_HEALTH" \
473 + -e HEALTH_TIMEOUT_S="$HEALTH_TIMEOUT_S" \
474 + -e DRY_RUN="$DRY_RUN" \
475 + -e MORIIO_READ_MODE="$MORIIO_READ_MODE" \
476 + -e ROUTER_TYPE="$ROUTER_TYPE" \
477 + -e ROUTER_PORT="$ROUTER_PORT" \
478 + -e ROUTER_POLICY="$ROUTER_POLICY" \
479 + -e GATEWAY_PORT="$GATEWAY_PORT" \
480 + --name "$DOCKER_CONT_NAME" \
481 + --entrypoint /bin/bash \
482 + "$IMAGE" -c "bash $CONTAINER_SCRIPTS/vllm_disagg.sh node" \
483 + 2>&1 | tee "$LOG_PATH/main_NODE${rank}.log"
484 + local node_rc=${PIPESTATUS[0]} # docker run exit code (not tee), for rank-0 pass/fail
485 + # Stop the external router (rank-0 node) once the main container returns.
486 + if [ "$ROUTER_TYPE" = "vllm-router" ] && [ "$rank" = "0" ]; then
487 + docker stop "$ROUTER_CONT_NAME" 2>/dev/null || true; docker rm "$ROUTER_CONT_NAME" 2>/dev/null || true;
488 + fi
489 + return "$node_rc"
490 +}
491 +
492 +# Teardown of the main + external router containers on every node ($1 = docker
493 +# stop timeout). `docker run` containers outlive the client that started them,
494 +# so on scancel the srun tasks die but the containers keep the GPUs until the
495 +# walltime expires unless they are stopped explicitly.
496 +teardown_containers() {
497 + export TEARDOWN_STOP_T="${1:-30}"
498 + if command -v srun >/dev/null 2>&1; then
499 + # shellcheck disable=SC2016 # must expand on the remote node, not here
500 + srun --nodelist="$SELECTED_NODELIST_STR" \
501 + --nodes="$NUM_NODES" --ntasks="$NUM_NODES" --ntasks-per-node=1 \
502 + 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
503 + [ "$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
504 + else
505 + docker stop -t "$TEARDOWN_STOP_T" "$DOCKER_CONT_NAME" >/dev/null 2>&1 || true
506 + docker rm "$DOCKER_CONT_NAME" >/dev/null 2>&1 || true
507 + if [ "$ROUTER_TYPE" = "vllm-router" ]; then
508 + docker stop -t "$TEARDOWN_STOP_T" "$ROUTER_CONT_NAME" >/dev/null 2>&1 || true
509 + docker rm "$ROUTER_CONT_NAME" >/dev/null 2>&1 || true
510 + fi
511 + fi
512 + return 0
513 +}
514 +
515 +# SLURM sends SIGTERM then SIGKILL after KillWait (30s by default), so the trap
516 +# uses a short stop timeout to fit inside that window.
517 +trap 'echo "[disagg-pd] signal received, tearing down containers"; teardown_containers 10; exit 143' TERM INT
518 +
519 +RC=0
520 +if command -v srun >/dev/null 2>&1; then
521 + # classic Slurm: fan out one task/node; ship the fn and call it with the
522 + # srun-supplied per-task SLURM_PROCID (escaped so the remote shell expands it).
523 + srun --label \
524 + --nodelist="$SELECTED_NODELIST_STR" \
525 + --nodes="$NUM_NODES" --ntasks="$NUM_NODES" --ntasks-per-node=1 \
526 + bash -c "$(declare -f preflight_node run_node_container); run_node_container \"\$SLURM_PROCID\"" || RC=$?
527 +else
528 + # Spur: the body already runs once per node; this node's rank is
529 + # SPUR_TASK_OFFSET (fall back to NODE_RANK, then 0 for a single-node run).
530 + run_node_container "${SPUR_TASK_OFFSET:-${NODE_RANK:-0}}" || RC=$?
531 +fi
532 +
533 +teardown_containers 30
534 +
535 +echo "[disagg-pd] launch finished rc=${RC}; logs in ${LOG_PATH}"
536 +exit "${RC}"
.buildkite/amd-disagg/vllm_disagg.sh new
+635
@@ -0,0 +1,635 @@
1 +#!/bin/bash
2 +# =============================================================================
3 +# vLLM disaggregated (P/D) launcher (rank-based / SLURM-native).
4 +# -----------------------------------------------------------------------------
5 +#
6 +# rank 0 .. xP-1 -> prefill (rank 0 is also the orchestrator + proxy)
7 +# rank xP .. xP+yD-1 -> decode
8 +#
9 +# Rank 0 additionally: starts the MoRIIO proxy, health-gates every server,
10 +# runs the post-health workload (RUN_AFTER_HEALTH=bench|accuracy|none), then
11 +# drops a shared-FS completion sentinel so the other ranks shut down and the
12 +# single `srun` returns. Rank 0's exit code is the run's pass/fail.
13 +#
14 +# Explicit per-role invocation is also supported (granular for manual use):
15 +# proxy | prefill | decode | bench | accuracy.
16 +#
17 +# ---- LAUNCH IT (both supported) --------------------------------------------
18 +# 1) Single srun (SLURM provides the rank) — what run_xPyD_disagg.slurm does:
19 +# srun --nodes=$((xP+yD)) --ntasks-per-node=1 \
20 +# bash .../vllm_disagg.sh node
21 +# (IPADDRS must be ordered to match ranks: prefill IPs first, then decode,
22 +# in the SAME order srun assigns PROCID — sort nodes alphabetically.)
23 +#
24 +# 2) Manual, one shell per node (rank from NODE_RANK; share LOG_PATH/FS):
25 +# IPADDRS=ipP0,ipD0 xP=1 yD=1 NODE_RANK=0 bash vllm_disagg.sh node
26 +# IPADDRS=ipP0,ipD0 xP=1 yD=1 NODE_RANK=1 bash vllm_disagg.sh node
27 +#
28 +# 3) Manual, fully granular (per-role):
29 +# bash vllm_disagg.sh proxy
30 +# bash vllm_disagg.sh prefill
31 +# NODE_RANK=1 bash vllm_disagg.sh decode
32 +# bash vllm_disagg.sh bench
33 +#
34 +# Cluster config: cluster.sh (sourced). Model flags: models.yaml.
35 +# Parallelism: WIDE_EP_MODE=0 tp (independent TP servers) | 1 ep (DP+EP groups).
36 +# =============================================================================
37 +set -euo pipefail
38 +
39 +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
40 +
41 +_PROXY_IP_OVERRIDE="${PROXY_IP:-}"
42 +
43 +log() { echo "[vllm_disagg] $*"; }
44 +die() { echo "ERROR: $*" >&2; exit 1; }
45 +
46 +usage() {
47 + echo "Usage: $0 [node|proxy|prefill|decode|bench|accuracy] [--wide-ep-mode 0|1]" >&2
48 + echo " (no role => 'node': derive role from \$SLURM_PROCID/\$NODE_RANK)" >&2
49 + exit 1
50 +}
51 +
52 +# ----------------------------------------------------------------- arg parsing
53 +# Sets globals: ROLE (default 'node'), _WIDE_EP_MODE_OVERRIDE.
54 +parse_args() {
55 + ROLE="${1:-node}"
56 + shift || true
57 + _WIDE_EP_MODE_OVERRIDE=""
58 + while [[ $# -gt 0 ]]; do
59 + case "$1" in
60 + --wide-ep-mode) _WIDE_EP_MODE_OVERRIDE="${2:-}"; shift 2 ;;
61 + --wide-ep-mode=*) _WIDE_EP_MODE_OVERRIDE="${1#*=}"; shift ;;
62 + -h|--help) usage ;;
63 + *) die "Unknown argument: $1" ;;
64 + esac
65 + done
66 +}
67 +
68 +# ----------------------------------------------------------------- load config
69 +# Sources cluster.sh and resolves WIDE_EP_MODE -> PARALLEL_MODE (tp|ep).
70 +load_config() {
71 + CLUSTER_ENV="${CLUSTER_ENV:-${SCRIPT_DIR}/cluster.sh}"
72 + [[ -f "${CLUSTER_ENV}" ]] || die "cluster env file not found: ${CLUSTER_ENV}"
73 + # shellcheck disable=SC1090
74 + source "${CLUSTER_ENV}"
75 +
76 + # Topology fanout normally arrives from cluster.sh; restate the defaults so
77 + # the rank arithmetic is defined even when it does not.
78 + : "${xP:=1}" "${yD:=1}"
79 +
80 + [[ -n "${_WIDE_EP_MODE_OVERRIDE}" ]] && WIDE_EP_MODE="${_WIDE_EP_MODE_OVERRIDE}"
81 + WIDE_EP_MODE="${WIDE_EP_MODE:-0}"
82 + case "${WIDE_EP_MODE}" in
83 + 0) PARALLEL_MODE="tp" ;;
84 + 1) PARALLEL_MODE="ep" ;;
85 + *) die "WIDE_EP_MODE must be 0 (tp) or 1 (ep); got '${WIDE_EP_MODE}'" ;;
86 + esac
87 +
88 + # node-mode orchestration knobs
89 + RUN_AFTER_HEALTH="${RUN_AFTER_HEALTH:-accuracy}" # bench | accuracy | none
90 + HEALTH_TIMEOUT_S="${HEALTH_TIMEOUT_S:-3600}"
91 +
92 + # MoRIIO KV transfer direction. 0 (default)
93 + MORIIO_READ_MODE="${MORIIO_READ_MODE:-0}"
94 +}
95 +
96 +# Emits the `read_mode` KV-config fragment (leading comma + newline) when read
97 +# mode is enabled, else nothing (leaving MoRIIO's write-mode default).
98 +moriio_read_mode_kv() {
99 + case "${MORIIO_READ_MODE:-0}" in
100 + 1|true|True|TRUE|yes|on) printf ',\n "read_mode": true' ;;
101 + *) : ;;
102 + esac
103 +}
104 +
105 +# ----------------------------------------------------------------- topology
106 +# Resolves MODEL_PATH, MODELS_YAML, the IP array, and the master/proxy addresses.
107 +# IP_ARRAY[0..xP-1] = prefill nodes, [xP..xP+yD-1] = decode nodes.
108 +resolve_topology() {
109 + MODEL_PATH="${MODEL_DIR%/}/${MODEL_NAME}"
110 + MODELS_YAML="${MODELS_YAML:-${SCRIPT_DIR}/models.yaml}"
111 + mkdir -p "${LOG_PATH}" 2>/dev/null || true
112 +
113 + [[ -z "${IPADDRS}" ]] && IPADDRS="${PREFILL_IP},${DECODE_IP}" # 1P1D fallback
114 + IFS=',' read -ra IP_ARRAY <<< "${IPADDRS}"
115 + PREFILL_MASTER_ADDR="${IP_ARRAY[0]}"
116 + DECODE_MASTER_ADDR="${IP_ARRAY[$xP]:-${IP_ARRAY[-1]}}"
117 + # Proxy is co-located on the prefill master (rank 0) unless the CALLER
118 + # explicitly set PROXY_IP. We use the pre-cluster.sh override (captured at the
119 + # top) — not the live PROXY_IP, which cluster.sh always populates with the
120 + # static PREFILL_IP fallback and would otherwise mask the real master.
121 + PROXY_IP="${_PROXY_IP_OVERRIDE:-${PREFILL_MASTER_ADDR}}"
122 +}
123 +
124 +# ============================================================================
125 +# Non-server roles
126 +# ============================================================================
127 +
128 +run_proxy() {
129 + [[ -f "${PROXY_SCRIPT}" ]] || die "proxy script not found: ${PROXY_SCRIPT}"
130 + log "proxy: http=:${PROXY_PORT} discovery=:${PROXY_PING_PORT} (${PROXY_SCRIPT})"
131 + exec python3 "${PROXY_SCRIPT}" --port "${PROXY_PORT}"
132 +}
133 +
134 +# Start the proxy in the BACKGROUND (node-mode rank 0). Sets PROXY_PID.
135 +start_proxy_bg() {
136 + [[ -f "${PROXY_SCRIPT}" ]] || die "proxy script not found: ${PROXY_SCRIPT}"
137 + local plog="${LOG_PATH}/proxy_$(date +%Y%m%d_%H%M%S).log"
138 + log "proxy(bg): http=:${PROXY_PORT} discovery=:${PROXY_PING_PORT} log=${plog}"
139 + python3 "${PROXY_SCRIPT}" --port "${PROXY_PORT}" >"${plog}" 2>&1 &
140 + PROXY_PID=$!
141 +}
142 +
143 +run_bench() {
144 + local base_url="http://127.0.0.1:${GATEWAY_PORT:-${PROXY_PORT}}"
145 + local result_dir="${LOG_PATH}/bench_$(date +%Y%m%d_%H%M%S)"
146 + mkdir -p "${result_dir}"
147 + log "bench -> ${base_url} model=${MODEL_PATH} wide_ep_mode=${WIDE_EP_MODE}(${PARALLEL_MODE})"
148 + log "combinations='${BENCHMARK_COMBINATIONS}' concurrency='${BENCHMARK_CON}'"
149 +
150 + local combo isl osl con nump logf
151 + for combo in ${BENCHMARK_COMBINATIONS}; do
152 + isl="${combo%%/*}"; osl="${combo##*/}"
153 + for con in ${BENCHMARK_CON}; do
154 + nump=$(( NUM_PROMPTS_FACTOR * con ))
155 + (( nump < BENCHMARK_MIN_PROMPTS )) && nump="${BENCHMARK_MIN_PROMPTS}"
156 + logf="${result_dir}/bench_${isl}_${osl}_con${con}.log"
157 + log "ISL=${isl} OSL=${osl} CON=${con} num_prompts=${nump} -> ${logf}"
158 + vllm bench serve \
159 + --backend openai \
160 + --base-url "${base_url}" \
161 + --endpoint /v1/completions \
162 + --model "${MODEL_PATH}" \
163 + --dataset-name random \
164 + --random-input-len "${isl}" \
165 + --random-output-len "${osl}" \
166 + --num-prompts "${nump}" \
167 + --max-concurrency "${con}" \
168 + --percentile-metrics ttft,tpot,itl,e2el \
169 + --ignore-eos \
170 + 2>&1 | tee "${logf}"
171 + done
172 + done
173 + log "bench complete. Results in ${result_dir}"
174 +}
175 +
176 +run_accuracy() {
177 + # lm_eval isn't in the stock vLLM image; install on demand (rank 0 only, and
178 + # only when accuracy is actually selected). Override with ACCURACY_PIP_SPEC.
179 + if ! command -v lm_eval >/dev/null 2>&1; then
180 + log "lm_eval not found — installing ${ACCURACY_PIP_SPEC:-lm_eval[api]}"
181 + python3 -m pip install --no-cache-dir "${ACCURACY_PIP_SPEC:-lm_eval[api]}"
182 + fi
183 + local base_url="http://127.0.0.1:${GATEWAY_PORT:-${PROXY_PORT}}/v1/completions"
184 + local ts; ts="$(date +%Y%m%d_%H%M%S)"
185 + local logf="${LOG_PATH}/accuracy_${MODEL_NAME}_${PARALLEL_MODE}_${ts}.log"
186 + local outdir="${LOG_PATH}/lm_eval_${MODEL_NAME}_${PARALLEL_MODE}_${ts}"
187 + log "accuracy -> ${base_url} model=${MODEL_PATH} tasks=${ACCURACY_TASKS} wide_ep_mode=${WIDE_EP_MODE}(${PARALLEL_MODE})"
188 + log "log=${logf} results=${outdir}"
189 +
190 + # --output_path makes lm_eval persist a results_*.json for scraping.
191 + local eval_rc=0
192 + ( set -o pipefail
193 + # Model weights are already served (loaded offline in a separate process);
194 + # allow the Hub only here so lm_eval can fetch the eval dataset (e.g. gsm8k).
195 + export HF_HUB_OFFLINE=0 HF_DATASETS_OFFLINE=0
196 + python3 -m lm_eval --model local-completions \
197 + --tasks "${ACCURACY_TASKS}" \
198 + --model_args "model=${MODEL_PATH},base_url=${base_url},num_concurrent=${ACCURACY_NUM_CONCURRENT},max_retries=${ACCURACY_MAX_RETRIES},tokenized_requests=False,trust_remote_code=True,timeout=${ACCURACY_TIMEOUT:-3600}" \
199 + --output_path "${outdir}" \
200 + 2>&1 | tee "${logf}"
201 + ) || eval_rc=$?
202 + if (( eval_rc != 0 )); then
203 + log "FAIL: lm_eval exited rc=${eval_rc} (no usable score) — Log: ${logf}"
204 + return "${eval_rc}"
205 + fi
206 +
207 + # Gate on the score. Find the results JSON lm_eval just wrote, pull the max
208 + # value of ACCURACY_METRIC across filters/tasks, and compare to the threshold.
209 + local results_json=""
210 + results_json="$(find "${outdir}" -name 'results_*.json' -type f 2>/dev/null | sort | tail -n1)"
211 + if [[ -z "${results_json}" ]]; then
212 + log "FAIL: no results_*.json under ${outdir} — cannot evaluate threshold"
213 + return 1
214 + fi
215 +
216 + local score=""
217 + score="$(python3 - "${results_json}" "${ACCURACY_METRIC}" <<'PY'
218 +import json, sys
219 +path, metric = sys.argv[1], sys.argv[2]
220 +with open(path) as f:
221 + data = json.load(f)
222 +best = None
223 +for task, metrics in (data.get("results") or {}).items():
224 + for key, val in metrics.items():
225 + if isinstance(val, bool) or not isinstance(val, (int, float)):
226 + continue
227 + if key.split(",")[0] == metric:
228 + best = val if best is None else max(best, val)
229 +print("" if best is None else "%.6f" % best)
230 +PY
231 +)" || score=""
232 +
233 + if [[ -z "${score}" ]]; then
234 + log "FAIL: could not parse metric '${ACCURACY_METRIC}' from ${results_json}"
235 + return 1
236 + fi
237 +
238 + if awk -v s="${score}" -v t="${ACCURACY_THRESHOLD}" 'BEGIN{exit !(s+0 >= t+0)}'; then
239 + log "PASS: ${ACCURACY_METRIC}=${score} >= threshold=${ACCURACY_THRESHOLD} (tasks=${ACCURACY_TASKS}). Log: ${logf}"
240 + return 0
241 + fi
242 + log "FAIL: ${ACCURACY_METRIC}=${score} < threshold=${ACCURACY_THRESHOLD} (tasks=${ACCURACY_TASKS}). Log: ${logf}"
243 + return 1
244 +}
245 +
246 +# ============================================================================
247 +# Server roles (prefill / decode) and their shared helpers
248 +# ============================================================================
249 +
250 +# Read model-specific flags from models.yaml (mode + role aware) into globals:
251 +# MODEL_BASE_FLAGS / MODEL_ROLE_FLAGS / MODEL_EXPERIMENTAL_FLAGS -> MODEL_CONFIG
252 +# Also exports the model's env: block (only if not already set; caller env wins).
253 +load_model_flags() {
254 + [[ -f "${MODELS_YAML}" ]] || die "models.yaml not found: ${MODELS_YAML}"
255 + export MODELS_YAML MODEL_NAME PARALLEL_MODE ROLE
256 + eval "$(python3 - <<'PY'
257 +import os, shlex, sys, yaml
258 +path = os.environ["MODELS_YAML"]; name = os.environ["MODEL_NAME"]
259 +mode = os.environ["PARALLEL_MODE"]; role = os.environ["ROLE"]
260 +with open(path, "r", encoding="utf-8") as f:
261 + doc = yaml.safe_load(f) or {}
262 +# Preferred form: top-level `models:` list of {model: <name>, ...} entries.
263 +# Back-compat: a top-level mapping keyed by model name.
264 +entries = doc.get("models") if isinstance(doc, dict) else doc
265 +cfg = None
266 +if isinstance(entries, list):
267 + cfg = next((e for e in entries if isinstance(e, dict) and e.get("model") == name), None)
268 +elif isinstance(doc, dict):
269 + cfg = doc.get(name)
270 +if cfg is None:
271 + print(f'echo "ERROR: model {name} not found in {path}" >&2; exit 1'); sys.exit(0)
272 +role_cfg = cfg.get(role, {}) or {}
273 +def q(v): return shlex.quote(str(v if v is not None else ""))
274 +exports = {
275 + "MODEL_BASE_FLAGS": cfg.get("base_flags", "") or "",
276 + "MODEL_ROLE_FLAGS": role_cfg.get(mode, "") or "",
277 + "MODEL_EXPERIMENTAL_FLAGS": cfg.get("experimental_flags", "") or "",
278 +}
279 +for k, v in exports.items(): print(f"{k}={q(v)}")
280 +# Model-specific env: exported only if not already set (caller env wins).
281 +for k, v in (cfg.get("env", {}) or {}).items():
282 + if k not in os.environ:
283 + print(f"export {k}={q(v)}")
284 +PY
285 +)"
286 + MODEL_CONFIG="${MODEL_BASE_FLAGS} ${MODEL_ROLE_FLAGS} ${MODEL_EXPERIMENTAL_FLAGS}"
287 +}
288 +
289 +# Build the TP-mode (WIDE_EP_MODE=0) serve command into the global CMD array.
290 +# Each node is an independent TP server registered to the proxy.
291 +build_tp_cmd() {
292 + local port="${TP_PORT}"
293 + [[ -n "${HOST_IP}" ]] && CMD+=(--host "${HOST_IP}")
294 + CMD+=(--port "${port}" --tensor-parallel-size "${TP_SIZE}")
295 + local _mc; read -ra _mc <<< "${MODEL_CONFIG}"; CMD+=("${_mc[@]}")
296 + local kv_cfg read_mode_kv
297 + read_mode_kv="$(moriio_read_mode_kv)"
298 + kv_cfg=$(cat <<EOF
299 +{
300 + "kv_connector": "MoRIIOConnector",
301 + "kv_role": "${KV_ROLE}",
302 + "kv_connector_extra_config": {
303 + "proxy_ip": "${PROXY_IP}",
304 + "proxy_ping_port": "${PROXY_PING_PORT}",
305 + "http_port": "${port}",
306 + "handshake_port": "${HANDSHAKE_PORT}",
307 + "notify_port": "${NOTIFY_PORT}"${read_mode_kv}
308 + }
309 +}
310 +EOF
311 +)
312 + CMD+=(--kv-transfer-config "${kv_cfg}")
313 + log "role=${ROLE} mode=tp host=${HOST_IP:-0.0.0.0}:${port} kv_role=${KV_ROLE} tp=${TP_SIZE} read_mode=${MORIIO_READ_MODE}"
314 +}
315 +
316 +# Build the EP-mode (WIDE_EP_MODE=1) serve command into the global CMD array.
317 +# DP + expert parallel; masters expose the API server + KV transfer, children
318 +# join headless.
319 +build_ep_cmd() {
320 + # prefill -> mori_high_throughput (paired with --enforce-eager in models.yaml)
321 + # decode -> mori_low_latency (paired with CUDA graphs in models.yaml)
322 + local default_backend="mori_low_latency"
323 + [[ "${ROLE}" == "prefill" ]] && default_backend="mori_high_throughput"
324 + local backend="${ALL2ALL_BACKEND:-${default_backend}}"
325 + export VLLM_ALL2ALL_BACKEND="${VLLM_ALL2ALL_BACKEND:-${backend}}"
326 + local port="${SERVE_PORT}"
327 + CMD+=(
328 + -tp 1
329 + --data-parallel-size "${DP_GROUP_SIZE}"
330 + --data-parallel-size-local "${GPUS_PER_NODE}"
331 + --data-parallel-address "${DP_MASTER_ADDR}"
332 + --data-parallel-rpc-port "${RPC_PORT}"
333 + --enable-expert-parallel
334 + --all2all-backend "${backend}"
335 + --port "${port}"
336 + --no-enable-prefix-caching
337 + --distributed-timeout-seconds "${DISTRIBUTED_TIMEOUT_SECONDS}"
338 + )
339 + local _mc; read -ra _mc <<< "${MODEL_CONFIG}"; CMD+=("${_mc[@]}")
340 +
341 + if [[ "${IS_MASTER}" == "1" ]]; then
342 + CMD+=(--api-server-count="${GPUS_PER_NODE}")
343 + local kv_cfg read_mode_kv
344 + read_mode_kv="$(moriio_read_mode_kv)"
345 + kv_cfg=$(cat <<EOF
346 +{
347 + "kv_connector": "MoRIIOConnector",
348 + "kv_role": "${KV_ROLE}",
349 + "kv_port": "${KV_PORT}",
350 + "kv_connector_extra_config": {
351 + "proxy_ip": "${PROXY_IP}",
352 + "proxy_port": "${PROXY_PORT}",
353 + "proxy_ping_port": "${PROXY_PING_PORT}",
354 + "http_port": "${port}",
355 + "local_ping_port": "${LOCAL_PING_PORT}",
356 + "handshake_port": "${HANDSHAKE_PORT}",
357 + "notify_port": "${NOTIFY_PORT}"${read_mode_kv}
358 + }
359 +}
360 +EOF
361 +)
362 + CMD+=(--kv-transfer-config "${kv_cfg}")
363 + log "role=${ROLE} mode=ep MASTER node_rank=${NODE_RANK} dp_size=${DP_GROUP_SIZE} dp_addr=${DP_MASTER_ADDR} kv_role=${KV_ROLE} all2all=${backend} read_mode=${MORIIO_READ_MODE}"
364 + else
365 + CMD+=(--data-parallel-start-rank "${DP_START_RANK}" --headless)
366 + log "role=${ROLE} mode=ep CHILD node_rank=${NODE_RANK} dp_size=${DP_GROUP_SIZE} dp_start_rank=${DP_START_RANK} dp_addr=${DP_MASTER_ADDR} all2all=${backend}"
367 + fi
368 +}
369 +
370 +# Set the prefill topology globals for the current NODE_RANK (no execution).
371 +configure_prefill() {
372 + : "${NODE_RANK:=0}" # manual default: prefill master
373 + KV_ROLE="kv_producer"
374 + DP_GROUP_SIZE=$(( xP * GPUS_PER_NODE ))
375 + DP_MASTER_ADDR="${PREFILL_MASTER_ADDR}"
376 + DP_START_RANK=$(( NODE_RANK * GPUS_PER_NODE ))
377 + [[ "${NODE_RANK}" -eq 0 ]] && IS_MASTER=1 || IS_MASTER=0
378 + TP_PORT="${PREFILL_PORT}"
379 +}
380 +
381 +# Set the decode topology globals for the current NODE_RANK (no execution).
382 +configure_decode() {
383 + : "${NODE_RANK:=${xP}}" # manual default: decode master
384 + KV_ROLE="kv_consumer"
385 + DP_GROUP_SIZE=$(( yD * GPUS_PER_NODE ))
386 + DP_MASTER_ADDR="${DECODE_MASTER_ADDR}"
387 + DP_START_RANK=$(( (NODE_RANK - xP) * GPUS_PER_NODE ))
388 + [[ "${NODE_RANK}" -eq "${xP}" ]] && IS_MASTER=1 || IS_MASTER=0
389 + TP_PORT="${DECODE_PORT}"
390 +}
391 +
392 +# Build the serve command (model flags + CMD array + LOGF). No exec.
393 +# Expects role globals set by configure_prefill/configure_decode.
394 +build_server_cmd() {
395 + HOST_IP="${IP_ARRAY[$NODE_RANK]:-}" # own IP for --host binding, if known
396 + load_model_flags
397 + LOGF="${LOG_PATH}/${ROLE}_${PARALLEL_MODE}_node${NODE_RANK}_$(date +%Y%m%d_%H%M%S).log"
398 + CMD=(vllm serve "${MODEL_PATH}")
399 + if [[ "${WIDE_EP_MODE}" == "0" ]]; then build_tp_cmd; else build_ep_cmd; fi
400 + log "model=${MODEL_PATH}"
401 +}
402 +
403 +# Foreground server run (explicit prefill/decode roles — original behavior).
404 +run_server_fg() {
405 + build_server_cmd
406 + log "log=${LOGF}"
407 + "${CMD[@]}" 2>&1 | tee "${LOGF}"
408 +}
409 +
410 +run_prefill() { configure_prefill; run_server_fg; }
411 +run_decode() { configure_decode; run_server_fg; }
412 +
413 +# ============================================================================
414 +# node mode — SLURM-native rank-based self-select + rank-0 orchestration
415 +# ============================================================================
416 +
417 +# Endpoints to health-check, one per line as "ip:port", depending on mode.
418 +# tp: every prefill node (:PREFILL_PORT) and every decode node (:DECODE_PORT)
419 +# ep: only the masters expose an API (prefill rank 0, decode rank xP) :SERVE_PORT
420 +health_endpoints() {
421 + local i
422 + if [[ "${WIDE_EP_MODE}" == "0" ]]; then
423 + for (( i=0; i<xP; i++ )); do echo "${IP_ARRAY[$i]}:${PREFILL_PORT}"; done
424 + for (( i=xP; i<xP+yD; i++ )); do echo "${IP_ARRAY[$i]}:${DECODE_PORT}"; done
425 + else
426 + echo "${IP_ARRAY[0]}:${SERVE_PORT}"
427 + echo "${IP_ARRAY[$xP]}:${SERVE_PORT}"
428 + fi
429 +}
430 +
431 +# Block until every endpoint's /health is OK, or timeout, or our local server dies.
432 +wait_all_healthy() {
433 + local deadline=$(( $(date +%s) + HEALTH_TIMEOUT_S ))
434 + local eps; mapfile -t eps < <(health_endpoints)
435 + log "health-gate: waiting on ${#eps[@]} endpoint(s): ${eps[*]}"
436 + local ep
437 + for ep in "${eps[@]}"; do
438 + until /usr/bin/curl -sf "http://${ep}/health" >/dev/null 2>&1; do
439 + (( $(date +%s) >= deadline )) && { log "TIMEOUT waiting for ${ep}"; return 1; }
440 + kill -0 "${SERVER_PID}" 2>/dev/null || { log "local server (pid ${SERVER_PID}) exited while waiting"; return 1; }
441 + sleep 10
442 + done
443 + log "healthy: ${ep}"
444 + done
445 + return 0
446 +}
447 +
448 +run_workload() {
449 + case "${RUN_AFTER_HEALTH}" in
450 + bench) run_bench ;;
451 + accuracy) run_accuracy ;;
452 + none|"") log "RUN_AFTER_HEALTH=none — skipping workload" ;;
453 + *) die "unknown RUN_AFTER_HEALTH='${RUN_AFTER_HEALTH}'" ;;
454 + esac
455 +}
456 +
457 +# rank 0: proxy + health-gate + workload, then write the completion sentinel.
458 +orchestrate_master() {
459 + local sentinel="$1" rc=0
460 + # Front door: the proxy runs in-container (started here); the vllm-router
461 + # runs as a SEPARATE container started by the SLURM job on this (rank-0) node,
462 + # so in that mode we don't start anything here.
463 + PROXY_PID=""
464 + if [[ "${ROUTER_TYPE:-vllm-router}" == "vllm-router" ]]; then
465 + log "ROUTER_TYPE=vllm-router: external router expected on gateway :${GATEWAY_PORT:-${ROUTER_PORT}} (not starting proxy)"
466 + else
467 + start_proxy_bg
468 + fi
469 + if wait_all_healthy; then
470 + # If using the external vllm-router, wait for it to become ready before
471 + # running the workload. The router starts as a separate container on this
472 + # (rank-0) node; give it up to 300s to bind port ROUTER_PORT.
473 + if [[ "${ROUTER_TYPE:-vllm-router}" == "vllm-router" ]]; then
474 + local router_deadline=$(( $(date +%s) + 300 ))
475 + log "waiting for vllm-router on :${ROUTER_PORT} (up to 300s)"
476 + until /usr/bin/curl -sf "http://127.0.0.1:${ROUTER_PORT}/health" >/dev/null 2>&1; do
477 + if (( $(date +%s) >= router_deadline )); then
478 + log "FAIL: vllm-router not ready on :${ROUTER_PORT} after 300s (bring-up)"
479 + rc=1
480 + break
481 + fi
482 + sleep 5
483 + done
484 + (( rc == 0 )) && log "vllm-router healthy on :${ROUTER_PORT}"
485 + fi
486 + if (( rc == 0 )); then
487 + set +e
488 + run_workload
489 + rc=$?
490 + set -e
491 + fi
492 + else
493 + log "FAIL: health-gate did not pass — server bring-up failed"
494 + rc=1
495 + fi
496 + echo "${rc}" > "${sentinel}" 2>/dev/null || true
497 + log "master: workload rc=${rc}; sentinel=${sentinel}; tearing down local proxy+server"
498 + [[ -n "${PROXY_PID:-}" ]] && kill "${PROXY_PID}" 2>/dev/null || true
499 + kill "${SERVER_PID}" 2>/dev/null || true
500 + wait "${SERVER_PID}" 2>/dev/null || true
501 + return "${rc}"
502 +}
503 +
504 +# non-master ranks: keep serving until the master signals completion (sentinel)
505 +# or our local server dies unexpectedly.
506 +watch_until_done() {
507 + local sentinel="$1" rc=0
508 + log "rank ${NODE_RANK} serving; waiting for completion sentinel (${sentinel})"
509 + while :; do
510 + [[ -f "${sentinel}" ]] && { log "completion sentinel seen; shutting down"; break; }
511 + if ! kill -0 "${SERVER_PID}" 2>/dev/null; then
512 + log "local server (pid ${SERVER_PID}) exited before completion"; rc=1; break
513 + fi
514 + sleep 10
515 + done
516 + kill "${SERVER_PID}" 2>/dev/null || true
517 + wait "${SERVER_PID}" 2>/dev/null || true
518 + return "${rc}"
519 +}
520 +
521 +run_node() {
522 + NODE_RANK="${NODE_RANK:-${SLURM_PROCID:-0}}"
523 + local total=$(( xP + yD ))
524 + (( NODE_RANK >= 0 && NODE_RANK < total )) \
525 + || die "NODE_RANK=${NODE_RANK} out of range [0,${total}) (xP=${xP} yD=${yD})"
526 +
527 + if (( NODE_RANK < xP )); then ROLE="prefill"; configure_prefill
528 + else ROLE="decode"; configure_decode
529 + fi
530 +
531 + # Shared-FS completion sentinel (LOG_PATH is shared & per-run/per-job).
532 + local sentinel="${LOG_PATH}/.disagg_done"
533 +
534 + log "node mode: rank=${NODE_RANK}/${total} role=${ROLE} master=${IS_MASTER} wide_ep=${WIDE_EP_MODE}(${PARALLEL_MODE})"
535 + log "topology: xP=${xP} yD=${yD} gpus/node=${GPUS_PER_NODE} IPADDRS=${IPADDRS}"
536 +
537 + # DRY_RUN: resolve+print the plan (role, serve cmd, endpoints) and exit 0.
538 + # Validates wiring without launching vLLM / the proxy / any workload.
539 + if [[ "${DRY_RUN:-0}" == "1" ]]; then
540 + build_server_cmd
541 + log "DRY_RUN: server cmd: ${CMD[*]}"
542 + if (( NODE_RANK == 0 )); then
543 + local eps; mapfile -t eps < <(health_endpoints)
544 + log "DRY_RUN: health endpoints: ${eps[*]}"
545 + log "DRY_RUN: post-health workload: ${RUN_AFTER_HEALTH}"
546 + log "DRY_RUN: sentinel: ${sentinel}"
547 + fi
548 + log "DRY_RUN complete (rank ${NODE_RANK}) — no processes launched"
549 + exit 0
550 + fi
551 +
552 + # aiter and flydsl keep runtime JIT/kernel caches inside the aiter install
553 + # tree (…/aiter/jit and …/aiter/jit/flydsl_cache), which is read-only for the
554 + # non-root (--user) container. AITER_JIT_DIR / FLYDSL_RUNTIME_CACHE_DIR
555 + # redirect those to node-local /tmp (both are sanctioned overrides: aiter only
556 + # points flydsl at its bundled cache when FLYDSL_RUNTIME_CACHE_DIR is unset).
557 + # Seed each once from the image's bundled copy so we keep the prebuilt
558 + # modules/kernels (fast, no NFS, no from-source rebuild / MLIR recompile).
559 + local _aiter_src
560 + _aiter_src="$(python3 -c 'import os,aiter; print(os.path.join(os.path.dirname(aiter.__file__),"jit"))' 2>/dev/null || true)"
561 + if [[ -n "${AITER_JIT_DIR:-}" && ! -e "${AITER_JIT_DIR}/.seeded" ]]; then
562 + if [[ -n "${_aiter_src}" && -d "${_aiter_src}" ]]; then
563 + mkdir -p "${AITER_JIT_DIR}"
564 + cp -a "${_aiter_src}/." "${AITER_JIT_DIR}/" 2>/dev/null || true
565 + touch "${AITER_JIT_DIR}/.seeded" 2>/dev/null || true
566 + log "seeded AITER_JIT_DIR=${AITER_JIT_DIR} from ${_aiter_src}"
567 + else
568 + log "WARN: could not locate aiter jit dir to seed AITER_JIT_DIR=${AITER_JIT_DIR}; aiter may rebuild from source"
569 + fi
570 + fi
571 + if [[ -n "${FLYDSL_RUNTIME_CACHE_DIR:-}" && ! -e "${FLYDSL_RUNTIME_CACHE_DIR}/.seeded" ]]; then
572 + mkdir -p "${FLYDSL_RUNTIME_CACHE_DIR}"
573 + if [[ -n "${_aiter_src}" && -d "${_aiter_src}/flydsl_cache" ]]; then
574 + cp -a "${_aiter_src}/flydsl_cache/." "${FLYDSL_RUNTIME_CACHE_DIR}/" 2>/dev/null || true
575 + log "seeded FLYDSL_RUNTIME_CACHE_DIR=${FLYDSL_RUNTIME_CACHE_DIR} from ${_aiter_src}/flydsl_cache"
576 + else
577 + log "FLYDSL_RUNTIME_CACHE_DIR=${FLYDSL_RUNTIME_CACHE_DIR} set; no bundled flydsl_cache to seed (will JIT-compile at runtime)"
578 + fi
579 + touch "${FLYDSL_RUNTIME_CACHE_DIR}/.seeded" 2>/dev/null || true
580 + fi
581 +
582 + # mori JIT-compiles its shmem/all2all kernels to ~/.mori/jit/<arch>_<nic>/ on
583 + # first EP use. HOME here is /workspace, bind-mounted to persistent+shared NFS
584 + # (/data/$USER), so a wrong-arch cache from an earlier run/build survives and
585 + # gets reloaded even after MORI_GPU_ARCHS changes -- an existing cache dir
586 + # wins over the arch setting. Reloading a gfx942 shmem_kernels.hsaco on gfx950
587 + # hardware fails with "device kernel image is invalid" and every EP worker
588 + # dies at init (surfacing async as a torch.tensor HIP error). Guard: on EP
589 + # runs, if the cache has no build for our target arch, purge it so mori
590 + # recompiles for MORI_GPU_ARCHS. A matching arch dir is kept (fast reload).
591 + if [[ "${WIDE_EP_MODE}" == "1" && -n "${MORI_GPU_ARCHS:-}" ]]; then
592 + local _mori_jit="${HOME:-/workspace}/.mori/jit"
593 + if [[ -d "${_mori_jit}" ]] && ! compgen -G "${_mori_jit}/${MORI_GPU_ARCHS}_*" >/dev/null 2>&1; then
594 + log "purging stale mori jit cache ${_mori_jit} (no ${MORI_GPU_ARCHS}_* build present; forcing recompile for ${MORI_GPU_ARCHS})"
595 + rm -rf "${_mori_jit}" 2>/dev/null || true
596 + fi
597 + fi
598 +
599 + # Start this node's server in the background (full server log -> LOGF; this
600 + # shell's stdout stays free for orchestration logs the srun/CI captures).
601 + build_server_cmd
602 + "${CMD[@]}" >"${LOGF}" 2>&1 &
603 + SERVER_PID=$!
604 + log "server started pid=${SERVER_PID} log=${LOGF}"
605 +
606 + local rc=0
607 + if (( NODE_RANK == 0 )); then
608 + orchestrate_master "${sentinel}" || rc=$?
609 + # Canonical, human-readable end-of-run verdict for the login-node poller
610 + # and CI logs (the .disagg_done sentinel remains the machine authority).
611 + log "VERDICT: $( (( rc == 0 )) && echo PASS || echo FAIL) rc=${rc}"
612 + else
613 + watch_until_done "${sentinel}" || rc=$?
614 + fi
615 + log "node rank=${NODE_RANK} exiting rc=${rc}"
616 + exit "${rc}"
617 +}
618 +
619 +# ============================================================================
620 +main() {
621 + parse_args "$@"
622 + load_config
623 + resolve_topology
624 + case "${ROLE}" in
625 + node) run_node ;;
626 + proxy) run_proxy ;;
627 + bench) run_bench ;;
628 + accuracy) run_accuracy ;;
629 + prefill) run_prefill ;;
630 + decode) run_decode ;;
631 + *) usage ;;
632 + esac
633 +}
634 +
635 +main "$@"