#!/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]`, semicolon-separated. The systemd unit is
restored at the end, so the endpoint comes back on whatever the env file says
whatever happens here.

TWO SHAPES, because a single one answers the wrong question:

  prose  a 1k-token prompt asking for a sentence — nothing to copy, so only a
         model-backed proposer can help. This is the shape that compares
         drafters.
  echo   a ~1.5k-token passage of unique lines with "repeat it back verbatim".
         The output is almost entirely a copy of the prompt, which is the only
         traffic where an n-gram matcher has anything to match. Measuring
         ngram-mod on prose alone would call it useless; measuring DFlash on
         echo alone would call it wasteful.

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

    prose 1k/512      none 17.7 | mtp n=3 32.4 | mtp n=7 26.7 | dflash=7 43.5 tok/s
    card 1 VRAM       none 26.07 | mtp n=3 29.35 | mtp n=7 29.65 | dflash 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. ngram-mod has no such
constraint — it is draftless and is added to the chain unconditionally.
"""
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|ngram_mod with"

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_INSTRUCTION = ("\n\nRepeat the passage above verbatim, every line, exactly as written, "
                    "with no commentary before or after it.")
PROSE_INSTRUCTION = "\n\nReply with a single short sentence."
IDENTITY_PROMPT = "List the first five prime numbers, comma separated, then stop."


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):
    out = []
    for chunk in spec.split(";"):
        chunk = chunk.strip()
        if not chunk:
            continue
        name, _, rest = chunk.partition("=")
        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):
    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 run_shape(base, prompt, max_tokens, label):
    t = post(base, {"model": MODEL,
                    "messages": [{"role": "user", "content": prompt}],
                    "max_tokens": max_tokens, "temperature": 0.0})
    tm = t["timings"]
    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": (t["choices"][0]["message"].get("content") or "").strip()}
    return row


def measure(base):
    run_shape(base, FILLER * (4000 // TOKENS_PER_REP) + PROSE_INSTRUCTION, 32, "warmup")
    rows = [
        run_shape(base, FILLER * (1000 // TOKENS_PER_REP) + PROSE_INSTRUCTION, 512, "prose 1k/512"),
        run_shape(base, echo_passage() + ECHO_INSTRUCTION, 400, "echo 1.5k/400"),
        run_shape(base, IDENTITY_PROMPT, 192, "identity"),
    ]
    return rows


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] 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
    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)
    stop_unit()
    for name, over in configs:
        env_path = write_env(args.env_file, name, over)
        logf = "/tmp/spec-ab-" + name.replace("/", "_") + ".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 -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 + " | head -6")}
        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 ""),
                logfile)
        log(name + " | vram " + str(results[name]["vram_gb"]), logfile)
        stop_unit()

    with open(args.json, "w") as fh:
        json.dump(results, fh, indent=2)
    for shape in ("prose 1k/512", "echo 1.5k/400", "identity"):
        texts = {v["rows"][i]["answer"] for v in results.values() if "rows" in v
                 for i, r in enumerate(v["rows"]) if r["shape"] == shape}
        if len(texts) > 1:
            log("WARNING: configurations disagree on greedy text for '" + shape +
                "' -- speculation is changing 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 " + args.json, logfile)
    return 0


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