feat: add native MLX inference benchmarks
This commit is contained in:
Christian Meyer
364
benchmarks/inference/coding_benchmark.py
Normal file
364
benchmarks/inference/coding_benchmark.py
Normal file
@@ -0,0 +1,364 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Measure end-to-end latency for deterministic, complex programming tasks."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import platform
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
from dataclasses import asdict, dataclass
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
TASKS = {
|
||||
"python_async_cache": """
|
||||
Act as a senior Python engineer. Implement a production-ready, single-file Python 3.12 module
|
||||
containing a generic AsyncTTLCache[K, V]. Requirements: weighted LRU capacity; per-entry TTL based
|
||||
on time.monotonic; single-flight get_or_create so concurrent misses for one key invoke the async
|
||||
factory exactly once; cancellation of one waiter must not cancel the shared factory; exceptions are
|
||||
never cached; invalidating a key while its factory is running must prevent stale publication; never
|
||||
hold the cache lock while awaiting user code; clear(), invalidate(), and immutable hit/miss/eviction
|
||||
statistics. Include concise pytest-asyncio tests covering races, cancellation, expiry, weighted
|
||||
capacity, factory failure, and invalidation during creation. Use only the standard library plus
|
||||
pytest. Output only the complete module, without Markdown fences or prose.
|
||||
""".strip(),
|
||||
"php_idempotent_transfer": """
|
||||
Act as a senior Laravel 12 engineer. Produce one self-contained PHP file that demonstrates an
|
||||
idempotent money-transfer service and its PHPUnit tests. Store money as integer cents. Requirements:
|
||||
validate positive amount and distinct accounts; use a database transaction; lock both account rows
|
||||
in deterministic ID order; reject insufficient funds; use a unique client idempotency key so retries
|
||||
return the original transfer; retry deadlocks without duplicating effects; persist an outbox event
|
||||
in the same transaction; and prove with tests that duplicate and competing requests preserve
|
||||
balances, create one transfer, and create one outbox event. Include minimal model/table assumptions
|
||||
as comments.
|
||||
Output only valid PHP, without Markdown fences or prose.
|
||||
""".strip(),
|
||||
"typescript_dag_scheduler": """
|
||||
Act as a senior TypeScript engineer. Implement a strict-mode, dependency-free DAG task scheduler in
|
||||
one TypeScript file. Each task has an ID, dependency IDs, and an async run(signal) function.
|
||||
Requirements: validate unknown dependencies and duplicate IDs; report an explicit cycle path; stable
|
||||
input-order scheduling; configurable maximum concurrency; AbortSignal support; fail fast after the
|
||||
first error while awaiting already-running tasks; never start dependents of failed tasks; return a
|
||||
Map of typed results on success; and avoid unhandled promise rejections. Include a compact test
|
||||
harness covering a diamond graph, concurrency enforcement, stable order, cycle detection, external
|
||||
cancellation, and rejection. Output only TypeScript, without Markdown fences or prose.
|
||||
""".strip(),
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Runtime:
|
||||
name: str
|
||||
kind: str
|
||||
base_url: str
|
||||
model: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class Result:
|
||||
task: str
|
||||
runtime: str
|
||||
model: str
|
||||
started_at: str
|
||||
prompt_chars: int
|
||||
max_tokens: int
|
||||
ttft_s: float | None = None
|
||||
total_s: float | None = None
|
||||
prompt_tokens: int | None = None
|
||||
output_tokens: int | None = None
|
||||
decode_tps: float | None = None
|
||||
finish_reason: str | None = None
|
||||
output_chars: int = 0
|
||||
output_sha256: str | None = None
|
||||
output: str = ""
|
||||
error: str | None = None
|
||||
|
||||
|
||||
def utc_now() -> str:
|
||||
return datetime.now(UTC).isoformat()
|
||||
|
||||
|
||||
def post_json(url: str, payload: dict[str, Any], timeout: float) -> dict[str, Any]:
|
||||
request = urllib.request.Request(
|
||||
url,
|
||||
data=json.dumps(payload).encode("utf-8"),
|
||||
headers={"Content-Type": "application/json"},
|
||||
method="POST",
|
||||
)
|
||||
with urllib.request.urlopen(request, timeout=timeout) as response:
|
||||
return json.load(response)
|
||||
|
||||
|
||||
def health(runtime: Runtime) -> None:
|
||||
path = "/api/tags" if runtime.kind == "ollama" else "/v1/models"
|
||||
with urllib.request.urlopen(f"{runtime.base_url.rstrip('/')}{path}", timeout=15):
|
||||
pass
|
||||
|
||||
|
||||
def consume_ollama(response: Any, started: float) -> tuple[float | None, str, dict[str, Any]]:
|
||||
first_at: float | None = None
|
||||
chunks: list[str] = []
|
||||
terminal: dict[str, Any] = {}
|
||||
for line in response:
|
||||
if not line.strip():
|
||||
continue
|
||||
event = json.loads(line)
|
||||
text = event.get("message", {}).get("content") or event.get("response") or ""
|
||||
if text:
|
||||
if first_at is None:
|
||||
first_at = time.perf_counter()
|
||||
chunks.append(text)
|
||||
if event.get("done"):
|
||||
terminal = event
|
||||
return (first_at - started if first_at else None, "".join(chunks), terminal)
|
||||
|
||||
|
||||
def consume_sse(
|
||||
response: Any, started: float
|
||||
) -> tuple[float | None, str, dict[str, Any], str | None]:
|
||||
first_at: float | None = None
|
||||
chunks: list[str] = []
|
||||
terminal: dict[str, Any] = {}
|
||||
finish_reason: str | None = None
|
||||
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
|
||||
for choice in event.get("choices") or []:
|
||||
text = choice.get("text") or choice.get("delta", {}).get("content") or ""
|
||||
if text:
|
||||
if first_at is None:
|
||||
first_at = time.perf_counter()
|
||||
chunks.append(text)
|
||||
finish_reason = choice.get("finish_reason") or finish_reason
|
||||
return (first_at - started if first_at else None, "".join(chunks), terminal, finish_reason)
|
||||
|
||||
|
||||
def request_payload(
|
||||
runtime: Runtime, prompt: str, max_tokens: int, seed: int
|
||||
) -> tuple[str, dict[str, Any]]:
|
||||
if runtime.kind == "ollama":
|
||||
return f"{runtime.base_url.rstrip('/')}/api/chat", {
|
||||
"model": runtime.model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"You are a senior software engineer. Follow every output constraint "
|
||||
"exactly and return a complete solution."
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
"think": False,
|
||||
"stream": True,
|
||||
"keep_alive": "30m",
|
||||
"options": {
|
||||
"temperature": 0,
|
||||
"seed": seed,
|
||||
"num_predict": max_tokens,
|
||||
"num_ctx": 8192,
|
||||
"draft_num_predict": 0,
|
||||
},
|
||||
}
|
||||
return f"{runtime.base_url.rstrip('/')}/v1/chat/completions", {
|
||||
"model": runtime.model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"You are a senior software engineer. Follow every output constraint exactly "
|
||||
"and return a complete solution."
|
||||
),
|
||||
},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True},
|
||||
"chat_template_kwargs": {"enable_thinking": False},
|
||||
"temperature": 0,
|
||||
"seed": seed,
|
||||
"max_tokens": max_tokens,
|
||||
}
|
||||
|
||||
|
||||
def run_task(
|
||||
runtime: Runtime,
|
||||
task_name: str,
|
||||
prompt: str,
|
||||
max_tokens: int,
|
||||
seed: int,
|
||||
timeout: float,
|
||||
) -> Result:
|
||||
result = Result(
|
||||
task=task_name,
|
||||
runtime=runtime.name,
|
||||
model=runtime.model,
|
||||
started_at=utc_now(),
|
||||
prompt_chars=len(prompt),
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
url, payload = request_payload(runtime, prompt, max_tokens, seed)
|
||||
request = urllib.request.Request(
|
||||
url,
|
||||
data=json.dumps(payload).encode("utf-8"),
|
||||
headers={"Content-Type": "application/json", "Accept": "text/event-stream"},
|
||||
method="POST",
|
||||
)
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
with urllib.request.urlopen(request, timeout=timeout) as response:
|
||||
if runtime.kind == "ollama":
|
||||
result.ttft_s, result.output, terminal = consume_ollama(response, started)
|
||||
result.prompt_tokens = terminal.get("prompt_eval_count")
|
||||
result.output_tokens = terminal.get("eval_count")
|
||||
result.finish_reason = terminal.get("done_reason")
|
||||
decode_ns = terminal.get("eval_duration")
|
||||
if decode_ns and result.output_tokens:
|
||||
result.decode_tps = result.output_tokens / (decode_ns / 1e9)
|
||||
else:
|
||||
result.ttft_s, result.output, terminal, result.finish_reason = consume_sse(
|
||||
response, started
|
||||
)
|
||||
usage = terminal.get("usage") or {}
|
||||
result.prompt_tokens = usage.get("prompt_tokens")
|
||||
result.output_tokens = usage.get("completion_tokens")
|
||||
result.total_s = time.perf_counter() - started
|
||||
if result.decode_tps is None and result.output_tokens and result.ttft_s is not None:
|
||||
decode_s = result.total_s - result.ttft_s
|
||||
if decode_s > 0:
|
||||
result.decode_tps = result.output_tokens / decode_s
|
||||
result.output_chars = len(result.output)
|
||||
result.output_sha256 = hashlib.sha256(result.output.encode()).hexdigest()
|
||||
if not result.output:
|
||||
result.error = "empty generated output"
|
||||
except urllib.error.HTTPError as error:
|
||||
detail = error.read().decode("utf-8", errors="replace")[:2000]
|
||||
result.error = f"HTTP {error.code}: {detail}"
|
||||
result.total_s = time.perf_counter() - started
|
||||
except (OSError, ValueError, json.JSONDecodeError) as error:
|
||||
result.error = f"{type(error).__name__}: {error}"
|
||||
result.total_s = time.perf_counter() - started
|
||||
return result
|
||||
|
||||
|
||||
def package_versions() -> dict[str, str | None]:
|
||||
script = (
|
||||
"import importlib.metadata as m,json; "
|
||||
"print(json.dumps({p:m.version(p) for p in "
|
||||
"['vllm','vllm-metal','mlx','mlx-lm','mlx-vlm']}))"
|
||||
)
|
||||
executable = Path.home() / ".venv-vllm-metal/bin/python"
|
||||
try:
|
||||
process = subprocess.run(
|
||||
[str(executable), "-c", script], capture_output=True, text=True, timeout=30, check=True
|
||||
)
|
||||
return json.loads(process.stdout)
|
||||
except (OSError, subprocess.SubprocessError, json.JSONDecodeError):
|
||||
return {}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--runtime", choices=("ollama", "vllm"), required=True)
|
||||
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("--max-tokens", type=int, default=1536)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--timeout", type=float, default=900)
|
||||
parser.add_argument("--output-dir", type=Path, default=Path(__file__).parent / "coding-results")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = parse_args()
|
||||
if args.max_tokens < 1:
|
||||
raise SystemExit("max-tokens must be positive")
|
||||
runtime = (
|
||||
Runtime("ollama", "ollama", args.ollama_url, args.ollama_model)
|
||||
if args.runtime == "ollama"
|
||||
else Runtime("vllm-metal", "openai", args.vllm_url, args.vllm_model)
|
||||
)
|
||||
try:
|
||||
health(runtime)
|
||||
except Exception as error:
|
||||
raise SystemExit(f"{runtime.name} is unavailable: {error}") from error
|
||||
|
||||
warmup = run_task(runtime, "warmup", "Output only the word READY.", 8, args.seed, args.timeout)
|
||||
if warmup.error:
|
||||
raise SystemExit(f"warmup failed: {warmup.error}")
|
||||
print(
|
||||
f"warmup runtime={runtime.name} total={warmup.total_s:.3f}s "
|
||||
f"tokens={warmup.output_tokens}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
results: list[Result] = []
|
||||
for index, (task_name, prompt) in enumerate(TASKS.items(), start=1):
|
||||
result = run_task(runtime, task_name, prompt, args.max_tokens, args.seed, args.timeout)
|
||||
results.append(result)
|
||||
print(
|
||||
f"[{index}/{len(TASKS)}] task={task_name} total={result.total_s:.3f}s "
|
||||
f"ttft={result.ttft_s:.3f}s tokens={result.output_tokens} "
|
||||
f"decode={result.decode_tps:.2f}tok/s finish={result.finish_reason} "
|
||||
f"error={result.error or 'none'}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
run_id = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
|
||||
output_dir = args.output_dir / f"{run_id}-{runtime.name}-{runtime.model.replace(':', '_')}"
|
||||
output_dir.mkdir(parents=True, exist_ok=False)
|
||||
raw = {
|
||||
"run_id": run_id,
|
||||
"created_at": utc_now(),
|
||||
"runtime": asdict(runtime),
|
||||
"host": {
|
||||
"platform": platform.platform(),
|
||||
"machine": platform.machine(),
|
||||
"python": platform.python_version(),
|
||||
},
|
||||
"package_versions": package_versions(),
|
||||
"settings": {
|
||||
"temperature": 0,
|
||||
"seed": args.seed,
|
||||
"max_tokens": args.max_tokens,
|
||||
"context_tokens": 8192,
|
||||
"raw_completion": False,
|
||||
"chat_template": True,
|
||||
"thinking": False,
|
||||
"ollama_speculative_drafting": False,
|
||||
},
|
||||
"warmup": asdict(warmup),
|
||||
"tasks": TASKS,
|
||||
"results": [asdict(result) for result in results],
|
||||
}
|
||||
(output_dir / "results.json").write_text(json.dumps(raw, indent=2), encoding="utf-8")
|
||||
summary = [
|
||||
{
|
||||
key: value
|
||||
for key, value in asdict(result).items()
|
||||
if key not in {"output"}
|
||||
}
|
||||
for result in results
|
||||
]
|
||||
(output_dir / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
|
||||
print(f"Results: {output_dir}")
|
||||
return 1 if any(result.error for result in results) else 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
Reference in New Issue
Block a user