iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
Chaitanya Sri Krishna Lolla
2026-08-10 18:39:26 +00:00
committed by GitHub
co-authored by Sheral Kumar avininjamay8 tej Cursor
parent 05f0a80016
commit 1482d2e015
6 changed files with 1974 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/amd-disagg/cluster.sh
+166
View File
@@ -0,0 +1,166 @@
#!/bin/bash
# =============================================================================
# cluster.sh — single, generic config for the vLLM disaggregated (P/D) launcher.
# -----------------------------------------------------------------------------
# This script is *sourced* by vllm_disagg.sh.
#
# Every value is `${VAR:-default}`, so the environment always wins:
# environment variable > built-in default below
#
# So you can override anything inline:
# PREFILL_IP=10.0.0.1 DECODE_IP=10.0.0.2 ./vllm_disagg.sh prefill
#
# Site-specific values (model dir, IPs, NIC list, partition) are the defaults
# in the "site defaults" sections below — edit those for a new cluster.
# Model-SPECIFIC perf flags live in models.yaml, NOT here.
# =============================================================================
_CLUSTER_SH_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# ----------------------------------------------------------------- model / mode
# MODEL_NAME indexes into models.yaml; MODEL_DIR is the parent dir holding it.
# MODEL_PATH (resolved by the launcher) = ${MODEL_DIR}/${MODEL_NAME}. [site]
export MODEL_NAME="${MODEL_NAME:-DeepSeek-V3}"
export MODEL_DIR="${MODEL_DIR:-/data/models2}"
# Shared NFS root (5 TB, visible on every node): model weights + per-run logs. [site]
export SHARED_MOUNT="${SHARED_MOUNT:-/data}"
# Parallelism mode (launcher derives PARALLEL_MODE tp|ep from this):
# WIDE_EP_MODE=0 tp : each node is an independent TP server (TP8, 1P1D)
# WIDE_EP_MODE=1 ep : data-parallel + expert-parallel across xP/yD nodes (wideep)
export WIDE_EP_MODE="${WIDE_EP_MODE:-0}"
# ----------------------------------------------------------------- topology
# xP prefill nodes + yD decode nodes. IPADDRS is the ordered, comma-separated
# node list (prefill IPs first, then decode IPs). NODE_RANK is this node's global
# 0-based rank; under SLURM it defaults to $SLURM_PROCID. Leave IPADDRS empty to
# use the PREFILL_IP/DECODE_IP fallback defaults below (1P1D only).
export xP="${xP:-1}"
export yD="${yD:-1}"
export IPADDRS="${IPADDRS:-}"
export NODE_RANK="${NODE_RANK:-${SLURM_PROCID:-}}"
export PREFILL_IP="${PREFILL_IP:-10.0.0.1}"
export DECODE_IP="${DECODE_IP:-10.0.0.2}"
# Per-node GPU count and TP degree
export GPUS_PER_NODE="${GPUS_PER_NODE:-8}"
export TP_SIZE="${TP_SIZE:-${GPUS_PER_NODE}}"
# ----------------------------------------------------------------- ports
# TP-mode (WIDE_EP_MODE=0) server ports.
export PREFILL_PORT="${PREFILL_PORT:-8100}"
export DECODE_PORT="${DECODE_PORT:-8200}"
# EP-mode (WIDE_EP_MODE=1) ports: API serve port, DP RPC port, KV transfer port, and
# per-node MoRIIO local ping port (must differ from PROXY_PING_PORT).
export SERVE_PORT="${SERVE_PORT:-20005}"
export RPC_PORT="${RPC_PORT:-13345}"
export KV_PORT="${KV_PORT:-9711}"
export LOCAL_PING_PORT="${LOCAL_PING_PORT:-61555}"
# MoRIIO proxy: HTTP port clients/benchmark hit, plus the connector control ports.
# PROXY_PING_PORT MUST be 36367 — the proxy hardcodes its zmq service-discovery
# socket on that port; prefill/decode register to PROXY_IP:PROXY_PING_PORT.
export PROXY_IP="${PROXY_IP:-${PREFILL_IP}}"
export PROXY_PORT="${PROXY_PORT:-10001}"
export PROXY_PING_PORT="${PROXY_PING_PORT:-36367}"
export HANDSHAKE_PORT="${HANDSHAKE_PORT:-6301}"
export NOTIFY_PORT="${NOTIFY_PORT:-61005}"
export PROXY_SCRIPT="${PROXY_SCRIPT:-/app/vllm/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py}"
# MoRIIO KV transfer direction (injected into --kv-transfer-config by the launcher):
# 0 -> omit read_mode (default; MoRIIO write mode: prefill pushes to decode)
# 1 -> "read_mode": true (decode pulls KV from prefill; matches upstream disagg)
export MORIIO_READ_MODE="${MORIIO_READ_MODE:-0}"
# ----------------------------------------------------------------- router / gateway
# Selection for client (bench/accuracy) traffic:
# proxy -> the in-container MoRIIO proxy started by the launcher
# vllm-router -> an external `vllm/vllm-router` container started by the SLURM job
# on the rank-0 node (default)
# Both use the SAME MoRIIO discovery mechanism (prefill/decode register to
# PROXY_IP:PROXY_PING_PORT=36367); only the client HTTP front door differs.
export ROUTER_TYPE="${ROUTER_TYPE:-vllm-router}"
export ROUTER_PORT="${ROUTER_PORT:-30000}"
export ROUTER_POLICY="${ROUTER_POLICY:-round_robin}"
export VLLM_ROUTER_IMAGE="${VLLM_ROUTER_IMAGE:-vllm/vllm-router:nightly}"
# Single client-facing port bench/accuracy target: the router port when routing,
# else the proxy port. Env override always wins.
if [[ "${ROUTER_TYPE}" == "vllm-router" ]]; then
export GATEWAY_PORT="${GATEWAY_PORT:-${ROUTER_PORT}}"
else
export GATEWAY_PORT="${GATEWAY_PORT:-${PROXY_PORT}}"
fi
# Where per-run logs / benchmark results are written. A $SLURM_JOB_ID subdir is
# appended so each CI run is self-scoped (falls back to 'local' off-SLURM). [site]
_LOG_BASE="${LOG_BASE:-/data/${USER:-$(id -un)}/disagg_logs}"
export LOG_PATH="${LOG_PATH:-${_LOG_BASE}/${SLURM_JOB_ID:-local}}"
# ----------------------------------------------------------------- vLLM runtime
# Engine/platform/transport-level env (NOT model-specific). Model-architecture
# AITER kernel toggles live in models.yaml under each model's `env:` block.
export VLLM_USE_V1="${VLLM_USE_V1:-1}"
export HSA_NO_SCRATCH_RECLAIM="${HSA_NO_SCRATCH_RECLAIM:-1}"
#export HF_HUB_OFFLINE="${HF_HUB_OFFLINE:-1}"
#export TRANSFORMERS_OFFLINE="${TRANSFORMERS_OFFLINE:-1}"
#
#export HOME=/tmp
export HF_HOME="${HF_HOME:-/tmp/hf_home}"
export XDG_CACHE_HOME="${XDG_CACHE_HOME:-/tmp/.cache}"
export VLLM_ENGINE_READY_TIMEOUT_S="${VLLM_ENGINE_READY_TIMEOUT_S:-3600}"
# ----------------------------------------------------------------- scale
# DP/EP group formation timeout across nodes (seconds).
export DISTRIBUTED_TIMEOUT_SECONDS="${DISTRIBUTED_TIMEOUT_SECONDS:-7200}"
# ----------------------------------------------------------------- RDMA / NCCL
# AMD Pensando AINIC RoCE fabric: 8 NICs exposed as ionic_0..7 (netdevs eth2..9),
# each rail on its own /24. GID index 1 + traffic class 104
_IB_DEVICES="${IB_DEVICES:-ionic_0,ionic_1,ionic_2,ionic_3,ionic_4,ionic_5,ionic_6,ionic_7}"
_IB_GID_INDEX="${NCCL_IB_GID_INDEX:-1}"
export IB_DEVICES="${IB_DEVICES:-${_IB_DEVICES}}"
export NCCL_IB_HCA="${NCCL_IB_HCA:-${_IB_DEVICES}}"
export NCCL_IB_GID_INDEX="${NCCL_IB_GID_INDEX:-${_IB_GID_INDEX}}"
export NCCL_IB_DISABLE="${NCCL_IB_DISABLE:-0}"
# In-box RCCL net transport (no external plugin); pin bootstrap to the VPC iface.
export NCCL_NET_PLUGIN="${NCCL_NET_PLUGIN:-none}"
export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-eth0}"
export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-${NCCL_SOCKET_IFNAME}}"
export NCCL_CROSS_NIC="${NCCL_CROSS_NIC:-0}"
export NCCL_PXN_DISABLE="${NCCL_PXN_DISABLE:-0}"
export NCCL_NET_DISABLE_INTRA="${NCCL_NET_DISABLE_INTRA:-1}"
export NCCL_IB_TC="${NCCL_IB_TC:-104}"
export NCCL_IB_FIFO_TC="${NCCL_IB_FIFO_TC:-192}"
export NCCL_IB_QPS_PER_CONNECTION="${NCCL_IB_QPS_PER_CONNECTION:-1}"
export NCCL_IB_TIMEOUT="${NCCL_IB_TIMEOUT:-22}"
export NCCL_IB_RETRY_CNT="${NCCL_IB_RETRY_CNT:-12}"
# MoRI uses the same NIC set as NCCL.
export MORI_RDMA_DEVICES="${MORI_RDMA_DEVICES:-${_IB_DEVICES}}"
export MORI_IB_GID_INDEX="${MORI_IB_GID_INDEX:-${_IB_GID_INDEX}}"
export MORI_SHMEM_HEAP_SIZE="${MORI_SHMEM_HEAP_SIZE:-16G}"
# Pin to gfx950 to avoid jit compilation failures with other archs on this cluster.
export MORI_GPU_ARCHS="gfx950"
# ----------------------------------------------------------------- benchmark
export BENCHMARK_COMBINATIONS="${BENCHMARK_COMBINATIONS:-1024/128 2048/128}"
export BENCHMARK_CON="${BENCHMARK_CON:-32 64}"
export NUM_PROMPTS_FACTOR="${NUM_PROMPTS_FACTOR:-2}"
export BENCHMARK_MIN_PROMPTS="${BENCHMARK_MIN_PROMPTS:-32}"
# ----------------------------------------------------------------- accuracy
# `accuracy` role runs lm_eval (local-completions backend) against the proxy.
export ACCURACY_TASKS="${ACCURACY_TASKS:-gsm8k}"
export ACCURACY_NUM_CONCURRENT="${ACCURACY_NUM_CONCURRENT:-64}"
export ACCURACY_MAX_RETRIES="${ACCURACY_MAX_RETRIES:-3}"
export ACCURACY_METRIC="${ACCURACY_METRIC:-exact_match}"
export ACCURACY_THRESHOLD="${ACCURACY_THRESHOLD:-0.90}"
# ----------------------------------------------------------------- SLURM (submit)
# Used by run-slurm-disagg-test.sh on the login node (harmless to export here). [site]
export SLURM_PARTITION="${SLURM_PARTITION:-default}"
.buildkite/amd-disagg/models.yaml
+97
View File
@@ -0,0 +1,97 @@
# =============================================================================
# Model catalog for the vLLM disaggregated (P/D) inference CI.
# -----------------------------------------------------------------------------
# Loaded by vllm_disagg.sh for the selected MODEL_NAME. WIDE_EP_MODE selects the
# flag set (0 -> tp keys, 1 -> ep keys):
# tp -> base_flags + prefill.tp / decode.tp
# ep -> base_flags + prefill.ep / decode.ep
#
# Keep ONLY model-specific perf flags here. The launcher owns everything
# topological / mode-structural:
# - tp: --host/--port, --tensor-parallel-size, --kv-transfer-config
# - ep: --tp 1, --data-parallel-size[-local], --data-parallel-address/-rpc-port,
# --data-parallel-start-rank/--headless (children), --enable-expert-parallel,
# --all2all-backend mori, --no-enable-prefix-caching,
# --api-server-count + --kv-transfer-config (masters only)
# Do NOT put any of those in this file.
#
# Format: a top-level `models:` list. The launcher selects the entry whose
# `model:` matches MODEL_NAME. Each entry:
# model : name (must match MODEL_NAME; also the dir under MODEL_DIR)
# env : model/arch-specific environment variables. The launcher
# exports these BEFORE `vllm serve`, but only if they are
# not already set, so precedence is:
# caller/inline env > models.yaml env:
# (Transport/fabric/engine env stays in cluster.sh,
# e.g. MORI_RDMA_DEVICES, NCCL_IB_HCA, VLLM_USE_V1.)
# base_flags : always applied (both roles, both modes)
# prefill/decode.<tp|ep> : role + mode specific perf flags
# experimental_flags : optional extra flags (omit if empty)
# =============================================================================
models:
- model: DeepSeek-V3
env:
VLLM_ROCM_USE_AITER: "1"
VLLM_ROCM_USE_AITER_MLA: "1"
VLLM_ROCM_USE_AITER_MOE: "1"
VLLM_ROCM_USE_AITER_RMSNORM: "1"
VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: "0"
base_flags: "--trust-remote-code --kv-cache-dtype fp8"
prefill:
tp: "--gpu-memory-utilization 0.85"
ep: "--gpu-memory-utilization 0.85 --enforce-eager"
decode:
tp: "--gpu-memory-utilization 0.85"
#ep: '--gpu-memory-utilization 0.75 --enforce-eager'
ep: '--gpu-memory-utilization 0.75 --compilation-config {"cudagraph_mode":"PIECEWISE","custom_ops":["+quant_fp8"]}'
- model: MiniMax-M3-MXFP8
env:
VLLM_ROCM_USE_AITER: "1"
VLLM_ROCM_USE_AITER_MOE: "1"
VLLM_ROCM_USE_AITER_RMSNORM: "1"
VLLM_USE_BREAKABLE_CUDAGRAPH: "0"
VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: "INT6"
base_flags: "--trust-remote-code --attention-backend TRITON_ATTN --block-size 128 --language-model-only --kv-cache-dtype fp8"
prefill:
tp: "--gpu-memory-utilization 0.85 --enforce-eager"
decode:
tp: "--gpu-memory-utilization 0.85"
- model: DeepSeek-R1-MXFP4
env:
VLLM_ROCM_USE_AITER: "1"
VLLM_ROCM_USE_AITER_MLA: "1"
VLLM_ROCM_USE_AITER_MOE: "1"
VLLM_ROCM_USE_AITER_RMSNORM: "1"
base_flags: "--trust-remote-code --kv-cache-dtype fp8"
prefill:
tp: "--gpu-memory-utilization 0.85"
decode:
tp: "--gpu-memory-utilization 0.85"
- model: Kimi-K2.5-MXFP4
env:
VLLM_ROCM_USE_AITER: "1"
VLLM_ROCM_USE_AITER_MOE: "1"
VLLM_ROCM_USE_AITER_RMSNORM: "1"
base_flags: "--trust-remote-code"
prefill:
tp: "--gpu-memory-utilization 0.85"
decode:
tp: "--gpu-memory-utilization 0.85"
- model: Kimi-K2.6-MXFP4
env:
VLLM_ROCM_USE_AITER: "1"
AMDGCN_USE_BUFFER_OPS: "1"
VLLM_ROCM_USE_AITER_MLA: "1"
VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: "INT4"
VLLM_ROCM_USE_SKINNY_GEMM: "0"
VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: "1"
base_flags: "--trust-remote-code --kv-cache-dtype fp8 --mm-encoder-tp-mode data --block-size 1 --attention-backend ROCM_AITER_MLA"
prefill:
tp: "--gpu-memory-utilization 0.9"
decode:
tp: "--gpu-memory-utilization 0.9"
.buildkite/amd-disagg/pipeline-disagg.yaml
+202
View File
@@ -0,0 +1,202 @@
# Native Buildkite pipeline — disaggregated PD SLURM tests (MoRIIO).
#
# ROUTER_TYPE must be set explicitly on every step: run-slurm-disagg-test.sh
# defaults to `vllm-router`
#
steps:
#-------------------------------------- mi350 · disagg 1P1D (TP8, 2 nodes) · SLURM --------------------------------------#
- label: "DeepSeek-V3-PD-1P1D-TP8-MoRIIO-proxy"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/2node"
commands:
- 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
- label: "DeepSeek-V3-PD-1P1D-TP8-MoRIIO-vllm-router"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/2node"
commands:
- 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
- label: "DeepSeek-R1-MXFP4-PD-1P1D-TP8-MoRIIO-proxy"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/2node"
commands:
- 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
- label: "DeepSeek-R1-MXFP4-PD-1P1D-TP8-MoRIIO-vllm-router"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/2node"
commands:
- 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
- label: "Kimi-K2.5-MXFP4-PD-1P1D-TP8-MoRIIO-proxy"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/2node"
commands:
- 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
- label: "Kimi-K2.5-MXFP4-PD-1P1D-TP8-MoRIIO-vllm-router"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/2node"
commands:
- 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
- label: "Kimi-K2.6-MXFP4-PD-1P1D-TP8-MoRIIO-proxy"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/2node"
commands:
- 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
- label: "Kimi-K2.6-MXFP4-PD-1P1D-TP8-MoRIIO-vllm-router"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/2node"
commands:
- 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
- label: "MiniMax-M3-MXFP8-PD-1P1D-TP8-MoRIIO-proxy"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/2node"
commands:
- 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
- label: "MiniMax-M3-MXFP8-PD-1P1D-TP8-MoRIIO-vllm-router"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/2node"
commands:
- 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
#-------------------------------------- mi350 · disagg 2P2D (TP8, 4 nodes) · SLURM --------------------------------------#
- label: "DeepSeek-V3-PD-2P2D-TP8-MoRIIO-proxy"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/4node"
commands:
- 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
- label: "DeepSeek-V3-PD-2P2D-TP8-MoRIIO-vllm-router"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/4node"
commands:
- 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
- label: "DeepSeek-R1-MXFP4-PD-2P2D-TP8-MoRIIO-proxy"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/4node"
commands:
- 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
- label: "DeepSeek-R1-MXFP4-PD-2P2D-TP8-MoRIIO-vllm-router"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/4node"
commands:
- 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
- label: "Kimi-K2.5-MXFP4-PD-2P2D-TP8-MoRIIO-proxy"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/4node"
commands:
- 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
- label: "Kimi-K2.5-MXFP4-PD-2P2D-TP8-MoRIIO-vllm-router"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/4node"
commands:
- 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
- label: "Kimi-K2.6-MXFP4-PD-2P2D-TP8-MoRIIO-proxy"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/4node"
commands:
- 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
- label: "Kimi-K2.6-MXFP4-PD-2P2D-TP8-MoRIIO-vllm-router"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/4node"
commands:
- 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
- label: "MiniMax-M3-MXFP8-PD-2P2D-TP8-MoRIIO-proxy"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/4node"
commands:
- 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
- label: "MiniMax-M3-MXFP8-PD-2P2D-TP8-MoRIIO-vllm-router"
agents:
queue: amd_mi350_ainic
timeout_in_minutes: 120
concurrency: 2
concurrency_group: "amd-disagg/mi350/4node"
commands:
- 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
#-------------------------------------- mi350 · disagg wide-EP (disabled) --------------------------------------#
# - label: "DeepSeek-V3-PD-1P1D-EP8/DP8-WideEP-MoRIIO-proxy"
# agents:
# queue: amd_mi350_ainic
# timeout_in_minutes: 120
# concurrency: 2
# concurrency_group: "amd-disagg/mi350/2node"
# commands:
# - 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
+338
View File
@@ -0,0 +1,338 @@
#!/bin/bash
# =============================================================================
# run-slurm-disagg-test.sh — foreground (login-node) submitter for the disagg
# P/D gate. The single command the Buildkite step runs.
# -----------------------------------------------------------------------------
# Thin glue: submits run_xPyD_disagg.slurm. By default (WAIT=0, Spur-safe) it
# fire-and-forgets and returns 0 on a good submit; with WAIT=1 it polls the
# scheduler and exits with the job's exit code so a CI step pass/fails.
# stdout carries ONLY the SLURM job info
# (image/nodes line, job id, log path); after sbatch nothing else is written to
# stdout — the full job log persists at the --output path and pass/fail
# diagnostics go to stderr. The actual work lives in run_xPyD_disagg.slurm: it
# selects nodes,
# fans out one container per node via a single srun, and hands off to
# vllm_disagg.sh (rank-based prefill/decode self-select).
#
# Default target: 1P1D TP8 (NODES=2), nightly image, accuracy gate.
#
# Spur usage (fire-and-forget; the default here):
# bash run-slurm-disagg-test.sh # 1P1D TP8
# WIDE_EP_MODE=1 bash run-slurm-disagg-test.sh # 1P1D wide-EP
# NODES=4 WIDE_EP_MODE=1 xP=2 yD=2 bash run-slurm-disagg-test.sh # 2P2D EP
# Prints the job id + NFS log paths and returns immediately (does NOT poll
# squeue/sacct, which hang on Spur). Track the run via the printed `tail -f`.
#
# Classic-Slurm / CI usage (block until done, exit with the gate's pass/fail):
# WAIT=1 bash run-slurm-disagg-test.sh
# =============================================================================
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
JOB_SCRIPT="${JOB_SCRIPT:-${SCRIPT_DIR}/run_xPyD_disagg.slurm}"
# ---- knobs (override from the Buildkite step env) --------------------------
# Defaults tuned for the Spur AMD MI350X cluster
IMAGE="${IMAGE:-vllm/vllm-openai-rocm:nightly}"
NODES="${NODES:-2}"
GPUS_PER_NODE="${GPUS_PER_NODE:-8}"
PARTITION="${SLURM_PARTITION:-}"
TIME_LIMIT="${SLURM_TIME_LIMIT:-02:00:00}"
WIDE_EP_MODE="${WIDE_EP_MODE:-0}" # 0 -> 1P1D TP8 (default); 1 -> wide-EP
xP="${xP:-1}"
yD="${yD:-1}"
RUN_AFTER_HEALTH="${RUN_AFTER_HEALTH:-accuracy}"
HEALTH_TIMEOUT_S="${HEALTH_TIMEOUT_S:-3600}" # P/D bring-up budget; +900s grace must fit the 2h wall
SHARED_MOUNT="${SHARED_MOUNT:-/data}"
LOG_ROOT="${LOG_ROOT:-${SHARED_MOUNT}/${USER:-$(whoami)}/disagg_logs}"
DRY_RUN="${DRY_RUN:-0}"
MORIIO_READ_MODE="${MORIIO_READ_MODE:-0}"
WAIT="${WAIT:-0}"
# ROUTER Type - defaults to vllm-router
ROUTER_TYPE="${ROUTER_TYPE:-vllm-router}"
ROUTER_PORT="${ROUTER_PORT:-30000}"
VLLM_ROUTER_IMAGE="${VLLM_ROUTER_IMAGE:-vllm/vllm-router:nightly}"
# Dry-run only validates wiring; cap its walltime low so it never holds the queue.
[[ "${DRY_RUN}" == "1" ]] && TIME_LIMIT="${SLURM_TIME_LIMIT:-00:10:00}"
mkdir -p "${LOG_ROOT}"
# Spur - sbatch scheduler.
DISAGG_SCRIPTS_STAGE="${DISAGG_SCRIPTS_STAGE:-/data/scratch/buildkite-agent}"
STAGED_DIR="${DISAGG_SCRIPTS_STAGE}/${BUILDKITE_COMMIT:-local}"
mkdir -p "${STAGED_DIR}"
cp -rL --no-preserve=ownership,timestamps "${SCRIPT_DIR}/." "${STAGED_DIR}/"
chmod -R u+rwX "${STAGED_DIR}" 2>/dev/null || true
export DISAGG_SCRIPTS_DIR="${STAGED_DIR}"
echo "[slurm-submit] staged scripts for compute nodes: ${DISAGG_SCRIPTS_DIR}" >&2
export IMAGE MODEL_NAME WIDE_EP_MODE xP yD GPUS_PER_NODE RUN_AFTER_HEALTH HEALTH_TIMEOUT_S
export SHARED_MOUNT LOG_ROOT DRY_RUN MORIIO_READ_MODE
export ROUTER_TYPE ROUTER_PORT VLLM_ROUTER_IMAGE
# Model selection.
[[ -n "${MODEL_NAME:-}" ]] && export MODEL_NAME
[[ -n "${MODEL_DIR:-}" ]] && export MODEL_DIR
[[ -n "${PARTITION}" ]] && echo "[slurm-submit] NOTE: PARTITION='${PARTITION}' ignored (Spur sbatch has no --partition)" >&2
SUBMIT_SCRIPT="${JOB_SCRIPT}"
FILE_NODES="$(grep -oE '^#SBATCH[[:space:]]+--nodes=[0-9]+' "${JOB_SCRIPT}" | grep -oE '[0-9]+' | head -n1)"
if [[ -n "${NODES}" && -n "${FILE_NODES}" && "${NODES}" != "${FILE_NODES}" ]]; then
SUBMIT_SCRIPT="/tmp/$(basename "${JOB_SCRIPT%.slurm}")-n${NODES}-$$.slurm"
sed -E "s/^#SBATCH([[:space:]]+)--nodes=[0-9]+/#SBATCH\\1--nodes=${NODES}/" "${JOB_SCRIPT}" > "${SUBMIT_SCRIPT}"
echo "[slurm-submit] node count ${FILE_NODES}->${NODES}: submitting patched copy ${SUBMIT_SCRIPT}" >&2
fi
# Job name comes from the script's #SBATCH --job-name (SLURM sets SLURM_JOB_NAME
JOB_NAME="$(grep -oE '^#SBATCH[[:space:]]+--job-name=[^[:space:]]+' "${JOB_SCRIPT}" | sed -E 's/.*--job-name=//' | head -n1)"
JOB_NAME="${JOB_NAME:-vllm-disagg-pd}"
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}"
SUBMIT_OUT="$(sbatch --time="${TIME_LIMIT}" "${SUBMIT_SCRIPT}")"
echo "${SUBMIT_OUT}"
# "Submitted batch job 114" -> 114 (last integer on the line).
JOB_ID="$(printf '%s\n' "${SUBMIT_OUT}" | grep -oE '[0-9]+' | tail -n1 || true)"
if [[ -z "${JOB_ID}" ]]; then
echo "[slurm-submit] ERROR: could not parse a job id from sbatch output above" >&2
exit 1
fi
echo "[slurm-submit] submitted job ${JOB_ID}"
# The real per-job log is written to NFS from inside the job body by
# run_xPyD_disagg.slurm (`exec > "$LOG_ROOT/${SLURM_JOB_NAME}-${SLURM_JOB_ID}.log"`).
# SLURM_JOB_NAME resolves to our --job-name, so we can compute the path here. The
# per-role server/proxy/bench logs land under ${LOG_ROOT}/${JOB_ID}/.
LOG_FILE="${LOG_ROOT}/${JOB_NAME}-${JOB_ID}.log"
LOG_DIR="${LOG_ROOT}/${JOB_ID}"
echo "[slurm-submit] job log: ${LOG_FILE}"
echo "[slurm-submit] role logs: ${LOG_DIR}/"
# --- Spur-safe by default: fire-and-forget --------------------------------------
if [[ "${WAIT}" != "1" ]]; then
echo "[slurm-submit] submitted (WAIT=0, not polling). Track with:" >&2
echo " tail -f ${LOG_FILE}" >&2
echo " grep -aE '(PASS|FAIL): |exact_match' ${LOG_FILE}" >&2
exit 0
fi
# --- WAIT=1: phase-aware poll of the NFS job log (scontrol-based) -------------
# Detection is phased so failures surface fast instead of waiting out the full
# walltime:
# 1) submitted -> running : scontrol catches infra/scheduler failures
# (NODE_FAIL/BOOT_FAIL/CANCELLED/TIMEOUT/...) within
# one poll, or a stuck-PENDING/never-started job.
# 2) running -> healthy : advance once every endpoint reports healthy (+ the
# vllm-router, when used); fail on bring-up errors or
# the health budget.
# 3) healthy -> verdict : PASS/FAIL from the accuracy gate, capped.
# scontrol is the scheduler authority on this cluster
echo "[slurm-submit] WAIT=1: phase-aware poll of ${LOG_FILE} (timeout ${TIME_LIMIT})" >&2
_h=0 _m=0 _s=0
IFS=: read -r _h _m _s <<< "${TIME_LIMIT}"
WAIT_DEADLINE=$(( $(date +%s) + 10#${_h}*3600 + 10#${_m}*60 + 10#${_s:-0} ))
unset -v _h _m _s
# Per-phase budgets (all overridable from the Buildkite step env).
POLL_INTERVAL="${POLL_INTERVAL:-20}"
SUBMIT_GRACE_S="${SUBMIT_GRACE_S:-900}" # reach RUNNING within 15m
PENDING_MAX_S="${PENDING_MAX_S:-1800}" # tolerate 30m queued
HEALTH_PHASE_TIMEOUT_S="${HEALTH_PHASE_TIMEOUT_S:-$(( HEALTH_TIMEOUT_S + 900 ))}"
WORKLOAD_TIMEOUT_S="${WORKLOAD_TIMEOUT_S:-1800}" # accuracy/bench cap
COMPLETED_GRACE_S="${COMPLETED_GRACE_S:-60}" # verdict lag after job exit
SENTINEL="${LOG_DIR}/.disagg_done"
# scontrol field extractor (authoritative here; timeout-guarded so a momentary
# scheduler stall can't wedge the poll). Returns the value or "".
job_field() { # $1=jobid $2=field -> value | ""
timeout 15 scontrol show job "$1" 2>/dev/null \
| grep -oE "$2=[^ ]+" | head -n1 | cut -d= -f2- || true
}
have() { grep -aqE "$1" "${LOG_FILE}" 2>/dev/null; }
# The job's own verdict sources, in authority order: the .disagg_done sentinel
# (rank-0 rc, written by a single writer) and then the gate line in the log.
# Sets STATE/RC/REASON and returns 0 when a verdict exists, 1 when it does not.
read_verdict() {
if [[ -f "${SENTINEL}" ]]; then
RC="$(tr -dc '0-9' < "${SENTINEL}" 2>/dev/null || true)"; RC="${RC:-1}"
if [[ "${RC}" == "0" ]]; then STATE="COMPLETED"; else STATE="FAILED"; fi
REASON="sentinel"; return 0
fi
if have '(PASS|FAIL): '; then
if have 'FAIL: '; then STATE="FAILED"; RC=1; else STATE="COMPLETED"; RC=0; fi
REASON="gate"; return 0
fi
return 1
}
# Never let the job outlive this poller. Cancelling a Buildkite build kills the
# agent's bootstrap, and without this the sbatch job keeps its whole allocation
# until the walltime expires. Only armed on the WAIT=1 path — under WAIT=0,
# leaving the job running is the point.
CANCEL_GRACE_S="${CANCEL_GRACE_S:-120}"
_CLEANED=0
# shellcheck disable=SC2329 # invoked from cleanup_job, which the traps below call
job_active() {
case "$(job_field "${JOB_ID}" JobState)" in
RUNNING|PENDING|COMPLETING|CONFIGURING|SUSPENDED|REQUEUED) return 0 ;;
*) return 1 ;;
esac
}
# shellcheck disable=SC2329 # invoked from the traps below
cleanup_job() {
[[ "${_CLEANED}" == "1" || -z "${JOB_ID:-}" ]] && return 0
_CLEANED=1
# A job that reported its own verdict is mid-teardown; give it a bounded
# moment to finish. Timeouts and infra failures get cancelled immediately —
# there is nothing to wait for and the nodes should come back now.
if [[ "${REASON:-}" == "sentinel" || "${REASON:-}" == "gate" || "${STATE:-}" == "COMPLETED" ]]; then
local deadline=$(( $(date +%s) + CANCEL_GRACE_S ))
while (( $(date +%s) < deadline )) && job_active; do sleep 5; done
fi
if job_active; then
echo "[slurm-submit] cleanup: cancelling job ${JOB_ID}" >&2
scancel "${JOB_ID}" 2>/dev/null \
|| echo "[slurm-submit] WARN: scancel ${JOB_ID} failed; job may still hold nodes" >&2
fi
return 0
}
trap cleanup_job EXIT
trap 'cleanup_job; exit 130' INT
trap 'cleanup_job; exit 143' TERM
trap 'cleanup_job; exit 129' HUP
STATE=""
RC=1
REASON=""
PHASE="submitted"
T_PHASE=$(date +%s)
while [[ $(date +%s) -lt ${WAIT_DEADLINE} ]]; do
NOW=$(date +%s)
# (1) Ultimate authority: terminal sentinel (holds rank-0 rc), then the
# explicit accuracy gate line. Honored regardless of phase.
read_verdict && break
# (2) Scheduler state via scontrol: drives phase transitions and catches
# infra/scheduler failures fast.
ST="$(job_field "${JOB_ID}" JobState)"
case "${ST}" in
RUNNING|COMPLETING)
if [[ "${PHASE}" == "submitted" ]]; then PHASE="bringup"; T_PHASE="${NOW}"; fi
;;
NODE_FAIL|BOOT_FAIL|CANCELLED|TIMEOUT|OUT_OF_MEMORY|DEADLINE|PREEMPTED)
STATE="infra-${ST}"; RC=1
REASON="scontrol Reason=$(job_field "${JOB_ID}" Reason)"; break
;;
FAILED)
# Ambiguous (infra crash vs a legit accuracy exit=1). A gate/sentinel
# line would have won above, so classify by the phase we reached.
case "${PHASE}" in
submitted) STATE="infra-FAILED" ;;
bringup) STATE="server-failed" ;;
*) STATE="workload-failed" ;;
esac
RC=1; REASON="scontrol JobState=FAILED phase=${PHASE}"; break
;;
COMPLETED)
_vd=$(( NOW + COMPLETED_GRACE_S ))
_found=0
while :; do
if read_verdict; then _found=1; break; fi
(( $(date +%s) >= _vd )) && break
sleep 5
done
if (( _found == 1 )); then
echo "[slurm-submit] verdict arrived after JobState=COMPLETED (${REASON})" >&2
else
STATE="completed-no-verdict"; RC=1
REASON="COMPLETED with no sentinel/gate within ${COMPLETED_GRACE_S}s (a walltime kill reports COMPLETED here)"
fi
break
;;
PENDING|CONFIGURING|RESV_DEL_HOLD|REQUEUED)
if (( NOW - T_PHASE > PENDING_MAX_S )); then
STATE="infra-stuck-${ST}"; RC=1; REASON="queued > ${PENDING_MAX_S}s"; break
fi
;;
"")
# scontrol lost the job (purged past MinJobAge) with no log verdict:
# rely on the terminal markers above + the per-phase deadlines below.
:
;;
esac
# (3) Log-driven phase progress + per-phase deadlines. Fast-fails hangs where
# the job is still RUNNING (so scontrol won't help) but bring-up/workload
# is stuck.
case "${PHASE}" in
submitted)
# Job body appearing in the log is a second "it started" signal for
# when scontrol is briefly empty.
if have 'Selected node IPs|health-gate:|\[disagg-pd\]'; then
PHASE="bringup"; T_PHASE="${NOW}"
elif [[ -z "${ST}" ]] && (( NOW - T_PHASE > SUBMIT_GRACE_S )); then
# scontrol can't confirm the job exists (empty state) AND no log
# output within the grace window -> treat as a lost/failed launch.
# NB: a genuinely-queued job reports PENDING via scontrol above and
# is governed by PENDING_MAX_S (default 30m), not this grace window.
STATE="infra-nostart"; RC=1; REASON="no scheduler state or log within ${SUBMIT_GRACE_S}s"; break
fi
;;
bringup)
if have 'FAIL:|TIMEOUT waiting for|exited while waiting|Traceback \(most recent'; then
STATE="server-failed"; RC=1; REASON="bring-up failure in log"; break
fi
NEED="$(grep -aoE 'waiting on [0-9]+' "${LOG_FILE}" 2>/dev/null | grep -oE '[0-9]+' | tail -n1 || true)"
GOT="$(grep -acE '\] healthy: ' "${LOG_FILE}" 2>/dev/null || true)"; GOT="${GOT:-0}"
ROUTER_OK=1
if [[ "${ROUTER_TYPE}" == "vllm-router" ]]; then
if ! have 'vllm-router healthy on'; then ROUTER_OK=0; fi
fi
if [[ -n "${NEED}" ]] && (( GOT >= NEED )) && (( ROUTER_OK == 1 )); then
PHASE="workload"; T_PHASE="${NOW}"
elif (( NOW - T_PHASE > HEALTH_PHASE_TIMEOUT_S )); then
STATE="bringup-timeout"; RC=1; REASON="no healthy within ${HEALTH_PHASE_TIMEOUT_S}s"; break
fi
;;
workload)
# PASS/FAIL handled at the top; only enforce the cap here.
if (( NOW - T_PHASE > WORKLOAD_TIMEOUT_S )); then
STATE="workload-timeout"; RC=1; REASON="no verdict within ${WORKLOAD_TIMEOUT_S}s"; break
fi
;;
esac
sleep "${POLL_INTERVAL}"
done
if [[ -z "${STATE}" ]]; then
echo "[slurm-submit] WARN: no verdict before ${TIME_LIMIT}; failing" >&2
STATE="deadline"; RC=1; REASON="poll deadline"
fi
# Pre-flight logic.
PF_LINES=()
if compgen -G "${LOG_DIR}/preflight_NODE*.log" >/dev/null 2>&1; then
cat -- "${LOG_DIR}"/preflight_NODE*.log >&2 || true
mapfile -t PF_LINES < <(grep -h '^PREFLIGHT-REJECTED: ' "${LOG_DIR}"/preflight_NODE*.log 2>/dev/null || true)
fi
if (( ${#PF_LINES[@]} > 0 )) && (( RC != 0 )); then
STATE="preflight-rejected"
REASON="${PF_LINES[0]#PREFLIGHT-REJECTED: }"
if (( ${#PF_LINES[@]} > 1 )); then
REASON="${REASON} [+$(( ${#PF_LINES[@]} - 1 )) more node(s), see ${LOG_DIR}/]"
fi
fi
# Surface the accuracy gate verdict (if any) from the job log — to stderr.
GATE_LINE=$(grep -aE '(PASS|FAIL): ' "${LOG_FILE}" 2>/dev/null | tail -n1 || true)
[[ -n "${GATE_LINE}" ]] && echo "[slurm-submit] gate: ${GATE_LINE}" >&2
echo "[slurm-submit] job ${JOB_ID} finished: state=${STATE:-unknown} phase=${PHASE} exit=${RC} reason=${REASON:-}" >&2
exit "${RC}"
.buildkite/amd-disagg/run_xPyD_disagg.slurm
+536
View File
@@ -0,0 +1,536 @@
#!/bin/bash
#SBATCH --job-name=vllm-disagg-pd
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=1
#SBATCH --gres=gpu:8
#SBATCH --exclusive
#SBATCH --time=02:00:00
#SBATCH --chdir=/tmp
#SBATCH --output=/tmp/spur-%j.out
#SBATCH --error=/tmp/spur-%j.err
set -euo pipefail
# ---- inputs (env > default; forwarded by run-slurm-disagg-test.sh via --export)
IMAGE="${IMAGE:-vllm/vllm-openai-rocm:nightly}"
WIDE_EP_MODE="${WIDE_EP_MODE:-0}"
xP="${xP:-1}"
yD="${yD:-1}"
GPUS_PER_NODE="${GPUS_PER_NODE:-8}"
SHARED_MOUNT="${SHARED_MOUNT:-/data}"
LOG_ROOT="${LOG_ROOT:-/data/${USER:-$(id -un)}/disagg_logs}"
RUN_AFTER_HEALTH="${RUN_AFTER_HEALTH:-accuracy}" # bench | accuracy
DRY_RUN="${DRY_RUN:-0}"
MORIIO_READ_MODE="${MORIIO_READ_MODE:-0}" # 0 -> prefill writes (default); 1 -> decode reads KV (read_mode)
# ---- router / gateway (host-side; also forwarded into the container) --------
# vllm-router -> external `vllm/vllm-router` container started on the rank-0 node (default)
# proxy -> in-container MoRIIO proxy (no extra container)
ROUTER_TYPE="${ROUTER_TYPE:-vllm-router}"
ROUTER_PORT="${ROUTER_PORT:-30000}"
ROUTER_POLICY="${ROUTER_POLICY:-round_robin}"
VLLM_ROUTER_IMAGE="${VLLM_ROUTER_IMAGE:-vllm/vllm-router:nightly}"
PROXY_PING_PORT="${PROXY_PING_PORT:-36367}" # must equal vllm-router --vllm-discovery-address port
PROXY_PORT="${PROXY_PORT:-10001}"
# Intra-node data-parallel size the router uses to route the MoRIIO KV-notify to
# the right DP rank. Required for wideEP DP (else decode hangs: "remote blocks
# never arrived"). EP -> GPUS_PER_NODE (matches --data-parallel-size-local); TP -> 1.
if [[ "${WIDE_EP_MODE}" == "1" ]]; then
ROUTER_DP_LOCAL="${ROUTER_DP_LOCAL:-${GPUS_PER_NODE}}"
else
ROUTER_DP_LOCAL="${ROUTER_DP_LOCAL:-1}"
fi
if [[ "${ROUTER_TYPE}" == "vllm-router" ]]; then
GATEWAY_PORT="${GATEWAY_PORT:-${ROUTER_PORT}}"
else
GATEWAY_PORT="${GATEWAY_PORT:-${PROXY_PORT}}"
fi
# This line is a recommended template from infra side.
exec > "$LOG_ROOT/${SLURM_JOB_NAME}-${SLURM_JOB_ID}.log" 2>&1
# Where the disagg scripts live ON THE HOST (bind-mounted into the container).
# Spur runs the job body from /tmp, so $(pwd)/SLURM_SUBMIT_DIR aren't reliable;
# default to the known host scripts dir. Override via env if you relocate it.
DISAGG_SCRIPTS_DIR="${DISAGG_SCRIPTS_DIR:-${SLURM_SUBMIT_DIR:-$(pwd)}}"
echo "DISAGG SCRIPTS DIR is ----"
echo "$DISAGG_SCRIPTS_DIR"
# In-container path the scripts are mounted to (and sourced from).
CONTAINER_SCRIPTS="${CONTAINER_SCRIPTS:-/vllm-workspace/.buildkite/amd-disagg}"
CLUSTER_ENV_IN="${CLUSTER_ENV_IN:-${CONTAINER_SCRIPTS}/cluster.sh}"
# Health endpoints (TP-mode ports from cluster.sh defaults).
PREFILL_PORT="${PREFILL_PORT:-8100}"
DECODE_PORT="${DECODE_PORT:-8200}"
HEALTH_TIMEOUT_S="${HEALTH_TIMEOUT_S:-3600}"
# EP-mode ports, mirrored from cluster.sh purely so the preflight below can test
# them host-side; the in-container values still come from cluster.sh.
SERVE_PORT="${SERVE_PORT:-20005}"
RPC_PORT="${RPC_PORT:-13345}"
KV_PORT="${KV_PORT:-9711}"
LOCAL_PING_PORT="${LOCAL_PING_PORT:-61555}"
# vllm-router binds this for Prometheus in addition to ROUTER_PORT, and panics
# at startup if it is taken.
ROUTER_METRICS_PORT="${ROUTER_METRICS_PORT:-29000}"
# Node preflight (see preflight_node): 0 disables, VRAM% above which a GPU counts
# as busy, and how long to let a reaped container's VRAM drain before judging.
PREFLIGHT="${PREFLIGHT:-1}"
PREFLIGHT_MAX_VRAM_PCT="${PREFLIGHT_MAX_VRAM_PCT:-10}"
PREFLIGHT_GPU_SETTLE_S="${PREFLIGHT_GPU_SETTLE_S:-30}"
NUM_NODES=$((xP + yD))
echo "Calculated NUM_NODES: $NUM_NODES (xP=$xP + yD=$yD) WIDE_EP_MODE=$WIDE_EP_MODE"
# ------------------------
# Extract NUM_NODES from SLURM allocation and select nodes
# ------------------------
echo "Original SLURM allocation:"
echo "SLURM_JOB_NODELIST: $SLURM_JOB_NODELIST"
echo "SLURM_NNODES: $SLURM_NNODES"
echo "SLURM_NTASKS: $SLURM_NTASKS"
# scontrol is available on classic Slurm but not in the Spur job body. Use it to
# expand compressed ranges when present, else split the (already comma-separated,
# full-hostname) SLURM_JOB_NODELIST on commas.
if command -v scontrol >/dev/null 2>&1; then
FULL_NODELIST=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
else
FULL_NODELIST=$(echo "$SLURM_JOB_NODELIST" | tr ',' '\n')
fi
AVAILABLE_COUNT=$(echo "$FULL_NODELIST" | wc -l)
if (( AVAILABLE_COUNT < NUM_NODES )); then
echo "ERROR: Need ${NUM_NODES} nodes (xP=${xP} + yD=${yD}) but only ${AVAILABLE_COUNT} allocated." >&2
exit 1
fi
# Pick the first NUM_NODES nodes and order them to match how IPs get assigned:
# * classic Slurm: sort alphabetically -- srun assigns PROCID=0 to the
# alphabetically-first hostname, so IPS[] (built per node below) lines up.
# * Spur: preserve RANK order (SPUR_NODELIST), so the hostnames pair 1:1 with
# the rank-ordered IPs in SPUR_PEER_NODES; an alphabetical sort here would
# desync MASTER_NODE / MASTER_ADDR / IPADDRS from the actual ranks.
# IPADDRS[0..xP-1] = prefill, IPADDRS[xP..] = decode, so this order must match.
if command -v srun >/dev/null 2>&1; then
SELECTED_NODES=$(echo "$FULL_NODELIST" | head -n "$NUM_NODES" | sort)
elif [[ -n "${SPUR_NODELIST:-}" ]]; then
SELECTED_NODES=$(echo "${SPUR_NODELIST}" | tr ',' '\n' | sed '/^$/d' | head -n "$NUM_NODES")
else
SELECTED_NODES=$(echo "$FULL_NODELIST" | head -n "$NUM_NODES")
fi
SELECTED_NODELIST_STR=$(echo "$SELECTED_NODES" | tr '\n' ',' | sed 's/,$//')
# Create new nodelist in SLURM format
# This is a simplified approach - for complex ranges, you might need more sophisticated parsing
NEW_SLURM_NODELIST=$(echo "$SELECTED_NODES" | paste -sd, | sed 's/,/,/g')
# Update SLURM environment variables
export SLURM_NNODES=$NUM_NODES
export SLURM_NTASKS=$NUM_NODES
export SLURM_JOB_NUM_NODES=$NUM_NODES
export SLURM_NPROCS=$NUM_NODES
export SLURM_JOB_NODELIST="$NEW_SLURM_NODELIST"
export SLURM_NODELIST="$NEW_SLURM_NODELIST"
# Keep other SLURM variables as they were or set defaults
export SLURM_TASKS_PER_NODE="1(x$NUM_NODES)"
export SLURM_SUBMIT_DIR="${SLURM_SUBMIT_DIR:-${HOME}}"
export SLURM_CLUSTER_NAME="${SLURM_CLUSTER_NAME:-vllm-di-ci}"
export SLURM_JOB_CPUS_PER_NODE="${SLURM_JOB_CPUS_PER_NODE:-1}"
export SLURM_JOB_PARTITION="${SLURM_JOB_PARTITION:-amd-rccl}"
export SLURM_JOBID="${SLURM_JOBID:-$SLURM_JOB_ID}"
export SLURM_JOB_QOS="${SLURM_JOB_QOS:-normal}"
export SLURM_JOB_ACCOUNT="${SLURM_JOB_ACCOUNT:-amd-rccl}"
export SLURM_NTASKS_PER_NODE=1
export SLURM_SUBMIT_HOST="${SLURM_SUBMIT_HOST:-$(hostname)}"
export SLURM_JOB_ID="${SLURM_JOB_ID:-}"
export SLURM_CONF="${SLURM_CONF:-/etc/slurm/slurm.conf}"
export SLURM_JOB_NAME="${SLURM_JOB_NAME:-${xP}p${yD}d_bench-serving}"
echo ""
echo "Updated SLURM Environment Variables:"
echo "SLURM_JOB_ID: $SLURM_JOB_ID"
echo "SLURM_JOB_NODELIST: $SLURM_JOB_NODELIST"
echo "SLURM_NNODES: $SLURM_NNODES"
echo "SLURM_NTASKS: $SLURM_NTASKS"
echo "SLURM_TASKS_PER_NODE: $SLURM_TASKS_PER_NODE"
echo "SLURM_JOB_CPUS_PER_NODE: $SLURM_JOB_CPUS_PER_NODE"
echo "SLURM_JOB_PARTITION: $SLURM_JOB_PARTITION"
echo "SLURM_JOB_NUM_NODES: $SLURM_JOB_NUM_NODES"
echo "SLURM_JOBID: $SLURM_JOBID"
echo "SLURM_JOB_QOS: $SLURM_JOB_QOS"
echo "SLURM_NODELIST: $SLURM_NODELIST"
echo "SLURM_JOB_ACCOUNT: $SLURM_JOB_ACCOUNT"
echo "SLURM_NPROCS: $SLURM_NPROCS"
echo "SLURM_SUBMIT_HOST: $SLURM_SUBMIT_HOST"
echo "SLURM_CONF: $SLURM_CONF"
echo "SLURM_JOB_NAME: $SLURM_JOB_NAME"
echo "SLURM_NTASKS_PER_NODE: $SLURM_NTASKS_PER_NODE"
echo "SLURM_SUBMIT_DIR: $SLURM_SUBMIT_DIR"
echo "SLURM_CLUSTER_NAME: $SLURM_CLUSTER_NAME"
echo "ulimit: $(ulimit -a)"
echo ""
echo "Selected nodes for execution:"
echo "$SELECTED_NODES"
echo ""
# Node information
MASTER_NODE=$(echo "$SELECTED_NODES" | head -n 1)
MASTER_PORT=39566 # Choose an open port
IPS=()
if command -v srun >/dev/null 2>&1; then
# Classic Slurm: resolve each node's IP via srun (original behavior).
for NODE in $SELECTED_NODES; do
IP=$(srun --nodes=1 --ntasks=1 --time=00:20:00 --nodelist="$NODE" bash -c 'hostname -I' | awk 'NR==1 {print $1}')
IPS+=("$IP")
done
else
# Spur job body (no srun): use the rank-ordered IPs Spur provides in
# SPUR_PEER_NODES ("ip:port,ip:port,..."). These are the cross-node-routable
# bootstrap addresses; hostnames resolve to public IPs and the local node to
# 127.0.1.1, so DNS is not usable here.
mapfile -t IPS < <(echo "${SPUR_PEER_NODES:-}" | tr ',' '\n' | cut -d: -f1 | sed '/^$/d')
if (( ${#IPS[@]} == 0 )); then
echo "ERROR: no srun and SPUR_PEER_NODES is empty; cannot determine node IPs." >&2
exit 1
fi
fi
MASTER_ADDR="${IPS[0]}"
IPADDRS="${IPS[*]}"
IPADDRS="${IPADDRS// /,}"
echo "Selected node IPs: ${IPADDRS}"
NNODES=$NUM_NODES
echo "MASTER_NODE is ${MASTER_NODE}"
echo "MASTER_ADDR is ${MASTER_ADDR}"
echo "MASTER_PORT is ${MASTER_PORT}"
echo "NNODES is ${NNODES}"
# Per-job log folder (MAD-style): all per-role server/proxy/bench logs
export LOG_PATH="${LOG_ROOT}/${SLURM_JOB_ID}"
mkdir -p "$LOG_PATH"
echo "[disagg-pd] LOG_PATH=$LOG_PATH"
# ----------------------------------------------------------------------------
# Launch one container per node via a single srun and hand off to
# vllm_disagg.sh node (rank-based self-select).
# ----------------------------------------------------------------------------
# IPADDRS is the ordered node IP list (prefill IPs first, then decode). IPS[]
# was built in SELECTED_NODES (alphabetical) order — the same order srun assigns
# SLURM_PROCID — so the rank->role mapping lines up:
# rank 0 -> prefill master + MoRIIO proxy + orchestrator
# rank 1 .. xP-1 -> prefill
# rank xP .. end -> decode
export IPADDRS
export WIDE_EP_MODE xP yD GPUS_PER_NODE RUN_AFTER_HEALTH HEALTH_TIMEOUT_S DRY_RUN MORIIO_READ_MODE
export IMAGE SHARED_MOUNT CONTAINER_SCRIPTS DISAGG_SCRIPTS_DIR
export CLUSTER_ENV="${CLUSTER_ENV_IN}"
export DOCKER_CONT_NAME="vllm_disagg_${MODEL_NAME:-model}_${SLURM_JOB_ID}"
# Router env: forwarded into the main container (cluster.sh honors these) and used
# host-side to launch the external router container on rank 0.
export ROUTER_TYPE ROUTER_PORT ROUTER_POLICY VLLM_ROUTER_IMAGE PROXY_PING_PORT GATEWAY_PORT ROUTER_DP_LOCAL
export ROUTER_CONT_NAME="vllm_router_${SLURM_JOB_ID}"
# Consumed by preflight_node on each node (srun forwards the environment).
export PREFILL_PORT DECODE_PORT PROXY_PORT ROUTER_METRICS_PORT
export SERVE_PORT RPC_PORT KV_PORT LOCAL_PING_PORT
export PREFLIGHT PREFLIGHT_MAX_VRAM_PCT PREFLIGHT_GPU_SETTLE_S
echo "[disagg-pd] IPADDRS=${IPADDRS}"
echo "[disagg-pd] launching ${NUM_NODES} containers (1/node) -> vllm_disagg.sh node"
# Refuse a node that is not actually free. $1 = this node's global rank.
#
# `#SBATCH --exclusive` only keeps out other SLURM jobs. A plain `docker run` on
# a node is invisible to the scheduler, so we get allocated nodes whose GPUs and
# ports are already in use. Without this the job pulls the image, loads the model
# and dies minutes later inside vLLM ("Free memory on device cuda:N ... is less
# than desired GPU memory utilization"); worse, on rank 0 a foreign router
# already listening on ROUTER_PORT can satisfy the health gate and take the
# accuracy traffic. Called after the stale-container reap, with a settle window
# so VRAM the reap is still releasing is not counted against us.
preflight_node() {
local rank="$1"
if [ "${PREFLIGHT:-1}" != "1" ]; then
echo "Rank $rank preflight: skipped (PREFLIGHT=0)"
return 0
fi
local host max_pct settle problems=""
local -a issues=() # one short tag per fault, for the one-line summary
host="$(hostname)"
max_pct="${PREFLIGHT_MAX_VRAM_PCT:-10}"
settle="${PREFLIGHT_GPU_SETTLE_S:-30}"
# GPUs: vLLM asks for a fraction of TOTAL, so memory already resident is
# subtracted from what it can get and pushes it under its own threshold.
if command -v rocm-smi >/dev/null 2>&1; then
local csv busy nbusy ncards deadline
deadline=$(( $(date +%s) + settle ))
while :; do
csv="$(rocm-smi --showmeminfo vram --csv 2>/dev/null || true)"
busy="$(printf '%s\n' "$csv" |
awk -F, -v m="$max_pct" '/^card/ && $2+0 > 0 {
pct = $3 * 100 / $2
if (pct > m) printf "%s=%d%%(%.0fGiB) ", $1, pct, $3/1073741824
}' || true)"
if [ -z "$busy" ]; then break; fi
if [ "$(date +%s)" -ge "$deadline" ]; then break; fi
sleep 5
done
if [ -n "$busy" ]; then
nbusy="$(printf '%s' "$busy" | tr ' ' '\n' | grep -c 'card' || true)"
ncards="$(printf '%s\n' "$csv" | grep -c '^card' || true)"
issues+=("${nbusy}/${ncards} GPUs hold VRAM above ${max_pct}%")
problems="${problems} GPU VRAM in use above ${max_pct}%: ${busy}"$'\n'
fi
else
echo "Rank $rank preflight: no rocm-smi on $host, skipping GPU check"
fi
# Ports this rank will bind. Every container here is --network host, so a
# foreign listener either collides or silently serves our traffic.
local -a want=()
if [ "${WIDE_EP_MODE:-0}" = "1" ]; then
want+=("$SERVE_PORT" "$RPC_PORT" "$KV_PORT" "$LOCAL_PING_PORT")
elif [ "$rank" -lt "${xP:-1}" ]; then
want+=("$PREFILL_PORT")
else
want+=("$DECODE_PORT")
fi
if [ "$rank" = "0" ]; then
want+=("$PROXY_PING_PORT")
if [ "$ROUTER_TYPE" = "vllm-router" ]; then
want+=("$ROUTER_PORT" "$ROUTER_METRICS_PORT")
else
want+=("$PROXY_PORT")
fi
fi
if command -v ss >/dev/null 2>&1; then
local bound taken="" p
bound="$(ss -lnt 2>/dev/null | awk '{ n = split($4, a, ":"); print a[n] }' | sort -u || true)"
for p in "${want[@]}"; do
if printf '%s\n' "$bound" | grep -qx -- "$p"; then
taken="${taken}${p} "
fi
done
if [ -n "$taken" ]; then
issues+=("ports already bound: ${taken% }")
problems="${problems} ports already bound: ${taken}"$'\n'
fi
else
echo "Rank $rank preflight: no ss on $host, skipping port check"
fi
if [ -n "$problems" ]; then
mkdir -p "$LOG_PATH" 2>/dev/null || true
local summary
summary="$(printf '%s; ' "${issues[@]}")"
summary="${summary%; }"
# Every node writes the shared job log concurrently, so this diagnostic can
# be overwritten there. Keep a single-writer copy per node as well, or the
# job fails in seconds with no surviving explanation. The submitter reads
# the PREFLIGHT-REJECTED line back out of it to use as the failure reason,
# so keep that prefix in sync with run-slurm-disagg-test.sh. It must not
# contain "FAIL:" — that is the accuracy gate's pattern.
{
echo "PREFLIGHT-REJECTED: rank ${rank} on ${host}: ${summary}"
printf '%s' "$problems"
echo " containers here: $(docker ps --format '{{.Names}} ({{.Image}})' 2>/dev/null | paste -sd'; ' - || true)"
echo " --exclusive does not exclude non-SLURM workloads; clear the node or exclude it from the allocation."
echo " Set PREFLIGHT=0 to submit anyway."
} | tee -a "$LOG_PATH/preflight_NODE${rank}.log" 2>/dev/null || true
# Unblock the other ranks: they park on this sentinel, and the submitter
# polls it. Without it a peer waits out HEALTH_TIMEOUT_S for an endpoint
# that is never coming.
echo 1 > "$LOG_PATH/.disagg_done" 2>/dev/null || true
return 1
fi
echo "Rank $rank preflight ok on $host: GPUs idle, ports free (${want[*]})"
return 0
}
# Per-node container launch, auto-dispatched by srun availability:
# Launch this node's disagg container and hand off to vllm_disagg.sh. $1 = this
# node's global rank. Relies on the env exported above (forwarded by srun on
# classic Slurm; already in scope when called directly on Spur).
run_node_container() {
local rank="$1"
set -o pipefail # propagate docker run exit code through the tee pipe
echo "Rank $rank on $(hostname)"
local -a stale=()
mapfile -t stale < <(docker ps -aq --filter name=vllm_disagg_ --filter name=vllm_router_ 2>/dev/null || true)
if (( ${#stale[@]} > 0 )); then
echo "Reaping stale containers on $(hostname): ${stale[*]}"
docker stop -t 30 "${stale[@]}" >/dev/null 2>&1 || true
docker rm "${stale[@]}" >/dev/null 2>&1 || true
fi
preflight_node "$rank" || return 1
# Bind-mount the host AMD Pensando (ionic) RoCE userspace provider into the
# container IF present: the image ships an ABI-incompatible copy and the NICs
# are invisible without the host libs. Each entry is added only when the host
# file exists, so this is a no-op on non-ionic (e.g. mlx5) clusters.
local IONIC_MOUNTS=()
local host_ionic
host_ionic=$(readlink -f /usr/lib/x86_64-linux-gnu/libionic.so.1 2>/dev/null || true)
[ -n "$host_ionic" ] && IONIC_MOUNTS+=(-v "$host_ionic:/usr/lib/x86_64-linux-gnu/libionic.so.1:ro")
[ -f /usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so ] && \
IONIC_MOUNTS+=(-v "/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so:/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so:ro")
[ -f /etc/libibverbs.d/ionic.driver ] && \
IONIC_MOUNTS+=(-v "/etc/libibverbs.d/ionic.driver:/etc/libibverbs.d/ionic.driver:ro")
echo "Rank $rank ionic mounts: ${IONIC_MOUNTS[*]:-<none>}"
# Pull latest container.
if [ "$DRY_RUN" != "1" ]; then
echo "Rank $rank pulling image on $(hostname): $IMAGE"
docker pull "$IMAGE" || echo "WARN: docker pull failed on $(hostname) for $IMAGE (using local copy if present)"
if [ "$ROUTER_TYPE" = "vllm-router" ] && [ "$rank" = "0" ]; then
echo "Rank $rank pulling router image on $(hostname): $VLLM_ROUTER_IMAGE"
docker pull "$VLLM_ROUTER_IMAGE" || echo "WARN: docker pull failed on $(hostname) for $VLLM_ROUTER_IMAGE (using local copy if present)"
fi
fi
# Start the external vLLM router as a SEPARATE container on the rank-0 node.
if [ "$ROUTER_TYPE" = "vllm-router" ] && [ "$rank" = "0" ] && [ "$DRY_RUN" = "1" ]; then
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)"
elif [ "$ROUTER_TYPE" = "vllm-router" ] && [ "$rank" = "0" ]; then
mkdir -p "$LOG_PATH" 2>/dev/null || true
{
echo "Starting vllm-router on $(hostname): $VLLM_ROUTER_IMAGE (discovery :$PROXY_PING_PORT, port :$ROUTER_PORT, policy $ROUTER_POLICY, dp_local=$ROUTER_DP_LOCAL)"
docker run -d \
--name "$ROUTER_CONT_NAME" \
--network host \
--user "$(id -u):$(id -g)" \
-v "$LOG_PATH":/run_logs \
-v "$SHARED_MOUNT":"$SHARED_MOUNT" \
-v "/data/$(id -un):/workspace" \
-e HOME=/workspace \
-e USER="$(id -un)" \
-e LOGNAME="$(id -un)" \
"$VLLM_ROUTER_IMAGE" \
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" \
&& echo "vllm-router container started ok" \
|| echo "WARN: failed to start vllm-router container on $(hostname)"
} 2>&1 | tee -a "$LOG_PATH/vllm_router_start.log"
fi
docker run --rm \
--user "$(id -u):$(id -g)" \
--device /dev/dri \
--device /dev/kfd \
--device /dev/infiniband \
--network host \
--ipc host \
--group-add "$(getent group video | cut -d: -f3)" \
--group-add "$(getent group render | cut -d: -f3)" \
--cap-add SYS_PTRACE \
--cap-add IPC_LOCK \
--security-opt seccomp=unconfined \
--shm-size 64G \
--ulimit nofile=1048576:1048576 \
-v "$SHARED_MOUNT":"$SHARED_MOUNT" \
-v "$DISAGG_SCRIPTS_DIR":"$CONTAINER_SCRIPTS" \
-v "/data/$(id -un):/workspace" \
"${IONIC_MOUNTS[@]}" \
--ulimit memlock=-1:-1 \
-e HF_TOKEN \
-e HOME=/workspace \
-e USER="$(id -un)" \
-e LOGNAME="$(id -un)" \
-e TRITON_CACHE_DIR=/tmp/triton_cache \
-e TORCHINDUCTOR_CACHE_DIR=/tmp/inductor_cache \
-e VLLM_CACHE_ROOT=/tmp/vllm_cache \
-e AITER_JIT_DIR=/tmp/aiter_jit \
-e FLYDSL_RUNTIME_CACHE_DIR=/tmp/flydsl_cache \
-e SLURM_JOB_ID="$SLURM_JOB_ID" \
-e SLURM_JOB_NODELIST="$SLURM_JOB_NODELIST" \
-e NODE_RANK="$rank" \
-e IPADDRS="$IPADDRS" \
-e WIDE_EP_MODE="$WIDE_EP_MODE" \
-e xP="$xP" -e yD="$yD" \
-e GPUS_PER_NODE="$GPUS_PER_NODE" \
-e CLUSTER_ENV="$CLUSTER_ENV" \
-e MODEL_NAME="${MODEL_NAME:-}" \
-e MODEL_DIR="${MODEL_DIR:-}" \
-e LOG_PATH="$LOG_PATH" \
-e RUN_AFTER_HEALTH="$RUN_AFTER_HEALTH" \
-e HEALTH_TIMEOUT_S="$HEALTH_TIMEOUT_S" \
-e DRY_RUN="$DRY_RUN" \
-e MORIIO_READ_MODE="$MORIIO_READ_MODE" \
-e ROUTER_TYPE="$ROUTER_TYPE" \
-e ROUTER_PORT="$ROUTER_PORT" \
-e ROUTER_POLICY="$ROUTER_POLICY" \
-e GATEWAY_PORT="$GATEWAY_PORT" \
--name "$DOCKER_CONT_NAME" \
--entrypoint /bin/bash \
"$IMAGE" -c "bash $CONTAINER_SCRIPTS/vllm_disagg.sh node" \
2>&1 | tee "$LOG_PATH/main_NODE${rank}.log"
local node_rc=${PIPESTATUS[0]} # docker run exit code (not tee), for rank-0 pass/fail
# Stop the external router (rank-0 node) once the main container returns.
if [ "$ROUTER_TYPE" = "vllm-router" ] && [ "$rank" = "0" ]; then
docker stop "$ROUTER_CONT_NAME" 2>/dev/null || true; docker rm "$ROUTER_CONT_NAME" 2>/dev/null || true;
fi
return "$node_rc"
}
# Teardown of the main + external router containers on every node ($1 = docker
# stop timeout). `docker run` containers outlive the client that started them,
# so on scancel the srun tasks die but the containers keep the GPUs until the
# walltime expires unless they are stopped explicitly.
teardown_containers() {
export TEARDOWN_STOP_T="${1:-30}"
if command -v srun >/dev/null 2>&1; then
# shellcheck disable=SC2016 # must expand on the remote node, not here
srun --nodelist="$SELECTED_NODELIST_STR" \
--nodes="$NUM_NODES" --ntasks="$NUM_NODES" --ntasks-per-node=1 \
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
[ "$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
else
docker stop -t "$TEARDOWN_STOP_T" "$DOCKER_CONT_NAME" >/dev/null 2>&1 || true
docker rm "$DOCKER_CONT_NAME" >/dev/null 2>&1 || true
if [ "$ROUTER_TYPE" = "vllm-router" ]; then
docker stop -t "$TEARDOWN_STOP_T" "$ROUTER_CONT_NAME" >/dev/null 2>&1 || true
docker rm "$ROUTER_CONT_NAME" >/dev/null 2>&1 || true
fi
fi
return 0
}
# SLURM sends SIGTERM then SIGKILL after KillWait (30s by default), so the trap
# uses a short stop timeout to fit inside that window.
trap 'echo "[disagg-pd] signal received, tearing down containers"; teardown_containers 10; exit 143' TERM INT
RC=0
if command -v srun >/dev/null 2>&1; then
# classic Slurm: fan out one task/node; ship the fn and call it with the
# srun-supplied per-task SLURM_PROCID (escaped so the remote shell expands it).
srun --label \
--nodelist="$SELECTED_NODELIST_STR" \
--nodes="$NUM_NODES" --ntasks="$NUM_NODES" --ntasks-per-node=1 \
bash -c "$(declare -f preflight_node run_node_container); run_node_container \"\$SLURM_PROCID\"" || RC=$?
else
# Spur: the body already runs once per node; this node's rank is
# SPUR_TASK_OFFSET (fall back to NODE_RANK, then 0 for a single-node run).
run_node_container "${SPUR_TASK_OFFSET:-${NODE_RANK:-0}}" || RC=$?
fi
teardown_containers 30
echo "[disagg-pd] launch finished rc=${RC}; logs in ${LOG_PATH}"
exit "${RC}"
.buildkite/amd-disagg/vllm_disagg.sh
+635
View File
@@ -0,0 +1,635 @@
#!/bin/bash
# =============================================================================
# vLLM disaggregated (P/D) launcher (rank-based / SLURM-native).
# -----------------------------------------------------------------------------
#
# rank 0 .. xP-1 -> prefill (rank 0 is also the orchestrator + proxy)
# rank xP .. xP+yD-1 -> decode
#
# Rank 0 additionally: starts the MoRIIO proxy, health-gates every server,
# runs the post-health workload (RUN_AFTER_HEALTH=bench|accuracy|none), then
# drops a shared-FS completion sentinel so the other ranks shut down and the
# single `srun` returns. Rank 0's exit code is the run's pass/fail.
#
# Explicit per-role invocation is also supported (granular for manual use):
# proxy | prefill | decode | bench | accuracy.
#
# ---- LAUNCH IT (both supported) --------------------------------------------
# 1) Single srun (SLURM provides the rank) — what run_xPyD_disagg.slurm does:
# srun --nodes=$((xP+yD)) --ntasks-per-node=1 \
# bash .../vllm_disagg.sh node
# (IPADDRS must be ordered to match ranks: prefill IPs first, then decode,
# in the SAME order srun assigns PROCID — sort nodes alphabetically.)
#
# 2) Manual, one shell per node (rank from NODE_RANK; share LOG_PATH/FS):
# IPADDRS=ipP0,ipD0 xP=1 yD=1 NODE_RANK=0 bash vllm_disagg.sh node
# IPADDRS=ipP0,ipD0 xP=1 yD=1 NODE_RANK=1 bash vllm_disagg.sh node
#
# 3) Manual, fully granular (per-role):
# bash vllm_disagg.sh proxy
# bash vllm_disagg.sh prefill
# NODE_RANK=1 bash vllm_disagg.sh decode
# bash vllm_disagg.sh bench
#
# Cluster config: cluster.sh (sourced). Model flags: models.yaml.
# Parallelism: WIDE_EP_MODE=0 tp (independent TP servers) | 1 ep (DP+EP groups).
# =============================================================================
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
_PROXY_IP_OVERRIDE="${PROXY_IP:-}"
log() { echo "[vllm_disagg] $*"; }
die() { echo "ERROR: $*" >&2; exit 1; }
usage() {
echo "Usage: $0 [node|proxy|prefill|decode|bench|accuracy] [--wide-ep-mode 0|1]" >&2
echo " (no role => 'node': derive role from \$SLURM_PROCID/\$NODE_RANK)" >&2
exit 1
}
# ----------------------------------------------------------------- arg parsing
# Sets globals: ROLE (default 'node'), _WIDE_EP_MODE_OVERRIDE.
parse_args() {
ROLE="${1:-node}"
shift || true
_WIDE_EP_MODE_OVERRIDE=""
while [[ $# -gt 0 ]]; do
case "$1" in
--wide-ep-mode) _WIDE_EP_MODE_OVERRIDE="${2:-}"; shift 2 ;;
--wide-ep-mode=*) _WIDE_EP_MODE_OVERRIDE="${1#*=}"; shift ;;
-h|--help) usage ;;
*) die "Unknown argument: $1" ;;
esac
done
}
# ----------------------------------------------------------------- load config
# Sources cluster.sh and resolves WIDE_EP_MODE -> PARALLEL_MODE (tp|ep).
load_config() {
CLUSTER_ENV="${CLUSTER_ENV:-${SCRIPT_DIR}/cluster.sh}"
[[ -f "${CLUSTER_ENV}" ]] || die "cluster env file not found: ${CLUSTER_ENV}"
# shellcheck disable=SC1090
source "${CLUSTER_ENV}"
# Topology fanout normally arrives from cluster.sh; restate the defaults so
# the rank arithmetic is defined even when it does not.
: "${xP:=1}" "${yD:=1}"
[[ -n "${_WIDE_EP_MODE_OVERRIDE}" ]] && WIDE_EP_MODE="${_WIDE_EP_MODE_OVERRIDE}"
WIDE_EP_MODE="${WIDE_EP_MODE:-0}"
case "${WIDE_EP_MODE}" in
0) PARALLEL_MODE="tp" ;;
1) PARALLEL_MODE="ep" ;;
*) die "WIDE_EP_MODE must be 0 (tp) or 1 (ep); got '${WIDE_EP_MODE}'" ;;
esac
# node-mode orchestration knobs
RUN_AFTER_HEALTH="${RUN_AFTER_HEALTH:-accuracy}" # bench | accuracy | none
HEALTH_TIMEOUT_S="${HEALTH_TIMEOUT_S:-3600}"
# MoRIIO KV transfer direction. 0 (default)
MORIIO_READ_MODE="${MORIIO_READ_MODE:-0}"
}
# Emits the `read_mode` KV-config fragment (leading comma + newline) when read
# mode is enabled, else nothing (leaving MoRIIO's write-mode default).
moriio_read_mode_kv() {
case "${MORIIO_READ_MODE:-0}" in
1|true|True|TRUE|yes|on) printf ',\n "read_mode": true' ;;
*) : ;;
esac
}
# ----------------------------------------------------------------- topology
# Resolves MODEL_PATH, MODELS_YAML, the IP array, and the master/proxy addresses.
# IP_ARRAY[0..xP-1] = prefill nodes, [xP..xP+yD-1] = decode nodes.
resolve_topology() {
MODEL_PATH="${MODEL_DIR%/}/${MODEL_NAME}"
MODELS_YAML="${MODELS_YAML:-${SCRIPT_DIR}/models.yaml}"
mkdir -p "${LOG_PATH}" 2>/dev/null || true
[[ -z "${IPADDRS}" ]] && IPADDRS="${PREFILL_IP},${DECODE_IP}" # 1P1D fallback
IFS=',' read -ra IP_ARRAY <<< "${IPADDRS}"
PREFILL_MASTER_ADDR="${IP_ARRAY[0]}"
DECODE_MASTER_ADDR="${IP_ARRAY[$xP]:-${IP_ARRAY[-1]}}"
# Proxy is co-located on the prefill master (rank 0) unless the CALLER
# explicitly set PROXY_IP. We use the pre-cluster.sh override (captured at the
# top) — not the live PROXY_IP, which cluster.sh always populates with the
# static PREFILL_IP fallback and would otherwise mask the real master.
PROXY_IP="${_PROXY_IP_OVERRIDE:-${PREFILL_MASTER_ADDR}}"
}
# ============================================================================
# Non-server roles
# ============================================================================
run_proxy() {
[[ -f "${PROXY_SCRIPT}" ]] || die "proxy script not found: ${PROXY_SCRIPT}"
log "proxy: http=:${PROXY_PORT} discovery=:${PROXY_PING_PORT} (${PROXY_SCRIPT})"
exec python3 "${PROXY_SCRIPT}" --port "${PROXY_PORT}"
}
# Start the proxy in the BACKGROUND (node-mode rank 0). Sets PROXY_PID.
start_proxy_bg() {
[[ -f "${PROXY_SCRIPT}" ]] || die "proxy script not found: ${PROXY_SCRIPT}"
local plog="${LOG_PATH}/proxy_$(date +%Y%m%d_%H%M%S).log"
log "proxy(bg): http=:${PROXY_PORT} discovery=:${PROXY_PING_PORT} log=${plog}"
python3 "${PROXY_SCRIPT}" --port "${PROXY_PORT}" >"${plog}" 2>&1 &
PROXY_PID=$!
}
run_bench() {
local base_url="http://127.0.0.1:${GATEWAY_PORT:-${PROXY_PORT}}"
local result_dir="${LOG_PATH}/bench_$(date +%Y%m%d_%H%M%S)"
mkdir -p "${result_dir}"
log "bench -> ${base_url} model=${MODEL_PATH} wide_ep_mode=${WIDE_EP_MODE}(${PARALLEL_MODE})"
log "combinations='${BENCHMARK_COMBINATIONS}' concurrency='${BENCHMARK_CON}'"
local combo isl osl con nump logf
for combo in ${BENCHMARK_COMBINATIONS}; do
isl="${combo%%/*}"; osl="${combo##*/}"
for con in ${BENCHMARK_CON}; do
nump=$(( NUM_PROMPTS_FACTOR * con ))
(( nump < BENCHMARK_MIN_PROMPTS )) && nump="${BENCHMARK_MIN_PROMPTS}"
logf="${result_dir}/bench_${isl}_${osl}_con${con}.log"
log "ISL=${isl} OSL=${osl} CON=${con} num_prompts=${nump} -> ${logf}"
vllm bench serve \
--backend openai \
--base-url "${base_url}" \
--endpoint /v1/completions \
--model "${MODEL_PATH}" \
--dataset-name random \
--random-input-len "${isl}" \
--random-output-len "${osl}" \
--num-prompts "${nump}" \
--max-concurrency "${con}" \
--percentile-metrics ttft,tpot,itl,e2el \
--ignore-eos \
2>&1 | tee "${logf}"
done
done
log "bench complete. Results in ${result_dir}"
}
run_accuracy() {
# lm_eval isn't in the stock vLLM image; install on demand (rank 0 only, and
# only when accuracy is actually selected). Override with ACCURACY_PIP_SPEC.
if ! command -v lm_eval >/dev/null 2>&1; then
log "lm_eval not found — installing ${ACCURACY_PIP_SPEC:-lm_eval[api]}"
python3 -m pip install --no-cache-dir "${ACCURACY_PIP_SPEC:-lm_eval[api]}"
fi
local base_url="http://127.0.0.1:${GATEWAY_PORT:-${PROXY_PORT}}/v1/completions"
local ts; ts="$(date +%Y%m%d_%H%M%S)"
local logf="${LOG_PATH}/accuracy_${MODEL_NAME}_${PARALLEL_MODE}_${ts}.log"
local outdir="${LOG_PATH}/lm_eval_${MODEL_NAME}_${PARALLEL_MODE}_${ts}"
log "accuracy -> ${base_url} model=${MODEL_PATH} tasks=${ACCURACY_TASKS} wide_ep_mode=${WIDE_EP_MODE}(${PARALLEL_MODE})"
log "log=${logf} results=${outdir}"
# --output_path makes lm_eval persist a results_*.json for scraping.
local eval_rc=0
( set -o pipefail
# Model weights are already served (loaded offline in a separate process);
# allow the Hub only here so lm_eval can fetch the eval dataset (e.g. gsm8k).
export HF_HUB_OFFLINE=0 HF_DATASETS_OFFLINE=0
python3 -m lm_eval --model local-completions \
--tasks "${ACCURACY_TASKS}" \
--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}" \
--output_path "${outdir}" \
2>&1 | tee "${logf}"
) || eval_rc=$?
if (( eval_rc != 0 )); then
log "FAIL: lm_eval exited rc=${eval_rc} (no usable score) — Log: ${logf}"
return "${eval_rc}"
fi
# Gate on the score. Find the results JSON lm_eval just wrote, pull the max
# value of ACCURACY_METRIC across filters/tasks, and compare to the threshold.
local results_json=""
results_json="$(find "${outdir}" -name 'results_*.json' -type f 2>/dev/null | sort | tail -n1)"
if [[ -z "${results_json}" ]]; then
log "FAIL: no results_*.json under ${outdir} — cannot evaluate threshold"
return 1
fi
local score=""
score="$(python3 - "${results_json}" "${ACCURACY_METRIC}" <<'PY'
import json, sys
path, metric = sys.argv[1], sys.argv[2]
with open(path) as f:
data = json.load(f)
best = None
for task, metrics in (data.get("results") or {}).items():
for key, val in metrics.items():
if isinstance(val, bool) or not isinstance(val, (int, float)):
continue
if key.split(",")[0] == metric:
best = val if best is None else max(best, val)
print("" if best is None else "%.6f" % best)
PY
)" || score=""
if [[ -z "${score}" ]]; then
log "FAIL: could not parse metric '${ACCURACY_METRIC}' from ${results_json}"
return 1
fi
if awk -v s="${score}" -v t="${ACCURACY_THRESHOLD}" 'BEGIN{exit !(s+0 >= t+0)}'; then
log "PASS: ${ACCURACY_METRIC}=${score} >= threshold=${ACCURACY_THRESHOLD} (tasks=${ACCURACY_TASKS}). Log: ${logf}"
return 0
fi
log "FAIL: ${ACCURACY_METRIC}=${score} < threshold=${ACCURACY_THRESHOLD} (tasks=${ACCURACY_TASKS}). Log: ${logf}"
return 1
}
# ============================================================================
# Server roles (prefill / decode) and their shared helpers
# ============================================================================
# Read model-specific flags from models.yaml (mode + role aware) into globals:
# MODEL_BASE_FLAGS / MODEL_ROLE_FLAGS / MODEL_EXPERIMENTAL_FLAGS -> MODEL_CONFIG
# Also exports the model's env: block (only if not already set; caller env wins).
load_model_flags() {
[[ -f "${MODELS_YAML}" ]] || die "models.yaml not found: ${MODELS_YAML}"
export MODELS_YAML MODEL_NAME PARALLEL_MODE ROLE
eval "$(python3 - <<'PY'
import os, shlex, sys, yaml
path = os.environ["MODELS_YAML"]; name = os.environ["MODEL_NAME"]
mode = os.environ["PARALLEL_MODE"]; role = os.environ["ROLE"]
with open(path, "r", encoding="utf-8") as f:
doc = yaml.safe_load(f) or {}
# Preferred form: top-level `models:` list of {model: <name>, ...} entries.
# Back-compat: a top-level mapping keyed by model name.
entries = doc.get("models") if isinstance(doc, dict) else doc
cfg = None
if isinstance(entries, list):
cfg = next((e for e in entries if isinstance(e, dict) and e.get("model") == name), None)
elif isinstance(doc, dict):
cfg = doc.get(name)
if cfg is None:
print(f'echo "ERROR: model {name} not found in {path}" >&2; exit 1'); sys.exit(0)
role_cfg = cfg.get(role, {}) or {}
def q(v): return shlex.quote(str(v if v is not None else ""))
exports = {
"MODEL_BASE_FLAGS": cfg.get("base_flags", "") or "",
"MODEL_ROLE_FLAGS": role_cfg.get(mode, "") or "",
"MODEL_EXPERIMENTAL_FLAGS": cfg.get("experimental_flags", "") or "",
}
for k, v in exports.items(): print(f"{k}={q(v)}")
# Model-specific env: exported only if not already set (caller env wins).
for k, v in (cfg.get("env", {}) or {}).items():
if k not in os.environ:
print(f"export {k}={q(v)}")
PY
)"
MODEL_CONFIG="${MODEL_BASE_FLAGS} ${MODEL_ROLE_FLAGS} ${MODEL_EXPERIMENTAL_FLAGS}"
}
# Build the TP-mode (WIDE_EP_MODE=0) serve command into the global CMD array.
# Each node is an independent TP server registered to the proxy.
build_tp_cmd() {
local port="${TP_PORT}"
[[ -n "${HOST_IP}" ]] && CMD+=(--host "${HOST_IP}")
CMD+=(--port "${port}" --tensor-parallel-size "${TP_SIZE}")
local _mc; read -ra _mc <<< "${MODEL_CONFIG}"; CMD+=("${_mc[@]}")
local kv_cfg read_mode_kv
read_mode_kv="$(moriio_read_mode_kv)"
kv_cfg=$(cat <<EOF
{
"kv_connector": "MoRIIOConnector",
"kv_role": "${KV_ROLE}",
"kv_connector_extra_config": {
"proxy_ip": "${PROXY_IP}",
"proxy_ping_port": "${PROXY_PING_PORT}",
"http_port": "${port}",
"handshake_port": "${HANDSHAKE_PORT}",
"notify_port": "${NOTIFY_PORT}"${read_mode_kv}
}
}
EOF
)
CMD+=(--kv-transfer-config "${kv_cfg}")
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}"
}
# Build the EP-mode (WIDE_EP_MODE=1) serve command into the global CMD array.
# DP + expert parallel; masters expose the API server + KV transfer, children
# join headless.
build_ep_cmd() {
# prefill -> mori_high_throughput (paired with --enforce-eager in models.yaml)
# decode -> mori_low_latency (paired with CUDA graphs in models.yaml)
local default_backend="mori_low_latency"
[[ "${ROLE}" == "prefill" ]] && default_backend="mori_high_throughput"
local backend="${ALL2ALL_BACKEND:-${default_backend}}"
export VLLM_ALL2ALL_BACKEND="${VLLM_ALL2ALL_BACKEND:-${backend}}"
local port="${SERVE_PORT}"
CMD+=(
-tp 1
--data-parallel-size "${DP_GROUP_SIZE}"
--data-parallel-size-local "${GPUS_PER_NODE}"
--data-parallel-address "${DP_MASTER_ADDR}"
--data-parallel-rpc-port "${RPC_PORT}"
--enable-expert-parallel
--all2all-backend "${backend}"
--port "${port}"
--no-enable-prefix-caching
--distributed-timeout-seconds "${DISTRIBUTED_TIMEOUT_SECONDS}"
)
local _mc; read -ra _mc <<< "${MODEL_CONFIG}"; CMD+=("${_mc[@]}")
if [[ "${IS_MASTER}" == "1" ]]; then
CMD+=(--api-server-count="${GPUS_PER_NODE}")
local kv_cfg read_mode_kv
read_mode_kv="$(moriio_read_mode_kv)"
kv_cfg=$(cat <<EOF
{
"kv_connector": "MoRIIOConnector",
"kv_role": "${KV_ROLE}",
"kv_port": "${KV_PORT}",
"kv_connector_extra_config": {
"proxy_ip": "${PROXY_IP}",
"proxy_port": "${PROXY_PORT}",
"proxy_ping_port": "${PROXY_PING_PORT}",
"http_port": "${port}",
"local_ping_port": "${LOCAL_PING_PORT}",
"handshake_port": "${HANDSHAKE_PORT}",
"notify_port": "${NOTIFY_PORT}"${read_mode_kv}
}
}
EOF
)
CMD+=(--kv-transfer-config "${kv_cfg}")
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}"
else
CMD+=(--data-parallel-start-rank "${DP_START_RANK}" --headless)
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}"
fi
}
# Set the prefill topology globals for the current NODE_RANK (no execution).
configure_prefill() {
: "${NODE_RANK:=0}" # manual default: prefill master
KV_ROLE="kv_producer"
DP_GROUP_SIZE=$(( xP * GPUS_PER_NODE ))
DP_MASTER_ADDR="${PREFILL_MASTER_ADDR}"
DP_START_RANK=$(( NODE_RANK * GPUS_PER_NODE ))
[[ "${NODE_RANK}" -eq 0 ]] && IS_MASTER=1 || IS_MASTER=0
TP_PORT="${PREFILL_PORT}"
}
# Set the decode topology globals for the current NODE_RANK (no execution).
configure_decode() {
: "${NODE_RANK:=${xP}}" # manual default: decode master
KV_ROLE="kv_consumer"
DP_GROUP_SIZE=$(( yD * GPUS_PER_NODE ))
DP_MASTER_ADDR="${DECODE_MASTER_ADDR}"
DP_START_RANK=$(( (NODE_RANK - xP) * GPUS_PER_NODE ))
[[ "${NODE_RANK}" -eq "${xP}" ]] && IS_MASTER=1 || IS_MASTER=0
TP_PORT="${DECODE_PORT}"
}
# Build the serve command (model flags + CMD array + LOGF). No exec.
# Expects role globals set by configure_prefill/configure_decode.
build_server_cmd() {
HOST_IP="${IP_ARRAY[$NODE_RANK]:-}" # own IP for --host binding, if known
load_model_flags
LOGF="${LOG_PATH}/${ROLE}_${PARALLEL_MODE}_node${NODE_RANK}_$(date +%Y%m%d_%H%M%S).log"
CMD=(vllm serve "${MODEL_PATH}")
if [[ "${WIDE_EP_MODE}" == "0" ]]; then build_tp_cmd; else build_ep_cmd; fi
log "model=${MODEL_PATH}"
}
# Foreground server run (explicit prefill/decode roles — original behavior).
run_server_fg() {
build_server_cmd
log "log=${LOGF}"
"${CMD[@]}" 2>&1 | tee "${LOGF}"
}
run_prefill() { configure_prefill; run_server_fg; }
run_decode() { configure_decode; run_server_fg; }
# ============================================================================
# node mode — SLURM-native rank-based self-select + rank-0 orchestration
# ============================================================================
# Endpoints to health-check, one per line as "ip:port", depending on mode.
# tp: every prefill node (:PREFILL_PORT) and every decode node (:DECODE_PORT)
# ep: only the masters expose an API (prefill rank 0, decode rank xP) :SERVE_PORT
health_endpoints() {
local i
if [[ "${WIDE_EP_MODE}" == "0" ]]; then
for (( i=0; i<xP; i++ )); do echo "${IP_ARRAY[$i]}:${PREFILL_PORT}"; done
for (( i=xP; i<xP+yD; i++ )); do echo "${IP_ARRAY[$i]}:${DECODE_PORT}"; done
else
echo "${IP_ARRAY[0]}:${SERVE_PORT}"
echo "${IP_ARRAY[$xP]}:${SERVE_PORT}"
fi
}
# Block until every endpoint's /health is OK, or timeout, or our local server dies.
wait_all_healthy() {
local deadline=$(( $(date +%s) + HEALTH_TIMEOUT_S ))
local eps; mapfile -t eps < <(health_endpoints)
log "health-gate: waiting on ${#eps[@]} endpoint(s): ${eps[*]}"
local ep
for ep in "${eps[@]}"; do
until /usr/bin/curl -sf "http://${ep}/health" >/dev/null 2>&1; do
(( $(date +%s) >= deadline )) && { log "TIMEOUT waiting for ${ep}"; return 1; }
kill -0 "${SERVER_PID}" 2>/dev/null || { log "local server (pid ${SERVER_PID}) exited while waiting"; return 1; }
sleep 10
done
log "healthy: ${ep}"
done
return 0
}
run_workload() {
case "${RUN_AFTER_HEALTH}" in
bench) run_bench ;;
accuracy) run_accuracy ;;
none|"") log "RUN_AFTER_HEALTH=none — skipping workload" ;;
*) die "unknown RUN_AFTER_HEALTH='${RUN_AFTER_HEALTH}'" ;;
esac
}
# rank 0: proxy + health-gate + workload, then write the completion sentinel.
orchestrate_master() {
local sentinel="$1" rc=0
# Front door: the proxy runs in-container (started here); the vllm-router
# runs as a SEPARATE container started by the SLURM job on this (rank-0) node,
# so in that mode we don't start anything here.
PROXY_PID=""
if [[ "${ROUTER_TYPE:-vllm-router}" == "vllm-router" ]]; then
log "ROUTER_TYPE=vllm-router: external router expected on gateway :${GATEWAY_PORT:-${ROUTER_PORT}} (not starting proxy)"
else
start_proxy_bg
fi
if wait_all_healthy; then
# If using the external vllm-router, wait for it to become ready before
# running the workload. The router starts as a separate container on this
# (rank-0) node; give it up to 300s to bind port ROUTER_PORT.
if [[ "${ROUTER_TYPE:-vllm-router}" == "vllm-router" ]]; then
local router_deadline=$(( $(date +%s) + 300 ))
log "waiting for vllm-router on :${ROUTER_PORT} (up to 300s)"
until /usr/bin/curl -sf "http://127.0.0.1:${ROUTER_PORT}/health" >/dev/null 2>&1; do
if (( $(date +%s) >= router_deadline )); then
log "FAIL: vllm-router not ready on :${ROUTER_PORT} after 300s (bring-up)"
rc=1
break
fi
sleep 5
done
(( rc == 0 )) && log "vllm-router healthy on :${ROUTER_PORT}"
fi
if (( rc == 0 )); then
set +e
run_workload
rc=$?
set -e
fi
else
log "FAIL: health-gate did not pass — server bring-up failed"
rc=1
fi
echo "${rc}" > "${sentinel}" 2>/dev/null || true
log "master: workload rc=${rc}; sentinel=${sentinel}; tearing down local proxy+server"
[[ -n "${PROXY_PID:-}" ]] && kill "${PROXY_PID}" 2>/dev/null || true
kill "${SERVER_PID}" 2>/dev/null || true
wait "${SERVER_PID}" 2>/dev/null || true
return "${rc}"
}
# non-master ranks: keep serving until the master signals completion (sentinel)
# or our local server dies unexpectedly.
watch_until_done() {
local sentinel="$1" rc=0
log "rank ${NODE_RANK} serving; waiting for completion sentinel (${sentinel})"
while :; do
[[ -f "${sentinel}" ]] && { log "completion sentinel seen; shutting down"; break; }
if ! kill -0 "${SERVER_PID}" 2>/dev/null; then
log "local server (pid ${SERVER_PID}) exited before completion"; rc=1; break
fi
sleep 10
done
kill "${SERVER_PID}" 2>/dev/null || true
wait "${SERVER_PID}" 2>/dev/null || true
return "${rc}"
}
run_node() {
NODE_RANK="${NODE_RANK:-${SLURM_PROCID:-0}}"
local total=$(( xP + yD ))
(( NODE_RANK >= 0 && NODE_RANK < total )) \
|| die "NODE_RANK=${NODE_RANK} out of range [0,${total}) (xP=${xP} yD=${yD})"
if (( NODE_RANK < xP )); then ROLE="prefill"; configure_prefill
else ROLE="decode"; configure_decode
fi
# Shared-FS completion sentinel (LOG_PATH is shared & per-run/per-job).
local sentinel="${LOG_PATH}/.disagg_done"
log "node mode: rank=${NODE_RANK}/${total} role=${ROLE} master=${IS_MASTER} wide_ep=${WIDE_EP_MODE}(${PARALLEL_MODE})"
log "topology: xP=${xP} yD=${yD} gpus/node=${GPUS_PER_NODE} IPADDRS=${IPADDRS}"
# DRY_RUN: resolve+print the plan (role, serve cmd, endpoints) and exit 0.
# Validates wiring without launching vLLM / the proxy / any workload.
if [[ "${DRY_RUN:-0}" == "1" ]]; then
build_server_cmd
log "DRY_RUN: server cmd: ${CMD[*]}"
if (( NODE_RANK == 0 )); then
local eps; mapfile -t eps < <(health_endpoints)
log "DRY_RUN: health endpoints: ${eps[*]}"
log "DRY_RUN: post-health workload: ${RUN_AFTER_HEALTH}"
log "DRY_RUN: sentinel: ${sentinel}"
fi
log "DRY_RUN complete (rank ${NODE_RANK}) — no processes launched"
exit 0
fi
# aiter and flydsl keep runtime JIT/kernel caches inside the aiter install
# tree (…/aiter/jit and …/aiter/jit/flydsl_cache), which is read-only for the
# non-root (--user) container. AITER_JIT_DIR / FLYDSL_RUNTIME_CACHE_DIR
# redirect those to node-local /tmp (both are sanctioned overrides: aiter only
# points flydsl at its bundled cache when FLYDSL_RUNTIME_CACHE_DIR is unset).
# Seed each once from the image's bundled copy so we keep the prebuilt
# modules/kernels (fast, no NFS, no from-source rebuild / MLIR recompile).
local _aiter_src
_aiter_src="$(python3 -c 'import os,aiter; print(os.path.join(os.path.dirname(aiter.__file__),"jit"))' 2>/dev/null || true)"
if [[ -n "${AITER_JIT_DIR:-}" && ! -e "${AITER_JIT_DIR}/.seeded" ]]; then
if [[ -n "${_aiter_src}" && -d "${_aiter_src}" ]]; then
mkdir -p "${AITER_JIT_DIR}"
cp -a "${_aiter_src}/." "${AITER_JIT_DIR}/" 2>/dev/null || true
touch "${AITER_JIT_DIR}/.seeded" 2>/dev/null || true
log "seeded AITER_JIT_DIR=${AITER_JIT_DIR} from ${_aiter_src}"
else
log "WARN: could not locate aiter jit dir to seed AITER_JIT_DIR=${AITER_JIT_DIR}; aiter may rebuild from source"
fi
fi
if [[ -n "${FLYDSL_RUNTIME_CACHE_DIR:-}" && ! -e "${FLYDSL_RUNTIME_CACHE_DIR}/.seeded" ]]; then
mkdir -p "${FLYDSL_RUNTIME_CACHE_DIR}"
if [[ -n "${_aiter_src}" && -d "${_aiter_src}/flydsl_cache" ]]; then
cp -a "${_aiter_src}/flydsl_cache/." "${FLYDSL_RUNTIME_CACHE_DIR}/" 2>/dev/null || true
log "seeded FLYDSL_RUNTIME_CACHE_DIR=${FLYDSL_RUNTIME_CACHE_DIR} from ${_aiter_src}/flydsl_cache"
else
log "FLYDSL_RUNTIME_CACHE_DIR=${FLYDSL_RUNTIME_CACHE_DIR} set; no bundled flydsl_cache to seed (will JIT-compile at runtime)"
fi
touch "${FLYDSL_RUNTIME_CACHE_DIR}/.seeded" 2>/dev/null || true
fi
# mori JIT-compiles its shmem/all2all kernels to ~/.mori/jit/<arch>_<nic>/ on
# first EP use. HOME here is /workspace, bind-mounted to persistent+shared NFS
# (/data/$USER), so a wrong-arch cache from an earlier run/build survives and
# gets reloaded even after MORI_GPU_ARCHS changes -- an existing cache dir
# wins over the arch setting. Reloading a gfx942 shmem_kernels.hsaco on gfx950
# hardware fails with "device kernel image is invalid" and every EP worker
# dies at init (surfacing async as a torch.tensor HIP error). Guard: on EP
# runs, if the cache has no build for our target arch, purge it so mori
# recompiles for MORI_GPU_ARCHS. A matching arch dir is kept (fast reload).
if [[ "${WIDE_EP_MODE}" == "1" && -n "${MORI_GPU_ARCHS:-}" ]]; then
local _mori_jit="${HOME:-/workspace}/.mori/jit"
if [[ -d "${_mori_jit}" ]] && ! compgen -G "${_mori_jit}/${MORI_GPU_ARCHS}_*" >/dev/null 2>&1; then
log "purging stale mori jit cache ${_mori_jit} (no ${MORI_GPU_ARCHS}_* build present; forcing recompile for ${MORI_GPU_ARCHS})"
rm -rf "${_mori_jit}" 2>/dev/null || true
fi
fi
# Start this node's server in the background (full server log -> LOGF; this
# shell's stdout stays free for orchestration logs the srun/CI captures).
build_server_cmd
"${CMD[@]}" >"${LOGF}" 2>&1 &
SERVER_PID=$!
log "server started pid=${SERVER_PID} log=${LOGF}"
local rc=0
if (( NODE_RANK == 0 )); then
orchestrate_master "${sentinel}" || rc=$?
# Canonical, human-readable end-of-run verdict for the login-node poller
# and CI logs (the .disagg_done sentinel remains the machine authority).
log "VERDICT: $( (( rc == 0 )) && echo PASS || echo FAIL) rc=${rc}"
else
watch_until_done "${sentinel}" || rc=$?
fi
log "node rank=${NODE_RANK} exiting rc=${rc}"
exit "${rc}"
}
# ============================================================================
main() {
parse_args "$@"
load_config
resolve_topology
case "${ROLE}" in
node) run_node ;;
proxy) run_proxy ;;
bench) run_bench ;;
accuracy) run_accuracy ;;
prefill) run_prefill ;;
decode) run_decode ;;
*) usage ;;
esac
}
main "$@"