diff --git a/scripts/bench_ollama.py b/scripts/bench_ollama.py new file mode 100755 index 0000000..b045864 --- /dev/null +++ b/scripts/bench_ollama.py @@ -0,0 +1,369 @@ +#!/usr/bin/env python3 +"""Benchmark Ollama models for FabledScribe chat / curator workloads. + +Measures time-to-first-token, total wall time, and tokens/sec for each +(model, scenario) pair. Designed to be runnable against both local and +remote Ollama servers so results from your two CPU servers are directly +comparable. + +Default forces CPU-only inference (num_gpu=0). Pass --num-gpu 99 for full +GPU offload, or any integer for partial. + +Scenarios: + chat — small input, small output, no thinking. Mirrors the chat-only + journal companion's expected load (short user message → + curious follow-up). + curator — longer transcript input, structured-output extraction, with + thinking enabled. Mirrors the curator's expected load + (read recent conversation → emit captures). + +Prerequisites: + - Ollama running and reachable at --server (default http://localhost:11434). + - The models named in --models must already be pulled + (`ollama pull qwen2.5:32b` etc). + +Usage examples: + # Curator candidate on CPU, 3 runs each (default), one model: + python scripts/bench_ollama.py --models qwen2.5:32b --scenario curator + + # Chat candidate on GPU, against a remote server: + python scripts/bench_ollama.py \\ + --server http://dock02:11434 \\ + --models llama3.2:3b \\ + --scenario chat --num-gpu 99 + + # Compare three curator candidates on CPU, 5 runs each, write markdown: + python scripts/bench_ollama.py \\ + --models qwen2.5:32b,qwen3:14b,gemma2:27b \\ + --scenario curator --runs 5 --out bench-cpu.md + +The first request to a (model, num_gpu) pair triggers a model load and is +excluded from the timing as a warm-up. Subsequent runs reflect warm-cache +performance — closest to how the running app will see the model in practice. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import sys +import time +from dataclasses import dataclass, field + +import httpx + + +SYSTEM_PROMPT_CHAT = ( + "You are the user's journal companion. Be warm and curious. Ask one " + "short follow-up question. Don't presume to help; let the user lead. " + "Match the user's length." +) + +USER_MESSAGE_CHAT = ( + "Got a lot done at Famous Supply today. Finished the network restage " + "on the Bedford swarm and there was one weird issue with a Windows " + "node. Tired but it feels good to be making progress." +) + +SYSTEM_PROMPT_CURATOR = ( + "You are reading a journal conversation. Extract the user's notable " + "moments — events, decisions, encounters, observations, plans. " + "Output a JSON array. Each entry has fields: content (1-2 sentence " + "summary in the user's voice), tags (list of strings), people (list " + "of names mentioned), projects (list of project names mentioned). " + "Skip meta-conversational chatter." +) + +USER_TRANSCRIPT_CURATOR = """\ +[12:04] User: morning. starting on the Famous Supply network restage today. +[12:04] Assistant: how big a job is that? +[12:05] User: bedford location swarm has 3 nodes. one of them is a Windows + host that's been flaky for a while. probably going to take all + afternoon. +[12:05] Assistant: anything you want to flag before you start in case + something goes sideways? +[12:06] User: yeah — the Windows node sometimes loses its bridge network + on reboot. if that happens i'll have to manually rejoin it. last + time it took an hour to figure out. +[12:06] Assistant: noted. how are you feeling going in? +[12:07] User: decent. didn't sleep great but i had coffee with Sarah this + morning which was nice. she mentioned she's hiring on her team if + i ever want to make a jump. +[12:07] Assistant: good to know. anything to capture about the work itself + before you dive in, or you good? +[12:08] User: nope, just wanted to externalize before starting. +[14:42] User: ok the Windows node DID lose its bridge. fixing it now. +[14:42] Assistant: any guess what triggers it? +[14:43] User: i think it's a stale dhcp lease that gets re-issued before + the bridge comes up. might write that down as a thing to fix + properly later. +[15:55] User: done. all three nodes restaged, the windows one too. wrote + a note about the dhcp issue for future me. +[15:55] Assistant: how do you feel about how the day went? +[15:56] User: tired but good. ate lunch standing up which was dumb. + watching a show with Victoria tonight to decompress. +""" + + +@dataclass +class RunResult: + ttft_ms: float + total_ms: float + prompt_tokens: int + output_tokens: int + tokens_per_sec: float + + +@dataclass +class ScenarioResult: + model: str + scenario: str + runs: list[RunResult] = field(default_factory=list) + error: str | None = None + + def summary(self) -> dict: + if self.error or not self.runs: + return { + "model": self.model, + "scenario": self.scenario, + "error": self.error or "no successful runs", + } + return { + "model": self.model, + "scenario": self.scenario, + "runs": len(self.runs), + "ttft_ms_p50": statistics.median(r.ttft_ms for r in self.runs), + "total_ms_p50": statistics.median(r.total_ms for r in self.runs), + "tokens_per_sec_p50": statistics.median( + r.tokens_per_sec for r in self.runs + ), + "output_tokens_mean": statistics.mean( + r.output_tokens for r in self.runs + ), + "prompt_tokens": self.runs[0].prompt_tokens, + } + + +def build_request( + scenario: str, model: str, num_gpu: int, keep_alive: str +) -> dict: + if scenario == "chat": + messages = [ + {"role": "system", "content": SYSTEM_PROMPT_CHAT}, + {"role": "user", "content": USER_MESSAGE_CHAT}, + ] + think = False + elif scenario == "curator": + messages = [ + {"role": "system", "content": SYSTEM_PROMPT_CURATOR}, + {"role": "user", "content": USER_TRANSCRIPT_CURATOR}, + ] + think = True + else: + raise ValueError(f"unknown scenario: {scenario}") + return { + "model": model, + "messages": messages, + "stream": True, + "think": think, + "keep_alive": keep_alive, + "options": { + "num_gpu": num_gpu, + "temperature": 0.3, + "num_ctx": 8192, + }, + } + + +def run_once(server: str, payload: dict) -> RunResult: + """Stream one chat request and time it. + + Uses Ollama-reported `eval_count` and `eval_duration` for tokens/sec + (authoritative; doesn't include client-side stream overhead). TTFT is + wall-clock from request send to first content chunk. + """ + url = f"{server.rstrip('/')}/api/chat" + t_start = time.monotonic() + ttft = None + prompt_tokens = 0 + output_tokens = 0 + eval_duration_ns = 0 + + with httpx.stream("POST", url, json=payload, timeout=600.0) as resp: + resp.raise_for_status() + for line in resp.iter_lines(): + if not line: + continue + chunk = json.loads(line) + if ttft is None and chunk.get("message", {}).get("content"): + ttft = time.monotonic() - t_start + if chunk.get("done"): + prompt_tokens = chunk.get("prompt_eval_count", 0) + output_tokens = chunk.get("eval_count", 0) + eval_duration_ns = chunk.get("eval_duration", 0) + + total = time.monotonic() - t_start + tps = ( + output_tokens / (eval_duration_ns / 1e9) + if eval_duration_ns + else 0.0 + ) + return RunResult( + ttft_ms=(ttft if ttft is not None else total) * 1000, + total_ms=total * 1000, + prompt_tokens=prompt_tokens, + output_tokens=output_tokens, + tokens_per_sec=tps, + ) + + +def benchmark( + *, + server: str, + models: list[str], + scenarios: list[str], + runs: int, + num_gpu: int, + keep_alive: str, +) -> list[ScenarioResult]: + results: list[ScenarioResult] = [] + for model in models: + for scenario in scenarios: + sr = ScenarioResult(model=model, scenario=scenario) + payload = build_request(scenario, model, num_gpu, keep_alive) + # Warm-up run loads the model into RAM/VRAM with the requested + # num_gpu setting. Excluded from the measured runs because it + # otherwise dominates TTFT with model-load time. + print( + f"[{model} :: {scenario}] warm-up (loading model)...", + flush=True, + ) + try: + run_once(server, payload) + except httpx.HTTPError as e: + sr.error = f"warm-up failed: {e}" + print(f" {sr.error}", file=sys.stderr) + results.append(sr) + continue + except Exception as e: + sr.error = f"warm-up exception: {e}" + print(f" {sr.error}", file=sys.stderr) + results.append(sr) + continue + for i in range(runs): + try: + r = run_once(server, payload) + except Exception as e: + print(f" run {i+1} failed: {e}", file=sys.stderr) + continue + sr.runs.append(r) + print( + f" run {i+1}/{runs}: ttft={r.ttft_ms:.0f}ms " + f"total={r.total_ms:.0f}ms tps={r.tokens_per_sec:.1f} " + f"out_tokens={r.output_tokens}", + flush=True, + ) + results.append(sr) + return results + + +def format_markdown(results: list[ScenarioResult], *, server: str, num_gpu: int) -> str: + mode = "CPU only" if num_gpu == 0 else ( + f"GPU offload ({num_gpu} layers)" if num_gpu > 0 else "Ollama default" + ) + lines = [ + "# Ollama benchmark", + "", + f"- Server: `{server}`", + f"- Mode: {mode} (`num_gpu={num_gpu}`)", + "", + "| Model | Scenario | Runs | Prompt tok | TTFT p50 (ms) " + "| Total p50 (ms) | tok/s p50 | Output tok (mean) |", + "|---|---|---|---|---|---|---|---|", + ] + for sr in results: + s = sr.summary() + if "error" in s: + lines.append( + f"| {s['model']} | {s['scenario']} | — | — | — | — | — " + f"| error: {s['error']} |" + ) + continue + lines.append( + f"| {s['model']} | {s['scenario']} | {s['runs']} " + f"| {s['prompt_tokens']} " + f"| {s['ttft_ms_p50']:.0f} | {s['total_ms_p50']:.0f} " + f"| {s['tokens_per_sec_p50']:.1f} " + f"| {s['output_tokens_mean']:.0f} |" + ) + return "\n".join(lines) + "\n" + + +def main(): + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument( + "--server", + default="http://localhost:11434", + help="Ollama server URL (default %(default)s)", + ) + parser.add_argument( + "--models", + required=True, + help="Comma-separated model tags (e.g. qwen2.5:32b,qwen3:14b)", + ) + parser.add_argument( + "--scenario", + choices=["chat", "curator", "both"], + default="both", + ) + parser.add_argument( + "--runs", + type=int, + default=3, + help="Runs per (model,scenario), excluding warm-up (default %(default)s)", + ) + parser.add_argument( + "--num-gpu", + type=int, + default=0, + help="0 = CPU only (default), 99 = full offload, -1 = Ollama default", + ) + parser.add_argument( + "--keep-alive", + default="10m", + help="Ollama keep_alive (default %(default)s)", + ) + parser.add_argument( + "--out", + help="Write markdown table to this file (also prints to stdout)", + ) + args = parser.parse_args() + + models = [m.strip() for m in args.models.split(",") if m.strip()] + scenarios = ( + ["chat", "curator"] if args.scenario == "both" else [args.scenario] + ) + + results = benchmark( + server=args.server, + models=models, + scenarios=scenarios, + runs=args.runs, + num_gpu=args.num_gpu, + keep_alive=args.keep_alive, + ) + + md = format_markdown(results, server=args.server, num_gpu=args.num_gpu) + print("\n" + md) + if args.out: + with open(args.out, "w") as f: + f.write(md) + print(f"Wrote results to {args.out}", file=sys.stderr) + + +if __name__ == "__main__": + main()