minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity
Files
b276a9acee8a02159a9a8d839da8a7b4dd898b58
sglang/scripts/ci/slurm/launch_mi355x.sh
T

493 lines
23 KiB
Bash
Executable File
Raw Blame History

#!/usr/bin/env bash
# Launch a 2-node 1P1D disaggregation benchmark on the AMD MI355X `amd-sglang`
# Slurm cluster, then emit per-concurrency result JSONs that
# scripts/ci/slurm/process_result.py aggregates.
#
# salloc's (prefill_workers + decode_workers) nodes -- one server per node --
# and runs the Docker harness: prefill server(s) on the first nodes, decode
# server(s) on the rest, a standalone load balancer on the prefill node, then an
# sglang.bench_serving concurrency sweep over MORI. Default recipe is 1P1D (2
# nodes); see the drive.sh note on reserving 2P2D / 1P3D / 3P1D.
#
# Required environment variables (set by the GitHub Actions workflow):
# MODEL - HuggingFace model id (table label / served model)
# MODEL_PREFIX - short prefix, e.g. dsv4flash
# PRECISION - fp8 / fp4
# ISL, OSL - input / output sequence lengths for the sweep
# CONFIG_FILE - path to the recipe YAML (relative to repo root)
# RESULT_FILENAME - prefix for the emitted result JSONs
# MATRIX_CONFIG_NAME - matrix entry name (used in filenames/tags)
# GITHUB_WORKSPACE - set by GitHub Actions; where result JSONs are written
# Optional:
# MODEL_PATH - local snapshot dir (preferred over downloading MODEL)
# SLURM_PARTITION - default: amd-sglang
# SLURM_NODELIST - optional explicit node pin (else scheduler chooses)
# SLURM_EXCLUDE - optional comma-separated nodes to keep the scheduler
# off (e.g. hosts with a broken RDMA driver)
# RUNNER_NAME - GitHub runner name (a built-in default env var)
# GITHUB_RUN_ID - GitHub Actions run id (a built-in default env var)
# The allocation is named
# mi355x-ci-<RUNNER_NAME>-<GITHUB_RUN_ID>-<config>
# so workflow cleanup can scancel exactly this leg's job
# (full name) or this runner's stale jobs (RUNNER_NAME
# prefix) -- never a blanket `squeue --me`. The run id +
# config make the name unique per matrix leg even if two
# runners happen to share a name.
# SLURM_EXCLUSIVE - request whole nodes (default 1); set 0 to disable
# TIME_LIMIT - salloc time limit, default 02:30:00 (covers server
# load + perf sweep + full GSM8K, under the 180m step cap)
set -euo pipefail
set -x
: "${MODEL_PREFIX:?}"
: "${PRECISION:?}"
: "${ISL:?}"
: "${OSL:?}"
: "${CONFIG_FILE:?}"
: "${RESULT_FILENAME:?}"
: "${MATRIX_CONFIG_NAME:?}"
: "${GITHUB_WORKSPACE:?}"
SLURM_PARTITION="${SLURM_PARTITION:-amd-sglang}"
TIME_LIMIT="${TIME_LIMIT:-02:30:00}"
MODEL_PATH="${MODEL_PATH:-${MODEL:-}}"
if [[ -z "$MODEL_PATH" ]]; then
echo "ERROR: set MODEL_PATH (local snapshot) or MODEL" >&2
exit 1
fi
# Resolve a HuggingFace cache dir (models--org--name) to its live snapshot dir.
# Lets nightly-configs point at the shared cache without hardcoding a snapshot
# hash; if MODEL_PATH is already a concrete snapshot (or plain dir), use as-is.
if [[ -f "$MODEL_PATH/refs/main" && -d "$MODEL_PATH/snapshots" ]]; then
SNAP_HASH="$(cat "$MODEL_PATH/refs/main")"
RESOLVED="$MODEL_PATH/snapshots/$SNAP_HASH"
if [[ -d "$RESOLVED" ]]; then
echo "resolved snapshot: $MODEL_PATH -> $RESOLVED"
MODEL_PATH="$RESOLVED"
else
echo "ERROR: refs/main=$SNAP_HASH but $RESOLVED missing" >&2
exit 1
fi
fi
# ---------------------------------------------------------------------------
# Parse the recipe (runtime + bench + topology) into shell vars.
# ---------------------------------------------------------------------------
# Ensure PyYAML is available to the host python used for parsing.
python3 -c 'import yaml' 2>/dev/null || pip install pyyaml -q 2>/dev/null \
|| pip install --user pyyaml -q 2>/dev/null || true
# Emit KEY=value lines and eval them (robust single-level command substitution;
# avoids a nested read<<EOF/$(<<PY) heredoc that misparses on some shells).
RECIPE_VARS="$(python3 - "$CONFIG_FILE" <<'PY'
import sys, yaml
r = yaml.safe_load(open(sys.argv[1]))
rt = r["runtime"]; b = r["backend"]["sglang_config"]; bn = r["bench"]
res = r.get("resources", {})
def emit(k, v): print(f"{k}={v}")
emit("IMAGE", rt["image"])
emit("ATTN", rt["attention_backend"])
emit("IB", rt["ib_devices"])
emit("PPORT", rt["prefill_port"])
emit("DPORT", rt["decode_port"])
emit("PBOOT", rt["prefill_bootstrap_port"])
emit("DBOOT", rt["decode_bootstrap_port"])
emit("LBPORT", rt["lb_port"])
emit("MEMFRAC", rt["mem_fraction_static"])
emit("PAGE", rt["page_size"])
emit("MAXREQ", rt["max_running_requests"])
emit("CHUNK", rt["chunked_prefill_size"])
emit("SWA", rt["swa_full_tokens_ratio"])
emit("PTP", b["prefill"]["tensor-parallel-size"])
emit("DTP", b["decode"]["tensor-parallel-size"])
emit("PEP", b["prefill"].get("expert-parallel-size", 1))
emit("PDP", b["prefill"].get("data-parallel-size", 1))
m = r.get("mtp", {}) or {}
emit("MTP_ENABLED", 1 if m.get("enabled") else 0)
emit("MTP_STEPS", m.get("num_steps", 3))
emit("MTP_TOPK", m.get("eagle_topk", 1))
emit("MTP_DRAFT", m.get("num_draft_tokens", 4))
# Worker counts double as node counts here: one server per node (TP == GPUs/node).
# 1P1D today; bumping these reserves 2P2D / 1P3D / 3P1D. Multi-node-per-worker
# (TP > GPUs/node, needs --dist-init-addr/--nnodes/--node-rank) is out of scope.
emit("PW", res.get("prefill_workers", 1))
emit("DW", res.get("decode_workers", 1))
emit("CONCS", ",".join(str(c) for c in bn["concurrencies"]))
emit("NPF", bn["num_prompts_factor"])
emit("RRR", bn["random_range_ratio"])
acc = bn.get("accuracy", {}) or {}
emit("ACC_ENABLED", 1 if acc.get("enabled") else 0)
emit("ACC_SHOTS", acc.get("num_shots", 8))
emit("ACC_NQ", acc.get("num_questions", 1319))
emit("ACC_THR", acc.get("threshold", 0.91))
PY
)"
if [[ -z "$RECIPE_VARS" ]]; then
echo "ERROR: failed to parse recipe $CONFIG_FILE (empty output from python3/yaml)" >&2
exit 1
fi
eval "$RECIPE_VARS"
# Optional image override from workflow_dispatch input.
if [[ -n "${IMAGE_OVERRIDE:-}" ]]; then
IMAGE="$IMAGE_OVERRIDE"
fi
echo "recipe: image=$IMAGE attn=$ATTN ib=$IB ptp=$PTP dtp=$DTP concs=$CONCS isl=$ISL osl=$OSL"
# ---------------------------------------------------------------------------
# Shared NFS scratch (visible to login node + compute nodes). Raw bench output
# lands here; the launcher normalizes it into GITHUB_WORKSPACE afterwards.
# ---------------------------------------------------------------------------
WORKDIR="$HOME/.mi355x_ci/${MATRIX_CONFIG_NAME}"
rm -rf "$WORKDIR"; mkdir -p "$WORKDIR"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# Accuracy-gate helpers (written when enabled). Pre-stage the GSM8K test set on
# shared NFS from the login node (which has internet) so the in-container eval
# doesn't depend on compute-node connectivity; fall back to in-container
# download if the pre-fetch fails.
if [[ "$ACC_ENABLED" == "1" ]]; then
GSM8K_URL="https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl"
curl -fsSL "$GSM8K_URL" -o "$WORKDIR/gsm8k_test.jsonl" 2>/dev/null \
&& echo "gsm8k dataset staged at $WORKDIR/gsm8k_test.jsonl" \
|| echo "WARN: gsm8k pre-stage failed; in-container download will be attempted"
cat > "$WORKDIR/check_acc.py" <<'PY'
import sys
acc, thr = float(sys.argv[1]), float(sys.argv[2])
print(f"[gsm8k] accuracy={acc:.3f} threshold={thr}")
sys.exit(0 if acc > thr else 1)
PY
fi
# DSV4 load-bearing env (see test/registered/amd/test_deepseek_v4_flash_fp8.py).
# SGLANG_DSV4_FP4_EXPERTS is precision-driven: true for fp4 weights, false for fp8.
if [[ "$PRECISION" == "fp4" ]]; then
FP4_EXPERTS=true
else
FP4_EXPERTS=false
fi
DSV4_ENV=(
-e SGLANG_DEFAULT_THINKING=1 -e SGLANG_DSV4_REASONING_EFFORT=max
-e SGLANG_OPT_DEEPGEMM_HC_PRENORM=false -e SGLANG_USE_AITER=1
-e SGLANG_USE_ROCM700A=1 -e SGLANG_OPT_USE_FUSED_COMPRESS=true
-e SGLANG_OPT_USE_FUSED_COMPRESS_TRITON=true
-e SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton
-e SGLANG_OPT_FP8_WO_A_GEMM=false -e SGLANG_OPT_USE_JIT_INDEXER_METADATA=false
-e SGLANG_OPT_USE_TOPK_V2=false -e SGLANG_OPT_USE_AITER_INDEXER=true
-e SGLANG_OPT_USE_TILELANG_INDEXER=false -e SGLANG_OPT_USE_TILELANG_MHC_PRE=false
-e SGLANG_OPT_USE_TILELANG_MHC_POST=false -e SGLANG_FP8_PAGED_MQA_LOGITS_TORCH=1
-e SGLANG_OPT_USE_MULTI_STREAM_OVERLAP=false -e SGLANG_ROCM_USE_MULTI_STREAM=false
-e AITER_BF16_FP8_MOE_BOUND=0 -e SGLANG_DSV4_FP4_EXPERTS=$FP4_EXPERTS
)
DSV4_ENV_STR="${DSV4_ENV[*]}"
MORI_ENV="-e MORI_DISABLE_AUTO_XGMI=1 -e NCCL_IB_HCA=ionic -e NCCL_IB_GID_INDEX=1 -e NCCL_CROSS_NIC=1"
# Optional topology / speculative-decode flags driven by the recipe. Base recipes
# (EP1/DP1, no mtp) leave EXTRA_FLAGS empty, preserving prior behavior exactly.
EXTRA_FLAGS=""
(( PDP > 1 )) && EXTRA_FLAGS="$EXTRA_FLAGS --enable-dp-attention --dp-size $PDP"
(( PEP > 1 )) && EXTRA_FLAGS="$EXTRA_FLAGS --ep-size $PEP"
if [[ "$MTP_ENABLED" == "1" ]]; then
EXTRA_FLAGS="$EXTRA_FLAGS --speculative-algorithm EAGLE \
--speculative-num-steps $MTP_STEPS --speculative-eagle-topk $MTP_TOPK \
--speculative-num-draft-tokens $MTP_DRAFT"
fi
echo "extra flags: ${EXTRA_FLAGS:-<none>} (pep=$PEP pdp=$PDP mtp=$MTP_ENABLED)"
COMMON_FLAGS="--trust-remote-code --tp $PTP --disable-radix-cache \
--attention-backend $ATTN --max-running-requests $MAXREQ --page-size $PAGE \
--mem-fraction-static $MEMFRAC --swa-full-tokens-ratio $SWA \
--chunked-prefill-size $CHUNK --disable-shared-experts-fusion \
--tool-call-parser deepseekv4 --reasoning-parser deepseek-v4 \
--disaggregation-transfer-backend mori --disaggregation-ib-device $IB$EXTRA_FLAGS"
DOCKER_COMMON="--rm --network host --ipc host --shm-size 32g --privileged \
--security-opt seccomp=unconfined \
--device /dev/kfd --device /dev/dri --device /dev/infiniband \
-v /it-share:/it-share:ro -v $HOME:/host_home"
# ---------------------------------------------------------------------------
# Write per-role scripts that srun dispatches to each compute node.
# ---------------------------------------------------------------------------
cat > "$WORKDIR/prefill.sh" <<EOF
#!/bin/bash
docker rm -f mi355x_prefill 2>/dev/null || true
docker run $DOCKER_COMMON --name mi355x_prefill \
-e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 $MORI_ENV $DSV4_ENV_STR \
$IMAGE python3 -m sglang.launch_server \
--model-path $MODEL_PATH --host 0.0.0.0 --port $PPORT \
$COMMON_FLAGS --disaggregation-mode prefill --disaggregation-bootstrap-port $PBOOT
EOF
cat > "$WORKDIR/decode.sh" <<EOF
#!/bin/bash
docker rm -f mi355x_decode 2>/dev/null || true
docker run $DOCKER_COMMON --name mi355x_decode \
-e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 $MORI_ENV $DSV4_ENV_STR \
$IMAGE python3 -m sglang.launch_server \
--model-path $MODEL_PATH --host 0.0.0.0 --port $DPORT \
$COMMON_FLAGS --disaggregation-mode decode --disaggregation-bootstrap-port $DBOOT
EOF
# Probe payload + validator (separate files to avoid quoting inside the
# bench.sh `bash -lc '...'` block). One real request exercises the full
# prefill->decode KV handoff before we commit to the whole sweep.
cat > "$WORKDIR/probe.json" <<'JSON'
{"text": "The capital of France is", "sampling_params": {"max_new_tokens": 16, "temperature": 0.0}}
JSON
cat > "$WORKDIR/assert_nonempty.py" <<'PY'
import sys, json
d = json.load(sys.stdin)
t = d.get("text", "") if isinstance(d, dict) else ""
if not (t and t.strip()):
print("[probe] empty/invalid output:", str(d)[:200])
sys.exit(1)
print("[probe] ok:", t[:80].replace("\n", " "))
PY
# Bench script runs on the prefill node; \$PIP/\$DIP injected at srun time.
cat > "$WORKDIR/bench.sh" <<EOF
#!/bin/bash
set -e
PIP=\$1; DIP=\$2
docker rm -f mi355x_bench 2>/dev/null || true
docker run $DOCKER_COMMON --name mi355x_bench \
-e PIP=\$PIP -e DIP=\$DIP \
$IMAGE bash -lc '
export PYTHONPATH=/sgl-workspace/sglang/python:\$PYTHONPATH
echo "[wait] prefill"; for i in \$(seq 1 600); do curl -sf http://\$PIP:$PPORT/health >/dev/null && break; sleep 5; done
echo "[wait] decode"; for i in \$(seq 1 600); do curl -sf http://\$DIP:$DPORT/health >/dev/null && break; sleep 5; done
python3 -m sglang_router.launch_router \
--pd-disaggregation \
--prefill http://\$PIP:$PPORT $PBOOT \
--decode http://\$DIP:$DPORT \
--host 0.0.0.0 --port $LBPORT \
--disable-circuit-breaker &
for i in \$(seq 1 30); do curl -sf http://127.0.0.1:$LBPORT/health >/dev/null && break; sleep 2; done
CIDIR=/host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}
echo "[probe] PD end-to-end check via LB"
curl -sf -X POST http://127.0.0.1:$LBPORT/generate \
-H "content-type: application/json" -d @\$CIDIR/probe.json > \$CIDIR/probe_out.json \
|| { echo "[probe] request failed -- PD path not serving; aborting before sweep"; exit 1; }
python3 \$CIDIR/assert_nonempty.py < \$CIDIR/probe_out.json \
|| { echo "[probe] empty/invalid generation; aborting before sweep"; exit 1; }
# Correctness gate runs BEFORE the perf sweep: if the model is wrong there
# is no point spending ~15min measuring how fast it is wrong, so a failure
# here exits immediately and the sweep never runs.
if [ "$ACC_ENABLED" = "1" ]; then
echo "=== GSM8K accuracy gate (num_questions=$ACC_NQ shots=$ACC_SHOTS) ==="
DP_ARG=""
[ -s \$CIDIR/gsm8k_test.jsonl ] && DP_ARG="--data-path \$CIDIR/gsm8k_test.jsonl"
python3 -m sglang.test.few_shot_gsm8k \
--num-shots $ACC_SHOTS --num-questions $ACC_NQ --parallel $MAXREQ \
--max-new-tokens 512 --host http://127.0.0.1 --port $LBPORT \
\$DP_ARG 2>&1 | tee \$CIDIR/gsm8k.log
ACC=\$(grep -oE "Accuracy: [0-9.]+" \$CIDIR/gsm8k.log | tail -1 | cut -d" " -f2)
[ -n "\$ACC" ] || { echo "[gsm8k] could not parse accuracy from harness output"; exit 1; }
python3 \$CIDIR/check_acc.py "\$ACC" "$ACC_THR" || { echo "[gsm8k] accuracy below threshold -- failing before sweep"; exit 1; }
fi
for C in ${CONCS//,/ }; do
echo "=== concurrency=\$C ==="
OUT=/host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}/raw_conc\${C}.json
rm -f \$OUT
python3 -m sglang.bench_serving --backend sglang \
--host 127.0.0.1 --port $LBPORT --model $MODEL_PATH \
--dataset-name random --random-input-len $ISL --random-output-len $OSL \
--random-range-ratio $RRR --max-concurrency \$C \
--num-prompts \$((C*$NPF)) --warmup-requests \$C \
--output-file \$OUT || true
done
'
EOF
chmod +x "$WORKDIR"/*.sh
# ---------------------------------------------------------------------------
# Orchestration drive (runs inside the salloc allocation on the login node).
# ---------------------------------------------------------------------------
# drive.sh splits the allocation into the first PW nodes (prefill) and the next
# DW nodes (decode), launches one server per node, then benches. For 1P1D
# (PW=DW=1) this is exactly prefill-on-node-A / decode-on-node-B. Larger PW/DW
# reserve 2P2D / 1P3D / 3P1D: all servers come up, but the load balancer and
# bench still target the first prefill + first decode (multi-P/D fan-out is the
# remaining LB piece), so a >1 topology logs an explicit NOTE rather than
# silently producing partial-coverage numbers.
cat > "$WORKDIR/drive.sh" <<'DRIVE'
#!/bin/bash
set -x
WORKDIR="$1"; PW="${2:-1}"; DW="${3:-1}"
mapfile -t NODES < <(scontrol show hostnames "$SLURM_JOB_NODELIST")
PNODES=("${NODES[@]:0:PW}")
DNODES=("${NODES[@]:PW:DW}")
PNODE="${PNODES[0]}"; DNODE="${DNODES[0]}"
PIP=$(getent ahostsv4 "$PNODE" | head -1 | awk '{print $1}')
DIP=$(getent ahostsv4 "$DNODE" | head -1 | awk '{print $1}')
echo "[drive] prefill nodes: ${PNODES[*]} ; decode nodes: ${DNODES[*]}"
echo "[drive] bench targets prefill=$PNODE($PIP) decode=$DNODE($DIP)"
if (( PW > 1 || DW > 1 )); then
echo "[drive] NOTE: router + bench use the first prefill and first decode only;"
echo "[drive] multi-prefill/multi-decode fan-out is not wired yet (LB work)."
fi
# Each server's srun runs here on the login node and returns exactly when its
# compute-node container exits. Wrap it so the return code lands in a marker
# file on shared NFS. The monitor then watches for markers instead of polling
# PIDs -- unambiguous (no zombie/kill -0 guesswork) and it records which role
# died and with what code. (A hung-but-alive server is NOT caught here; that is
# bounded by bench.sh's health-wait timeout.)
rm -f "$WORKDIR"/server_exit_* "$WORKDIR/bench_exit"
for n in "${PNODES[@]}"; do
( srun --overlap -N1 --nodelist="$n" bash "$WORKDIR/prefill.sh" > "$WORKDIR/prefill_$n.log" 2>&1
echo "prefill@$n rc=$?" > "$WORKDIR/server_exit_prefill_$n" ) &
done
for n in "${DNODES[@]}"; do
( srun --overlap -N1 --nodelist="$n" bash "$WORKDIR/decode.sh" > "$WORKDIR/decode_$n.log" 2>&1
echo "decode@$n rc=$?" > "$WORKDIR/server_exit_decode_$n" ) &
done
sleep 5
# Bench in the background with its own marker, so the wait loop is purely file
# based: finish when bench writes its marker, abort if any server marker shows up
# first (a server died before the sweep completed).
( srun --overlap -N1 --nodelist="$PNODE" bash "$WORKDIR/bench.sh" "$PIP" "$DIP" > "$WORKDIR/bench.log" 2>&1
echo $? > "$WORKDIR/bench_exit" ) &
BENCH_BG=$!
# Stream bench output live and poll the markers with xtrace OFF, so the console
# shows clean benchmark/accuracy output instead of a compgen/sleep trace every
# 10s. (Mirrors NVIDIA's launch_gb200.sh, which set +x around its log stream.)
touch "$WORKDIR/bench.log"
tail -n +1 -F "$WORKDIR/bench.log" 2>/dev/null &
TAIL_PID=$!
set +x
RC=0
while [[ ! -f "$WORKDIR/bench_exit" ]]; do
if compgen -G "$WORKDIR/server_exit_*" > /dev/null; then
echo "[drive] ERROR: a server exited early before bench finished:"
cat "$WORKDIR"/server_exit_* || true
kill "$BENCH_BG" 2>/dev/null || true
RC=1
break
fi
sleep 10
done
set -x
kill "$TAIL_PID" 2>/dev/null || true
[[ "$RC" -eq 0 ]] && RC=$(cat "$WORKDIR/bench_exit" 2>/dev/null || echo 1)
echo "[drive] bench finished (rc=$RC), tearing down"
for n in "${PNODES[@]}"; do srun --overlap -N1 --nodelist="$n" docker kill mi355x_prefill >/dev/null 2>&1 || true; done
for n in "${DNODES[@]}"; do srun --overlap -N1 --nodelist="$n" docker kill mi355x_decode >/dev/null 2>&1 || true; done
exit "$RC"
DRIVE
chmod +x "$WORKDIR/drive.sh"
NODELIST_ARG=()
[[ -n "${SLURM_NODELIST:-}" ]] && NODELIST_ARG=(--nodelist="$SLURM_NODELIST")
# Request whole nodes so a co-scheduled job can't share a node and skew the
# benchmark numbers. Toggle off with SLURM_EXCLUSIVE=0 on partitions that
# disallow --exclusive.
EXCLUSIVE_ARG=()
[[ "${SLURM_EXCLUSIVE:-1}" == "1" ]] && EXCLUSIVE_ARG=(--exclusive)
# Keep the scheduler off known-bad nodes (e.g. a host whose ionic RDMA driver
# ABI mismatches the container, where MORI reports "no active RDMA device" and
# the disagg server dies on init). Comma-separated node list.
EXCLUDE_ARG=()
[[ -n "${SLURM_EXCLUDE:-}" ]] && EXCLUDE_ARG=(--exclude="$SLURM_EXCLUDE")
# One node per prefill/decode worker (TP == GPUs/node). 1P1D -> 2 nodes.
TOTAL_NODES=$((PW + DW))
# Name the allocation <RUNNER_NAME>-<GITHUB_RUN_ID>-<config> so the workflow's
# cleanup steps can scancel precisely instead of a blanket `squeue --me` that
# would kill a concurrent matrix leg. RUNNER_NAME alone is not assumed unique;
# GITHUB_RUN_ID + config make the name unique per matrix leg regardless.
JOB_NAME="mi355x-ci-${RUNNER_NAME:-norunner}-${GITHUB_RUN_ID:-0}-${MATRIX_CONFIG_NAME}"
set +e
salloc -p "$SLURM_PARTITION" -N"$TOTAL_NODES" "${NODELIST_ARG[@]}" "${EXCLUDE_ARG[@]}" "${EXCLUSIVE_ARG[@]}" \
--job-name "$JOB_NAME" -t "$TIME_LIMIT" \
bash "$WORKDIR/drive.sh" "$WORKDIR" "$PW" "$DW"
SALLOC_RC=$?
set -e
# bench output already streamed live from drive.sh (tail -F). drive.sh exits
# non-zero when a server died or bench failed; on failure dump bench.log + the
# server logs (the actual root cause). We still fall through to normalize
# whatever raw results the completed concurrencies produced -- partial perf data
# is worth uploading -- and propagate the failure via the exit code at the end.
if [[ "$SALLOC_RC" -ne 0 ]]; then
echo "ERROR: allocation/bench failed (rc=$SALLOC_RC); bench + server logs:" >&2
echo "--- bench.log (tail) ---"; tail -40 "$WORKDIR/bench.log" 2>/dev/null || true
for f in "$WORKDIR"/prefill_*.log "$WORKDIR"/decode_*.log; do
[[ -f "$f" ]] && { echo "--- $f (tail) ---"; tail -30 "$f"; }
done
fi
# Surface the GSM8K accuracy in the job summary -- it scrolls past in the live
# log, and the perf table (collect-results/summarize.py) doesn't include it.
if [[ "$ACC_ENABLED" == "1" && -n "${GITHUB_STEP_SUMMARY:-}" ]]; then
ACC_LINE=$(grep -aoE "Accuracy: [0-9.]+" "$WORKDIR/bench.log" 2>/dev/null | tail -1 || true)
{
echo "### GSM8K accuracy gate — ${MATRIX_CONFIG_NAME}"
echo '```'
echo "${ACC_LINE:-Accuracy: <not found in bench.log>} (threshold > ${ACC_THR})"
echo '```'
} >> "$GITHUB_STEP_SUMMARY"
fi
# ---------------------------------------------------------------------------
# Normalize raw bench_serving output -> process_result.py schema.
#
# bench_serving and process_result.py disagree on field names, so we remap the
# last JSON line of each raw file. If bench_serving ever renames an output
# field, the KeyError raised here (rather than a silently wrong table) is the
# signal to update this mapping. Field-by-field:
#
# bench_serving key -> process_result.py key (purpose)
# -------------------------- ------------------------ -------------------------
# max_concurrency -> max_concurrency (sweep point; falls back to $C)
# total_throughput -> total_token_throughput (in+out tok/s, tput_per_gpu)
# output_throughput -> output_throughput (out tok/s, output_tput_per_gpu)
# median_ttft_ms -> median_ttft_ms (TTFT; /1000 -> s)
# median_tpot_ms -> median_tpot_ms (TPOT; -> interactivity)
# median_e2e_latency_ms -> median_e2el_ms (E2E latency; /1000 -> s)
# (none; injected here) -> model_id (served model, from $MODEL_PATH)
# ---------------------------------------------------------------------------
TOTAL_GPUS=$((PTP + DTP))
PROCESSED=0
for C in ${CONCS//,/ }; do
RAW="$WORKDIR/raw_conc${C}.json"
[[ -f "$RAW" ]] || { echo "WARN: missing $RAW"; continue; }
DEST="$GITHUB_WORKSPACE/${RESULT_FILENAME}_${MATRIX_CONFIG_NAME}_conc${C}_gpus_${TOTAL_GPUS}_ctx_${PTP}_gen_${DTP}.json"
MODEL_ID="$MODEL_PATH" python3 - "$RAW" "$DEST" "$C" <<'PY'
import json, os, sys
raw_path, dest, conc = sys.argv[1], sys.argv[2], int(sys.argv[3])
line = [l for l in open(raw_path).read().splitlines() if l.strip()][-1]
r = json.loads(line)
norm = {
"max_concurrency": r.get("max_concurrency") or conc,
"model_id": os.environ["MODEL_ID"],
"total_token_throughput": r["total_throughput"],
"output_throughput": r["output_throughput"],
"median_ttft_ms": r["median_ttft_ms"],
"median_tpot_ms": r["median_tpot_ms"],
"median_e2el_ms": r["median_e2e_latency_ms"],
}
json.dump(norm, open(dest, "w"), indent=2)
print("normalized ->", dest)
PY
PROCESSED=$((PROCESSED + 1))
done
# Propagate a benchmark/allocation failure even though we emitted partial
# results above (the workflow uploads them with `always()`).
if [[ "$SALLOC_RC" -ne 0 ]]; then
echo "ERROR: benchmark failed (rc=$SALLOC_RC); emitted $PROCESSED partial result file(s)." >&2
exit "$SALLOC_RC"
fi
if [[ "$PROCESSED" -eq 0 ]]; then
echo "ERROR: no result files produced" >&2
exit 1
fi
echo "Done. $PROCESSED result file(s) in $GITHUB_WORKSPACE."
Reference in New Issue View Git Blame Copy Permalink