#!/usr/bin/env bash
# arm3_window.sh — ONE bench arm of the Jev bench's arm 3, start to finish, on inferencebox.
#
# WHAT THIS IS, for a reader who opened it cold. Arm 3 runs inside a 45-minute planned
# maintenance window on a LIVE box (README §A3.8). Everything that can be decided before
# the window is decided here rather than typed at 04:10Z: which snapshot, which flags,
# what a refusal is, when to give up on a server, and what the receipt line says. The
# operator's paste is one line per arm:
#
#     ssh inferencebox '~/bench-arm3/arm3_window.sh bf16'
#     ssh inferencebox '~/bench-arm3/arm3_window.sh fp8'
#
# It starts a vLLM server on 127.0.0.1:8000, waits for it, records the kernel line the
# pre-registration requires, runs the arm's three tasks in the pre-registered order
# (c first — the gated one — then a, then b), stops the server, and prints ONE receipt
# line the operator can paste back.
#
# IT NEVER TOUCHES A SEAT. No systemctl, no sudo, nothing outside ~/bench-arm3. The
# stops and the starts are the runbook's own numbered steps, in the operator's hands.
#
# IT REFUSES RATHER THAN IMPROVISES. Every exit below 0 names the exact output that
# produced it (README §A3.9). A refusal inside the window is a result; a workaround
# invented at 04:20Z is not.
set -uo pipefail

ARM_KIND="${1:-}"
LARGECARD=GPU-25bc3288-319a-790c-a14e-82acdb8d15b8
HOME_DIR="${HOME}"
BENCH="${HOME_DIR}/bench-arm3"
TREE="${BENCH}/jev"
VENV="${HOME_DIR}/bench-vllm-venv/bin"
PORT=8000
#: The reranker seat stays resident on largecard by the operator's ruling (README §A3.3.4).
#: This is the ONLY process the pre-flight tolerates on that board, and the ceiling is its
#: measured 954 MiB with room for its own allocator, not a round number picked to pass.
PREFLIGHT_CEILING_MIB=2000

case "$ARM_KIND" in
  bf16) REPO=models--openjev--openjev      ; REV=5ec9e5fd2f80a6fff386779b1e5ac7e389971889
        ARM=openjev-bf16-largecard           ; EXTRA=(--dtype bfloat16)
        START_DEADLINE=540 ;;   # 9 min — 54.73 GB of weights plus a cold compile
  fp8)  REPO=models--openjev--openjev-FP8  ; REV=4ec320f267401e67c9be04d5df1be4d2b6b64f10
        ARM=openjev-fp8-largecard            ; EXTRA=(--quantization fp8)
        START_DEADLINE=420 ;;   # 7 min — 30.41 GB, the same compile
  *)    echo "usage: arm3_window.sh bf16|fp8" >&2; exit 64 ;;
esac

SNAP="${HOME_DIR}/hf-cache/hub/${REPO}/snapshots/${REV}"
LOG="${TREE}/receipts/vllm-serve-arm3-${ARM_KIND}.log"
RUNLOG="${TREE}/receipts/run-arm3-${ARM_KIND}.log"
stamp() { date -u +%Y-%m-%dT%H:%M:%SZ; }
say() { echo "[$(stamp)] $*"; }
mib_on_largecard() { nvidia-smi -i "$LARGECARD" --query-gpu=memory.used --format=csv,noheader,nounits; }

mkdir -p "${TREE}/receipts" "${TREE}/rows"
say "ARM3 ${ARM} starting — box inferencebox, card largecard ${LARGECARD:0:16}, one card, no link"

# ── REFUSAL 6 of README §A3.9, checked before anything else: the snapshot must be whole ──
[ -d "$SNAP" ] || { say "REFUSE: no snapshot at $SNAP"; exit 65; }
SHARDS=$(ls -L "$SNAP"/model-*.safetensors 2>/dev/null | wc -l)
[ "$SHARDS" -eq 12 ] || { say "REFUSE: $SHARDS safetensors shards at $SNAP, expected 12"; exit 65; }
say "snapshot ok: $SNAP ($SHARDS shards, rev ${REV:0:8})"

# ── REFUSAL 1: the free-VRAM pre-flight. Only the reranker may hold largecard. ─────────────
APPS=$(nvidia-smi --query-compute-apps=gpu_uuid,pid,process_name,used_memory --format=csv,noheader \
       | grep -F "$LARGECARD" | grep -v "rerank-env" || true)
USED=$(mib_on_largecard)
if [ -n "$APPS" ]; then
  say "REFUSE: largecard is held by something that is not the reranker:"; echo "$APPS"; exit 66
fi
if [ "$USED" -gt "$PREFLIGHT_CEILING_MIB" ]; then
  say "REFUSE: largecard reads ${USED} MiB used, ceiling ${PREFLIGHT_CEILING_MIB} MiB"; exit 66
fi
say "pre-flight ok: largecard ${USED} MiB used, no compute app but the reranker"

# ── the server ───────────────────────────────────────────────────────────────────────────
export HF_HOME="${HOME_DIR}/hf-cache" HF_HUB_OFFLINE=1
export CUDA_HOME=/usr/local/cuda PATH="/usr/local/cuda/bin:${PATH}"
export CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES="$LARGECARD"
#: Set deliberately, not for want of a toolchain: inferencebox HAS nvcc 13.1 and gcc 15.2. It keeps
#: the sampling path byte-identical to arm 2 and keeps a cold flashinfer JIT out of a
#: 45-minute window. It cannot move a number — every arm is greedy at temperature 0 and the
#: readout reads raw logprobs (README §A3.3.5).
export VLLM_USE_FLASHINFER_SAMPLER=0

say "starting vLLM 0.29.0 on 127.0.0.1:${PORT} — log ${LOG}"
nohup "${VENV}/vllm" serve "$SNAP" --served-model-name "$ARM" \
  --host 127.0.0.1 --port "$PORT" --tensor-parallel-size 1 "${EXTRA[@]}" \
  --max-model-len 16384 --gpu-memory-utilization 0.90 --enable-prefix-caching \
  --max-num-seqs 1 --max-logprobs 64 --limit-mm-per-prompt '{"image":0}' \
  --trust-remote-code --gdn-prefill-backend triton \
  > "$LOG" 2>&1 < /dev/null &
SERVER_PID=$!

stop_server() {
  [ -n "${SERVER_PID:-}" ] || return 0
  say "stopping the bench server (pid ${SERVER_PID})"
  kill "$SERVER_PID" 2>/dev/null
  for _ in $(seq 1 60); do kill -0 "$SERVER_PID" 2>/dev/null || break; sleep 1; done
  kill -0 "$SERVER_PID" 2>/dev/null && { say "server did not exit on TERM — KILL"; kill -9 "$SERVER_PID"; }
  for _ in $(seq 1 30); do [ "$(mib_on_largecard)" -le "$PREFLIGHT_CEILING_MIB" ] && break; sleep 2; done
  say "largecard after the stop: $(mib_on_largecard) MiB used"
}
trap stop_server EXIT

# ── REFUSAL 2: the server answers inside its slot, or the arm gives the window back ──────
BOOT_START=$(date +%s)
READY=0
for _ in $(seq 1 "$START_DEADLINE"); do
  kill -0 "$SERVER_PID" 2>/dev/null || { say "REFUSE: the server died during startup — tail:"; tail -40 "$LOG"; exit 67; }
  if curl -sf --max-time 3 "http://127.0.0.1:${PORT}/v1/models" > /dev/null 2>&1; then READY=1; break; fi
  sleep 1
done
BOOT_S=$(( $(date +%s) - BOOT_START ))
[ "$READY" -eq 1 ] || { say "REFUSE: no /v1/models after ${START_DEADLINE}s — tail:"; tail -40 "$LOG"; exit 67; }
say "server answered /v1/models after ${BOOT_S}s"

# ── THE KERNEL LINE — a required receipt of the pre-registration (README §A3.4). ─────────
#: Arm 2 on sm86 read `Selected MarlinFP8ScaledMMLinearKernel` — weight-only, FP8's memory
#: without FP8's arithmetic. On sm_120 this should read a native block-scaled FP8 kernel, and
#: a Marlin line HERE is a finding rather than a detail. Copied verbatim, never paraphrased.
say "---- the kernel and backend lines, verbatim ----"
grep -aE "Selected .*Kernel|attention backend|GDN (prefill|decode)|Using Flash|quantization|dtype" "$LOG" \
  | sed 's/^/    /' | head -20
say "---- end kernel lines ----"

# ── the arm ──────────────────────────────────────────────────────────────────────────────
#: c FIRST: it is the only gated task, so a window that overruns loses the tail and not the
#: verdict. a before b keeps arm 2's adjacency and its prefix-cache warmth. b is droppable
#: and is declared skipped rather than hurried (README §A3.7).
say "running ${ARM}: tasks c, a, b — log ${RUNLOG}"
( cd "$TREE" && "${VENV}/python" run.py --arm "$ARM" --task c --task a --task b --out rows ) \
  > "$RUNLOG" 2>&1
RC=$?
tail -5 "$RUNLOG"

C_REPORT="${TREE}/rows/${ARM}.c.report.json"
ACC=$("${VENV}/python" - "$C_REPORT" <<'PY' 2>/dev/null || echo "unreadable"
import json,sys
try:
    s=json.load(open(sys.argv[1]))["summary"]
    print("%.1f%% (%d/%d)"%(100*s["accuracy"],s["correct"],s["n_labelled"]))
except Exception:
    print("unreadable")
PY
)
say "RECEIPT  arm=${ARM}  rc=${RC}  boot=${BOOT_S}s  task_c=${ACC}  gate_floor=91.5%  rows=${TREE}/rows"
exit "$RC"