| 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:-03: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 CI step timeout |
| 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}" |