501 lines
18 KiB
Python
501 lines
18 KiB
Python
#!/usr/bin/env python3
|
|
"""Reproducible HTTP benchmark for Ollama and OpenAI-compatible runtimes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import concurrent.futures
|
|
import csv
|
|
import json
|
|
import math
|
|
import platform
|
|
import statistics
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
import uuid
|
|
from dataclasses import asdict, dataclass
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
PROMPTS = {
|
|
"short": (
|
|
"In exactly four concise bullet points, explain why database transactions are useful. "
|
|
"Do not use a heading."
|
|
),
|
|
"medium": (
|
|
"You are reviewing a Laravel service that transfers money between two accounts. "
|
|
"Explain the failure modes involving concurrency, partial writes, retries, and audit logs. "
|
|
"Then provide a robust implementation plan. Keep the response technical and concise."
|
|
),
|
|
"code": (
|
|
"Write a Python 3.12 function named merge_intervals that accepts a list of integer pairs, "
|
|
"validates each pair, merges overlapping or adjacent intervals, and returns sorted tuples. "
|
|
"Include type hints, a short docstring, and five assert-based examples. Output only code."
|
|
),
|
|
}
|
|
|
|
CSV_FIELDS = (
|
|
"run_id",
|
|
"runtime",
|
|
"model",
|
|
"phase",
|
|
"concurrency",
|
|
"round",
|
|
"slot",
|
|
"request_id",
|
|
"prompt_name",
|
|
"prompt_chars",
|
|
"max_tokens",
|
|
"started_at",
|
|
"ttft_s",
|
|
"total_s",
|
|
"prompt_tokens",
|
|
"output_tokens",
|
|
"prompt_tps",
|
|
"decode_tps",
|
|
"request_tps",
|
|
"group_elapsed_s",
|
|
"group_output_tps",
|
|
"error",
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RuntimeConfig:
|
|
name: str
|
|
kind: str
|
|
base_url: str
|
|
model: str
|
|
|
|
|
|
@dataclass
|
|
class Measurement:
|
|
run_id: str
|
|
runtime: str
|
|
model: str
|
|
phase: str
|
|
concurrency: int
|
|
round: int
|
|
slot: int
|
|
request_id: str
|
|
prompt_name: str
|
|
prompt_chars: int
|
|
max_tokens: int
|
|
started_at: str
|
|
ttft_s: float | None = None
|
|
total_s: float | None = None
|
|
prompt_tokens: int | None = None
|
|
output_tokens: int | None = None
|
|
prompt_tps: float | None = None
|
|
decode_tps: float | None = None
|
|
request_tps: float | None = None
|
|
group_elapsed_s: float | None = None
|
|
group_output_tps: float | None = None
|
|
error: str | None = None
|
|
|
|
|
|
def utc_now() -> str:
|
|
return datetime.now(UTC).isoformat()
|
|
|
|
|
|
def percentile(values: list[float], fraction: float) -> float | None:
|
|
if not values:
|
|
return None
|
|
ordered = sorted(values)
|
|
position = (len(ordered) - 1) * fraction
|
|
lower = math.floor(position)
|
|
upper = math.ceil(position)
|
|
if lower == upper:
|
|
return ordered[lower]
|
|
return ordered[lower] + (ordered[upper] - ordered[lower]) * (position - lower)
|
|
|
|
|
|
def safe_version(command: list[str]) -> str | None:
|
|
try:
|
|
result = subprocess.run(command, capture_output=True, text=True, timeout=15, check=False)
|
|
except (OSError, subprocess.SubprocessError):
|
|
return None
|
|
value = (result.stdout or result.stderr).strip()
|
|
return value or None
|
|
|
|
|
|
def http_json(url: str, timeout: float = 15) -> dict[str, Any]:
|
|
request = urllib.request.Request(url, headers={"Accept": "application/json"})
|
|
with urllib.request.urlopen(request, timeout=timeout) as response:
|
|
return json.load(response)
|
|
|
|
|
|
def read_ndjson(response: Any, started: float) -> tuple[float | None, dict[str, Any], int]:
|
|
first_token_at: float | None = None
|
|
output_chars = 0
|
|
terminal: dict[str, Any] = {}
|
|
for raw_line in response:
|
|
if not raw_line.strip():
|
|
continue
|
|
event = json.loads(raw_line)
|
|
content = event.get("response", "")
|
|
if content:
|
|
output_chars += len(content)
|
|
if first_token_at is None:
|
|
first_token_at = time.perf_counter()
|
|
if event.get("done"):
|
|
terminal = event
|
|
return (first_token_at - started if first_token_at else None, terminal, output_chars)
|
|
|
|
|
|
def read_sse(response: Any, started: float) -> tuple[float | None, dict[str, Any], int]:
|
|
first_token_at: float | None = None
|
|
output_chars = 0
|
|
terminal: dict[str, Any] = {}
|
|
for raw_line in response:
|
|
line = raw_line.decode("utf-8").strip()
|
|
if not line.startswith("data:"):
|
|
continue
|
|
payload = line[5:].strip()
|
|
if payload == "[DONE]":
|
|
break
|
|
event = json.loads(payload)
|
|
terminal = event
|
|
choices = event.get("choices") or []
|
|
if choices:
|
|
text = choices[0].get("text") or choices[0].get("delta", {}).get("content") or ""
|
|
if text:
|
|
output_chars += len(text)
|
|
if first_token_at is None:
|
|
first_token_at = time.perf_counter()
|
|
return (first_token_at - started if first_token_at else None, terminal, output_chars)
|
|
|
|
|
|
def benchmark_request(
|
|
runtime: RuntimeConfig,
|
|
prompt_name: str,
|
|
prompt: str,
|
|
max_tokens: int,
|
|
seed: int,
|
|
timeout: float,
|
|
run_id: str,
|
|
phase: str,
|
|
concurrency: int,
|
|
round_number: int,
|
|
slot: int,
|
|
) -> Measurement:
|
|
measurement = Measurement(
|
|
run_id=run_id,
|
|
runtime=runtime.name,
|
|
model=runtime.model,
|
|
phase=phase,
|
|
concurrency=concurrency,
|
|
round=round_number,
|
|
slot=slot,
|
|
request_id=str(uuid.uuid4()),
|
|
prompt_name=prompt_name,
|
|
prompt_chars=len(prompt),
|
|
max_tokens=max_tokens,
|
|
started_at=utc_now(),
|
|
)
|
|
headers = {"Content-Type": "application/json", "Accept": "text/event-stream"}
|
|
if runtime.kind == "ollama":
|
|
url = f"{runtime.base_url.rstrip('/')}/api/generate"
|
|
payload = {
|
|
"model": runtime.model,
|
|
"prompt": prompt,
|
|
"stream": True,
|
|
"raw": True,
|
|
"keep_alive": "30m",
|
|
"options": {
|
|
"temperature": 0,
|
|
"seed": seed,
|
|
"num_predict": max_tokens,
|
|
"num_ctx": 8192,
|
|
"draft_num_predict": 0,
|
|
},
|
|
}
|
|
else:
|
|
url = f"{runtime.base_url.rstrip('/')}/v1/completions"
|
|
payload = {
|
|
"model": runtime.model,
|
|
"prompt": prompt,
|
|
"stream": True,
|
|
"stream_options": {"include_usage": True},
|
|
"temperature": 0,
|
|
"seed": seed,
|
|
"max_tokens": max_tokens,
|
|
}
|
|
|
|
started = time.perf_counter()
|
|
try:
|
|
request = urllib.request.Request(
|
|
url,
|
|
data=json.dumps(payload).encode("utf-8"),
|
|
headers=headers,
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(request, timeout=timeout) as response:
|
|
if runtime.kind == "ollama":
|
|
ttft, terminal, output_chars = read_ndjson(response, started)
|
|
measurement.prompt_tokens = terminal.get("prompt_eval_count")
|
|
measurement.output_tokens = terminal.get("eval_count")
|
|
prompt_ns = terminal.get("prompt_eval_duration")
|
|
decode_ns = terminal.get("eval_duration")
|
|
if prompt_ns and measurement.prompt_tokens:
|
|
measurement.prompt_tps = measurement.prompt_tokens / (prompt_ns / 1e9)
|
|
if decode_ns and measurement.output_tokens:
|
|
measurement.decode_tps = measurement.output_tokens / (decode_ns / 1e9)
|
|
else:
|
|
ttft, terminal, output_chars = read_sse(response, started)
|
|
usage = terminal.get("usage") or {}
|
|
measurement.prompt_tokens = usage.get("prompt_tokens")
|
|
measurement.output_tokens = usage.get("completion_tokens")
|
|
measurement.ttft_s = ttft
|
|
measurement.total_s = time.perf_counter() - started
|
|
if measurement.output_tokens and measurement.total_s:
|
|
measurement.request_tps = measurement.output_tokens / measurement.total_s
|
|
if measurement.decode_tps is None and measurement.ttft_s is not None:
|
|
decode_time = measurement.total_s - measurement.ttft_s
|
|
if decode_time > 0:
|
|
measurement.decode_tps = measurement.output_tokens / decode_time
|
|
if output_chars == 0:
|
|
measurement.error = "stream contained no generated text"
|
|
except urllib.error.HTTPError as error:
|
|
detail = error.read().decode("utf-8", errors="replace")[:1000]
|
|
measurement.error = f"HTTP {error.code}: {detail}"
|
|
measurement.total_s = time.perf_counter() - started
|
|
except (OSError, ValueError, json.JSONDecodeError) as error:
|
|
measurement.error = f"{type(error).__name__}: {error}"
|
|
measurement.total_s = time.perf_counter() - started
|
|
return measurement
|
|
|
|
|
|
def run_round(
|
|
runtime: RuntimeConfig,
|
|
prompt_items: list[tuple[str, str]],
|
|
max_tokens: int,
|
|
seed: int,
|
|
timeout: float,
|
|
run_id: str,
|
|
phase: str,
|
|
concurrency: int,
|
|
round_number: int,
|
|
) -> tuple[list[Measurement], float]:
|
|
barrier = threading.Barrier(concurrency)
|
|
|
|
def execute(slot: int) -> Measurement:
|
|
prompt_name, prompt = prompt_items[(round_number * concurrency + slot) % len(prompt_items)]
|
|
barrier.wait()
|
|
return benchmark_request(
|
|
runtime,
|
|
prompt_name,
|
|
prompt,
|
|
max_tokens,
|
|
seed,
|
|
timeout,
|
|
run_id,
|
|
phase,
|
|
concurrency,
|
|
round_number,
|
|
slot,
|
|
)
|
|
|
|
group_started = time.perf_counter()
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as executor:
|
|
measurements = list(executor.map(execute, range(concurrency)))
|
|
group_elapsed = time.perf_counter() - group_started
|
|
successful_tokens = sum(item.output_tokens or 0 for item in measurements if not item.error)
|
|
group_tps = successful_tokens / group_elapsed if group_elapsed else 0.0
|
|
for item in measurements:
|
|
item.group_elapsed_s = group_elapsed
|
|
item.group_output_tps = group_tps
|
|
return measurements, group_elapsed
|
|
|
|
|
|
def summarize(measurements: list[Measurement]) -> list[dict[str, Any]]:
|
|
groups: dict[tuple[str, int], list[Measurement]] = {}
|
|
for measurement in measurements:
|
|
if measurement.phase != "measured":
|
|
continue
|
|
groups.setdefault((measurement.runtime, measurement.concurrency), []).append(measurement)
|
|
|
|
summaries: list[dict[str, Any]] = []
|
|
for (runtime, concurrency), items in sorted(groups.items()):
|
|
successful = [item for item in items if not item.error]
|
|
summary: dict[str, Any] = {
|
|
"runtime": runtime,
|
|
"model": items[0].model,
|
|
"concurrency": concurrency,
|
|
"requests": len(items),
|
|
"successes": len(successful),
|
|
"failures": len(items) - len(successful),
|
|
}
|
|
for field in ("ttft_s", "total_s", "prompt_tps", "decode_tps", "request_tps"):
|
|
values = [
|
|
float(value)
|
|
for item in successful
|
|
if (value := getattr(item, field)) is not None
|
|
]
|
|
summary[f"{field}_mean"] = statistics.fmean(values) if values else None
|
|
summary[f"{field}_p50"] = percentile(values, 0.50)
|
|
summary[f"{field}_p95"] = percentile(values, 0.95)
|
|
group_tps = [
|
|
float(item.group_output_tps)
|
|
for item in successful
|
|
if item.slot == 0 and item.group_output_tps is not None
|
|
]
|
|
summary["aggregate_output_tps_mean"] = statistics.fmean(group_tps) if group_tps else None
|
|
summary["aggregate_output_tps_p50"] = percentile(group_tps, 0.50)
|
|
summary["aggregate_output_tps_p95"] = percentile(group_tps, 0.95)
|
|
summaries.append(summary)
|
|
return summaries
|
|
|
|
|
|
def write_csv(path: Path, measurements: list[Measurement]) -> None:
|
|
with path.open("w", newline="", encoding="utf-8") as handle:
|
|
writer = csv.DictWriter(handle, fieldnames=CSV_FIELDS, lineterminator="\n")
|
|
writer.writeheader()
|
|
for measurement in measurements:
|
|
writer.writerow(asdict(measurement))
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--output-dir", type=Path, default=Path(__file__).parent / "results")
|
|
parser.add_argument("--runtime", choices=("all", "ollama", "vllm"), default="all")
|
|
parser.add_argument("--ollama-url", default="http://127.0.0.1:11434")
|
|
parser.add_argument("--ollama-model", default="qwen3.8:27b-mlx")
|
|
parser.add_argument("--vllm-url", default="http://127.0.0.1:8001")
|
|
parser.add_argument("--vllm-model", default="qwen3.8-27b-nvfp4")
|
|
parser.add_argument("--warmups", type=int, default=2)
|
|
parser.add_argument("--rounds", type=int, default=5)
|
|
parser.add_argument("--concurrency", type=int, nargs="+", default=(1, 4))
|
|
parser.add_argument("--max-tokens", type=int, default=128)
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
parser.add_argument("--timeout", type=float, default=600)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
if args.warmups < 0 or args.rounds < 1 or any(value < 1 for value in args.concurrency):
|
|
raise SystemExit("warmups must be >= 0; rounds and concurrency must be >= 1")
|
|
|
|
runtimes = {
|
|
"ollama": RuntimeConfig("ollama", "ollama", args.ollama_url, args.ollama_model),
|
|
"vllm": RuntimeConfig("vllm-metal", "openai", args.vllm_url, args.vllm_model),
|
|
}
|
|
selected = list(runtimes.values()) if args.runtime == "all" else [runtimes[args.runtime]]
|
|
for runtime in selected:
|
|
health_url = (
|
|
f"{runtime.base_url.rstrip('/')}/api/tags"
|
|
if runtime.kind == "ollama"
|
|
else f"{runtime.base_url.rstrip('/')}/v1/models"
|
|
)
|
|
try:
|
|
http_json(health_url)
|
|
except Exception as error:
|
|
raise SystemExit(f"{runtime.name} is not ready at {health_url}: {error}") from error
|
|
|
|
run_id = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
|
|
output_dir = args.output_dir / run_id
|
|
output_dir.mkdir(parents=True, exist_ok=False)
|
|
metadata = {
|
|
"run_id": run_id,
|
|
"created_at": utc_now(),
|
|
"host": {
|
|
"platform": platform.platform(),
|
|
"machine": platform.machine(),
|
|
"python": platform.python_version(),
|
|
},
|
|
"arguments": vars(args) | {"output_dir": str(args.output_dir)},
|
|
"runtimes": [asdict(runtime) for runtime in selected],
|
|
"versions": {
|
|
"ollama": safe_version(["ollama", "--version"]),
|
|
"vllm": safe_version([str(Path.home() / ".venv-vllm-metal/bin/vllm"), "--version"]),
|
|
},
|
|
"model_metadata": {},
|
|
"prompts": PROMPTS,
|
|
"measurement_notes": {
|
|
"ttft": (
|
|
"Client wall time from request start until first non-empty generated text chunk."
|
|
),
|
|
"total": "Client wall time from request start until stream completion.",
|
|
"decode_tps": "Ollama server eval metric; otherwise output tokens / (total - TTFT).",
|
|
"aggregate_output_tps": "Sum of successful output tokens / concurrent group wall time.",
|
|
},
|
|
}
|
|
for runtime in selected:
|
|
url = (
|
|
f"{runtime.base_url.rstrip('/')}/api/show"
|
|
if runtime.kind == "ollama"
|
|
else f"{runtime.base_url.rstrip('/')}/v1/models"
|
|
)
|
|
try:
|
|
if runtime.kind == "ollama":
|
|
request = urllib.request.Request(
|
|
url,
|
|
data=json.dumps({"model": runtime.model}).encode(),
|
|
headers={"Content-Type": "application/json"},
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(request, timeout=30) as response:
|
|
metadata["model_metadata"][runtime.name] = json.load(response)
|
|
else:
|
|
metadata["model_metadata"][runtime.name] = http_json(url)
|
|
except Exception as error:
|
|
metadata["model_metadata"][runtime.name] = {"metadata_error": str(error)}
|
|
|
|
prompt_items = list(PROMPTS.items())
|
|
measurements: list[Measurement] = []
|
|
total_groups = len(selected) * sum(args.warmups + args.rounds for _ in args.concurrency)
|
|
completed_groups = 0
|
|
for runtime in selected:
|
|
for concurrency in args.concurrency:
|
|
for phase, rounds in (("warmup", args.warmups), ("measured", args.rounds)):
|
|
for round_number in range(rounds):
|
|
items, elapsed = run_round(
|
|
runtime,
|
|
prompt_items,
|
|
args.max_tokens,
|
|
args.seed,
|
|
args.timeout,
|
|
run_id,
|
|
phase,
|
|
concurrency,
|
|
round_number,
|
|
)
|
|
measurements.extend(items)
|
|
completed_groups += 1
|
|
errors = sum(bool(item.error) for item in items)
|
|
print(
|
|
f"[{completed_groups}/{total_groups}] {runtime.name} phase={phase} "
|
|
f"concurrency={concurrency} round={round_number + 1}/{rounds} "
|
|
f"elapsed={elapsed:.3f}s errors={errors}",
|
|
flush=True,
|
|
)
|
|
|
|
summaries = summarize(measurements)
|
|
metadata["completed_at"] = utc_now()
|
|
metadata["summary"] = summaries
|
|
(output_dir / "raw.json").write_text(
|
|
json.dumps([asdict(item) | {
|
|
"group_elapsed_s": getattr(item, "group_elapsed_s", None),
|
|
"group_output_tps": getattr(item, "group_output_tps", None),
|
|
} for item in measurements], indent=2),
|
|
encoding="utf-8",
|
|
)
|
|
(output_dir / "metadata.json").write_text(
|
|
json.dumps(metadata, indent=2, default=str), encoding="utf-8"
|
|
)
|
|
(output_dir / "summary.json").write_text(json.dumps(summaries, indent=2), encoding="utf-8")
|
|
write_csv(output_dir / "raw.csv", measurements)
|
|
print(json.dumps(summaries, indent=2))
|
|
print(f"Results: {output_dir}")
|
|
return 1 if any(item.error for item in measurements if item.phase == "measured") else 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|