365 lines
14 KiB
Python
365 lines
14 KiB
Python
#!/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())
|