#!/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)
#   SGLANG_USE_CHECKOUT_RUNTIME
#                      - default 1. Reinstall this workflow checkout's Python
#                        sglang package inside each runtime container, and the
#                        checkout sglang-router package inside the bench
#                        container, before launching servers/bench. Set 0 to
#                        use the image's baked-in packages.
#   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:-}}"

# Scheduler profile. `slurm` (default) is the amd-sglang MI355X cluster and is
# the original code path, unchanged. `spur` is the TW PIT cluster, which runs
# Spur -- a Slurm-compatible scheduler that differs in three ways that matter
# here: salloc takes no trailing command, `scontrol show hostnames` does not
# exist (its nodelist env var is already comma-expanded), and `srun --overlap`
# requires an explicit --jobid instead of inheriting one from the environment.
CLUSTER="${CLUSTER:-slurm}"
case "$CLUSTER" in
    slurm|spur) ;;
    *) echo "ERROR: CLUSTER must be 'slurm' or 'spur', got '$CLUSTER'" >&2; exit 1 ;;
esac

# Spur reads a two-field time limit as HH:MM where Slurm reads MM:SS, so "30:00"
# silently becomes 30 HOURS. Require the unambiguous HH:MM:SS / D-HH:MM:SS form
# there rather than let a typo book a node for a day.
if [[ "$CLUSTER" == "spur" && ! "$TIME_LIMIT" =~ ^([0-9]+-)?[0-9]+:[0-9]{2}:[0-9]{2}$ ]]; then
    echo "ERROR: on spur TIME_LIMIT must be HH:MM:SS or D-HH:MM:SS (got '$TIME_LIMIT');" >&2
    echo "       a two-field value like '30:00' is parsed as 30 hours, not 30 minutes." >&2
    exit 1
fi

# Optional account / QoS. Empty on the mi355x cluster (which gates on partition
# alone); required on spur, where every job must name an account and its QoS.
# Both flags are spelled the same in real Slurm, so passing them is portable.
SLURM_ACCOUNT="${SLURM_ACCOUNT:-}"
SLURM_QOS="${SLURM_QOS:-}"

# Relocate the shared HuggingFace cache root. nightly-configs.yaml pins
# model_path under /it-share/model_coverage (the mi355x cluster's NFS); a
# cluster that mirrors the same models elsewhere sets MODEL_ROOT and every
# recipe/config path is rewritten, so no recipe or config file needs editing.
MODEL_ROOT="${MODEL_ROOT:-}"
# Root used to RESOLVE the snapshot hash; defaults to MODEL_ROOT. Set it when
# MODEL_ROOT is node-local and therefore unreadable from the driver node.
MODEL_RESOLVE_ROOT="${MODEL_RESOLVE_ROOT:-$MODEL_ROOT}"
MODEL_ROOT_FROM="${MODEL_ROOT_FROM:-/it-share/model_coverage}"
relocate_model_root() {
    local p="$1"
    if [[ -n "$MODEL_ROOT" && "$p" == "$MODEL_ROOT_FROM"/* ]]; then
        printf '%s\n' "$MODEL_ROOT/${p#"$MODEL_ROOT_FROM"/}"
    else
        printf '%s\n' "$p"
    fi
}
SGLANG_USE_CHECKOUT_RUNTIME="${SGLANG_USE_CHECKOUT_RUNTIME:-1}"
case "${SGLANG_USE_CHECKOUT_RUNTIME,,}" in
    0|false|no|off) SGLANG_USE_CHECKOUT_RUNTIME=0 ;;
    *) SGLANG_USE_CHECKOUT_RUNTIME=1 ;;
esac

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 / recipes point at the shared cache without hardcoding a
# snapshot hash; a concrete snapshot dir (or plain dir) is returned unchanged.
# Used for both MODEL_PATH and an optional speculative draft model path.
resolve_snapshot() {
    local p="$1"
    if [[ -f "$p/refs/main" && -d "$p/snapshots" ]]; then
        local hash resolved
        hash="$(cat "$p/refs/main")"
        resolved="$p/snapshots/$hash"
        if [[ -n "$hash" && -d "$resolved" ]]; then
            echo "resolved snapshot: $p -> $resolved" >&2
            echo "$resolved"
            return 0
        fi
        echo "ERROR: refs/main=$hash but $resolved missing" >&2
        return 1
    fi
    echo "$p"
}
# Resolve the snapshot against the shared cache first, then relocate. MODEL_ROOT
# may name a node-local copy that does not exist on the node running this
# script; resolve_snapshot would then find no refs/main and pass the bare cache
# dir to the container. Three steps, since the resolve root and the destination
# root can differ: rewrite onto a readable root, resolve refs/main, then rewrite
# onto the root the container will see.
MODEL_PATH="$(MODEL_ROOT="$MODEL_RESOLVE_ROOT" relocate_model_root "$MODEL_PATH")"
MODEL_PATH="$(resolve_snapshot "$MODEL_PATH")" || exit 1
MODEL_PATH="$(MODEL_ROOT_FROM="$MODEL_RESOLVE_ROOT" relocate_model_root "$MODEL_PATH")"

# ---------------------------------------------------------------------------
# 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, shlex
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}")
# Quoted emit for values that may contain spaces (per-role wide-EP flag/env
# strings). eval treats the RHS as a single shell word. Existing emits stay
# space-free so the EP<=8 recipes are byte-identical.
def emitq(k, v): print(f"{k}={shlex.quote(str(v))}")
emit("IMAGE", rt["image"])
# Attention backend: single (`attention_backend`) or split
# (`prefill_attention_backend`/`decode_attention_backend`). Empty when absent so
# the flag is dropped for a model that omits it.
emit("ATTN", rt.get("attention_backend", ""))
emit("PATTN", rt.get("prefill_attention_backend", ""))
emit("DATTN", rt.get("decode_attention_backend", ""))
emit("IB", rt["ib_devices"])
# Wide-EP (EP > GPUs/node) options; empty for the single-node EP<=8 recipes so
# the flags/env below are dropped and their argv stays byte-identical.
#   moe_a2a_backend    -> --moe-a2a-backend <x> --deepep-mode normal (MoE all-to-all)
#   dist_socket_ifname -> NCCL_/GLOO_SOCKET_IFNAME for cross-node torch-dist init
emit("A2A", rt.get("moe_a2a_backend", ""))
emit("DIST_SOCK", rt.get("dist_socket_ifname", ""))
# KV transfer backend for the P->D handoff. Defaults to mori (the pre-wide
# hardcoded value) so EP<=8 recipes stay byte-identical; wide-EP spur recipes
# set mooncake, which is the validated cross-node KV path there.
emit("XFER", rt.get("kv_transfer_backend", "mori"))
# SGLANG_USE_ROCM700A toggles the ROCm-7.0.0-alpha codepath. Default 1 (the
# pre-wide hardcoded value) keeps EP<=8 recipes byte-identical; the validated
# wide-EP run on the rocm720 0715 image needs 0, set via runtime.rocm700a.
emit("ROCM700A", rt.get("rocm700a", 1))
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("MAXTOK", rt.get("max_total_tokens", ""))
emit("CHUNK", rt["chunked_prefill_size"])
# swa is DSV4-specific; emit empty when a model omits it so the flag is dropped.
emit("SWA", rt.get("swa_full_tokens_ratio", ""))
# 1 when the recipe carries a `model:` block (env + server_args written to
# model_flags.sh); 0 for the DSV4 recipes, which keep the hardcoded DSV4 path.
emit("HAS_MODEL", 1 if r.get("model") else 0)
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))
# Decode-side EP/DP: default to the prefill values so a recipe that omits them
# (every EP<=8 recipe today, where both roles are identical) is byte-identical.
# Oren's wide-EP recipes set decode EP/DP=16 while prefill stays EP8.
emit("DEP", b["decode"].get("expert-parallel-size", b["prefill"].get("expert-parallel-size", 1)))
emit("DDP", b["decode"].get("data-parallel-size", b["prefill"].get("data-parallel-size", 1)))
# Wide-EP per-role overrides (optional `runtime.wide_ep` block). Absent for
# EP<=8 recipes -> all defaults collapse to the existing single-value knobs, so
# the generated prefill/decode argv is unchanged. Present only in the asymmetric
# narrow-prefill/wide-decode (Oren) EP16 recipes.
we = rt.get("wide_ep", {}) or {}
emit("KVDTYPE", we.get("kv_cache_dtype", ""))
emit("PMEMFRAC", we.get("prefill_mem_fraction_static", rt["mem_fraction_static"]))
emit("DMEMFRAC", we.get("decode_mem_fraction_static", rt["mem_fraction_static"]))
emit("PCHUNK", we.get("prefill_chunked_prefill_size", rt["chunked_prefill_size"]))
emit("PMAXREQ", we.get("prefill_max_running_requests", rt["max_running_requests"]))
emit("DMAXREQ", we.get("decode_max_running_requests", rt["max_running_requests"]))
emitq("PEXTRA", we.get("prefill_extra_flags", ""))
emitq("DEXTRA", we.get("decode_extra_flags", ""))
emitq("WECOMMON", we.get("common_extra_flags", ""))
def render_env(d):
    return " ".join(f"-e {k}={v}" for k, v in (d or {}).items())
emitq("PENV", render_env(we.get("prefill_extra_env")))
emitq("DENV", render_env(we.get("decode_extra_env")))
m = r.get("mtp", {}) or {}
emit("MTP_ENABLED", 1 if m.get("enabled") else 0)
emit("MTP_ALGO", m.get("algorithm", "EAGLE"))
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))
# External draft checkpoint (EAGLE3 etc.); empty for DSV4's built-in EAGLE head.
emit("MTP_DRAFT_PATH", m.get("draft_model_path", ""))
# 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

# Nodes per engine: an engine whose TP exceeds one node's GPU count spans
# ceil(TP/GPUS_PER_NODE) nodes and needs torch-dist multi-node init. EP<=8
# recipes give 1 (single node) so all downstream multi-node logic no-ops.
GPUS_PER_NODE="${GPUS_PER_NODE:-8}"
PN_PER=$(( (PTP + GPUS_PER_NODE - 1) / GPUS_PER_NODE ))
DN_PER=$(( (DTP + GPUS_PER_NODE - 1) / GPUS_PER_NODE ))
# torch-dist rendezvous port for a multi-node engine. Default 29500 (torch's
# conventional MASTER_PORT); avoids :5000, which a node-local daemon holds on
# some clusters (spur). Overridable per environment.
DIST_PORT="${DIST_PORT:-29500}"
# A role may have multiple single-node engines (PW>1, each PN_PER=1: the router
# fans out over them) OR one multi-node wide engine (PN_PER>1, workers==1). What
# is NOT wired is *multiple copies of a multi-node engine* (a wide engine with
# workers>1), because the drive split assumes one contiguous node block per wide
# engine and the router only knows one endpoint per wide engine.
if (( (PN_PER > 1 && PW > 1) || (DN_PER > 1 && DW > 1) )); then
    echo "ERROR: a wide engine (nodes/engine>1) cannot have >1 worker of that role (got PW=$PW PN_PER=$PN_PER DW=$DW DN_PER=$DN_PER)" >&2
    exit 1
fi
echo "recipe: image=$IMAGE attn=${ATTN:-$PATTN/$DATTN} ib=$IB ptp=$PTP dtp=$DTP pn_per=$PN_PER dn_per=$DN_PER a2a=${A2A:-<none>} 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}"
if [[ "$CLUSTER" == "spur" ]]; then
  # `rm -rf "$WORKDIR"` is not safe to gate a leg on over NFS: if anything still
  # holds a descriptor inside (an orphaned `tail -F` on bench.log), the server
  # renames the file to .nfsXXXX rather than unlinking it, so rm sees a
  # non-empty directory and exits non-zero. Clear the contents instead, retry,
  # and fall back to moving the directory aside. Housekeeping must not fail a
  # leg.
  for _try in 1 2 3; do
    rm -rf "$WORKDIR"/* "$WORKDIR"/.[!.]* 2>/dev/null || true
    # Anything left is a silent .nfsXXXX handle; give the server a moment.
    [[ -z "$(ls -A "$WORKDIR" 2>/dev/null)" ]] && break
    sleep 5
  done
  mkdir -p "$WORKDIR"
  if [[ -n "$(ls -A "$WORKDIR" 2>/dev/null)" ]]; then
    _stale="$WORKDIR.stale.$$"
    if mv "$WORKDIR" "$_stale" 2>/dev/null; then
      echo "WARN: $WORKDIR would not clear (open handles); moved to $_stale" >&2
    else
      WORKDIR="$WORKDIR.$$"
      echo "WARN: could not clear or move workdir; using fresh $WORKDIR" >&2
    fi
    mkdir -p "$WORKDIR"
  fi
else
rm -rf "$WORKDIR"; mkdir -p "$WORKDIR"
fi
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Stage the workflow checkout on shared NFS so Slurm compute-node containers can
# reinstall the same code SHA the workflow checked out. The container gets a
# read-only mount and copies it to /tmp before mutating pyproject.toml.
CHECKOUT_DOCKER_ARGS="-e SGLANG_USE_CHECKOUT_RUNTIME=$SGLANG_USE_CHECKOUT_RUNTIME"
if [[ "$SGLANG_USE_CHECKOUT_RUNTIME" == "1" ]]; then
    CHECKOUT_STAGE="$WORKDIR/checkout"
    CHECKOUT_SHA="$(git -C "$GITHUB_WORKSPACE" rev-parse HEAD)"
    echo "Staging checkout runtime: sha=$CHECKOUT_SHA -> $CHECKOUT_STAGE"
    rm -rf "$CHECKOUT_STAGE"
    mkdir -p "$CHECKOUT_STAGE"
    # See install_checkout_sglang.sh: tar exit 1 is a warning ("file changed as
    # we read it", raised spuriously by NFS attribute caching), 2+ is an error.
    set +e
    tar --exclude='__pycache__' --exclude='*.pyc' --exclude='.git/config' \
        -C "$GITHUB_WORKSPACE" -cf - . | tar -C "$CHECKOUT_STAGE" -xf -
    STAGE_TAR_RC=("${PIPESTATUS[@]}")
    set -e
    if (( STAGE_TAR_RC[0] > 1 || STAGE_TAR_RC[1] != 0 )); then
        echo "ERROR: staging tar failed (create=${STAGE_TAR_RC[0]} extract=${STAGE_TAR_RC[1]})" >&2
        exit 1
    fi
    # Tolerating create-rc=1 above means a warning cannot abort the copy -- but a
    # TRUNCATED copy must still not reach a container, where it only surfaces
    # ~15 minutes later as "invalid checkout mount". Assert the same file the
    # in-container installer checks, right where the copy was made.
    if [[ ! -f "$CHECKOUT_STAGE/python/sglang/version.py" ]]; then
        echo "ERROR: staged checkout is incomplete (no python/sglang/version.py in $CHECKOUT_STAGE);" >&2
        echo "       tar rc were create=${STAGE_TAR_RC[0]} extract=${STAGE_TAR_RC[1]}" >&2
        exit 1
    fi
    CHECKOUT_DOCKER_ARGS="$CHECKOUT_DOCKER_ARGS -e SGLANG_CHECKOUT_SHA=$CHECKOUT_SHA -v $CHECKOUT_STAGE:/sglang-checkout:ro"
else
    echo "SGLANG_USE_CHECKOUT_RUNTIME=0; using sglang package baked into image."
fi

# 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=$ROCM700A
  -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[*]}"
# A recipe carrying a `model:` block supplies its OWN docker env (below), so the
# DSV4 env must not leak into it; the DSV4 recipes keep the string above.
[[ "$HAS_MODEL" == "1" ]] && DSV4_ENV_STR=""
# NCCL_IB_HCA must name real devices. "ionic" is the mi355x cluster's spelling;
# on pit2 the HCAs are rdma0..7 (the recipe's $IB) and "ionic" matches nothing,
# which fails TP group init outright with
#   RuntimeError: NCCL error: remote process exited or there was a network error
# from ncclCommInitRank. The wide-EP block below already overrides this for
# EP>8; EP<=8 recipes need it too.
MORI_NCCL_HCA=ionic
[[ "$CLUSTER" == "spur" ]] && MORI_NCCL_HCA="$IB"
MORI_ENV="-e MORI_DISABLE_AUTO_XGMI=1 -e NCCL_IB_HCA=$MORI_NCCL_HCA -e NCCL_IB_GID_INDEX=1 -e NCCL_CROSS_NIC=1"
# Whether NCCL may use IB depends on the role, not on the recipe. A single-node
# engine does all its NCCL traffic intra-node (cross-node KV goes over mori),
# and on pit2 letting it reach for IB hangs ncclCommInitRank. Gating on the
# whole recipe is wrong for 2P1D-EP16 (PN_PER=1, DN_PER=2), whose prefill
# engines are single-node TP8 but whose recipe-wide test left IB on for both
# roles.
#
# Both cases still need a socket interface: NCCL uses one for its bootstrap
# handshake even on an intra-node data path, and pit2 exposes eight tw-eth* /31
# links that auto-selection picks and hangs on. Wide recipes set it in the block
# below; pin it here for any role that block does not cover.
SPUR_PNCCL=""
SPUR_DNCCL=""
if [[ "$CLUSTER" == "spur" ]]; then
    SPUR_SOCK="${DIST_SOCK:-${ADDR_NIC_OVERRIDE:-eno0}}"
    if (( PN_PER == 1 )); then
        SPUR_PNCCL=" -e NCCL_IB_DISABLE=1 -e NCCL_SOCKET_IFNAME=$SPUR_SOCK -e GLOO_SOCKET_IFNAME=$SPUR_SOCK"
    fi
    if (( DN_PER == 1 )); then
        SPUR_DNCCL=" -e NCCL_IB_DISABLE=1 -e NCCL_SOCKET_IFNAME=$SPUR_SOCK -e GLOO_SOCKET_IFNAME=$SPUR_SOCK"
    fi
fi
# Wide-EP (engine spans >1 node) adds mori all-to-all MoE tuning + the cross-node
# torch-dist socket NIC. Gated on nodes-per-engine>1 so EP<=8 recipes are untouched.
if (( PN_PER > 1 || DN_PER > 1 )); then
    # TC=104 (SL=3) = the ionic lossless RoCE queue (DSCP26/pri3); TC=96 is the
    # lossy pri0 queue (~1% BW) and wedges cross-node a2a under load. bf16
    # dispatch/combine matches the validated wide-EP run (job 13196).
    # The base MORI_ENV sets NCCL_IB_HCA=ionic (a spur-ism); on this fabric the
    # IB device names are the recipe's $IB (rdma0..7), and a wide engine's
    # cross-node TP/attention collectives ride NCCL, so point NCCL at the real
    # HCAs. Docker last-wins => this overrides the base value for wide recipes.
    MORI_ENV="$MORI_ENV \
-e NCCL_IB_HCA=$IB \
-e MORI_IB_GID_INDEX=1 \
-e SGLANG_MORI_DISPATCH_DTYPE=bf16 -e SGLANG_MORI_COMBINE_DTYPE=bf16 \
-e SGLANG_MORI_QP_PER_TRANSFER=4 -e SGLANG_MORI_NUM_WORKERS=4 \
-e MORI_IO_SQ_BACKOFF_TIMEOUT_US=50000 -e MORI_IO_QP_MAX_SEND_WR=16384 \
-e MORI_IO_QP_MAX_CQE=32768 -e MORI_IO_QP_MAX_SGE=4 \
-e MORI_SHMEM_MODE=ISOLATION -e MORI_EP_LAUNCH_CONFIG_MODE=AUTO -e MORI_APP_LOG_LEVEL=INFO \
-e MORI_RDMA_SL=3 -e MORI_RDMA_TC=104 -e MORI_IO_SL=3 -e MORI_IO_TC=104 -e MORI_IO_TC_DISABLE=0 \
-e SGLANG_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD=4096 \
-e SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=3600 -e SGLANG_DISAGGREGATION_WAITING_TIMEOUT=3600 \
-e SGLANG_DISAGGREGATION_NUM_PRE_ALLOCATE_REQS=32 -e SGLANG_EAGER_INPUT_NO_COPY=true \
-e MORI_BOOTSTRAP_TIMEOUT=300"
    # The overlap plan stream is a wide-EP perf knob, but it makes
    # EAGLEWorkerV2.verify call attn_backend.update_verify_buffers_to_fill_after_draft,
    # which the DSV4 MLA backend does not implement (NotImplementedError). Only
    # enable it for non-MTP wide-EP; MTP wide-EP uses the plain verify path (as EP8 does).
    if [[ "$MTP_ENABLED" != "1" ]]; then
        MORI_ENV="$MORI_ENV -e SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1"
    fi
    if [[ -n "$DIST_SOCK" ]]; then
        MORI_ENV="$MORI_ENV -e GLOO_SOCKET_IFNAME=$DIST_SOCK -e NCCL_SOCKET_IFNAME=$DIST_SOCK -e MORI_SOCKET_IFNAME=$DIST_SOCK"
    fi
fi

# Model-specific docker `-e` env + sglang server args from the recipe's optional
# `model:` block, written as bash arrays to model_flags.sh (sourced by
# prefill.sh/decode.sh). DSV4 recipes have no `model:` block -> empty arrays, so
# their generated docker argv is unchanged. Each server arg + its value MUST be a
# separate YAML list item so shlex.quote keeps "--foo" and "bar" as two tokens.
python3 - "$CONFIG_FILE" "$WORKDIR/model_flags.sh" <<'PY'
import shlex, sys, yaml
r = yaml.safe_load(open(sys.argv[1]))
model = r.get("model", {}) or {}
env = model.get("env", {}) or {}
server_args = model.get("server_args", []) or []
# YAML true/false parse to Python bool; render lowercase so env values stay
# byte-identical to shell (`=false`, not `=False`) -- SGLang parsing is
# case-sensitive for some of these.
def fmt(v):
    if isinstance(v, bool):
        return "true" if v else "false"
    return str(v)
env_args = []
for k, v in env.items():
    env_args += ["-e", f"{k}={fmt(v)}"]
def q(items):
    return " ".join(shlex.quote(fmt(x)) for x in items)
with open(sys.argv[2], "w") as f:
    f.write(f"MODEL_ENV_ARGS=({q(env_args)})\n")
    f.write(f"MODEL_SERVER_ARGS=({q(server_args)})\n")
PY

# Resolve host ionic userspace mounts on each compute node. This keeps server,
# benchmark, and wide-EP containers compatible with the host RDMA ABI.
cat > "$WORKDIR/ionic_mounts.sh" <<'IONIC_EOF'
IONIC_MOUNTS=()
_ionic_provider="/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so"
_ionic_soname="/usr/lib/x86_64-linux-gnu/libionic.so.1"
if [ ! -e "$_ionic_provider" ]; then
    _ionic_provider=$(find /usr/lib/x86_64-linux-gnu -maxdepth 1 \
        -name "libionic.so.*" -print -quit 2>/dev/null)
fi
if [ -n "$_ionic_provider" ] && [ -e "$_ionic_provider" ]; then
    _ionic_real=$(readlink -f "$_ionic_provider" 2>/dev/null)
    _ionic_soname_real=$(readlink -f "$_ionic_soname" 2>/dev/null)
    [ -f "$_ionic_real" ] && IONIC_MOUNTS+=( -v "$_ionic_real:$_ionic_real:ro" )
    [ -f "$_ionic_soname_real" ] && \
        IONIC_MOUNTS+=( -v "$_ionic_soname_real:$_ionic_soname:ro" )
    [ -d /usr/lib/x86_64-linux-gnu/libibverbs ] && \
        IONIC_MOUNTS+=( -v /usr/lib/x86_64-linux-gnu/libibverbs:/usr/lib/x86_64-linux-gnu/libibverbs:ro )
    [ -d /etc/libibverbs.d ] && \
        IONIC_MOUNTS+=( -v /etc/libibverbs.d:/etc/libibverbs.d:ro )
    for _pattern in "libnl-3.so*" "libnl-route-3.so*" "libmnl.so*"; do
        _lib=$(find /lib/x86_64-linux-gnu /usr/lib/x86_64-linux-gnu /lib64 /usr/lib64 \
            -maxdepth 1 -name "$_pattern" -print -quit 2>/dev/null)
        [ -n "$_lib" ] && IONIC_MOUNTS+=( -v "$_lib:$_lib:ro" )
    done
fi
IONIC_EOF

# Optional topology / speculative-decode flags driven by the recipe. Base recipes
# (EP1/DP1, no mtp) leave the extra strings empty, preserving prior behavior.
#
# EP/DP are now PER ROLE: the DP-attention + ep-size flags come from PDP/PEP for
# prefill and DDP/DEP for decode. For every EP<=8 recipe DDP==PDP and DEP==PEP
# (decode inherits prefill), so both role strings equal the old single EXTRA_FLAGS
# and the generated argv is byte-identical. Oren's wide-EP recipes set decode
# EP/DP=16 while prefill stays EP8.
PREFILL_DPEP=""
(( PDP > 1 )) && PREFILL_DPEP="$PREFILL_DPEP --enable-dp-attention --dp-size $PDP"
(( PEP > 1 )) && PREFILL_DPEP="$PREFILL_DPEP --ep-size $PEP"
DECODE_DPEP=""
(( DDP > 1 )) && DECODE_DPEP="$DECODE_DPEP --enable-dp-attention --dp-size $DDP"
(( DEP > 1 )) && DECODE_DPEP="$DECODE_DPEP --ep-size $DEP"
# Flags shared by both roles (a2a backend, mtp). --max-total-tokens stays here for
# non-wide recipes; wide recipes carry a per-role prefill_max_total via wide_ep.
EXTRA_COMMON=""
[[ -n "$A2A" ]] && EXTRA_COMMON="$EXTRA_COMMON --moe-a2a-backend $A2A --deepep-mode normal"
[[ -n "$MAXTOK" ]] && EXTRA_COMMON="$EXTRA_COMMON --max-total-tokens $MAXTOK"
if [[ "$MTP_ENABLED" == "1" ]]; then
    EXTRA_COMMON="$EXTRA_COMMON --speculative-algorithm $MTP_ALGO \
--speculative-num-steps $MTP_STEPS --speculative-eagle-topk $MTP_TOPK \
--speculative-num-draft-tokens $MTP_DRAFT"
    # EAGLE3 (and other draft-model algos) need an external draft checkpoint;
    # built-in EAGLE (DSV4) omits draft_model_path and this stays unset.
    if [[ -n "$MTP_DRAFT_PATH" ]]; then
        _d="$(MODEL_ROOT="$MODEL_RESOLVE_ROOT" relocate_model_root "$MTP_DRAFT_PATH")"
        _d="$(resolve_snapshot "$_d")" || exit 1
        DRAFT_RESOLVED="$(MODEL_ROOT_FROM="$MODEL_RESOLVE_ROOT" relocate_model_root "$_d")"
        EXTRA_COMMON="$EXTRA_COMMON --speculative-draft-model-path $DRAFT_RESOLVED"
    fi
fi
# Prefix a leading space only when the arg is non-empty (keeps EP<=8 argv byte-
# identical: the wide-only strings are empty and contribute nothing).
sp() { [[ -n "$1" ]] && printf ' %s' "$1"; return 0; }
# --kv-cache-dtype is emitted only when the recipe sets wide_ep.kv_cache_dtype;
# the pre-wide DSV4 path had no such flag, so EP<=8 recipes omit it.
KV_FLAG=""
[[ -n "$KVDTYPE" ]] && KV_FLAG=" --kv-cache-dtype $KVDTYPE"
# Assemble the per-role tail: role DP/EP + shared + wide common + role-specific
# wide extras. All wide pieces (WECOMMON/PEXTRA/DEXTRA) are empty for EP<=8.
PREFILL_TAIL="$PREFILL_DPEP$EXTRA_COMMON$(sp "$WECOMMON")$(sp "$PEXTRA")"
DECODE_TAIL="$DECODE_DPEP$EXTRA_COMMON$(sp "$WECOMMON")$(sp "$DEXTRA")"
echo "prefill tail:${PREFILL_TAIL:-<none>} | decode tail:${DECODE_TAIL:-<none>} (pep=$PEP pdp=$PDP dep=$DEP ddp=$DDP mtp=$MTP_ENABLED)"

if [[ "$HAS_MODEL" == "1" ]]; then
    # Generic path (e.g. Kimi): attention + swa from the recipe, model parsers /
    # quirks ride MODEL_SERVER_ARGS. Single `--attention-backend` when the recipe
    # sets `attention_backend`; split `--prefill-/--decode-attention-backend` when
    # it sets the per-role keys. swa dropped when the recipe omits it.
    ATTN_FLAGS=""
    [[ -n "$ATTN" ]]  && ATTN_FLAGS="$ATTN_FLAGS --attention-backend $ATTN"
    [[ -n "$PATTN" ]] && ATTN_FLAGS="$ATTN_FLAGS --prefill-attention-backend $PATTN"
    [[ -n "$DATTN" ]] && ATTN_FLAGS="$ATTN_FLAGS --decode-attention-backend $DATTN"
    SWA_FLAG=""
    [[ -n "$SWA" ]] && SWA_FLAG=" --swa-full-tokens-ratio $SWA"
    PREFILL_COMMON_FLAGS="--trust-remote-code --tp $PTP --disable-radix-cache \
$ATTN_FLAGS --max-running-requests $PMAXREQ --page-size $PAGE \
--mem-fraction-static $PMEMFRAC$SWA_FLAG \
--chunked-prefill-size $PCHUNK \
--disaggregation-transfer-backend $XFER --disaggregation-ib-device $IB$KV_FLAG$PREFILL_TAIL"
    DECODE_COMMON_FLAGS="--trust-remote-code --tp $DTP --disable-radix-cache \
$ATTN_FLAGS --max-running-requests $DMAXREQ --page-size $PAGE \
--mem-fraction-static $DMEMFRAC$SWA_FLAG \
--chunked-prefill-size $CHUNK \
--disaggregation-transfer-backend $XFER --disaggregation-ib-device $IB$KV_FLAG$DECODE_TAIL"
else
    # DSV4 path: for EP<=8 recipes (PTP==DTP, no wide_ep) both role strings equal
    # the pre-Kimi launcher's COMMON_FLAGS exactly.
    PREFILL_COMMON_FLAGS="--trust-remote-code --tp $PTP --disable-radix-cache \
--attention-backend $ATTN --max-running-requests $PMAXREQ --page-size $PAGE \
--mem-fraction-static $PMEMFRAC --swa-full-tokens-ratio $SWA \
--chunked-prefill-size $PCHUNK --disable-shared-experts-fusion \
--tool-call-parser deepseekv4 --reasoning-parser deepseek-v4 \
--disaggregation-transfer-backend $XFER --disaggregation-ib-device $IB$KV_FLAG$PREFILL_TAIL"
    DECODE_COMMON_FLAGS="--trust-remote-code --tp $DTP --disable-radix-cache \
--attention-backend $ATTN --max-running-requests $DMAXREQ --page-size $PAGE \
--mem-fraction-static $DMEMFRAC --swa-full-tokens-ratio $SWA \
--chunked-prefill-size $CHUNK --disable-shared-experts-fusion \
--tool-call-parser deepseekv4 --reasoning-parser deepseek-v4 \
--disaggregation-transfer-backend $XFER --disaggregation-ib-device $IB$KV_FLAG$DECODE_TAIL"
fi

# /it-share is the mi355x cluster's model NFS. Mount it only where it exists --
# docker would otherwise silently create a root-owned empty /it-share on every
# node of a cluster that keeps its models elsewhere (see MODEL_ROOT).
IT_SHARE_MOUNT="-v /it-share:/it-share:ro "
[[ -d /it-share ]] || IT_SHARE_MOUNT=""
# A cluster that relocates the model cache must bind that root in too, at the
# same path: --model-path is resolved inside the container, and a missing
# directory does not fail as "no such file" -- transformers falls back to
# treating it as a HuggingFace repo id and dies with a confusing
# "Repo id must be in the form 'repo_name' or 'namespace/repo_name'".
if [[ -n "$MODEL_ROOT" && "$MODEL_ROOT" != /it-share/* ]]; then
    IT_SHARE_MOUNT="$IT_SHARE_MOUNT-v $MODEL_ROOT:$MODEL_ROOT:ro "
fi
# Slurm hosts each server as a long-lived job step and reaps the container when
# the step ends, so --rm is right there. Spur has no equivalent step (its
# --overlap path is an interactive exec-into-job), so the container runs
# DETACHED under dockerd and drive.sh polls it -- which means it must NOT be
# --rm, or the exit code is gone before the poller can read it. Teardown removes
# it explicitly instead.
DOCKER_LIFECYCLE="--rm"
[[ "$CLUSTER" == "spur" ]] && DOCKER_LIFECYCLE="-d"
# pit2 runs rootless docker, whose spur-authz plugin refuses --privileged
# outright ("denied [P1]"). Spell out what --privileged was actually buying us:
# the device nodes (already listed below), the groups that own kfd/dri, and two
# capabilities -- IPC_LOCK, which RDMA memory registration needs, and SYS_PTRACE
# for profiling. Verified on g11: ibv_devinfo sees all 8 HCAs PORT_ACTIVE and
# --ipc/--network host and --shm-size are all permitted under this set.
PRIV_ARGS="--privileged"
if [[ "$CLUSTER" == "spur" ]]; then
    PRIV_ARGS="--group-add video --group-add render --cap-add IPC_LOCK --cap-add SYS_PTRACE"
fi
DOCKER_COMMON="$DOCKER_LIFECYCLE --network host --ipc host --shm-size 32g $PRIV_ARGS \
--security-opt seccomp=unconfined \
--device /dev/kfd --device /dev/dri --device /dev/infiniband \
${IT_SHARE_MOUNT}-v $HOME:/host_home $CHECKOUT_DOCKER_ARGS"
# Optional extra docker args (e.g. bind-mounting a locally-built lib for
# validation). Empty by default so the docker argv is byte-identical otherwise.
[[ -n "${EXTRA_DOCKER_ARGS:-}" ]] && DOCKER_COMMON="$DOCKER_COMMON ${EXTRA_DOCKER_ARGS}"

# Spur compute nodes run prolog/epilog hooks that reap any container not tagged
# with the owning job, so an untagged server is SIGKILLed (rc=137) seconds after
# it starts. Single-quoted so the literal $SPUR_JOB_ID survives into the
# generated per-role script and expands on the compute node, where the job
# environment actually defines it.
if [[ "$CLUSTER" == "spur" ]]; then
    DOCKER_COMMON="$DOCKER_COMMON "'--label spur_job_id=$SPUR_JOB_ID'
fi

# Per-role wide-EP docker env (MORI dispatch-token tuning etc.). Empty for EP<=8
# recipes; carries its own leading space so an empty value leaves the docker argv
# byte-identical (no stray double space).
PENV_ARG=""; [[ -n "$PENV" ]] && PENV_ARG=" $PENV"
DENV_ARG=""; [[ -n "$DENV" ]] && DENV_ARG=" $DENV"

# Pin the torch-dist rendezvous port for single-node engines. Without
# --nccl-port, get_free_port() binds port 0, closes the socket, and rank 0
# re-binds the same number later for the TCPStore; anything handed that port in
# between makes rank 0 die with EADDRINUSE. It reproduces rather than being a
# lottery, since the kernel walks its ephemeral range in order.
#
# Only single-node engines need this -- a multi-node engine gets
# --dist-init-addr, which _resolve_dist_init_method() prefers over the port.
# The value is the recipe's own per-role port +1000, which clears the server
# port and stays under the 32768 ephemeral floor. These recipes place at most
# one single-node engine per node, so two engines cannot collide.
#
# Pinning is safe here because --dp-size is only emitted together with
# --enable-dp-attention, and that path gives every DP rank the same nccl port.
# launch_dp_schedulers(), which needs a distinct port per worker, is never
# reached by these recipes.
PNCCL_ARG=""
DNCCL_ARG=""
if [[ "$CLUSTER" == "spur" ]]; then
    PNCCL_ARG=" --nccl-port $((PPORT + 1000))"
    DNCCL_ARG=" --nccl-port $((DPORT + 1000))"
fi
# Per-role NCCL IB/socket policy (spur only; empty under slurm, so the argv
# there is unchanged). Appended after $PENV/$DENV so docker's last-wins applies.
PENV_ARG="$PENV_ARG$SPUR_PNCCL"
DENV_ARG="$DENV_ARG$SPUR_DNCCL"

# ---------------------------------------------------------------------------
# Write per-role scripts that srun dispatches to each compute node.
# ---------------------------------------------------------------------------
# These are UNQUOTED `<<EOF` heredocs, so $MORI_ENV/$DSV4_ENV_STR/$COMMON_FLAGS/
# $IMAGE/$MODEL_PATH expand now (at generation). model_flags.sh is sourced at
# runtime for the model's env/server arrays, so the `${MODEL_ENV_ARGS[@]}` /
# `${MODEL_SERVER_ARGS[@]}` refs are backslash-escaped to survive into the script
# and expand after `source`. For DSV4 those arrays are empty and $DSV4_ENV_STR is
# set, so the resulting docker argv is byte-identical to the pre-Kimi launcher.
cat > "$WORKDIR/install_checkout_sglang.sh" <<'EOF'
#!/bin/bash
set -euo pipefail

case "${SGLANG_USE_CHECKOUT_RUNTIME:-1}" in
  0|false|False|FALSE|no|No|NO|off|Off|OFF)
    echo "[checkout-sglang] disabled; using image-baked sglang"
    exit 0
    ;;
esac

CHECKOUT_SRC="${CHECKOUT_SRC:-/sglang-checkout}"
RUNTIME_CHECKOUT="${RUNTIME_CHECKOUT:-/tmp/sglang-checkout-runtime}"

if [[ ! -f "$CHECKOUT_SRC/python/sglang/version.py" ]]; then
  echo "[checkout-sglang] ERROR: invalid checkout mount: $CHECKOUT_SRC" >&2
  exit 1
fi

echo "[checkout-sglang] reinstalling sglang from $CHECKOUT_SRC"
rm -rf "$RUNTIME_CHECKOUT"
mkdir -p "$RUNTIME_CHECKOUT"
# tar exits 1 for warnings -- notably "file changed as we read it", which NFS
# attribute caching raises spuriously on a tree nothing is writing (seen on the
# pit2 shared-NFS checkout). Only 2+ is a real error, and the extracting tar
# must still succeed outright. Without this, set -o pipefail aborts the install
# on a warning and the server container exits before it ever starts.
set +e
tar --exclude='__pycache__' --exclude='*.pyc' \
  -C "$CHECKOUT_SRC" -cf - . | tar -C "$RUNTIME_CHECKOUT" -xf -
_tar_rc=("${PIPESTATUS[@]}")
set -e
if (( _tar_rc[0] > 1 || _tar_rc[1] != 0 )); then
  echo "[checkout-sglang] ERROR: staging tar failed (create=${_tar_rc[0]} extract=${_tar_rc[1]})" >&2
  exit 1
fi

git config --global --add safe.directory "$RUNTIME_CHECKOUT" || true

# The ROCm pyproject variant is the one used by AMD CI. Mutate only the private
# /tmp copy so prefill/decode/bench never race on the read-only checkout mount.
rm -f "$RUNTIME_CHECKOUT/python/pyproject.toml"
cp "$RUNTIME_CHECKOUT/python/pyproject_other.toml" "$RUNTIME_CHECKOUT/python/pyproject.toml"
for f in README.md LICENSE; do
  if [[ -f "$RUNTIME_CHECKOUT/$f" && ! -e "$RUNTIME_CHECKOUT/python/$f" ]]; then
    cp "$RUNTIME_CHECKOUT/$f" "$RUNTIME_CHECKOUT/python/$f"
  fi
done

python3 -m pip uninstall -y sglang || true
python3 -m pip install --no-deps --no-build-isolation -e "$RUNTIME_CHECKOUT/python"

export RUNTIME_CHECKOUT
export PYTHONPATH="$RUNTIME_CHECKOUT/python:${PYTHONPATH:-}"
python3 - <<'PY'
import importlib.metadata
import os
import subprocess
import sglang

checkout = os.environ["RUNTIME_CHECKOUT"]
expected = os.path.realpath(os.path.join(checkout, "python", "sglang")) + os.sep
actual = os.path.realpath(os.path.dirname(sglang.__file__)) + os.sep
try:
    sha = subprocess.check_output(
        ["git", "-C", checkout, "rev-parse", "HEAD"], text=True
    ).strip()
except Exception:
    sha = os.environ.get("SGLANG_CHECKOUT_SHA", "unknown")

print(f"[checkout-sglang] sha={sha}")
print(f"[checkout-sglang] sglang_file={sglang.__file__}")
print(f"[checkout-sglang] sglang_version={importlib.metadata.version('sglang')}")
if not actual.startswith(expected):
    raise SystemExit(f"sglang did not import from checkout: {sglang.__file__}")
PY
EOF

cat > "$WORKDIR/install_checkout_router.sh" <<'EOF'
#!/bin/bash
set -euo pipefail

case "${SGLANG_USE_CHECKOUT_RUNTIME:-1}" in
  0|false|False|FALSE|no|No|NO|off|Off|OFF)
    echo "[checkout-router] disabled; using image-baked sglang-router"
    python3 - <<'PY' || true
import importlib.metadata
import sglang_router

print(f"[checkout-router] sglang_router_file={sglang_router.__file__}")
print(
    "[checkout-router] sglang_router_version="
    f"{importlib.metadata.version('sglang-router')}"
)
try:
    import sglang_router.sglang_router_rs as rs

    print(f"[checkout-router] sglang_router_rs_file={rs.__file__}")
except Exception as exc:
    print(f"[checkout-router] sglang_router_rs_import_error={exc}")
PY
    exit 0
    ;;
esac

RUNTIME_CHECKOUT="${RUNTIME_CHECKOUT:-/tmp/sglang-checkout-runtime}"
ROUTER_SRC="$RUNTIME_CHECKOUT/sgl-model-gateway/bindings/python"
WHEEL_DIR="${SGLANG_ROUTER_WHEEL_DIR:-/tmp/sglang-router-wheels}"

if [[ ! -f "$ROUTER_SRC/pyproject.toml" ]]; then
  echo "[checkout-router] ERROR: invalid router checkout: $ROUTER_SRC" >&2
  exit 1
fi

echo "[checkout-router] building sglang-router from $ROUTER_SRC"
export CARGO_BUILD_JOBS="${CARGO_BUILD_JOBS:-4}"
python3 -m maturin --version >/dev/null 2>&1 \
  || python3 -m pip install --no-cache-dir "maturin<1.14"

# Match the ROCm image build recipe when compiling from the checkout copy.
if [[ -f "$RUNTIME_CHECKOUT/sgl-model-gateway/Cargo.toml" ]]; then
  sed -i -E 's|^(smg-[a-zA-Z-]+)\s*=\s*"~1\.0\.0"|\1 = "=1.0.0"|' \
    "$RUNTIME_CHECKOUT/sgl-model-gateway/Cargo.toml"
fi

rm -rf "$WHEEL_DIR"
mkdir -p "$WHEEL_DIR"
(
  cd "$ROUTER_SRC"
  ulimit -n 65536 || true
  python3 -m maturin build --release --features vendored-openssl --out "$WHEEL_DIR"
)

python3 -m pip uninstall -y sglang-router || true
python3 -m pip install --force-reinstall --no-deps "$WHEEL_DIR"/*.whl

python3 - <<'PY'
import importlib.metadata
import sglang_router
import sglang_router.sglang_router_rs as rs
from sglang_router.sglang_router_rs import Router

print(f"[checkout-router] sglang_router_file={sglang_router.__file__}")
print(
    "[checkout-router] sglang_router_version="
    f"{importlib.metadata.version('sglang-router')}"
)
print(f"[checkout-router] sglang_router_rs_file={rs.__file__}")
print(f"[checkout-router] Router={Router}")
PY
EOF

if (( PN_PER > 1 || DN_PER > 1 )); then
  # Wide-EP path: an engine spans >1 node. Entry scripts add torch-dist
  # rendezvous args from NODE_RANK/NNODES/DIST_ADDR; launch scripts take the
  # per-node rank as $1..$3 and forward it into the container. Only reached when
  # a recipe sets TP>GPUS_PER_NODE, so EP<=8 recipes never take this branch.
  cat > "$WORKDIR/prefill_entry.sh" <<EOF
#!/bin/bash
set -euo pipefail
CIDIR=/host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}
source "\$CIDIR/model_flags.sh"
bash "\$CIDIR/install_checkout_sglang.sh"
if [[ "\${SGLANG_USE_CHECKOUT_RUNTIME:-1}" != "0" ]]; then
  export PYTHONPATH=/tmp/sglang-checkout-runtime/python:\${PYTHONPATH:-}
fi
DIST_ARGS=""
if [[ "\${NNODES:-1}" != "1" ]]; then
  DIST_ARGS="--nnodes \$NNODES --node-rank \$NODE_RANK --dist-init-addr \$DIST_ADDR:\${DIST_PORT:-29500}"
else
  # Single-node engine: no --dist-init-addr, so the rendezvous port is the one
  # sglang would otherwise pick with get_free_port(). Pin it. Empty on slurm.
  DIST_ARGS="$PNCCL_ARG"
fi
exec python3 -m sglang.launch_server \
  --model-path $MODEL_PATH --host 0.0.0.0 --port $PPORT \
  $PREFILL_COMMON_FLAGS "\${MODEL_SERVER_ARGS[@]}" \$DIST_ARGS \
  --disaggregation-mode prefill --disaggregation-bootstrap-port $PBOOT
EOF

  cat > "$WORKDIR/decode_entry.sh" <<EOF
#!/bin/bash
set -euo pipefail
CIDIR=/host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}
source "\$CIDIR/model_flags.sh"
bash "\$CIDIR/install_checkout_sglang.sh"
if [[ "\${SGLANG_USE_CHECKOUT_RUNTIME:-1}" != "0" ]]; then
  export PYTHONPATH=/tmp/sglang-checkout-runtime/python:\${PYTHONPATH:-}
fi
DIST_ARGS=""
if [[ "\${NNODES:-1}" != "1" ]]; then
  DIST_ARGS="--nnodes \$NNODES --node-rank \$NODE_RANK --dist-init-addr \$DIST_ADDR:\${DIST_PORT:-29500}"
else
  # Single-node engine: no --dist-init-addr, so the rendezvous port is the one
  # sglang would otherwise pick with get_free_port(). Pin it. Empty on slurm.
  DIST_ARGS="$DNCCL_ARG"
fi
exec python3 -m sglang.launch_server \
  --model-path $MODEL_PATH --host 0.0.0.0 --port $DPORT \
  $DECODE_COMMON_FLAGS "\${MODEL_SERVER_ARGS[@]}" \$DIST_ARGS \
  --disaggregation-mode decode --disaggregation-bootstrap-port $DBOOT
EOF

  cat > "$WORKDIR/prefill.sh" <<EOF
#!/bin/bash
source "$WORKDIR/model_flags.sh"
source "$WORKDIR/ionic_mounts.sh"
NODE_RANK="\${1:-0}"; NNODES="\${2:-1}"; DIST_ADDR="\${3:-}"
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 \
  -e NODE_RANK="\$NODE_RANK" -e NNODES="\$NNODES" -e DIST_ADDR="\$DIST_ADDR" -e DIST_PORT=${DIST_PORT} \
  "\${IONIC_MOUNTS[@]}" \
  $MORI_ENV$PENV_ARG $DSV4_ENV_STR "\${MODEL_ENV_ARGS[@]}" \
  $IMAGE bash /host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}/prefill_entry.sh
EOF

  cat > "$WORKDIR/decode.sh" <<EOF
#!/bin/bash
source "$WORKDIR/model_flags.sh"
source "$WORKDIR/ionic_mounts.sh"
NODE_RANK="\${1:-0}"; NNODES="\${2:-1}"; DIST_ADDR="\${3:-}"
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 \
  -e NODE_RANK="\$NODE_RANK" -e NNODES="\$NNODES" -e DIST_ADDR="\$DIST_ADDR" -e DIST_PORT=${DIST_PORT} \
  "\${IONIC_MOUNTS[@]}" \
  $MORI_ENV$DENV_ARG $DSV4_ENV_STR "\${MODEL_ENV_ARGS[@]}" \
  $IMAGE bash /host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}/decode_entry.sh
EOF

else
  # Single-node-per-engine path (all EP<=8 recipes): byte-identical to the
  # pre-wide launcher. No NODE_RANK/NNODES/DIST_* plumbing is emitted.
  cat > "$WORKDIR/prefill_entry.sh" <<EOF
#!/bin/bash
set -euo pipefail
CIDIR=/host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}
source "\$CIDIR/model_flags.sh"
bash "\$CIDIR/install_checkout_sglang.sh"
if [[ "\${SGLANG_USE_CHECKOUT_RUNTIME:-1}" != "0" ]]; then
  export PYTHONPATH=/tmp/sglang-checkout-runtime/python:\${PYTHONPATH:-}
fi
exec python3 -m sglang.launch_server \
  --model-path $MODEL_PATH --host 0.0.0.0 --port $PPORT \
  $PREFILL_COMMON_FLAGS "\${MODEL_SERVER_ARGS[@]}" \
  --disaggregation-mode prefill --disaggregation-bootstrap-port $PBOOT$PNCCL_ARG
EOF

  cat > "$WORKDIR/decode_entry.sh" <<EOF
#!/bin/bash
set -euo pipefail
CIDIR=/host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}
source "\$CIDIR/model_flags.sh"
bash "\$CIDIR/install_checkout_sglang.sh"
if [[ "\${SGLANG_USE_CHECKOUT_RUNTIME:-1}" != "0" ]]; then
  export PYTHONPATH=/tmp/sglang-checkout-runtime/python:\${PYTHONPATH:-}
fi
exec python3 -m sglang.launch_server \
  --model-path $MODEL_PATH --host 0.0.0.0 --port $DPORT \
  $DECODE_COMMON_FLAGS "\${MODEL_SERVER_ARGS[@]}" \
  --disaggregation-mode decode --disaggregation-bootstrap-port $DBOOT$DNCCL_ARG
EOF

  cat > "$WORKDIR/prefill.sh" <<EOF
#!/bin/bash
source "$WORKDIR/model_flags.sh"
source "$WORKDIR/ionic_mounts.sh"
docker rm -f mi355x_prefill 2>/dev/null || true
docker run $DOCKER_COMMON "\${IONIC_MOUNTS[@]}" --name mi355x_prefill \
  -e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 $MORI_ENV$PENV_ARG $DSV4_ENV_STR "\${MODEL_ENV_ARGS[@]}" \
  $IMAGE bash /host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}/prefill_entry.sh
EOF

  cat > "$WORKDIR/decode.sh" <<EOF
#!/bin/bash
source "$WORKDIR/model_flags.sh"
source "$WORKDIR/ionic_mounts.sh"
docker rm -f mi355x_decode 2>/dev/null || true
docker run $DOCKER_COMMON "\${IONIC_MOUNTS[@]}" --name mi355x_decode \
  -e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 $MORI_ENV$DENV_ARG $DSV4_ENV_STR "\${MODEL_ENV_ARGS[@]}" \
  $IMAGE bash /host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}/decode_entry.sh
EOF
fi

# 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

# Prefill health-wait + router --prefill args. For PW=1 this is a single endpoint
# (byte-identical to the pre-fan-out launcher); for PW>1 it iterates the comma-
# separated engine node0 IPs in PCSV ($3), health-waits each, and passes one
# --prefill per engine so the router fans requests across all prefill engines.
if (( PW > 1 )); then
  PREFILL_WAIT_ROUTER="    IFS=',' read -ra PIPS <<< \"\${PCSV:-\$PIP}\"
    PREFILL_ARGS=\"\"
    for pip in \"\${PIPS[@]}\"; do
      echo \"[wait] prefill \$pip\"; for i in \$(seq 1 600); do curl -sf http://\$pip:$PPORT/health >/dev/null && break; sleep 5; done
      PREFILL_ARGS=\"\$PREFILL_ARGS --prefill http://\$pip:$PPORT $PBOOT\"
    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_ARGS \\
      --decode http://\$DIP:$DPORT \\
      --host 0.0.0.0 --port $LBPORT \\
      --disable-circuit-breaker &"
else
  PREFILL_WAIT_ROUTER="    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 &"
fi

# 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; PCSV=\${3:-\$PIP}
source "$WORKDIR/ionic_mounts.sh"
docker rm -f mi355x_bench 2>/dev/null || true
docker run $DOCKER_COMMON "\${IONIC_MOUNTS[@]}" --name mi355x_bench \
  -e PIP=\$PIP -e DIP=\$DIP -e PCSV=\$PCSV \
  $IMAGE bash -lc '
    CIDIR=/host_home/.mi355x_ci/${MATRIX_CONFIG_NAME}
    bash \$CIDIR/install_checkout_sglang.sh
    if [ "\${SGLANG_USE_CHECKOUT_RUNTIME:-1}" != "0" ]; then
      export PYTHONPATH=/tmp/sglang-checkout-runtime/python:\${PYTHONPATH:-}
    else
      export PYTHONPATH=/sgl-workspace/sglang/python:\${PYTHONPATH:-}
    fi
    bash \$CIDIR/install_checkout_router.sh
$PREFILL_WAIT_ROUTER
    for i in \$(seq 1 30); do curl -sf http://127.0.0.1:$LBPORT/health >/dev/null && break; sleep 2; done
    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}"; PN_PER="${4:-1}"; DN_PER="${5:-1}"; DIST_NIC="${6:-}"
# Interface resolve_ip() falls back to when the recipe names no dist NIC (every
# EP<=8 recipe). Empty on Slurm, where drive.sh runs on the login node and never
# resolves a node it is running on; set on Spur, where it runs on an allocated
# node and that node's own name resolves to 127.0.1.1 via Ubuntu's stock
# /etc/hosts -- useless as an address to advertise to the other engine.
ADDR_NIC="${7:-}"
# Scheduler seam (see the CLUSTER block at the top of launch_mi355x.sh). Slurm
# hands out a bracket-compressed hostlist that only `scontrol show hostnames`
# can expand, and srun inherits the allocation from the environment. Spur's
# nodelist env var is already comma-expanded, has no `show hostnames`, and its
# `srun --overlap` requires an explicit --jobid.
CLUSTER="${CLUSTER:-slurm}"
if [[ "$CLUSTER" == "spur" ]]; then
  SRUN=(srun --jobid="$SPUR_JOB_ID" --overlap)
  sched_hostnames() { tr ',' '\n' <<< "$SPUR_JOB_NODELIST" | sed '/^$/d'; }
else
  SRUN=(srun --overlap)
  sched_hostnames() { scontrol show hostnames "$SLURM_JOB_NODELIST"; }
fi
mapfile -t NODES < <(sched_hostnames)
# Dispatch one long-running, log-producing step to a node.
#
# Spur allocates a pty for every --overlap step. drive.sh runs the servers and
# the bench as BACKGROUND subshells, and a background process group writing
# through that pty is stopped by SIGTTOU -- the step then dies (marker rc=137)
# and everything it printed after the first few bytes is lost, which is why a
# failing leg used to leave a log containing nothing but the srun banner. Having
# srun write the file itself keeps the step's output off the pty entirely.
# Slurm has no pty here, so it keeps the original shell redirect.
srun_log() {  # srun_log <logfile> <node> <cmd...>
  local log="$1" node="$2"; shift 2
  # On spur drive.sh runs on one of the allocated nodes, so a dispatch to that
  # node would be srun exec-ing into its own job -- the one spur case that has
  # repeatedly misbehaved. Run it directly instead; the result is identical and
  # it removes a whole class of scheduler interaction.
  if [[ "$CLUSTER" == "spur" && "$node" == "$(hostname)" ]]; then
    "$@" > "$log" 2>&1
    return
  fi
  "${SRUN[@]}" -N1 --nodelist="$node" "$@" > "$log" 2>&1
}

srun_local_or_step() {  # same rule, for short commands whose output we capture
  local node="$1"; shift
  if [[ "$CLUSTER" == "spur" && "$node" == "$(hostname)" ]]; then
    "$@"
    return
  fi
  "${SRUN[@]}" -N1 --nodelist="$node" "$@"
}

# --- spur: containers are detached, liveness comes from dockerd ---------------
# Spur cannot host a long-lived process in a job step, so on spur the servers
# and the bench run as DETACHED containers and these helpers derive the same
# server_exit_*/bench_exit markers the monitor loop already consumes -- from
# real container state rather than from srun's status. Every srun issued here is
# a sub-second foreground command, which is the part of spur that behaves.
declare -A SPUR_SEEN=()   # container -> 1 once observed running, so a container
                          # that has not started yet is not mistaken for a crash
declare -A SPUR_ABSENT=() # container -> consecutive polls that found it gone
SPUR_SERVERS=()           # "role:node:container"
SPUR_POLLS=0
SPUR_ABSENT_LIMIT=3       # polls (10s apart) before "gone" is believed

container_state() {  # <node> <name> -> "running" | "<exit code>" | absent | unknown
  # The obvious `docker inspect | tail -1` cannot be trusted on spur. drive.sh
  # runs on one allocated node, so every poll of a container on the OTHER node
  # is an `srun --overlap` step, and those intermittently fail to spawn
  # ("step command failed to spawn: Permission denied (os error 13)"). A failed
  # step and a genuinely absent container are the same observation -- nonzero
  # exit, empty stdout -- so one flaky poll was enough to declare a healthy
  # server dead. Observed on dsv4flash-fp8-1k1k-1p1d: prefill marked
  # "rc=missing" while `docker ps` on that node still showed it Up.
  #
  # The __ok__ sentinel is emitted only if the remote command actually ran, so a
  # dispatch failure is reported as "unknown" and the caller ignores it rather
  # than counting it as evidence.
  local out
  out="$(srun_local_or_step "$1" bash -c \
      "docker inspect -f '{{if .State.Running}}running{{else}}{{.State.ExitCode}}{{end}}' $2 2>/dev/null || echo absent; echo __ok__" \
      2>/dev/null | tr -d '\r')"
  [[ "$out" == *__ok__* ]] || { echo unknown; return; }
  grep -vxE '__ok__|' <<< "$out" | tail -1
}

container_log() {  # <node> <name> <logfile>
  # Fetch to a temp file and never let a failed fetch shrink the log: the same
  # flaky step that produced the false "missing" also truncated the 5-line
  # prefill log to 0 bytes, destroying the only evidence of what happened.
  local tmp="$3.fetch"
  srun_local_or_step "$1" docker logs "$2" > "$tmp" 2>/dev/null || true
  # Publish in place (truncate + rewrite the SAME inode), never by renaming over
  # the target. drive.sh keeps a `tail -F` on bench.log, and swapping the inode
  # every 10s made tail reopen and re-emit the whole file each poll -- the
  # driver log reached 177 MB and legs took 2 h instead of 36 min.
  if [[ -s "$tmp" || ! -f "$3" ]]; then cat "$tmp" > "$3"; fi
  rm -f "$tmp"
}

spur_rm() {  # <node> <container...> -- remove, retrying past a flaky dispatch
  local n="$1"; shift
  local i
  for i in 1 2 3; do
    srun_local_or_step "$n" bash -c "docker rm -f $* >/dev/null 2>&1; echo __ok__" \
      2>/dev/null | grep -q __ok__ && return 0
    sleep 3
  done
  return 1
}

# Refresh one container's log and write its marker once it is no longer running.
spur_track() {  # <node> <name> <logfile> <markerfile> <marker prefix>
  local node="$1" name="$2" log="$3" marker="$4" label="$5" st
  container_log "$node" "$name" "$log"
  st="$(container_state "$node" "$name")"
  if [[ "$st" == "running" ]]; then
    SPUR_SEEN["$name"]=1
    SPUR_ABSENT["$name"]=0
    return
  fi
  # A step that never ran says nothing about the container.
  [[ "$st" == "unknown" ]] && return
  if [[ "$st" == "absent" || -z "$st" ]]; then
    # Not created yet, or one poll that came back empty. Only a container that
    # is absent on SPUR_ABSENT_LIMIT polls in a row is treated as a crash, so
    # neither a slow `docker run -d` nor a single flaky poll ends the leg.
    SPUR_ABSENT["$name"]=$(( ${SPUR_ABSENT[$name]:-0} + 1 ))
    if [[ -z "${SPUR_SEEN[$name]:-}" && "$SPUR_POLLS" -lt 6 ]]; then return; fi
    [[ "${SPUR_ABSENT[$name]}" -lt "$SPUR_ABSENT_LIMIT" ]] && return
    st="missing"
  fi
  echo "$label rc=$st" > "$marker"
}

spur_refresh_markers() {
  local entry role node cname st
  SPUR_POLLS=$((SPUR_POLLS + 1))
  for entry in "${SPUR_SERVERS[@]}"; do
    IFS=: read -r role node cname <<< "$entry"
    [[ -f "$WORKDIR/server_exit_${role}_${node}" ]] && continue
    spur_track "$node" "$cname" "$WORKDIR/${role}_${node}.log" \
               "$WORKDIR/server_exit_${role}_${node}" "$role@$node"
  done
  # The bench marker holds a bare exit code, not the "role@node rc=" form.
  container_log "$PNODE" mi355x_bench "$WORKDIR/bench.log"
  st="$(container_state "$PNODE" mi355x_bench)"
  if [[ "$st" == "running" ]]; then
    SPUR_SEEN[mi355x_bench]=1
    SPUR_ABSENT[mi355x_bench]=0
  elif [[ "$st" == "unknown" ]]; then
    :
  elif [[ "$st" == "absent" || -z "$st" ]]; then
    SPUR_ABSENT[mi355x_bench]=$(( ${SPUR_ABSENT[mi355x_bench]:-0} + 1 ))
    if [[ -n "${SPUR_SEEN[mi355x_bench]:-}" \
          && "${SPUR_ABSENT[mi355x_bench]}" -ge "$SPUR_ABSENT_LIMIT" ]]; then
      echo 1 > "$WORKDIR/bench_exit"
    fi
  else
    echo "$st" > "$WORKDIR/bench_exit"
  fi
}
# Resolve a node's IP. When a cross-node dist NIC is named (wide engines set
# $DIST_NIC to the recipe's dist_socket_ifname), read the address off that
# interface on the node itself -- the multi-node dist-init-addr must be a
# LOCALLY-BINDABLE IP (torch-dist + tokenizer ZMQ bind to it), and a node's
# forward DNS can be stale/point at a non-local mgmt alias (observed: a decode
# root whose hostname resolved to an unpingable IP, ZMQ bind => "Cannot assign
# requested address"). Fall back to Slurm's NodeAddr, then DNS.
resolve_ip() {
  local n="$1" ip="" addr="" field nic="${DIST_NIC:-$ADDR_NIC}"
  if [[ -n "$nic" ]]; then
    # Go through srun_local_or_step, not a bare srun: on spur drive.sh runs on an
    # allocated node, and resolving its OWN node through `srun --overlap` is the
    # dispatch that has repeatedly misbehaved. Here it did not fail, it HUNG --
    # a leg sat 30 min on `srun ... ip -4 -o addr show eno0` for the node it was
    # already running on, before any container was started. The helper runs it
    # directly in that case. The timeout bounds the genuinely remote call so a
    # wedged step costs 60s and a DNS fallback instead of the whole leg.
    if [[ "$CLUSTER" == "spur" ]]; then
      if [[ "$n" == "$(hostname)" ]]; then
        ip=$(command ip -4 -o addr show "$nic" 2>/dev/null \
               | awk '{print $4}' | cut -d/ -f1 | head -1)
      else
        ip=$(timeout 60 "${SRUN[@]}" -N1 --nodelist="$n" ip -4 -o addr show "$nic" 2>/dev/null \
               | awk '{print $4}' | cut -d/ -f1 | head -1)
      fi
    else
      ip=$("${SRUN[@]}" -N1 --nodelist="$n" ip -4 -o addr show "$nic" 2>/dev/null \
             | awk '{print $4}' | cut -d/ -f1 | head -1)
    fi
  fi
  if [[ -z "$ip" ]]; then
    # Spur's `scontrol show node` does not emit Slurm's NodeAddr= field, so the
    # loop below would just fall through; skip it there and go straight to DNS.
    if [[ "$CLUSTER" != "spur" ]]; then
      for field in $(scontrol show node "$n" -o); do
        case "$field" in
          NodeAddr=*) addr="${field#NodeAddr=}"; break ;;
        esac
      done
    fi
    [[ -n "$addr" ]] || addr="$n"
    read -r ip _ < <(getent ahostsv4 "$addr")
  fi
  # A loopback or empty answer is never a usable advertised address. Fail here
  # rather than let a server bind 127.x (or nothing) and have the peer engine
  # time out connecting to it with no obvious cause. Callers must use
  # `x=$(resolve_ip n) || exit 1` -- the status is visible at the assignment,
  # but an `exit` inside the substitution would only leave the subshell.
  if [[ -z "$ip" ]]; then
    echo "[drive] ERROR: no address found for $n" >&2
    return 1
  fi
  if [[ "$ip" == 127.* ]]; then
    echo "[drive] ERROR: $n resolved to loopback $ip; set ADDR_NIC to a real interface" >&2
    return 1
  fi
  printf '%s\n' "$ip"
}
# SLURM canonicalizes (sorts) SLURM_JOB_NODELIST, so the requested --nodelist
# order is already lost here. To keep a "slow-root" node out of any engine's
# rank0 slot -- e.g. mia1-p01-g20, whose MORI EP bootstrap reaches the connect
# stage ~16s late and loses the hardcoded 10s worker-connect race whenever it is
# rank0 -- push every node named in SLURM_DIST_TAIL to the END of the list. It
# then lands in a trailing decode-worker position, where its late ranks connect
# into the decode root's 30s accept window (harmless). Comma-separated; unmatched
# names ignored; empty (every EP<=8 recipe) => no-op, order unchanged.
if [[ -n "${SLURM_DIST_TAIL:-}" ]]; then
  IFS=',' read -ra _TAIL <<< "$SLURM_DIST_TAIL"
  _HEAD=(); _TL=()
  for n in "${NODES[@]}"; do
    _skip=0; for t in "${_TAIL[@]}"; do [[ "$n" == "$t" ]] && { _skip=1; break; }; done
    if (( _skip )); then _TL+=("$n"); else _HEAD+=("$n"); fi
  done
  NODES=("${_HEAD[@]}" "${_TL[@]}")
  echo "[drive] SLURM_DIST_TAIL=$SLURM_DIST_TAIL -> node order: ${NODES[*]}"
fi
# Each engine may span PN_PER/DN_PER nodes (ceil(TP/GPUs-per-node)); PN_PER=1 for
# EP<=8 so this reduces to the original one-node-per-worker split.
PN_TOTAL=$((PW * PN_PER)); DN_TOTAL=$((DW * DN_PER))
PNODES=("${NODES[@]:0:PN_TOTAL}")
DNODES=("${NODES[@]:PN_TOTAL:DN_TOTAL}")
PNODE="${PNODES[0]}"; DNODE="${DNODES[0]}"
PIP=$(resolve_ip "$PNODE") || exit 1
DIP=$(resolve_ip "$DNODE") || exit 1
echo "[drive] prefill nodes: ${PNODES[*]} ; decode nodes: ${DNODES[*]}"
echo "[drive] bench targets prefill=$PNODE($PIP) decode=$DNODE($DIP)"
if (( DW > 1 )); then
  echo "[drive] NOTE: router + bench use the first decode engine only;"
  echo "[drive]       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"
if [[ "$CLUSTER" == "spur" ]]; then
  # Detached containers are not --rm, and a container whose docker client was
  # killed can survive as an exited husk. Clear all three names on every
  # allocated node first: otherwise `docker run` fails with a name conflict and,
  # worse, the poller reads the STALE container's exit code and reports a crash
  # that never happened.
  for n in "${NODES[@]}"; do
    srun_local_or_step "$n" \
      docker rm -f mi355x_prefill mi355x_decode mi355x_bench >/dev/null 2>&1 || true
  done
fi
# Launch PW prefill engines; each spans PN_PER nodes as its own torch-dist group
# (engine node0 = dist-init addr; node_rank = position WITHIN the engine, so an
# engine's ranks are 0..PN_PER-1, not a global index). PN_PER=1 => NNODES=1 in
# the entry script => dist args dropped => byte-identical single-node launch.
# Collect each engine's node0 IP into PCSV for the router's prefill fan-out.
PCSV=""
for ((k=0; k<PW; k++)); do
  e0="${PNODES[k*PN_PER]}"
  eip=$(resolve_ip "$e0") || exit 1
  PCSV="${PCSV:+$PCSV,}$eip"
  for ((j=0; j<PN_PER; j++)); do
    n="${PNODES[k*PN_PER + j]}"
    if [[ "$CLUSTER" == "spur" ]]; then
      # Returns as soon as the container is detached; liveness is polled later.
      srun_log "$WORKDIR/prefill_start_$n.log" "$n" \
        bash "$WORKDIR/prefill.sh" "$j" "$PN_PER" "$eip" || true
      SPUR_SERVERS+=("prefill:$n:mi355x_prefill")
    else
    ( srun_log "$WORKDIR/prefill_$n.log" "$n" bash "$WORKDIR/prefill.sh" "$j" "$PN_PER" "$eip"
      echo "prefill@$n rc=$?" > "$WORKDIR/server_exit_prefill_$n" ) &
    fi
  done
done
for ((k=0; k<DW; k++)); do
  e0="${DNODES[k*DN_PER]}"
  eip=$(resolve_ip "$e0") || exit 1
  for ((j=0; j<DN_PER; j++)); do
    n="${DNODES[k*DN_PER + j]}"
    if [[ "$CLUSTER" == "spur" ]]; then
      srun_log "$WORKDIR/decode_start_$n.log" "$n" \
        bash "$WORKDIR/decode.sh" "$j" "$DN_PER" "$eip" || true
      SPUR_SERVERS+=("decode:$n:mi355x_decode")
    else
    ( srun_log "$WORKDIR/decode_$n.log" "$n" bash "$WORKDIR/decode.sh" "$j" "$DN_PER" "$eip"
      echo "decode@$n rc=$?" > "$WORKDIR/server_exit_decode_$n" ) &
    fi
  done
done
echo "[drive] prefill engine endpoints (fan-out): $PCSV"
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).
BENCH_BG=""
if [[ "$CLUSTER" == "spur" ]]; then
  srun_log "$WORKDIR/bench_start.log" "$PNODE" \
    bash "$WORKDIR/bench.sh" "$PIP" "$DIP" "$PCSV" || true
else
( srun_log "$WORKDIR/bench.log" "$PNODE" bash "$WORKDIR/bench.sh" "$PIP" "$DIP" "$PCSV"
  echo $? > "$WORKDIR/bench_exit" ) &
BENCH_BG=$!
fi
# 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
  # On spur the markers are produced here, from container state, instead of by a
  # backgrounded srun. Everything below is unchanged.
  [[ "$CLUSTER" == "spur" ]] && spur_refresh_markers
  if compgen -G "$WORKDIR/server_exit_*" > /dev/null; then
    echo "[drive] ERROR: a server exited early before bench finished:"
    cat "$WORKDIR"/server_exit_* || true
    [[ -n "$BENCH_BG" ]] && 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"
if [[ "$CLUSTER" == "spur" ]]; then
  # Detached containers are not --rm, so kill is not enough -- remove them, and
  # the bench container too, or the next leg on this node inherits them.
  # Retry the removals: a step that fails to spawn used to leave the containers
  # running, and the next leg on that node then inherited engines holding all
  # eight GPUs.
  for n in "${PNODES[@]}"; do spur_rm "$n" mi355x_prefill || true; done
  for n in "${DNODES[@]}"; do spur_rm "$n" mi355x_decode  || true; done
  spur_rm "$PNODE" mi355x_bench || true
else
for n in "${PNODES[@]}"; do "${SRUN[@]}" -N1 --nodelist="$n" docker kill mi355x_prefill >/dev/null 2>&1 || true; done
for n in "${DNODES[@]}"; do "${SRUN[@]}" -N1 --nodelist="$n" docker kill mi355x_decode  >/dev/null 2>&1 || true; done
fi
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)

# Optional comma-separated nodes to keep the scheduler off.
EXCLUDE_ARG=()
[[ -n "${SLURM_EXCLUDE:-}" ]] && EXCLUDE_ARG=(--exclude="$SLURM_EXCLUDE")

# Optional account / QoS (spur requires both; mi355x sets neither).
ACCT_ARG=()
[[ -n "$SLURM_ACCOUNT" ]] && ACCT_ARG+=(-A "$SLURM_ACCOUNT")
[[ -n "$SLURM_QOS" ]] && ACCT_ARG+=(-q "$SLURM_QOS")

# Nodes = sum over engines of nodes-per-engine. EP<=8 (PN_PER=DN_PER=1) gives the
# original PW+DW (1P1D -> 2 nodes); wide EP16 1P1D gives 2+2 = 4 nodes.
TOTAL_NODES=$(( PW * PN_PER + DW * DN_PER ))

# Keep g20 out of a wide engine's root position, where its slower MORI
# bootstrap can miss the worker-connect timeout.
if [[ "$MATRIX_CONFIG_NAME" == *-2p1d-ep16* ]]; then
    export SLURM_DIST_TAIL="${SLURM_DIST_TAIL:-mia1-p01-g20}"
fi

# 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}"

# Address-resolution fallback NIC for drive.sh (see ADDR_NIC there). Only spur
# needs one, because only there does drive.sh run on an allocated node.
ADDR_NIC=""
[[ "$CLUSTER" == "spur" ]] && ADDR_NIC="${ADDR_NIC_OVERRIDE:-eno0}"

set +e
if [[ "$CLUSTER" == "spur" ]]; then
    # Spur's salloc takes no trailing command, so drive.sh is submitted as a
    # batch job instead. That runs it on the first allocated node rather than
    # the login node, which is fine: $WORKDIR and $GITHUB_WORKSPACE are on
    # shared NFS, and drive.sh only ever reaches the other nodes through srun.
    # We stream the batch output back live and take the exit code from a file
    # drive_batch.sh writes, since sbatch returns as soon as the job is queued.
    cat > "$WORKDIR/drive_batch.sh" <<EOF
#!/bin/bash
# Set explicitly rather than relying on sbatch to propagate the submit
# environment; drive.sh needs it to pick the spur scheduler seam.
export CLUSTER=spur
export PATH="$PATH"
# Slurm runs a batch script once, on the first allocated node. Spur runs it as
# one task PER NODE (scontrol reports NumTasks=NumNodes), so without this guard
# every node gets its own drive.sh: the copies race, and each one's pre-launch
# container cleanup kills the servers the other just started. Only rank 0
# drives; the rest exit immediately and simply hold their node in the
# allocation, which is all drive.sh needs them for.
# Elect the driver with an atomic mkdir on shared NFS rather than by rank:
# spur does not consistently place the batch script on rank 0 (it has been
# observed running only on SPUR_NODEID=1), so a rank test would sometimes elect
# nobody and the leg would never start.
if ! mkdir "$WORKDIR/driver.lock" 2>/dev/null; then
  # Must BLOCK, not exit: spur starts tearing the job down as soon as one of its
  # tasks finishes, and a job in Completing rejects further steps
  # ("CreateJobStep failed: job N is not running") -- which killed the decode
  # launch. Holding here keeps the node in the allocation until rank 0 is done.
  echo "[drive_batch] \$(hostname) standing by; another task drives" >> "$WORKDIR/standby.log"
  while [[ ! -f "$WORKDIR/drive_exit" ]]; do sleep 10; done
  exit 0
fi
echo "[drive_batch] \$(hostname) rank \${SPUR_NODEID:-?} elected driver"
# Send the driver's output to a private per-node file, not to the shared
# --output. All tasks of the job write sbatch.out concurrently, so it loses
# records: a leg that failed in the monitor loop left an sbatch.out that simply
# stopped mid-run, with the teardown lines never visible.
#
# A plain redirect, deliberately NOT `| tee`. With a pipeline, drive_batch waits
# for the whole pipeline, and that does not finish when drive.sh does -- the
# backgrounded `tail -F` on bench.log inherits drive.sh's stdout and holds the
# pipe's write end open. drive.sh exited 0 and drive_exit was still unwritten
# 11 minutes later, so the standby tasks kept the allocation alive and the leg
# looked hung after it had actually passed.
bash "$WORKDIR/drive.sh" "$WORKDIR" "$PW" "$DW" "$PN_PER" "$DN_PER" "$DIST_SOCK" "$ADDR_NIC" \
  > "$WORKDIR/drive_\$(hostname).log" 2>&1
echo \$? > "$WORKDIR/drive_exit"
EOF
    chmod +x "$WORKDIR/drive_batch.sh"
    rm -f "$WORKDIR/drive_exit"
    SBATCH_OUT="$WORKDIR/sbatch.out"
    : > "$SBATCH_OUT"
    # Node count is what matters here; do NOT pass -n1 to make the task count
    # match. Spur sizes the allocation from the task count, so -n1 collapses a
    # 2-node request to a single node and the decode engine has nowhere to run.
    # Uniqueness of the driver is handled in drive_batch.sh by an atomic lock,
    # not by the task count.
    # Spur's controller is a Raft cluster, so submission can fail for reasons
    # that have nothing to do with the request: during a leader election it
    # answers `The service is currently unavailable ... "not the Raft leader"`.
    # That is transient and clears in seconds, but an unretried submit turns it
    # into a failed leg. Worse, the message is multi-line, so the ${VAR##* }
    # job-id parse below extracted `leader"` from it and the leg was reported as
    # a benchmark failure with an empty bench.log -- indistinguishable, at a
    # glance, from a model bug. Observed on glm52-fp4-1k1k-2p1d-ep16 2026-09-17.
    # Retry only on that signature; a genuinely bad sbatch request must still
    # fail immediately rather than being retried six times.
    SPUR_LEG_JOB_ID=""
    for _attempt in 1 2 3 4 5 6; do
        SBATCH_MSG=$(sbatch -p "$SLURM_PARTITION" -N"$TOTAL_NODES" "${NODELIST_ARG[@]}" \
            "${EXCLUDE_ARG[@]}" "${EXCLUSIVE_ARG[@]}" "${ACCT_ARG[@]}" \
            --job-name "$JOB_NAME" -t "$TIME_LIMIT" \
            --output "$SBATCH_OUT" --error "$SBATCH_OUT" \
            "$WORKDIR/drive_batch.sh" 2>&1)
        echo "$SBATCH_MSG"
        SPUR_LEG_JOB_ID="${SBATCH_MSG##* }"
        [[ "$SPUR_LEG_JOB_ID" =~ ^[0-9]+$ ]] && break
        if [[ "$SBATCH_MSG" == *"not the Raft leader"* \
           || "$SBATCH_MSG" == *"service is currently unavailable"* ]]; then
            echo "[launch] spur controller unavailable (attempt $_attempt/6); retrying in 20s" >&2
            SPUR_LEG_JOB_ID=""
            sleep 20
            continue
        fi
        break
    done
    if [[ ! "$SPUR_LEG_JOB_ID" =~ ^[0-9]+$ ]]; then
        echo "ERROR: could not parse a job id out of: $SBATCH_MSG" >&2
        SALLOC_RC=1
    else
        echo "[launch] spur job $SPUR_LEG_JOB_ID submitted; streaming $SBATCH_OUT"
        tail -F "$SBATCH_OUT" 2>/dev/null &
        SPUR_TAIL_PID=$!
        # drive_batch.sh writes drive_exit last, so that file is the primary
        # completion signal; the squeue check catches a job that died without
        # ever running it (node failure, scheduler kill, time limit).
        while :; do
            [[ -f "$WORKDIR/drive_exit" ]] && break
            if ! squeue -h -o "%i" 2>/dev/null | grep -qx "$SPUR_LEG_JOB_ID"; then
                # The job has left the queue. drive_exit is written on a compute
                # node and read here on the login node, so NFS close-to-open
                # visibility can delay it well past a single short sleep. A five
                # second grace produced three false reds on 2026-09-16, each on a
                # leg that had finished every concurrency point and passed the
                # accuracy gate; one of them was the whole GLM EP16 row.
                for _ in $(seq 1 30); do
                    [[ -f "$WORKDIR/drive_exit" ]] && break
                    sleep 3
                done
                if [[ ! -f "$WORKDIR/drive_exit" ]]; then
                    # Still not visible. Do not invent a failure we did not
                    # observe -- decide from what the run actually produced.
                    # bench_exit is written earlier and is usually visible by now;
                    # when even that is missing (seen on glm52-fp4-1k1k-2p1d-ep16,
                    # which had all seven results at 0.945), a passed accuracy
                    # gate with results on disk is the stronger evidence.
                    #
                    # "Some results plus a passed gate" is not enough. A leg
                    # whose decode ran at ~9 tok/s got through conc=1, 8 and 16
                    # of seven points before the job ended, and this branch
                    # called it green: the gate had passed and raw files
                    # existed. Require every expected concurrency point, so a
                    # sweep that stopped early is red rather than a green row
                    # standing on someone else's numbers.
                    WANT=0; HAVE=0
                    for _c in ${CONCS//,/ }; do
                        WANT=$((WANT + 1))
                        [[ -f "$WORKDIR/raw_conc${_c}.json" ]] && HAVE=$((HAVE + 1))
                    done
                    if [[ "$(cat "$WORKDIR/bench_exit" 2>/dev/null)" == "0" ]]; then
                        echo 0 > "$WORKDIR/drive_exit"
                    elif [[ "$HAVE" -eq "$WANT" ]] \
                         && awk '/\[gsm8k\] accuracy=/{
                                   match($0,/accuracy=([0-9.]+)/,a)
                                   match($0,/threshold=([0-9.]+)/,t)
                                   if (a[1]+0 >= t[1]+0) ok=1
                                 } END{ exit ok?0:1 }' \
                                "$WORKDIR/bench.log" 2>/dev/null; then
                        echo 0 > "$WORKDIR/drive_exit"
                    else
                        echo "[launch] incomplete sweep: $HAVE/$WANT concurrency points" >&2
                        echo 1 > "$WORKDIR/drive_exit"
                    fi
                fi
                break
            fi
            sleep 10
        done
        kill "$SPUR_TAIL_PID" 2>/dev/null || true
        SALLOC_RC=$(cat "$WORKDIR/drive_exit" 2>/dev/null || echo 1)
        # Release the allocation. On slurm the job ends when salloc's command
        # returns; on spur drive.sh is an sbatch job, and the loop above exits as
        # soon as drive_exit appears -- while the batch job, and every container
        # it started, is still running. Nothing else ever cancels it, so the
        # allocation is held until the time limit. The next leg then cannot get
        # nodes and dies with `JobLaunchFailure (dispatch confirmation failed
        # (2/4 confirmed))`, which reads as an unrelated infrastructure fault and
        # cascades through every leg after it. Observed 2026-09-17: jobs
        # 1774/1775 held all four nodes long after their legs had been recorded
        # rc=1. Unconditional, because reaching here means this leg is done with
        # its nodes either way.
        if squeue -h -o "%i" 2>/dev/null | grep -qx "$SPUR_LEG_JOB_ID"; then
            echo "[launch] releasing spur job $SPUR_LEG_JOB_ID"
            scancel "$SPUR_LEG_JOB_ID" 2>/dev/null || true
        fi
    fi
else
    salloc -p "$SLURM_PARTITION" -N"$TOTAL_NODES" "${NODELIST_ARG[@]}" "${EXCLUDE_ARG[@]}" "${EXCLUSIVE_ARG[@]}" \
        "${ACCT_ARG[@]}" \
        --job-name "$JOB_NAME" -t "$TIME_LIMIT" \
        bash "$WORKDIR/drive.sh" "$WORKDIR" "$PW" "$DW" "$PN_PER" "$DN_PER" "$DIST_SOCK"
    SALLOC_RC=$?
fi
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)
# ---------------------------------------------------------------------------
# GPU counts are per-ROLE totals across all engines: PW prefill engines of PTP
# GPUs each, DW decode engines of DTP each. For EP<=8 (PW=DW=1) this is PTP/DTP
# exactly, so the filename fields are unchanged. process_result.py reads the
# _ctx_/_gen_ fields as PREFILL_GPUS/DECODE_GPUS for its per-GPU throughput math,
# so they must be the role totals (Oren EP16: ctx=2*8=16, gen=1*16=16, gpus=32).
PREFILL_GPUS_TOTAL=$((PW * PTP)); DECODE_GPUS_TOTAL=$((DW * DTP))
TOTAL_GPUS=$((PREFILL_GPUS_TOTAL + DECODE_GPUS_TOTAL))
# Clear this leg's results from any previous run before writing new ones.
# The filenames are a pure function of the config, so a run that produces
# fewer concurrency points than the last one leaves the missing slots filled
# by the old run's files and the published table silently mixes two runs.
# Seen on kimik26-mxfp4-1k1k-2p1d-ep16-mxfp4: 2 points measured, 7 published,
# 5 of them a day old. Deleting first means a short run looks short.
rm -f "$GITHUB_WORKSPACE/${RESULT_FILENAME}_${MATRIX_CONFIG_NAME}_conc"*"_gpus_"*".json"

PROCESSED=0
for C in ${CONCS//,/ }; do
    RAW="$WORKDIR/raw_conc${C}.json"
    # raw_conc*.json is written by the bench container on a compute node; this
    # loop runs on the driver. NFS close-to-open means the last file written is
    # routinely not visible here yet, and skipping it outright silently dropped
    # one or two concurrency points from four legs of a 34-leg sweep -- the runs
    # were fine, only the published results were short. Wait a bounded 60s.
    for _ in $(seq 1 20); do
        [[ -f "$RAW" ]] && break
        sleep 3
    done
    [[ -f "$RAW" ]] || { echo "WARN: missing $RAW (not visible after 60s)"; continue; }
    DEST="$GITHUB_WORKSPACE/${RESULT_FILENAME}_${MATRIX_CONFIG_NAME}_conc${C}_gpus_${TOTAL_GPUS}_ctx_${PREFILL_GPUS_TOTAL}_gen_${DECODE_GPUS_TOTAL}.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."