#!/usr/bin/env python3 """ fleet-bench.py — Full fleet benchmark: gen speed, prompt processing, stress test. Runs 3 benchmarks per model across the defined fleet, saves structured JSON results, and prints clean summary tables. Handles cold-load skew by separating run 1 from warm runs 2+. Usage: python3 fleet-bench.py # benchmark all models python3 fleet-bench.py --quick # 1 gen run + 1 stress run per model python3 fleet-bench.py --output /tmp/myb.json """ import json, subprocess, sys, time, os # ── Config ────────────────────────────────────────────────────────────── LLAMA_SWAP_URL = "http://127.0.0.1:9292/v1" MODELS = [ ("nemotron-term-14b", "Fast dense 14B"), ("gpt-oss-20b", "Lightweight MoE"), ("glm-4.7-flash", "Agentic MoE"), ("qwen36-35b-mtp", "Primary + MTP"), ("gemma-26b-200k", "Deep reasoning"), ] GEN_PROMPT = ( "Write a detailed Python implementation of a binary search tree with " "insert, delete, search, and traversal methods. Include docstrings and comments." ) STRESS_PROMPT = "Count from 1 to 10, one number per line." LONG_PROMPT = "The quick brown fox jumps over the lazy dog. " * 100 + "\nSummarize this text in 3 bullet points." OUTPUT_FILE = "/tmp/bench_results.json" # ── Helpers ───────────────────────────────────────────────────────────── def shq(s): """Shell-quote a string for safe single-quoted insertion.""" return "'" + s.replace("'", "'\\''") + "'" def make_payload(model_id, prompt_text, max_tokens): return json.dumps({ "model": model_id, "messages": [{"role": "user", "content": prompt_text}], "max_tokens": max_tokens, "temperature": 0.2, "stream": False, }) def curl_bench(payload_json, timeout=120): """Run one bench request, return parsed dict with timings + wall clock.""" fname = "/tmp/bench_out.json" cmd = ( f'curl -s --max-time {timeout} -X POST {LLAMA_SWAP_URL}/chat/completions ' f'-H "Content-Type: application/json" ' f'-d {shq(payload_json)} ' f'-o {fname} -w "HTTP_CODE:%{{http_code}}\\nTIME_TOTAL:%{{time_total}}\\n"' ) start = time.time() subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout) elapsed = time.time() - start # Parse wall clock from curl -w wall_time = round(elapsed, 2) # Parse response JSON try: with open(fname) as f: data = json.load(f) except (json.JSONDecodeError, FileNotFoundError) as e: return {"error": str(e)} # Parse timings from response t = data.get("timings", data) # fallback to top-level for older builds usage = data.get("usage", {}) # VRAM smi = subprocess.run( "nvidia-smi --query-gpu=memory.used --format=csv,noheader", shell=True, capture_output=True, text=True, timeout=5 ) pt = usage.get("prompt_tokens", 0) ct = usage.get("completion_tokens", 0) finish = data.get("choices", [{}])[0].get("finish_reason", "?") draft_n = t.get("draft_n", 0) draft_accepted = t.get("draft_n_accepted", 0) draft_rate = round(draft_accepted / draft_n * 100, 1) if draft_n > 0 else 0 return { "http_code": "200", "prompt_tokens": pt, "completion_tokens": ct, "finish_reason": finish, "gen_tps": round(t.get("predicted_per_second", 0), 1), "prompt_tps": round(t.get("prompt_per_second", 0), 0) if t.get("prompt_per_second") else 0, "gen_mspt": round(t.get("predicted_per_token_ms", 0), 2), "prompt_mspt": round(t.get("prompt_per_token_ms", 0), 2), "cache_hit_tokens": t.get("cache_n", 0), "draft_accepted": draft_accepted, "draft_total": draft_n, "draft_accept_pct": draft_rate, "vram_mib": smi.stdout.strip(), "wall_time_s": wall_time, } def print_banner(text): width = 60 print(f"\n{'=' * width}") print(f" {text}") print(f"{'=' * width}\n") # ── Benchmarks ────────────────────────────────────────────────────────── def run_gen_bench(quick=False): """Generation speed benchmark (max_tokens=512).""" print_banner("Generation Speed (max_tokens=512)") all_runs = [] for model_id, desc in MODELS: print(f"\n {model_id} ({desc})") payload = make_payload(model_id, GEN_PROMPT, 512) runs = 1 if quick else 3 for i in range(runs): print(f" Run {i+1}/{runs}...", end=" ", flush=True) r = curl_bench(payload) r["model"] = model_id r["run"] = f"gen-{i+1}" all_runs.append(r) if "error" in r: print(f"ERROR: {r['error'][:60]}") else: draft = f", draft={r['draft_accept_pct']}%/{r['draft_total']}" if r['draft_total'] > 0 else "" print(f"gen={r['gen_tps']} t/s, prompt={r['prompt_tps']} t/s, tok={r['completion_tokens']}, wall={r['wall_time_s']}s{draft}") return all_runs def run_prompt_bench(quick=False): """Prompt processing speed benchmark (~3K tokens input).""" print_banner("Prompt Processing Speed (~3K tokens input)") all_runs = [] for model_id, desc in MODELS: print(f"\n {model_id}") payload = make_payload(model_id, LONG_PROMPT, 128) runs = 1 if quick else 2 for i in range(runs): print(f" Run {i+1}/{runs}...", end=" ", flush=True) r = curl_bench(payload, timeout=180) r["model"] = model_id r["run"] = f"prompt-{i+1}" all_runs.append(r) if "error" in r: print(f"ERROR: {r['error'][:60]}") else: print(f"prompt={r['prompt_tps']} t/s, gen={r['gen_tps']} t/s, wall={r['wall_time_s']}s, cache={r['cache_hit_tokens']}t") return all_runs def run_stress_bench(quick=False): """Stress test: 5 rapid requests per model.""" print_banner("Stress Test (5 rapid requests, max_tokens=50)") all_runs = [] for model_id, desc in MODELS: print(f"\n {model_id} ({desc})") payload = make_payload(model_id, STRESS_PROMPT, 50) runs = 5 times = [] for i in range(runs): print(f" Req {i+1}/{runs}...", end=" ", flush=True) r = curl_bench(payload, timeout=60) r["model"] = model_id r["run"] = f"stress-{i+1}" all_runs.append(r) if "error" not in r: times.append(r["wall_time_s"]) print(f"wall={r['wall_time_s']}s, gen={r['gen_tps']} t/s") else: print(f"ERROR: {r['error'][:60]}") if times: print(f" → Cold: {times[0]}s, Warm avg: {round(sum(times[1:]) / len(times[1:]), 2)}s, Min: {min(times)}s") return all_runs # ── Summary ───────────────────────────────────────────────────────────── def print_summary(all_results): print_banner("SUMMARY") # Generation (best warm run) print("Generation Speed (warm, best of runs 2-3)") for mid, desc in MODELS: runs = [r for r in all_results if r.get("model") == mid and "gen-" in r.get("run", "") and "error" not in r] warm = [r for r in runs if r.get("run") not in ("gen-1",)] if warm: best = max(warm, key=lambda x: x.get("gen_tps", 0)) d = f", draft={best['draft_accept_pct']}%" if best['draft_total'] > 0 else "" print(f" {mid:25s} {best['gen_tps']:>5.1f} t/s {best['prompt_tps']:>4.0f} ppt/s VRAM {best['vram_mib']}{d}") print() # Stress (cold vs warm) print("Stress Test (cold vs warm)") for mid, desc in MODELS: runs = [r for r in all_results if r.get("model") == mid and "stress-" in r.get("run", "") and "error" not in r] if runs: cold = runs[0]["wall_time_s"] warm_avg = round(sum(r["wall_time_s"] for r in runs[1:]) / len(runs[1:]), 3) if len(runs) > 1 else 0 best_tps = max(r.get("gen_tps", 0) for r in runs) print(f" {mid:25s} cold={cold:>5.1f}s warm_avg={warm_avg:>5.3f}s best_tps={best_tps}") print(f"\nRaw results saved to {OUTPUT_FILE}") # ── Main ──────────────────────────────────────────────────────────────── if __name__ == "__main__": quick_mode = "--quick" in sys.argv if quick_mode: sys.argv.remove("--quick") out_path = OUTPUT_FILE if "--output" in sys.argv: idx = sys.argv.index("--output") out_path = sys.argv[idx + 1] all_results = [] all_results += run_gen_bench(quick_mode) all_results += run_prompt_bench(quick_mode) all_results += run_stress_bench(quick_mode) print_summary(all_results) # Save all raw data with open(out_path, "w") as f: json.dump(all_results, f, indent=2) print(f"\nDone — {len(all_results)} total runs saved to {out_path}")