scripts: add bench_ollama.py for CPU/GPU model benchmarking
Standalone tool to measure Ollama model performance under the two workload shapes the chat+curator architecture would impose: - chat scenario: short user message, short reply, no thinking. Mirrors the no-tools chat companion's expected load. - curator scenario: ~700-token journal transcript with an extraction prompt, thinking enabled. Mirrors the curator's expected load. Defaults to CPU-only inference (num_gpu=0). Streams responses; reports TTFT, total wall time, tokens/sec (from Ollama's eval_count/eval_duration so it excludes client-side stream overhead), and prompt token count. First request per (model, num_gpu) is a warm-up to load the model into memory; not counted in the measured runs. Designed for cross-server comparison: --server points at any Ollama instance, --out writes a markdown table. Comparing the two CPU servers becomes a matter of running the same command on each and diffing the output. Lives outside the chat/curator architecture commitment — measurement tool only. Tells us "is qwen2.5:32b on CPU fast enough for a 10-20 min curator cadence?" without writing any of the architecture code yet. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Executable
+369
@@ -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()
|
||||
Reference in New Issue
Block a user