feat: add native MLX inference benchmarks
This commit is contained in:
Christian Meyer
500
benchmarks/inference/benchmark.py
Normal file
500
benchmarks/inference/benchmark.py
Normal file
@@ -0,0 +1,500 @@
|
||||
#!/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())
|
||||
Reference in New Issue
Block a user