#!/usr/bin/env python3 """ Full test + benchmark suite for the local model fleet. Perf : per-model generation tok/s, prompt tok/s, cold load time, GPU placement (from Ollama's own eval metrics - objective, not wall-clock guesses) Function : every model returns non-empty output; vision models pass an OCR check Routing : representative prompts hit the router and must land on the right model Run AFTER all models are downloaded and the router is up: ./.venv/bin/python benchmark.py # or: python3 benchmark.py Writes bench_report.md and prints a summary table. """ import json, time, base64, io, re, subprocess, urllib.request, urllib.error, datetime, sys OLLAMA = "http://127.0.0.1:11434" ROUTER = "http://127.0.0.1:8080" NUM_PREDICT = 200 PERF_PROMPT = ("Explain how TCP congestion control works (slow start, congestion " "avoidance, fast retransmit) in about 150 words.") VISION_MODELS = ("qwen3-vl", "gemma4", "gemma3") # accept image input def user(txt): return {"role": "user", "content": txt} def img_user(txt, b64): # OpenAI format — for the router (LiteLLM) path return {"role": "user", "content": [ {"type": "text", "text": txt}, {"type": "image_url", "image_url": {"url": "data:image/png;base64," + b64}}]} def img_user_native(txt, b64): # Ollama-native format — for direct /api/chat return {"role": "user", "content": txt, "images": [b64]} def ollama_chat(model, messages, num_predict=NUM_PREDICT, timeout=1800, think=None): body = {"model": model, "messages": messages, "stream": False, "keep_alive": "30m", "options": {"temperature": 0, "num_predict": num_predict}} if think is not None: body["think"] = think req = urllib.request.Request(OLLAMA + "/api/chat", data=json.dumps(body).encode(), headers={"content-type": "application/json"}) with urllib.request.urlopen(req, timeout=timeout) as r: return json.loads(r.read()) def metrics(d): sec = lambda ns: (ns or 0) / 1e9 ec, ed = d.get("eval_count", 0), sec(d.get("eval_duration")) pc, pd = d.get("prompt_eval_count", 0), sec(d.get("prompt_eval_duration")) return {"gen_tps": ec / ed if ed else 0, "prompt_tps": pc / pd if pd else 0, "load_s": sec(d.get("load_duration")), "ttft_s": sec(d.get("load_duration")) + pd, "gen_tokens": ec, "text": (d.get("message") or {}).get("content", "")} def ps_info(model): try: out = subprocess.run(["ollama", "ps"], capture_output=True, text=True, timeout=15).stdout for line in out.splitlines()[1:]: cols = re.split(r"\s{2,}", line.strip()) if cols and cols[0] == model: size = cols[2] if len(cols) > 2 else "" proc = cols[3] if len(cols) > 3 else "" return size, proc except Exception: pass return "", "" def make_ocr_image(text): from PIL import Image, ImageDraw, ImageFont img = Image.new("RGB", (760, 220), "white") d = ImageDraw.Draw(img) font = None for path in ("/System/Library/Fonts/Supplemental/Arial.ttf", "/System/Library/Fonts/Helvetica.ttc", "/Library/Fonts/Arial.ttf"): try: font = ImageFont.truetype(path, 72); break except Exception: pass d.text((40, 70), text, fill="black", font=font) buf = io.BytesIO(); img.save(buf, format="PNG") return base64.b64encode(buf.getvalue()).decode() def list_models(): out = subprocess.run(["ollama", "list"], capture_output=True, text=True).stdout return [ln.split()[0] for ln in out.splitlines()[1:] if ln.strip()] def router_call(messages, model="auto", timeout=1800): body = {"model": model, "messages": messages, "stream": False, "options": {"num_predict": 32}} req = urllib.request.Request(ROUTER + "/v1/chat/completions", data=json.dumps(body).encode(), headers={"content-type": "application/json"}) try: with urllib.request.urlopen(req, timeout=timeout) as r: # compare the routing DECISION, not the post-fallback served model picked = r.headers.get("x-router-initial-model") or r.headers.get("x-router-model", "?") data = json.loads(r.read()) content = data["choices"][0]["message"]["content"] return 200, picked, content except urllib.error.HTTPError as e: return e.code, "-", e.read().decode(errors="replace")[:120] except Exception as e: return 0, "-", str(e)[:120] def main(): models = list_models() ocr_text = "INVOICE-7X42" img = None try: img = make_ocr_image(ocr_text) except Exception as e: print(f"(pillow unavailable, OCR check skipped: {e})") print(f"Benchmarking {len(models)} models ...\n") perf = [] for m in models: vis = any(v in m for v in VISION_MODELS) try: ollama_chat(m, [user("hi")], num_predict=4) # cold-load / warm d = ollama_chat(m, [user(PERF_PROMPT)]) # measured run mm = metrics(d) size, proc = ps_info(m) ocr = "-" if vis and img: try: r = ollama_chat(m, [img_user_native( "What text is shown in this image?", img)], 256) msg = r.get("message") or {} got = (msg.get("content") or "") + " " + (msg.get("thinking") or "") ocr = "PASS" if ocr_text.replace("-", "").lower() in \ got.replace("-", "").replace(" ", "").lower() else "miss" except Exception: ocr = "err" ok = "ok" if mm["gen_tokens"] > 0 else "EMPTY" perf.append((m, mm, size, proc, ocr, ok)) print(f" {m:32s} {mm['gen_tps']:6.1f} tok/s load {mm['load_s']:5.1f}s " f"{proc or '?':>9s} ocr={ocr}") except Exception as e: perf.append((m, None, "", "", "-", f"ERROR {e}")) print(f" {m:32s} ERROR: {e}") # routing correctness print("\nRouting tests:") cases = [ ("short chat", [user("hey, how's it going?")], "qwen3:8b"), ("code", [user("Refactor this function and fix the bug in app.py")], "glm-4.7-flash"), ("repo code", [user("Refactor auth across the whole codebase, multiple files")], "qwen3-coder:30b"), ("reasoning", [user("Prove step by step why sqrt(2) is irrational")], "qwen3.6:35b-a3b"), ] if img: cases += [ ("vision", [img_user("what is in this image?", img)], "qwen3-vl:8b"), ("ocr/doc", [img_user("extract the text from this document table", img)], "qwen3-vl:30b-a3b-instruct"), ] routing = [] for name, msgs, expect in cases: code, picked, _ = router_call(msgs) verdict = "PASS" if picked == expect else "DIFF" routing.append((name, expect, picked, code, verdict)) print(f" {name:12s} expect {expect:26s} got {picked:26s} [{verdict}]") # write report ts = datetime.datetime.now().strftime("%Y-%m-%d %H:%M") with open("bench_report.md", "w") as f: f.write(f"# Local fleet benchmark — {ts}\n\n") f.write("## Performance (sorted by generation tok/s)\n\n") f.write("| Model | Gen tok/s | Prompt tok/s | Cold load | TTFT | Placement | OCR | Status |\n") f.write("|---|--:|--:|--:|--:|:--|:--:|:--|\n") for m, mm, size, proc, ocr, ok in sorted( perf, key=lambda x: (x[1]["gen_tps"] if x[1] else -1), reverse=True): if mm: f.write(f"| `{m}` | {mm['gen_tps']:.1f} | {mm['prompt_tps']:.0f} | " f"{mm['load_s']:.1f}s | {mm['ttft_s']:.1f}s | {proc or '?'} | {ocr} | {ok} |\n") else: f.write(f"| `{m}` | - | - | - | - | - | - | {ok} |\n") f.write("\n## Routing correctness\n\n") f.write("| Case | Expected | Routed to | HTTP | Verdict |\n|---|---|---|--:|:--:|\n") for name, expect, picked, code, verdict in routing: f.write(f"| {name} | `{expect}` | `{picked}` | {code} | {verdict} |\n") print("\nWrote bench_report.md") if __name__ == "__main__": main()