serverAppConfigs/scripts/.AGENTTRASH/bench-spec-decode.py.ik-llama.cpp

#!/usr/bin/env python3
"""Benchmark decode throughput across speculative decoding configs.

Sweeps 4 configs: none, ngram-mod, draft-mtp, draft-mtp+ngram-mod.
Each config generates 16k tokens from a minimal prompt at temp 0.

Runs ON the target host (Dankclaw). Service management uses local sudo.
Usage:
    python3 bench-spec-decode.py [--max-tokens N] [--host H] [--port P]
"""
import json
import subprocess
import sys
import time
import urllib.request
import urllib.error

SERVICE = "frank-dragoon"
ENV_FILE = "/home/dank/serverAppConfigs/scripts/frank-dragoon.env"
ALIAS = "frank-dragoon"
DEFAULT_HOST = "localhost"
DEFAULT_PORT = 7770
DEFAULT_MAX_TOKENS = 16384
HEALTH_TIMEOUT = 1200  # 20 min for cold load
HEALTH_INTERVAL = 5

PROMPT = (
    "Write a comprehensive, detailed essay about the complete history of "
    "computing, from the earliest mechanical calculators through modern AI. "
    "Cover every major era in detail, including key figures, machines, and "
    "breakthroughs. Be as thorough as possible."
)

CONFIGS = [
    ("none", "none"),
    ("ngram-mod", "ngram-mod"),
    ("draft-mtp", "draft-mtp"),
    ("draft-mtp+ngram-mod", "draft-mtp,ngram-mod"),
]


def local(cmd):
    """Run a command locally (the script runs on the target host)."""
    r = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=120)
    if r.returncode != 0:
        print(f"  WARNING: {cmd} -> exit {r.returncode}: {r.stderr.strip()}",
              file=sys.stderr)
    return r


def sed_spec_type(value):
    local(f"sudo sed -i 's/^SPEC_TYPE=.*/SPEC_TYPE={value}/' {ENV_FILE}")


def wait_health(host, port, timeout=HEALTH_TIMEOUT):
    start = time.time()
    while time.time() - start < timeout:
        try:
            resp = urllib.request.urlopen(
                f"http://{host}:{port}/health", timeout=10
            )
            if resp.status == 200:
                return True
        except (urllib.error.URLError, urllib.error.HTTPError, ConnectionError,
                OSError):
            pass
        time.sleep(HEALTH_INTERVAL)
    return False


def benchmark(host, port, alias, max_tokens):
    body = json.dumps({
        "model": alias,
        "messages": [{"role": "user", "content": PROMPT}],
        "max_tokens": max_tokens,
        "temperature": 0,
        "stream": False,
    }).encode()
    req = urllib.request.Request(
        f"http://{host}:{port}/v1/chat/completions",
        data=body,
        headers={"Content-Type": "application/json"},
    )
    t0 = time.time()
    resp = urllib.request.urlopen(req, timeout=1200)
    wall = time.time() - t0
    data = json.loads(resp.read())
    return data, wall


def main():
    import argparse
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--host", default=DEFAULT_HOST)
    ap.add_argument("--port", type=int, default=DEFAULT_PORT)
    ap.add_argument("--max-tokens", type=int, default=DEFAULT_MAX_TOKENS)
    args = ap.parse_args()

    # Read original SPEC_TYPE from the env file
    r = local(f"grep '^SPEC_TYPE=' {ENV_FILE}")
    original_spec = r.stdout.strip()
    orig_val = (original_spec.split("=", 1)[1]
                if "=" in original_spec else "draft-mtp,ngram-mod")
    print(f"Original: {original_spec}")
    print(f"Max tokens: {args.max_tokens}")
    print(f"Configs: {[c[0] for c in CONFIGS]}")
    print()

    # Stop + mask the production service
    print("Stopping + masking production service...")
    local(f"sudo systemctl stop {SERVICE}")
    local(f"sudo systemctl mask {SERVICE}")
    time.sleep(2)

    results = []
    for name, spec_type in CONFIGS:
        print(f"\n{'='*60}")
        print(f"Config: {name} (SPEC_TYPE={spec_type})")
        print(f"{'='*60}")

        # Set SPEC_TYPE
        sed_spec_type(spec_type)
        print(f"Set SPEC_TYPE={spec_type}")

        # Start the service
        local(f"sudo systemctl unmask {SERVICE}")
        local(f"sudo systemctl start {SERVICE}")
        print("Starting service...")

        # Wait for health
        print(f"Waiting for health (up to {HEALTH_TIMEOUT}s)...")
        if not wait_health(args.host, args.port):
            print(f"FAILED: service did not become healthy for {name}")
            results.append({
                "config": name, "spec_type": spec_type,
                "error": "health timeout",
            })
            local(f"sudo systemctl stop {SERVICE}")
            local(f"sudo systemctl mask {SERVICE}")
            continue

        print("Healthy. Running benchmark...")

        # Run benchmark
        try:
            data, wall = benchmark(
                args.host, args.port, ALIAS, args.max_tokens
            )
        except Exception as e:
            print(f"FAILED: {e}")
            results.append({
                "config": name, "spec_type": spec_type,
                "error": str(e),
            })
            local(f"sudo systemctl stop {SERVICE}")
            local(f"sudo systemctl mask {SERVICE}")
            continue

        timings = data.get("timings", {})
        usage = data.get("usage", {})
        pred_n = timings.get("predicted_n", usage.get("completion_tokens", 0))
        pred_ms = timings.get("predicted_ms", 0)
        pred_ps = timings.get("predicted_per_second", 0)
        prompt_n = timings.get("prompt_n", usage.get("prompt_tokens", 0))
        prompt_ms = timings.get("prompt_ms", 0)

        print(f"  Generated:   {pred_n} tokens")
        print(f"  Wall time:   {wall:.1f}s")
        print(f"  Decode:      {pred_ps:.2f} tok/s")
        print(f"  Prefill:     {prompt_n} tokens in {prompt_ms:.0f}ms"
              f" ({prompt_n / max(prompt_ms, 1) * 1000:.1f} tok/s)")

        results.append({
            "config": name,
            "spec_type": spec_type,
            "predicted_tokens": pred_n,
            "predicted_ms": pred_ms,
            "predicted_per_second": pred_ps,
            "prompt_tokens": prompt_n,
            "prompt_ms": prompt_ms,
            "wall_time": wall,
        })

        # Stop + mask between configs
        local(f"sudo systemctl stop {SERVICE}")
        local(f"sudo systemctl mask {SERVICE}")
        time.sleep(2)

    # Restore original SPEC_TYPE
    sed_spec_type(orig_val)
    print(f"\nRestoring SPEC_TYPE={orig_val}")
    print("Unmasking + starting production service...")
    local(f"sudo systemctl unmask {SERVICE}")
    local(f"sudo systemctl start {SERVICE}")

    # Health check
    if wait_health(args.host, args.port):
        print("Production service healthy.")
    else:
        print("WARNING: production service not healthy after restore!")

    # Print results table
    print(f"\n{'='*70}")
    print("RESULTS")
    print(f"{'='*70}")
    print(f"{'Config':<30} {'Tokens':>8} {'Wall(s)':>10} {'Decode tok/s':>14} "
          f"{'Error':>20}")
    print(f"{'-'*30} {'-'*8} {'-'*10} {'-'*14} {'-'*20}")
    for r in results:
        if "error" in r:
            print(f"{r['config']:<30} {'':>8} {'':>10} {'':>14} "
                  f"{r['error']:>20}")
        else:
            print(f"{r['config']:<30} {r['predicted_tokens']:>8} "
                  f"{r['wall_time']:>10.1f} {r['predicted_per_second']:>14.2f} "
                  f"{'':>20}")
    print(f"{'='*70}")

    # Save JSON results
    ts = time.strftime("%Y%m%d-%H%M%S")
    out = f"/tmp/bench-spec-decode-{ts}.json"
    with open(out, "w") as f:
        json.dump({"results": results, "max_tokens": args.max_tokens,
                   "timestamp": ts}, f, indent=2)
    print(f"Results saved: {out}")


if __name__ == "__main__":
    main()