#!/usr/bin/env python3
"""Warm decode tok/s + deep-prefill OOM check for the fv-ml1 LLM seats.

Complements scripts/seat-inventory.py: the inventory reads STATIC state
(placement, KV, concurrency, quant) with zero load; this applies LOAD to
measure warm decode throughput and to prove each seat survives a near-max-context
prefill without OOM. The numbers it prints are what docs/pfi/llm-seat-catalog.md
records, with the harness stated so they travel honestly.

    scripts/seat-bench.py                 # bench every generative seat, serially
    scripts/seat-bench.py --host fv-ml1

⚠ SERIAL BY DESIGN. Two deep prefills at once contend for GPU memory/compute and
would both confound the OOM result and depress tok/s. One seat at a time.

HARNESS (state it with any number this prints):
  warm decode = greedy, temperature 0, conc=1 (single stream), n=3 reps, median,
    a fixed ~40-word prompt generating 300 tokens (decode throughput; generation
    is never prefix-cached, so the reps are valid).
  deep prefill = one non-repeating random prompt at ~0.97x max-model-len, 8 output
    tokens; PASS = returns with no error AND the seat's allocator log shows no
    OOM / CUBLAS / illegal-memory across the probe window.

These are clean, uncontended, single-stream figures — a ceiling, not a loaded
number. Aggregate throughput under real concurrency is higher per-GPU and lower
per-request.
"""
import argparse, json, time, random, statistics, subprocess, urllib.request, urllib.error

# (label, served-model-name, port, max_model_len). Refresh from seat-inventory if seats change.
SEATS = [
    ("cyberprev (sec)", "cyberprev-27b", 8025, 262144),
    ("gen-small",       "gen-small",     8026, 262144),
    ("char-rp",         "char-rp",       8016, 262144),
    ("char-rp-fast",    "G4-MeroMero-26B-A4B-it-uncensored-heretic-NVFP4A16", 8021, 262144),
    ("gen (flash-next)","qwen3.8-flash-next-uncensored", 8022, 262144),
    ("coder",           "qwen2.5-coder-1.5b", 8020, 16384),
]
DECODE_PROMPT = ("Write a detailed technical explanation of how a modern CPU branch predictor "
                 "works, covering the pattern history table, the branch target buffer, and "
                 "misprediction cost.")

def run(host, cmd):
    r = subprocess.run(["ssh","-o","BatchMode=yes","-o","ConnectTimeout=10",
                        f"infra-ops@{host}", cmd], capture_output=True, text=True, timeout=1900)
    return r.stdout

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--host", default="10.251.50.54")
    a = ap.parse_args()
    # This driver runs the HTTP calls ON the host (loopback to each seat) via a pushed helper.
    helper = "/tmp/_seat_bench_http.py"
    open("/tmp/_seat_bench_http.py","w").write(_HTTP_DRIVER)
    subprocess.run(["scp","-o","BatchMode=yes","/tmp/_seat_bench_http.py",
                    f"infra-ops@{a.host}:{helper}"], check=True)
    print(run(a.host, f"python3 {helper}"))

_HTTP_DRIVER = r'''
import json,time,random,statistics,urllib.request,urllib.error
WORDS=[f"{random.Random(i).randint(0,1<<30):x}" for i in range(320000)]
SEATS=[("cyberprev (sec)","cyberprev-27b",8025,262144),
 ("gen-small","gen-small",8026,262144),
 ("char-rp","char-rp",8016,262144),
 ("char-rp-fast","G4-MeroMero-26B-A4B-it-uncensored-heretic-NVFP4A16",8021,262144),
 ("gen (flash-next)","qwen3.8-flash-next-uncensored",8022,262144),
 ("coder","qwen2.5-coder-1.5b",8020,16384)]
P=("Write a detailed technical explanation of how a modern CPU branch predictor works, "
   "covering the pattern history table, the branch target buffer, and misprediction cost.")
def call(port,model,prompt,mx,to=1800):
    b=json.dumps({"model":model,"prompt":prompt,"max_tokens":mx,"temperature":0}).encode()
    r=urllib.request.Request(f"http://127.0.0.1:{port}/v1/completions",data=b,headers={"Content-Type":"application/json"})
    t0=time.time()
    try: d=json.loads(urllib.request.urlopen(r,timeout=to).read())
    except urllib.error.HTTPError as e: return None,time.time()-t0,f"HTTP {e.code}: {e.read().decode()[:120]}"
    except Exception as e: return None,time.time()-t0,f"{type(e).__name__}: {str(e)[:100]}"
    return d.get("usage",{}),time.time()-t0,None
for name,model,port,mc in SEATS:
    print(f"\n===== {name} (:{port}, max {mc}) =====",flush=True)
    call(port,model,P,64,to=120)
    rates=[]
    for i in range(3):
        u,dt,err=call(port,model,P,300,to=180)
        if err: print(f"  rep{i+1} ERR {err}",flush=True); continue
        t=u.get("completion_tokens",0)
        if t: rates.append(t/dt); print(f"  rep{i+1}: {t}/{dt:.2f}s = {t/dt:.1f} tok/s",flush=True)
    if rates: print(f"  >> warm decode MEDIAN {statistics.median(rates):.1f} tok/s (n={len(rates)})",flush=True)
    nw=int(mc*0.97/7.9); r=random.Random(port*13+int(time.time())%997)
    dp=" ".join(r.choice(WORDS) for _ in range(nw))
    u,dt,err=call(port,model,dp,8,to=1800)
    print(("  >> DEEP PREFILL FAIL "+err) if err else f"  >> DEEP PREFILL: {u.get('prompt_tokens'):,} tok in {dt:.1f}s OK",flush=True)
print("\nBENCH DONE",flush=True)
'''

if __name__ == "__main__":
    main()