#!/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 is compared across them — speculation may buy speed, never change
the answer.

    bench-spec-ab.py <env-file> <serve-script> [options]
    bench-spec-ab.py /home/dank/serverAppConfigs/scripts/frank-medium.env \
                     /home/dank/serverAppConfigs/scripts/serve-frank-medium.sh \
                     --configs "dflash=draft-dflash@7;ngram=ngram-mod@64;chain=draft-dflash,ngram-mod@7"

Config syntax: `name=SPEC_TYPE[@n_max][|ENV=VAL|ENV=VAL]`, entries joined by `;`.
The systemd unit is restored in a finally block, so it comes back on whatever the
env file says even if a configuration crashes the harness.

THREE SHAPES, because a single one answers the wrong question. A proposer can
only draft text it has seen, so the shape decides the verdict:

  prose       a 1k prompt asking for one sentence. Nothing to copy and nothing
              repetitive, so only a model-backed proposer can help: this is the
              shape that compares drafters.
  echo        a ~1.5k passage of unique lines, "repeat it back verbatim". The
              output IS the prompt, which is ngram-mod's home turf.
  structured  "25 JSON records, one per line" — repetitive OUTPUT with no prompt
              overlap, which is what an ngram matcher is really for in an
              agentic fleet (config edits, tables, record lists).

echo and structured disable thinking (`chat_template_kwargs`), otherwise the
model thinks the answer into `reasoning_content`, `content` comes back empty and
the cross-configuration identity check passes vacuously — which is exactly how
this harness first "verified" an echo it never looked at.

Results on Dankserver (frank-medium payload, 2026-09-16, temperature 0). Decode:

  shape            dflash=7   ngram-mod(48/64)   chain        (tok/s)
  prose 1k/512        43.8        17.7 (0 drafts)   43.7
  echo 1.5k/400       76.8       102.5             200.6
  structured/400      91.1        17.8 (0 drafts)   91.2
  identity            70.4        17.8 (0 drafts)   70.4

`draft-mtp,draft-dflash` fails at load (a chain has ONE draft context and MTP
builds against whatever model that context holds); ngram-mod has no such
constraint — it is draftless and joins the chain unconditionally.
"""
import argparse
import json
import os
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|"
             "draft acceptance")

DEFAULT_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"}),
]

WORDS = ["latency", "throughput", "cache", "adapter", "scheduler", "allocator",
         "queue", "kernel", "buffer", "shard", "replica", "window"]


def echo_passage(lines=90):
    """Unique lines: no long span repeats inside the prompt, so an n-gram can
    only fire when the OUTPUT copies the INPUT."""
    return "\n".join(
        "Record %03d: the %s of subsystem %d reports state %d with checksum %05d."
        % (i, WORDS[i % len(WORDS)], i % 7, i % 3, (i * 7919) % 100000)
        for i in range(lines))


ECHO = (echo_passage() + "\n\nRepeat the passage above verbatim, every line, exactly as "
        "written, with no commentary before or after it.")
STRUCTURED = ("Output exactly 25 JSON records, one per line, no prose and no code fence. "
              'Each line must look like {"id": N, "state": "active", "region": "us-west-2"} '
              "with N running from 1 to 25.")
PROSE = FILLER * (1000 // TOKENS_PER_REP) + "\n\nReply with a single short sentence."
IDENTITY = "List the first five prime numbers, comma separated, then stop."
NO_THINK = {"chat_template_kwargs": {"enable_thinking": False}}


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 parse_configs(spec):
    """`name=SPEC_TYPE[@n_max][|KEY=VALUE]...`, entries joined by ';'.

    KEY=VALUE parts override any env file key, which is how an impl's own window
    (SPEC_NGRAM_N_MIN/MAX, for instance) gets swept without hand-editing env
    files."""
    out = []
    for chunk in spec.split(";"):
        chunk = chunk.strip()
        if not chunk:
            continue
        name, _, rest = chunk.partition("=")
        over = {}
        if "|" in rest:
            rest, _, extras = rest.partition("|")
            for item in extras.split("|"):
                if "=" in item:
                    k, _, v = item.partition("=")
                    over[k.strip()] = v.strip()
        n_max = None
        if "@" in rest:
            rest, _, n_max = rest.partition("@")
        over["SPEC_TYPE"] = rest.strip()
        if n_max:
            over["SPEC_DRAFT_N_MAX"] = n_max.strip()
        out.append((name.strip(), over))
    return out


def write_env(src, name, overrides):
    """The env file with this configuration's overrides applied."""
    path = "/tmp/spec-ab-" + name.replace("/", "_") + ".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=1800):
    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 restore_unit(base, logfile):
    """Bring the endpoint back on the deployed configuration. Runs in a
    finally, because a harness that can leave production down is a hazard."""
    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":
                    log(UNIT + " healthy again", logfile)
                    return True
        except Exception:
            time.sleep(5)
    log("WARNING: " + UNIT + " did not come back healthy -- check it by hand", logfile)
    return False


def run_shape(base, prompt, max_tokens, label, extra=None):
    body = {"model": MODEL, "messages": [{"role": "user", "content": prompt}],
            "max_tokens": max_tokens, "temperature": 0.0}
    if extra:
        body.update(extra)
    d = post(base, body)
    tm = d["timings"]
    msg = d["choices"][0]["message"]
    content = (msg.get("content") or "").strip()
    row = {"shape": label,
           "prompt_n": tm["prompt_n"],
           "prefill_tps": round(1000 * tm["prompt_n"] / tm["prompt_ms"], 1),
           "predicted_n": tm["predicted_n"],
           "decode_tps": round(tm["predicted_per_second"], 2),
           "draft_n": tm.get("draft_n", 0),
           "accepted": tm.get("draft_n_accepted", 0),
           "answer": content,
           "reasoning_len": len(msg.get("reasoning_content") or "")}
    if label.startswith("echo"):
        # How much of the answer really is the passage: a configuration that
        # "wins" by answering something else is not a win.
        want = {l.strip() for l in echo_passage().splitlines()}
        got = [l.strip() for l in content.splitlines() if l.strip()]
        row["copy_ratio"] = round(sum(1 for l in got if l in want) / len(got), 3) if got else 0.0
    return row


def measure(base):
    run_shape(base, FILLER * (4000 // TOKENS_PER_REP) +
              "\n\nReply with a single short sentence.", 32, "warmup")
    return [
        run_shape(base, PROSE, 512, "prose 1k/512"),
        run_shape(base, ECHO, 400, "echo 1.5k/400", NO_THINK),
        run_shape(base, STRUCTURED, 400, "structured/400", NO_THINK),
        run_shape(base, IDENTITY, 192, "identity"),
    ]


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")
    ap.add_argument("--configs", default=None,
                    help="name=SPEC_TYPE[@n_max][|ENV=VAL] joined by ';' "
                         "(default: the four configs measured on 2026-09-16)")
    ap.add_argument("--json", default="/tmp/bench-spec-ab.json")
    args = ap.parse_args()

    configs = parse_configs(args.configs) if args.configs else DEFAULT_CONFIGS
    # Absolute paths: the launcher must not depend on the caller's cwd, and a
    # relative path plus a `cd` to its own directory double-counts the prefix.
    env_file = os.path.abspath(args.env_file)
    serve_script = os.path.abspath(args.serve_script)
    base, logfile = args.base_url, args.log
    results = {}
    log("spec A/B start: " + ", ".join(n for n, _ in configs) +
        " (unit " + UNIT + " stopped for the duration)", logfile)
    try:
        stop_unit()
        for name, over in configs:
            env_path = write_env(env_file, name, over)
            logf = "/tmp/spec-ab-" + name.replace("/", "_") + ".log"
            sh("FRANK_MEDIUM_ENV_FILE=" + env_path + " setsid bash " + serve_script +
               " > " + logf + " 2>&1 < /dev/null &")
            if not wait_health(base):
                log(name + ": FAILED to become healthy -- " + sh("tail -3 " + logf), logfile)
                results[name] = {"error": "not healthy", "tail": sh("tail -6 " + logf)}
                stop_unit()
                continue
            rows = measure(base)
            results[name] = {"vram_gb": vram(), "rows": rows,
                             "spec_log": sh("grep -E '" + SPEC_GREP + "' " + logf + " | tail -4")}
            for r in rows:
                log(name + " | " + r["shape"] + " prefill " + str(r["prefill_tps"]) +
                    " tok/s | decode " + str(r["decode_tps"]) + " tok/s | draft " +
                    str(r["draft_n"]) + "/" + str(r["accepted"]) +
                    (" (" + str(round(100.0 * r["accepted"] / r["draft_n"])) + " %)"
                     if r["draft_n"] else "") +
                    (" | copy " + str(r["copy_ratio"]) if "copy_ratio" in r else "") +
                    (" | EMPTY ANSWER" if not r["answer"] else ""), logfile)
            log(name + " | vram " + str(results[name]["vram_gb"]), logfile)
            stop_unit()
    finally:
        restore_unit(base, logfile)

    with open(args.json, "w") as fh:
        json.dump(results, fh, indent=2)
    for shape in ("prose 1k/512", "echo 1.5k/400", "structured/400", "identity"):
        texts = [r["answer"] for v in results.values() if "rows" in v
                 for r in v["rows"] if r["shape"] == shape]
        if any(not t for t in texts):
            log("WARNING: '" + shape + "' produced an empty answer in some configuration -- "
                "its identity check proves nothing", logfile)
        elif len(set(texts)) > 1:
            log("WARNING: configurations disagree on greedy text for '" + shape +
                "' -- speculation is changing answers", logfile)
    log("done; results in " + args.json, logfile)
    return 0


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