#!/usr/bin/env bash
# autoresearch-frank-primary.sh — RUNS ON DANKSERVER, driven over ssh by the
# dev-box autoresearch.sh. It is the canonical frank-primary benchmark entry
# point for the autoresearch loop.
#
# Method (mandated by the Benchmarking Policy): the production service serves
# real users and shares the two R9700s, so a bench beside it starves the GPUs
# and corrupts results. This script therefore:
#   1. stops the production frank-primary unit (root, relink symlink: stop,
#      NOT mask — mask refuses to replace a symlinked unit file),
#   2. boots a scratch radiance container named `arbench` on the SAME two
#      cards with the candidate engine argv (the env file the loop edits),
#   3. polls /health until the listener is live (it binds AFTER the full
#      load, so a bare curl is an honest gate),
#   4. runs ONE self-contained completion request locally (127.0.0.1, no LAN
#      jitter): a fixed ~140k-char real-text prompt (~32k tokens) decoded for
#      256 tokens, greedy. The engine's `timings` object reports the honest
#      prefill and decode rates for that single request.
#   5. prints METRIC lines and tears the scratch down,
#   6. restores production (systemctl start), via a trap so a failed bench
#      never leaves frank-primary down (the stale-probe-tail lesson).
#
# Determinism: the prompt is built from a fixed, sorted corpus of repo files
# with a fixed random seed, so prompt_n is stable across runs. Sampling is
# greedy (temperature 0), so the decode length is fixed; the only residual
# nondeterminism is the wht6 lossy TP all-reduce, which is a constant across
# runs and cancels in A/B comparisons.
#
# The scratch mirrors the production container contract exactly (serve-frank-
# primary.sh): --network host, --init, seccomp=unconfined (io_uring for the
# n-gram table), memlock unbounded (the ~132 GiB pinned ledger), /dev/kfd +
# /dev/dri, the .rad ro, the state dir rw, the tune cache mount. The ONLY
# differences: the container name (arbench, so prod's identity is untouched)
# and a FRESH prefix-cache disk tier (arbench/, wiped per run) so the cold
# prefill is cold — the prod tier at kvcache/frank-primary is left intact.
set -uo pipefail

ENV_FILE="${1:-/home/dank/serverAppConfigs/scripts/frank-primary.env}"
SCRATCH=arbench
PORT=7770
HEALTH="http://127.0.0.1:$PORT/health"
COMPLETE="http://127.0.0.1:$PORT/v1/completions"
BODY=/tmp/arbench-body.json
LOG=/home/dank/logs/arbench.log

restore() {
  podman rm -f "$SCRATCH" >/dev/null 2>&1 || true
  sudo systemctl start frank-primary >/dev/null 2>&1 || true
}
trap restore EXIT

# --- source the candidate config (same file the production unit reads) ------
if [ ! -f "$ENV_FILE" ]; then
  echo "no env file at $ENV_FILE" >&2; exit 1
fi
# shellcheck disable=SC1090
set -a; . "$ENV_FILE"; set +a

# --- stop production; the scratch needs both cards to itself ----------------
sudo systemctl stop frank-primary >/dev/null 2>&1 || true
podman rm -f "$SCRATCH" >/dev/null 2>&1 || true

# --- fresh prefix-cache tier for a cold prefill (prod's tier untouched) ------
FRESH_KV="$FRANK_PRIMARY_STATE_DIR/kvcache/$SCRATCH"
rm -rf "$FRESH_KV"; mkdir -p "$FRESH_KV"
mkdir -p "$FRANK_PRIMARY_STATE_DIR/tune"

# --- boot the scratch: the production run line, scratch name + fresh KV ------
podman run -d --name "$SCRATCH" \
  --network host \
  --init \
  --security-opt seccomp=unconfined \
  --ulimit memlock=-1:-1 \
  --device /dev/kfd --device /dev/dri \
  -v "$FRANK_PRIMARY_MODEL_DIR":/models:ro \
  -v "$FRANK_PRIMARY_STATE_DIR":/data \
  -v "$FRANK_PRIMARY_STATE_DIR/tune":/opt/radiance/share/radiance/tune \
  -v "$FRANK_PRIMARY_TEMPLATE":/templates/template.jinja:ro \
  "$FRANK_PRIMARY_IMAGE" \
  --model "/models/$FRANK_PRIMARY_MODEL_FILE" \
  --host :: --port "$PORT" \
  --served-model-name "$FRANK_PRIMARY_SERVED_NAME" \
  --override-chat-template /templates/template.jinja \
  --tp "$FRANK_PRIMARY_TP" --tp-wire "$FRANK_PRIMARY_TP_WIRE" \
  --max-num-seqs "$FRANK_PRIMARY_MAX_SEQS" \
  --max-model-len "$FRANK_PRIMARY_MAX_MODEL_LEN" \
  --placement "$FRANK_PRIMARY_PLACEMENT" \
  --host-pool-mib "$FRANK_PRIMARY_HOST_POOL_MIB" \
  --gpu-headroom-mib "$FRANK_PRIMARY_GPU_HEADROOM_MIB" \
  --vram-weights-mib "$FRANK_PRIMARY_VRAM_WEIGHTS_MIB" \
  --vram-kv-mib "$FRANK_PRIMARY_VRAM_KV_MIB" \
  --kv-cache-dtype "$FRANK_PRIMARY_KV_DTYPE" \
  --prefix-cache-host-mib "$FRANK_PRIMARY_PREFIX_HOST_MIB" \
  --prefix-cache-dir "/data/kvcache/$SCRATCH" \
  --prefix-cache-disk-mib "$FRANK_PRIMARY_PREFIX_DISK_MIB" \
  --ngram-placement "$FRANK_PRIMARY_NGRAM_PLACEMENT" \
  --num-speculative-tokens "$FRANK_PRIMARY_SPEC_TOKENS" \
  --max-num-batched-tokens "$FRANK_PRIMARY_MAX_BATCHED_TOKENS" \
  ${FRANK_PRIMARY_EXTRA_ARGS:+$FRANK_PRIMARY_EXTRA_ARGS} \
  > "$LOG" 2>&1
if [ $? -ne 0 ]; then echo "scratch boot failed; see $LOG" >&2; exit 1; fi

# --- poll health (listener binds after the full load; cold boot ~110 s) ------
ok=0
for _ in $(seq 1 360); do
  code=$(curl -s -o /dev/null -w "%{http_code}" --max-time 3 "$HEALTH" 2>/dev/null || echo 000)
  if [ "$code" = "200" ]; then ok=1; break; fi
  sleep 2
done
if [ "$ok" -ne 1 ]; then
  echo "scratch never reached health 200; tail of $LOG:" >&2
  tail -30 "$LOG" >&2; exit 1
fi

# --- build the workload: a fixed ~140k-char real-text prompt (repo corpus) ---
python3 - "$BODY" <<'PY'
import json, sys, random, glob
random.seed(1234)
root = "/home/dank/serverAppConfigs"
files = sorted(
    glob.glob(f"{root}/serverAppConfigs-docs/**/*.md", recursive=True)
    + glob.glob(f"{root}/scripts/*.sh")
    + glob.glob(f"{root}/scripts/*.py")
    + glob.glob(f"{root}/AGENTS.md")
)
corpus = []
total = 0
for f in files:
    try:
        t = open(f, encoding="utf-8", errors="ignore").read()
    except Exception:
        continue
    if not t.strip():
        continue
    corpus.append(t)
    total += len(t)
    if total > 240000:
        break
corpus = "\n".join(corpus)
TGT = 140000
parts, n = [], 0
while n < TGT:
    start = random.randint(0, max(0, len(corpus) - 500))
    seg = corpus[start:start + random.randint(200, 500)]
    parts.append(seg)
    n += len(seg) + 1
prompt = " ".join(parts)
json.dump(
    {"model": "frank-archon", "prompt": prompt,
     "max_tokens": 256, "temperature": 0.0, "stream": False,
     "stop": ["<|im_end|>"]},
    open(sys.argv[1], "w"),
)
PY
if [ ! -f "$BODY" ]; then echo "prompt build failed" >&2; exit 1; fi

# --- the bench request (prefill ~32k + decode 256, greedy) -------------------
RESP=$(curl -s --max-time 600 -X POST "$COMPLETE" \
  -H 'Content-Type: application/json' -d "@$BODY")
echo "$RESP" > /tmp/arbench-resp.json

# --- parse the timings into METRIC lines -------------------------------------
python3 - <<'PY'
import json, sys
try:
    r = json.load(open("/tmp/arbench-resp.json"))
except Exception as e:
    sys.stderr.write(f"unparseable response: {e}\n"); sys.exit(1)
if "timings" not in r:
    sys.stderr.write(f"no timings in response: {json.dumps(r)[:500]}\n"); sys.exit(1)
t = r["timings"]
dtps = t.get("predicted_per_second")
ptps = t.get("prompt_per_second")
pms  = t.get("prompt_ms")
dn   = t.get("draft_n", 0)
da   = t.get("draft_n_accepted", 0)
if dtps is None:
    sys.stderr.write("predicted_per_second null (no decode produced)\n"); sys.exit(1)
print(f"METRIC decode_tps={dtps:.4f}")
if ptps is not None:
    print(f"METRIC prefill_tps={ptps:.4f}")
if pms is not None:
    print(f"METRIC ttft_ms={pms:.4f}")
if dn:
    print(f"METRIC draft_accept={da/dn:.4f}")
print(f"INFO prompt_n={t.get('prompt_n')} cache_n={t.get('cache_n')} predicted_n={t.get('predicted_n')}")
PY