#!/bin/bash
# Full acceptance under one gpu-lock window. Boots the candidate config,
# runs: 391 check, validate.py (quality dump), compare vs r2 + base,
# short-ctx decode bench, 3x >=500k-token rounds (longctx.py), teardown.
# Usage: accept-512k.sh [MAXLEN] [IMAGE] [EXTRA vllm args...]   (env: SEQS BATCHED UTIL MM)
set -u
D="$(cd "$(dirname "$0")" && pwd)"
LOG=/tmp/glm-512k-accept.log
MODEL=glm-5.3-flash-fp8attn
PORT=18996

docker rm -f glm-512k-fit >/dev/null 2>&1
# teardown race guard: another worker's dying container may still hold VRAM
# for a few seconds after the lock frees; wait for both cards to drain.
for i in $(seq 1 60); do
  MAXUSED=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | sort -n | tail -1)
  [ "$MAXUSED" -lt 200 ] && break
  sleep 3
done
echo "GPU drain: max used ${MAXUSED} MiB before boot"
setsid "$D/boot-512k-fit.sh" "$@" >"$LOG" 2>&1 &
BOOT=$!
READY=0
for i in $(seq 1 240); do
  sleep 5
  if curl -sf -m 2 http://127.0.0.1:$PORT/health >/dev/null 2>&1; then READY=1; break; fi
  kill -0 "$BOOT" 2>/dev/null || break
done
echo "=== LEDGER ==="
grep -aiE "consumed|non-torch|activation|KV cache size|kv cache memory|Maximum concurrency|graph|needs|available|ERROR|OOM" "$LOG" | tail -25
if [ "$READY" != "1" ]; then echo "BOOT FAILED"; docker stop glm-512k-fit >/dev/null 2>&1; exit 2; fi

# VRAM watermark sampler (min free MiB seen during the run)
( MINFREE=999999; while docker ps --format '{{.Names}}' | grep -q glm-512k-fit; do
    F=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits | sort -n | head -1)
    [ -n "$F" ] && [ "$F" -lt "$MINFREE" ] && MINFREE=$F && echo "$MINFREE" > /tmp/glm-512k-minfree
    sleep 4
  done ) &
WATCH=$!

echo "=== 391 ==="
curl -s -m 300 http://127.0.0.1:$PORT/v1/chat/completions -H 'Content-Type: application/json' \
  -d "{\"model\":\"$MODEL\",\"messages\":[{\"role\":\"user\",\"content\":\"What is 17*23? Reply with just the number.\"}],\"max_tokens\":200,\"temperature\":0}" \
  | python3 -c 'import json,sys; r=json.load(sys.stdin); m=r["choices"][0]["message"]; print((m.get("content") or "")[:100], "| reasoning:", (m.get("reasoning_content") or "")[-80:])'

echo "=== validate.py ==="
python3 "$D/validate.py" $PORT $MODEL "$D/results-512k.json" || echo "validate FAILED"

echo "=== compare vs r2 ==="
python3 "$D/compare.py" "$D/results-fp8attn-r2.json" "$D/results-512k.json" || true
echo "=== compare vs base ==="
python3 "$D/compare.py" "$D/results-base.json" "$D/results-512k.json" || true

echo "=== decode short-ctx ==="
python3 "$D/bench_decode.py" $PORT $MODEL 512 || true

echo "=== longctx 3 rounds ==="
python3 "$D/longctx.py" $PORT $MODEL 3 505000
LC=$?

echo "=== post-run VRAM ==="
nvidia-smi --query-gpu=index,memory.used --format=csv,noheader
kill "$WATCH" 2>/dev/null
echo "=== min free watermark (MiB): $(cat /tmp/glm-512k-minfree 2>/dev/null || echo n/a) ==="

if [ "${KEEP:-0}" != "1" ]; then docker stop glm-512k-fit >/dev/null 2>&1; wait "$BOOT" 2>/dev/null; fi
echo "=== ACCEPT EXIT: longctx=$LC ==="
exit $LC