#!/usr/bin/env python3
"""bench-llm-soak.py — sustained multi-family load, and what it reveals.

Purpose is not a headline number. It is to catch what single measurements miss:
a pinned replica that degrades, KV pressure that appears only after many large
prefills, a routing table that grows, a proxy failover storm. Each family keeps
its own large shared prefix (the agent-loop shape) and runs turn after turn for
the requested duration.

What it found on frank-medium (2026-09-16): with three concurrent families at
~30k tokens each, ~75 % of turns were cold re-prefills at a 13.9 s median,
because two families plus their cached prefixes exceed one card's 122,636-token
KV pool — a capacity wall that placement and large-request admission were both
measured against and neither moved. Zero requests failed: the cost of the wall
is latency, not errors.

Usage:
    bench-llm-soak.py <base-url> <minutes> <families> [salt]

The salt rides at the front of every prompt so a run never inherits a warm
prefix cache, and so the router's identity window sees a new conversation.
"""

from __future__ import annotations

import json
import statistics
import sys
import threading
import time
import urllib.request

BASE = sys.argv[1].rstrip("/") if len(sys.argv) > 1 else "http://127.0.0.1:7770"
MINUTES = float(sys.argv[2]) if len(sys.argv) > 2 else 10
FAMILIES = int(sys.argv[3]) if len(sys.argv) > 3 else 3
SALT = sys.argv[4] if len(sys.argv) > 4 else "SOAK"

DOCS = [
    "/home/dank/bench/bigdoc-a.txt",
    "/home/dank/bench/bigdoc-b.txt",
    "/home/dank/serverAppConfigs/scripts/deploy-remote.sh",
]

QUESTIONS = [
    "Which parts of this document would break if the host rebooted?",
    "Summarise the ordering guarantees this document relies on.",
    "Which commands here need root, and what fails without it?",
    "Point out the two riskiest assumptions in this document.",
    "What would you change first if you had ten minutes?",
]

deadline = time.time() + MINUTES * 60
lock = threading.Lock()
rows: list[tuple[str, float, bool]] = []


def ask(kind: str, prompt: str) -> None:
    body = {
        "model": "frank-medium",
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": 48,
        "stream": True,
    }
    req = urllib.request.Request(
        BASE + "/v1/chat/completions",
        data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json"},
    )
    start = time.time()
    ok = False
    try:
        with urllib.request.urlopen(req, timeout=900) as resp:
            for line in resp:
                if line.startswith(b"data: ") and b'"content"' in line:
                    ok = True
                    break
    except Exception as exc:
        print(f"    {kind} error: {type(exc).__name__}: {str(exc)[:70]}", flush=True)
    elapsed = (time.time() - start) * 1000.0
    with lock:
        rows.append((kind, elapsed, ok))
    print(f"  {kind}: {elapsed:9.1f} ms {'ok' if ok else 'FAIL'}", flush=True)


def family(index: int) -> None:
    try:
        with open(DOCS[index % len(DOCS)], errors="replace") as fh:
            doc = (fh.read() * 3)[:120_000]
    except OSError:
        # Fall back to any available source so the soak can run on a host that
        # has not had the big documents generated.
        with open(DOCS[-1], errors="replace") as fh:
            doc = (fh.read() * 3)[:120_000]
    history: list[dict] = []
    turn = 0
    while time.time() < deadline:
        question = QUESTIONS[turn % len(QUESTIONS)]
        ask(f"fam{index}", f"{SALT}-{index}.\n{doc}\n\nQuestion: {question}")
        history += [{"role": "assistant", "content": "Acknowledged."},
                    {"role": "user", "content": question}]
        turn += 1


def main() -> int:
    threads = [threading.Thread(target=family, args=(i,), daemon=True)
               for i in range(FAMILIES)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    with lock:
        snapshot = list(rows)
    ok = [r for r in snapshot if r[2]]
    print(f"\nSOAK salt={SALT} families={FAMILIES} minutes={MINUTES} "
          f"requests={len(snapshot)} ok={len(ok)} failed={len(snapshot) - len(ok)}")
    for fam in sorted({r[0] for r in snapshot}):
        times = [r[1] for r in ok if r[0] == fam]
        if not times:
            continue
        cold = [t for t in times if t >= 5000]
        print(f"  {fam}: n={len(times)} median={statistics.median(times):.0f} ms "
              f"re-prefills={len(cold)}/{len(times)}")
    return 0 if len(ok) == len(snapshot) else 1


if __name__ == "__main__":
    sys.exit(main())