calib-push: drive calib.jsonl through scratch prefill until the tap writes Hessians

scripts/calib-push.py

@@ -0,0 +1,110 @@
+#!/usr/bin/env python3
+"""calib-push.py — drive a calibration corpus through a radiance scratch
+server's prefill so its tap writes layer Hessians (docs/MOE-W4.md).
+
+Runs ON DANKSERVER next to the scratch server (the one autoresearch- /
+serve-style scratch container booted with RADIANCE_CALIB_DIR set).
+Reads the corpus JSONL produced by make-calib-corpus.py
+({"prompt": "..."} a line), POSTs each prompt to /v1/completions with
+max_tokens=1 (prefill-only cost, one token out), at N-way concurrency
+matching the scratch server's --max-num-seqs budget.
+
+Corpus rounds repeat until the engine prints the tap's completion line
+("calibration wrote layer Hessians") in `podman logs <container>` — that
+is the ground truth that RADIANCE_CALIB_TOKENS was reached. After that
+line, further prompts are pure waste, so the driver exits.
+
+Mechanics notes:
+- requests carry no sampling params that matter (max_tokens=1, greedy).
+- a per-request timeout of 30 min: at ~1k tok/s prefill a 12k-token
+  prompt is a few seconds; the timeout only catches a wedged engine.
+- failures print and continue: the tap pools over tokens, a few lost
+  prompts do not bias the Hessians (the corpus is drawn round-robin over
+  source classes, so loss is class-neutral as long as failures are rare;
+  a systematic 4xx is loud in the log and must be fixed, not ignored).
+
+Usage: calib-push.py [--port 7770] [--parallel 64] [--container arbench]
+       [--corpus /home/dank/corpora/out/calib.jsonl] [--max-rounds 4]
+"""
+import argparse
+import concurrent.futures
+import json
+import subprocess
+import sys
+import time
+import urllib.request
+
+DONE_MARK = "calibration wrote"
+
+
+def tap_done(container):
+    """True if the engine logged the tap's completion line."""
+    try:
+        out = subprocess.run(
+            ["podman", "logs", "--tail", "50", container],
+            capture_output=True, text=True, timeout=30,
+        )
+    except Exception:
+        return False
+    return DONE_MARK in (out.stdout + out.stderr)
+
+
+def push_one(url, prompt):
+    body = json.dumps({
+        "model": "frank-primary",
+        "prompt": prompt,
+        "max_tokens": 1,
+        "temperature": 0,
+        "stream": False,
+    }).encode()
+    req = urllib.request.Request(
+        url, data=body, headers={"Content-Type": "application/json"})
+    try:
+        with urllib.request.urlopen(req, timeout=1800) as r:
+            r.read()
+        return True
+    except Exception as e:
+        print(f"FAIL {str(e)[:120]}", flush=True)
+        return False
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--port", type=int, default=7770)
+    ap.add_argument("--parallel", type=int, default=64)
+    ap.add_argument("--container", default="arbench")
+    ap.add_argument("--corpus", default="/home/dank/corpora/out/calib.jsonl")
+    ap.add_argument("--max-rounds", type=int, default=4)
+    a = ap.parse_args()
+    url = f"http://127.0.0.1:{a.port}/v1/completions"
+
+    prompts = []
+    with open(a.corpus) as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                prompts.append(json.loads(line)["prompt"])
+    print(f"corpus: {len(prompts)} prompts", flush=True)
+
+    for rnd in range(1, a.max_rounds + 1):
+        if tap_done(a.container):
+            break
+        t0 = time.time()
+        ok = 0
+        with concurrent.futures.ThreadPoolExecutor(a.parallel) as ex:
+            for good in ex.map(lambda p: push_one(url, p), prompts):
+                ok += 1 if good else 0
+        dt = time.time() - t0
+        print(f"round {rnd}: {ok}/{len(prompts)} ok in {dt:.0f}s "
+              f"({sum(len(p) for p in prompts)/4/dt:.0f} chars/s est)",
+              flush=True)
+        if tap_done(a.container):
+            break
+    else:
+        print("max rounds exhausted, tap not done", flush=True)
+        sys.exit(1)
+    print("TAP_DONE: calibration wrote layer Hessians", flush=True)
+
+
+if __name__ == "__main__":
+    main()