290 lines
11 KiB
Python
290 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""Benchmark Ollama cloud models: TTFT, throughput, thinking vs response phases.
|
|
|
|
Measures:
|
|
- TTFT-thinking: time to first thinking token
|
|
- TTFT-response: time to first response token (visible output)
|
|
- Thinking phase: duration and tok/s of reasoning
|
|
- Response phase: duration and tok/s of visible output
|
|
- Total wall time and overall tok/s
|
|
- Token counts (prompt, thinking, response, total eval)
|
|
|
|
Usage:
|
|
python3 benchmark_ollama.py [--host nvidia.hell] [--port 11434]
|
|
python3 benchmark_ollama.py --models glm-5.1,glm-5.2,kimi-k2.6 --runs 3
|
|
python3 benchmark_ollama.py --list
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import time
|
|
import sys
|
|
import urllib.request
|
|
import urllib.error
|
|
|
|
DEFAULT_PROMPT = "Write a detailed 500-word essay about the history of computing, from Babbage to modern AI. Include key milestones, people, and technologies."
|
|
DEFAULT_HOST = "nvidia.hell"
|
|
DEFAULT_PORT = 11434
|
|
DEFAULT_NUM_PREDICT = 1000
|
|
|
|
|
|
def api_get(host, port, path):
|
|
url = f"http://{host}:{port}{path}"
|
|
req = urllib.request.Request(url, method="GET")
|
|
with urllib.request.urlopen(req, timeout=10) as resp:
|
|
return json.loads(resp.read())
|
|
|
|
|
|
def api_post_stream(host, port, path, payload):
|
|
"""POST with streaming response. Yields (line_dict, elapsed_since_start)."""
|
|
url = f"http://{host}:{port}{path}"
|
|
data = json.dumps(payload).encode("utf-8")
|
|
req = urllib.request.Request(
|
|
url, data=data, method="POST",
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
t0 = time.monotonic()
|
|
resp = urllib.request.urlopen(req, timeout=300)
|
|
buffer = b""
|
|
while True:
|
|
chunk = resp.read(1)
|
|
if not chunk:
|
|
break
|
|
buffer += chunk
|
|
if chunk == b"\n":
|
|
line = buffer.strip()
|
|
buffer = b""
|
|
if not line:
|
|
continue
|
|
try:
|
|
obj = json.loads(line)
|
|
yield obj, time.monotonic() - t0
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
|
|
def list_models(host, port):
|
|
try:
|
|
data = api_get(host, port, "/api/tags")
|
|
return [m["name"] for m in data.get("models", [])]
|
|
except Exception as e:
|
|
print(f"Error listing models: {e}", file=sys.stderr)
|
|
return []
|
|
|
|
|
|
def benchmark_model(host, port, model, prompt, num_predict):
|
|
"""Run a single benchmark against one model."""
|
|
payload = {
|
|
"model": model,
|
|
"prompt": prompt,
|
|
"stream": True,
|
|
"options": {"temperature": 0.0, "num_predict": num_predict},
|
|
}
|
|
|
|
# Track timing for thinking and response phases
|
|
first_thinking_time = None
|
|
last_thinking_time = None
|
|
first_response_time = None
|
|
last_response_time = None
|
|
thinking_chars = 0
|
|
response_chars = 0
|
|
|
|
# Stats from final chunk
|
|
prompt_eval_count = 0
|
|
eval_count = 0
|
|
total_duration_ns = 0
|
|
done_reason = ""
|
|
|
|
try:
|
|
for obj, elapsed in api_post_stream(host, port, "/api/generate", payload):
|
|
if obj.get("error"):
|
|
return {"model": model, "error": obj["error"]}
|
|
|
|
thinking = obj.get("thinking", "") or ""
|
|
response = obj.get("response", "") or ""
|
|
|
|
if thinking:
|
|
if first_thinking_time is None:
|
|
first_thinking_time = elapsed
|
|
last_thinking_time = elapsed
|
|
thinking_chars += len(thinking)
|
|
|
|
if response:
|
|
if first_response_time is None:
|
|
first_response_time = elapsed
|
|
last_response_time = elapsed
|
|
response_chars += len(response)
|
|
|
|
if obj.get("done"):
|
|
prompt_eval_count = obj.get("prompt_eval_count", 0)
|
|
eval_count = obj.get("eval_count", 0)
|
|
total_duration_ns = obj.get("total_duration", 0)
|
|
done_reason = obj.get("done_reason", "")
|
|
break
|
|
except urllib.error.HTTPError as e:
|
|
body = e.read().decode("utf-8", errors="replace")[:200]
|
|
return {"model": model, "error": f"HTTP {e.code}: {body}"}
|
|
except Exception as e:
|
|
return {"model": model, "error": str(e)}
|
|
|
|
# Calculate metrics
|
|
# Estimate token split: thinking vs response by char ratio
|
|
total_gen_chars = thinking_chars + response_chars
|
|
if total_gen_chars > 0:
|
|
thinking_tokens_est = round(eval_count * thinking_chars / total_gen_chars)
|
|
response_tokens_est = round(eval_count * response_chars / total_gen_chars)
|
|
else:
|
|
thinking_tokens_est = 0
|
|
response_tokens_est = 0
|
|
|
|
# Phase durations
|
|
thinking_duration = 0
|
|
if first_thinking_time and last_thinking_time:
|
|
thinking_duration = last_thinking_time - first_thinking_time
|
|
|
|
response_duration = 0
|
|
if first_response_time and last_response_time:
|
|
response_duration = last_response_time - first_response_time
|
|
|
|
# Total generation time (from first token to last token, whether thinking or response)
|
|
first_token_time = None
|
|
last_token_time = None
|
|
if first_thinking_time is not None or first_response_time is not None:
|
|
first_token_time = min(
|
|
t for t in [first_thinking_time, first_response_time] if t is not None
|
|
)
|
|
if last_thinking_time is not None or last_response_time is not None:
|
|
last_token_time = max(
|
|
t for t in [last_thinking_time, last_response_time] if t is not None
|
|
)
|
|
|
|
total_gen_time = 0
|
|
if first_token_time and last_token_time:
|
|
total_gen_time = last_token_time - first_token_time
|
|
|
|
# Throughput calculations
|
|
# Overall: eval_count / total_gen_time (all tokens including thinking)
|
|
overall_tps = (eval_count / total_gen_time) if total_gen_time > 0 else 0
|
|
|
|
# Thinking phase throughput
|
|
thinking_tps = (thinking_tokens_est / thinking_duration) if thinking_duration > 0 else 0
|
|
|
|
# Response phase throughput
|
|
response_tps = (response_tokens_est / response_duration) if response_duration > 0 else 0
|
|
|
|
# Wall-clock total (from request start to last token)
|
|
wall_total = last_token_time or 0
|
|
|
|
# TTFT metrics
|
|
ttft_thinking = first_thinking_time
|
|
ttft_response = first_response_time
|
|
|
|
return {
|
|
"model": model,
|
|
"done_reason": done_reason,
|
|
"ttft_thinking_s": round(ttft_thinking, 3) if ttft_thinking else None,
|
|
"ttft_response_s": round(ttft_response, 3) if ttft_response else None,
|
|
"wall_total_s": round(wall_total, 3),
|
|
"prompt_tokens": prompt_eval_count,
|
|
"eval_count": eval_count,
|
|
"thinking_tokens_est": thinking_tokens_est,
|
|
"response_tokens_est": response_tokens_est,
|
|
"thinking_duration_s": round(thinking_duration, 3),
|
|
"response_duration_s": round(response_duration, 3),
|
|
"thinking_tps": round(thinking_tps, 1),
|
|
"response_tps": round(response_tps, 1),
|
|
"overall_tps": round(overall_tps, 1),
|
|
"thinking_chars": thinking_chars,
|
|
"response_chars": response_chars,
|
|
}
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Benchmark Ollama models (thinking + response)")
|
|
parser.add_argument("--host", default=DEFAULT_HOST)
|
|
parser.add_argument("--port", type=int, default=DEFAULT_PORT)
|
|
parser.add_argument("--prompt", default=DEFAULT_PROMPT)
|
|
parser.add_argument("--models", help="Comma-separated model list (default: all cloud)")
|
|
parser.add_argument("--num-predict", type=int, default=DEFAULT_NUM_PREDICT)
|
|
parser.add_argument("--runs", type=int, default=1)
|
|
parser.add_argument("--list", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
print(f"Connecting to http://{args.host}:{args.port} ...")
|
|
models = list_models(args.host, args.port)
|
|
if not models:
|
|
print("No models found or connection failed.", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
# Filter to cloud models by default for relevance
|
|
cloud_models = [m for m in models if ":cloud" in m]
|
|
all_models = models
|
|
|
|
if args.list:
|
|
print(f"Available models ({len(models)}):")
|
|
for m in all_models:
|
|
print(f" - {m}")
|
|
return
|
|
|
|
# Select models
|
|
if args.models:
|
|
wanted = [m.strip() for m in args.models.split(",")]
|
|
selected = []
|
|
for w in wanted:
|
|
matches = [m for m in all_models if w.lower() in m.lower()]
|
|
if matches:
|
|
selected.extend(matches)
|
|
else:
|
|
print(f" Warning: '{w}' not found", file=sys.stderr)
|
|
seen = set()
|
|
selected = [m for m in selected if not (m in seen or seen.add(m))]
|
|
else:
|
|
selected = cloud_models
|
|
|
|
if not selected:
|
|
print("No matching models selected.", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
print(f"\nBenchmarking {len(selected)} model(s), {args.runs} run(s) each, max {args.num_predict} tokens")
|
|
print(f"Prompt: \"{args.prompt[:80]}...\"")
|
|
print(f"{'='*120}")
|
|
|
|
results = []
|
|
for model in selected:
|
|
for run in range(args.runs):
|
|
run_label = f"run {run+1}/{args.runs}" if args.runs > 1 else ""
|
|
print(f"\n[{model}] {run_label}")
|
|
r = benchmark_model(args.host, args.port, model, args.prompt, args.num_predict)
|
|
results.append(r)
|
|
|
|
if "error" in r:
|
|
print(f" ERROR: {r['error']}")
|
|
continue
|
|
|
|
print(f" Done reason: {r['done_reason']}")
|
|
tt = f"{r['ttft_thinking_s']:.3f}s" if r['ttft_thinking_s'] else "N/A"
|
|
tr = f"{r['ttft_response_s']:.3f}s" if r['ttft_response_s'] else "N/A"
|
|
print(f" TTFT thinking: {tt}")
|
|
print(f" TTFT response: {tr}")
|
|
print(f" Wall total: {r['wall_total_s']:.3f}s")
|
|
print(f" Tokens: prompt={r['prompt_tokens']} total_gen={r['eval_count']} (think~{r['thinking_tokens_est']} resp~{r['response_tokens_est']})")
|
|
print(f" Think phase: {r['thinking_duration_s']:.3f}s @ {r['thinking_tps']:.1f} tok/s")
|
|
print(f" Resp phase: {r['response_duration_s']:.3f}s @ {r['response_tps']:.1f} tok/s")
|
|
print(f" Overall tps: {r['overall_tps']:.1f} tok/s ({r['eval_count']} tok in {r['wall_total_s']:.1f}s)")
|
|
|
|
# Summary table
|
|
print(f"\n{'='*120}")
|
|
print("SUMMARY")
|
|
print(f"{'Model':<22} {'Reason':>8} {'TTFT-t':>7} {'TTFT-r':>7} {'Wall s':>7} {'Tokens':>7} {'Think t':>7} {'Think tps':>9} {'Resp t':>7} {'Resp tps':>9} {'Overall':>8}")
|
|
print("-" * 120)
|
|
for r in results:
|
|
if "error" in r:
|
|
print(f"{r['model']:<22} ERROR: {r['error'][:50]}")
|
|
continue
|
|
tt = f"{r['ttft_thinking_s']:.3f}" if r['ttft_thinking_s'] else " N/A"
|
|
tr = f"{r['ttft_response_s']:.3f}" if r['ttft_response_s'] else " N/A"
|
|
print(f"{r['model']:<22} {r['done_reason']:>8} {tt:>7} {tr:>7} {r['wall_total_s']:>7.3f} {r['eval_count']:>7} {r['thinking_duration_s']:>7.3f} {r['thinking_tps']:>9.1f} {r['response_duration_s']:>7.3f} {r['response_tps']:>9.1f} {r['overall_tps']:>8.1f}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |