#!/usr/bin/env python3
"""bench-spec-ab.py — settle "which proposer should this endpoint use?" with data.

Runs one payload through several speculative-decoding configurations of the same
binary, one at a time, with the unit stopped between runs: on a single-GPU host
there is no room for two copies of a 29 GiB model, and the benchmarking policy
applies. Every configuration gets identical prompts at temperature 0, and the
greedy text of a fixed prompt is compared across them — speculation may buy
speed, never change the answer.

    bench-spec-ab.py <env-file> <serve-script> [--base-url URL]
    bench-spec-ab.py /home/dank/serverAppConfigs/scripts/frank-medium.env \
                     /home/dank/serverAppConfigs/scripts/frank-medium.sh

The configurations are the ones that exist for THIS payload (frank-medium,
2026-09-16): the DFlash2 drafter, the model's own MTP head at two depths, and
autoregressive. The systemd unit is restored at the end, so the endpoint comes
back on whatever the env file says whatever happens here.

Result on this box (1k prompt / 512 out, temperature 0):

    none            17.7 tok/s                     card1 26.07 GB
    draft-mtp  n=3  32.4 tok/s  58 % accepted       card1 29.35 GB
    draft-mtp  n=7  26.7 tok/s  33 % accepted       card1 29.65 GB
    draft-dflash 7  43.5 tok/s  41 % accepted       card1 29.42 GB

...and `draft-mtp,draft-dflash` fails at load: a chain has ONE draft context and
MTP builds against whatever model that context holds.
"""
import argparse
import json
import re
import subprocess
import sys
import time
import urllib.request

FILLER = ("The quick brown fox jumps over the lazy dog near the river bank. "
          "A cache is a component that stores data so future requests are faster. ")
TOKENS_PER_REP = 28
UNIT = "frank-medium.service"
MODEL = "frank-medium"
SPEC_GREP = "adding speculative implementation|- n_max=|- block_size=|creating MTP"

CONFIGS = [
    ("none", {"SPEC_TYPE": "none"}),
    ("mtp-n3", {"SPEC_TYPE": "draft-mtp", "SPEC_DRAFT_N_MAX": "3"}),
    ("mtp-n7", {"SPEC_TYPE": "draft-mtp", "SPEC_DRAFT_N_MAX": "7"}),
    ("dflash-n7", {"SPEC_TYPE": "draft-dflash", "SPEC_DRAFT_N_MAX": "7"}),
]


def log(msg, logfile):
    line = "[" + time.strftime("%H:%M:%S") + "] " + str(msg)
    print(line, flush=True)
    with open(logfile, "a") as fh:
        fh.write(line + "\n")


def sh(cmd):
    return subprocess.run(cmd, shell=True, capture_output=True, text=True).stdout.strip()


def write_env(src, name, overrides):
    path = "/tmp/spec-ab-" + name + ".env"
    with open(src) as fh:
        lines = fh.read().splitlines()
    out = []
    for line in lines:
        for k, v in overrides.items():
            if line.startswith(k + "="):
                line = k + "=" + v
                break
        out.append(line)
    with open(path, "w") as fh:
        fh.write("\n".join(out) + "\n")
    return path


def post(base, body, timeout=1200):
    req = urllib.request.Request(
        base + "/v1/chat/completions", data=json.dumps(body).encode(),
        headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.load(r)


def vram():
    out = sh("rocm-smi --showmeminfo vram 2>/dev/null | grep 'Used Memory'")
    return [round(int(x) / 1e9, 2) for x in re.findall(r":\s+(\d+)\s*$", out, re.M)][:2]


def wait_health(base, timeout=420):
    t0 = time.time()
    while time.time() - t0 < timeout:
        if not sh("pgrep -x llama-server >/dev/null && echo up"):
            return False
        try:
            with urllib.request.urlopen(base + "/health", timeout=5) as r:
                if r.status == 200 and json.load(r).get("status") == "ok":
                    return True
        except Exception:
            time.sleep(3)
    return False


def stop_unit():
    sh("sudo systemctl stop " + UNIT)
    sh("pkill -x llama-server")
    for _ in range(40):
        if not sh("pgrep -x llama-server || true"):
            return
        time.sleep(1)


def shape(base, prompt_tokens, max_tokens):
    t = post(base, {"model": MODEL,
                    "messages": [{"role": "user",
                                  "content": FILLER * (prompt_tokens // TOKENS_PER_REP) +
                                             "\n\nReply with a single short sentence."}],
                    "max_tokens": max_tokens, "temperature": 0.0})["timings"]
    return {"shape": "decode " + str(prompt_tokens) + "/" + str(max_tokens),
            "prefill_tps": round(1000 * t["prompt_n"] / t["prompt_ms"], 1),
            "predicted_n": t["predicted_n"],
            "decode_tps": round(t["predicted_per_second"], 2),
            "draft_n": t.get("draft_n", 0),
            "accepted": t.get("draft_n_accepted", 0)}


def measure(base):
    shape(base, 4000, 32)                                   # warmup
    rows = [shape(base, 1000, 512), shape(base, 16000, 16)]  # decode, prefill
    d = post(base, {"model": MODEL,
                    "messages": [{"role": "user",
                                  "content": "List the first five prime numbers, comma separated, then stop."}],
                    "max_tokens": 192, "temperature": 0.0})
    msg = d["choices"][0]["message"]
    return rows, (msg.get("content") or "").strip()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("env_file")
    ap.add_argument("serve_script")
    ap.add_argument("--base-url", default="http://127.0.0.1:7770")
    ap.add_argument("--log", default="/tmp/bench-spec-ab.log")
    args = ap.parse_args()

    base, logfile = args.base_url, args.log
    results = {}
    log("spec A/B start (unit " + UNIT + " stopped for the duration)", logfile)
    stop_unit()
    for name, over in CONFIGS:
        env_path = write_env(args.env_file, name, over)
        logf = "/tmp/spec-ab-" + name + ".log"
        sh("cd " + sh("dirname " + args.serve_script) + " && FRANK_MEDIUM_ENV_FILE=" + env_path +
           " setsid bash " + args.serve_script + " > " + logf + " 2>&1 < /dev/null &")
        if not wait_health(base):
            log(name + ": FAILED to become healthy -- " + sh("tail -2 " + logf), logfile)
            results[name] = {"error": "not healthy", "tail": sh("tail -4 " + logf)}
            stop_unit()
            continue
        rows, text = measure(base)
        results[name] = {"vram_gb": vram(), "rows": rows, "answer": text[:240],
                         "spec_log": sh("grep -E '" + SPEC_GREP + "' " + logf + " | head -4")}
        log(name + ": " + " | ".join(
            r["shape"] + " prefill " + str(r["prefill_tps"]) + " decode " + str(r["decode_tps"]) +
            " draft " + str(r["draft_n"]) + "/" + str(r["accepted"]) for r in rows) +
            " | vram " + str(vram()) + " | answer=" + repr(text[:40]), logfile)
        stop_unit()

    with open("/tmp/bench-spec-ab.json", "w") as fh:
        json.dump(results, fh, indent=2)
    answers = {v.get("answer") for v in results.values() if v.get("answer")}
    if len(answers) > 1:
        log("WARNING: configurations disagree on greedy text -- speculation is "
            "changing answers: " + repr(answers), logfile)
    log("restoring " + UNIT, logfile)
    sh("sudo systemctl start " + UNIT)
    for _ in range(120):
        try:
            with urllib.request.urlopen(base + "/health", timeout=5) as r:
                if json.load(r).get("status") == "ok":
                    break
        except Exception:
            time.sleep(5)
    log("done; results in /tmp/bench-spec-ab.json", logfile)
    return 0


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