#!/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-default-large" ENV_FILE = "/home/dank/serverAppConfigs/scripts/frank-default-large.env" ALIAS = "frank-default-large" 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()