feat: add automated MTP benchmark runner for llama-server via podman containers
This commit is contained in:
Donato Capitella
Executable
+431
@@ -0,0 +1,431 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
MTP Benchmark Runner — Automated Multi-Token Prediction benchmarking.
|
||||||
|
|
||||||
|
Auto-discovers MTP GGUF models, starts llama-server in podman containers
|
||||||
|
(with and without MTP), runs mtp-bench.py against each configuration,
|
||||||
|
and collects structured JSON results.
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
python run_mtp_bench.py # run everything
|
||||||
|
python run_mtp_bench.py --model "Qwen3.6-35B" # filter by model name
|
||||||
|
python run_mtp_bench.py --toolbox vulkan-radv-mtp # filter by toolbox
|
||||||
|
python run_mtp_bench.py --models-dir /path/to/models
|
||||||
|
python run_mtp_bench.py --port 8081
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import datetime
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import platform
|
||||||
|
import subprocess
|
||||||
|
import sys
|
||||||
|
import time
|
||||||
|
from pathlib import Path
|
||||||
|
from urllib import request
|
||||||
|
from urllib.error import URLError
|
||||||
|
|
||||||
|
# ── Toolbox definitions ──────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
TOOLBOXES = {
|
||||||
|
"rocm-7.2.3-mtp": {
|
||||||
|
"image": "docker.io/kyuz0/amd-strix-halo-toolboxes:rocm-7.2.3-mtp",
|
||||||
|
"engine_args": [
|
||||||
|
"--device", "/dev/dri",
|
||||||
|
"--device", "/dev/kfd",
|
||||||
|
"--group-add", "video",
|
||||||
|
"--group-add", "render",
|
||||||
|
"--security-opt", "seccomp=unconfined",
|
||||||
|
],
|
||||||
|
},
|
||||||
|
"vulkan-radv-mtp": {
|
||||||
|
"image": "docker.io/kyuz0/amd-strix-halo-toolboxes:vulkan-radv-mtp",
|
||||||
|
"engine_args": [
|
||||||
|
"--device", "/dev/dri",
|
||||||
|
"--group-add", "video",
|
||||||
|
"--security-opt", "seccomp=unconfined",
|
||||||
|
],
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
# ── Benchmark modes ──────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
MODES = {
|
||||||
|
"baseline": [],
|
||||||
|
"mtp-2": ["--spec-type", "draft-mtp", "--spec-draft-n-max", "2", "-np", "1"],
|
||||||
|
"mtp-3": ["--spec-type", "draft-mtp", "--spec-draft-n-max", "3", "-np", "1"],
|
||||||
|
}
|
||||||
|
|
||||||
|
# ── Constants ────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
CONTAINER_NAME = "mtp-bench-server"
|
||||||
|
HEALTH_TIMEOUT = 180 # seconds to wait for server readiness
|
||||||
|
HEALTH_INTERVAL = 3 # seconds between health polls
|
||||||
|
COOLDOWN = 5 # seconds between runs
|
||||||
|
BENCH_SCRIPT = Path(__file__).parent / "mtp-bench.py"
|
||||||
|
|
||||||
|
|
||||||
|
# ── Model discovery ──────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def discover_models(models_dir: Path) -> list[dict]:
|
||||||
|
"""Scan models_dir for GGUF files with 'MTP' in their path."""
|
||||||
|
models = []
|
||||||
|
if not models_dir.is_dir():
|
||||||
|
print(f"❌ Models directory not found: {models_dir}")
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
for gguf in sorted(models_dir.rglob("*.gguf")):
|
||||||
|
rel = gguf.relative_to(models_dir)
|
||||||
|
# Must have MTP somewhere in the path (case-insensitive)
|
||||||
|
if "mtp" not in str(rel).lower():
|
||||||
|
continue
|
||||||
|
# Skip non-first shards of multi-shard models
|
||||||
|
name = gguf.name
|
||||||
|
if "-000" in name and "-00001-of-" not in name:
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Derive display name from parent directory or filename
|
||||||
|
if gguf.parent != models_dir:
|
||||||
|
display_name = gguf.parent.name
|
||||||
|
else:
|
||||||
|
display_name = gguf.stem
|
||||||
|
|
||||||
|
models.append({
|
||||||
|
"name": display_name,
|
||||||
|
"gguf": str(rel),
|
||||||
|
})
|
||||||
|
|
||||||
|
return models
|
||||||
|
|
||||||
|
|
||||||
|
# ── System info ──────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def capture_system_info(results_dir: Path):
|
||||||
|
"""Write system_info.json if it doesn't exist."""
|
||||||
|
path = results_dir / "system_info.json"
|
||||||
|
if path.exists():
|
||||||
|
return
|
||||||
|
|
||||||
|
def get_distro():
|
||||||
|
try:
|
||||||
|
with open("/etc/os-release") as f:
|
||||||
|
for line in f:
|
||||||
|
if line.startswith("PRETTY_NAME="):
|
||||||
|
return line.split("=", 1)[1].strip().strip('"')
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
return "Linux"
|
||||||
|
|
||||||
|
def get_linux_firmware():
|
||||||
|
try:
|
||||||
|
r = subprocess.run(["rpm", "-q", "linux-firmware"],
|
||||||
|
capture_output=True, text=True)
|
||||||
|
if r.returncode == 0:
|
||||||
|
return r.stdout.strip()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
|
return "unknown"
|
||||||
|
|
||||||
|
info = {
|
||||||
|
"distro": get_distro(),
|
||||||
|
"kernel": platform.release(),
|
||||||
|
"linux_firmware": get_linux_firmware(),
|
||||||
|
"timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||||||
|
}
|
||||||
|
path.write_text(json.dumps(info, indent=2) + "\n")
|
||||||
|
print(f"📋 Captured system info → {path}")
|
||||||
|
|
||||||
|
|
||||||
|
# ── Container lifecycle ──────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def stop_container():
|
||||||
|
"""Stop and remove the benchmark container if it exists."""
|
||||||
|
subprocess.run(["podman", "stop", "-t", "5", CONTAINER_NAME],
|
||||||
|
capture_output=True)
|
||||||
|
subprocess.run(["podman", "rm", "-f", CONTAINER_NAME],
|
||||||
|
capture_output=True)
|
||||||
|
|
||||||
|
|
||||||
|
def start_server(toolbox: dict, gguf: str, models_dir: Path,
|
||||||
|
spec_flags: list[str], port: int) -> bool:
|
||||||
|
"""Start llama-server in a podman container. Returns True on success."""
|
||||||
|
stop_container()
|
||||||
|
|
||||||
|
cmd = [
|
||||||
|
"podman", "run", "--rm", "-d",
|
||||||
|
"--name", CONTAINER_NAME,
|
||||||
|
"--security-opt", "label=disable",
|
||||||
|
"--userns=keep-id",
|
||||||
|
*toolbox["engine_args"],
|
||||||
|
"-v", f"{models_dir}:/models:ro",
|
||||||
|
"-p", f"127.0.0.1:{port}:{port}",
|
||||||
|
toolbox["image"],
|
||||||
|
"llama-server",
|
||||||
|
"-m", f"/models/{gguf}",
|
||||||
|
"-c", "12288",
|
||||||
|
"-ngl", "999",
|
||||||
|
"--host", "0.0.0.0",
|
||||||
|
"--port", str(port),
|
||||||
|
"--no-mmap",
|
||||||
|
"-fa", "1",
|
||||||
|
"--jinja",
|
||||||
|
*spec_flags,
|
||||||
|
]
|
||||||
|
|
||||||
|
print(f" 🐳 Starting container...")
|
||||||
|
print(f" {' '.join(cmd)}")
|
||||||
|
result = subprocess.run(cmd, capture_output=True, text=True)
|
||||||
|
if result.returncode != 0:
|
||||||
|
print(f" ❌ Failed to start container: {result.stderr.strip()}")
|
||||||
|
return False
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
def wait_for_health(port: int) -> bool:
|
||||||
|
"""Poll /health until the server is ready."""
|
||||||
|
url = f"http://127.0.0.1:{port}/health"
|
||||||
|
deadline = time.time() + HEALTH_TIMEOUT
|
||||||
|
print(f" ⏳ Waiting for server health ({HEALTH_TIMEOUT}s timeout)...", end="", flush=True)
|
||||||
|
|
||||||
|
while time.time() < deadline:
|
||||||
|
try:
|
||||||
|
req = request.Request(url, method="GET")
|
||||||
|
with request.urlopen(req, timeout=5) as r:
|
||||||
|
if r.status == 200:
|
||||||
|
data = json.loads(r.read())
|
||||||
|
if data.get("status") == "ok":
|
||||||
|
print(f" ✅ ready")
|
||||||
|
return True
|
||||||
|
except (URLError, OSError, json.JSONDecodeError):
|
||||||
|
pass
|
||||||
|
print(".", end="", flush=True)
|
||||||
|
time.sleep(HEALTH_INTERVAL)
|
||||||
|
|
||||||
|
print(f" ❌ timeout")
|
||||||
|
# Dump container logs for debugging
|
||||||
|
logs = subprocess.run(["podman", "logs", "--tail", "30", CONTAINER_NAME],
|
||||||
|
capture_output=True, text=True)
|
||||||
|
if logs.stdout:
|
||||||
|
print(f" 📝 Last 30 lines of server logs:")
|
||||||
|
for line in logs.stdout.strip().split("\n")[-15:]:
|
||||||
|
print(f" {line}")
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
# ── Benchmark execution ──────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def run_benchmark(port: int, out_path: Path) -> dict | None:
|
||||||
|
"""Run mtp-bench.py as subprocess. Returns parsed JSON or None."""
|
||||||
|
cmd = [
|
||||||
|
sys.executable, str(BENCH_SCRIPT),
|
||||||
|
"--url", f"http://127.0.0.1:{port}",
|
||||||
|
"--out", str(out_path),
|
||||||
|
]
|
||||||
|
print(f" 🔬 Running benchmark...")
|
||||||
|
result = subprocess.run(cmd, capture_output=False)
|
||||||
|
if result.returncode != 0:
|
||||||
|
print(f" ❌ Benchmark failed (exit {result.returncode})")
|
||||||
|
return None
|
||||||
|
|
||||||
|
try:
|
||||||
|
return json.loads(out_path.read_text())
|
||||||
|
except Exception as e:
|
||||||
|
print(f" ❌ Failed to read results: {e}")
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
# ── Results handling ─────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def save_result(out_path: Path, model: dict, toolbox_name: str,
|
||||||
|
mode: str, spec_flags: list[str], bench_data: dict):
|
||||||
|
"""Wrap mtp-bench output with run metadata and save."""
|
||||||
|
wrapped = {
|
||||||
|
"model": model["name"],
|
||||||
|
"gguf": model["gguf"],
|
||||||
|
"toolbox": toolbox_name,
|
||||||
|
"mode": mode,
|
||||||
|
"spec_flags": " ".join(spec_flags) if spec_flags else "(none)",
|
||||||
|
"timestamp": datetime.datetime.now().isoformat(),
|
||||||
|
**bench_data,
|
||||||
|
}
|
||||||
|
out_path.write_text(json.dumps(wrapped, indent=2) + "\n")
|
||||||
|
|
||||||
|
|
||||||
|
def print_summary(results_dir: Path):
|
||||||
|
"""Read all result JSONs and print a summary table."""
|
||||||
|
results = []
|
||||||
|
for f in sorted(results_dir.glob("*.json")):
|
||||||
|
if f.name == "system_info.json" or f.name == "summary.json":
|
||||||
|
continue
|
||||||
|
try:
|
||||||
|
data = json.loads(f.read_text())
|
||||||
|
if "aggregate" in data:
|
||||||
|
results.append(data)
|
||||||
|
except Exception:
|
||||||
|
continue
|
||||||
|
|
||||||
|
if not results:
|
||||||
|
print("\n📊 No results found.")
|
||||||
|
return
|
||||||
|
|
||||||
|
# Compute average tok/s per result
|
||||||
|
for r in results:
|
||||||
|
prompts = r.get("results", [])
|
||||||
|
if prompts:
|
||||||
|
r["_avg_toks"] = sum(p.get("predicted_per_second", 0) for p in prompts) / len(prompts)
|
||||||
|
else:
|
||||||
|
r["_avg_toks"] = 0
|
||||||
|
|
||||||
|
# Build baseline lookup for speedup calculation
|
||||||
|
baselines = {}
|
||||||
|
for r in results:
|
||||||
|
if r["mode"] == "baseline":
|
||||||
|
key = (r["model"], r["toolbox"])
|
||||||
|
baselines[key] = r["_avg_toks"]
|
||||||
|
|
||||||
|
# Print table
|
||||||
|
print("\n" + "=" * 100)
|
||||||
|
print(f"{'Model':<30} {'Toolbox':<20} {'Mode':<10} {'Avg tok/s':>10} {'Accept%':>9} {'Wall(s)':>8} {'Speedup':>8}")
|
||||||
|
print("-" * 100)
|
||||||
|
|
||||||
|
for r in results:
|
||||||
|
agg = r.get("aggregate", {})
|
||||||
|
accept = agg.get("aggregate_accept_rate")
|
||||||
|
wall = agg.get("wall_s_total", 0)
|
||||||
|
accept_str = f"{accept * 100:.1f}%" if accept is not None else "—"
|
||||||
|
avg_toks = r["_avg_toks"]
|
||||||
|
|
||||||
|
# Speedup relative to baseline
|
||||||
|
baseline_key = (r["model"], r["toolbox"])
|
||||||
|
baseline_toks = baselines.get(baseline_key)
|
||||||
|
if baseline_toks and baseline_toks > 0:
|
||||||
|
speedup = f"{avg_toks / baseline_toks:.2f}×"
|
||||||
|
else:
|
||||||
|
speedup = "—"
|
||||||
|
|
||||||
|
print(f"{r['model']:<30} {r['toolbox']:<20} {r['mode']:<10} {avg_toks:>10.1f} {accept_str:>9} {wall:>8.1f} {speedup:>8}")
|
||||||
|
|
||||||
|
print("=" * 100)
|
||||||
|
|
||||||
|
# Write summary.json
|
||||||
|
summary_data = []
|
||||||
|
for r in results:
|
||||||
|
agg = r.get("aggregate", {})
|
||||||
|
summary_data.append({
|
||||||
|
"model": r["model"],
|
||||||
|
"toolbox": r["toolbox"],
|
||||||
|
"mode": r["mode"],
|
||||||
|
"avg_tok_s": round(r["_avg_toks"], 1),
|
||||||
|
"accept_rate": agg.get("aggregate_accept_rate"),
|
||||||
|
"wall_s_total": agg.get("wall_s_total"),
|
||||||
|
})
|
||||||
|
|
||||||
|
summary_path = results_dir / "summary.json"
|
||||||
|
summary_path.write_text(json.dumps(summary_data, indent=2) + "\n")
|
||||||
|
print(f"\n📄 Summary written to {summary_path}")
|
||||||
|
|
||||||
|
|
||||||
|
# ── Main ─────────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
def main():
|
||||||
|
ap = argparse.ArgumentParser(description="MTP Benchmark Runner")
|
||||||
|
ap.add_argument("--models-dir", type=Path, default=Path.home() / "models",
|
||||||
|
help="Directory containing GGUF models (default: ~/models)")
|
||||||
|
ap.add_argument("--model", type=str, default=None,
|
||||||
|
help="Filter: only run models whose name contains this string")
|
||||||
|
ap.add_argument("--toolbox", type=str, default=None,
|
||||||
|
help="Filter: only run this toolbox (e.g. 'vulkan-radv-mtp')")
|
||||||
|
ap.add_argument("--port", type=int, default=8080,
|
||||||
|
help="Port for llama-server (default: 8080)")
|
||||||
|
args = ap.parse_args()
|
||||||
|
|
||||||
|
models_dir = args.models_dir.expanduser().resolve()
|
||||||
|
results_dir = Path(__file__).parent / "results-mtp"
|
||||||
|
results_dir.mkdir(exist_ok=True)
|
||||||
|
|
||||||
|
# Discover models
|
||||||
|
models = discover_models(models_dir)
|
||||||
|
if args.model:
|
||||||
|
models = [m for m in models if args.model.lower() in m["name"].lower()]
|
||||||
|
|
||||||
|
if not models:
|
||||||
|
print(f"❌ No MTP models found in {models_dir}")
|
||||||
|
sys.exit(1)
|
||||||
|
|
||||||
|
# Filter toolboxes
|
||||||
|
toolboxes = TOOLBOXES
|
||||||
|
if args.toolbox:
|
||||||
|
if args.toolbox not in TOOLBOXES:
|
||||||
|
print(f"❌ Unknown toolbox: {args.toolbox}")
|
||||||
|
print(f" Available: {', '.join(TOOLBOXES.keys())}")
|
||||||
|
sys.exit(1)
|
||||||
|
toolboxes = {args.toolbox: TOOLBOXES[args.toolbox]}
|
||||||
|
|
||||||
|
# Print run plan
|
||||||
|
print(f"\n🔍 Discovered {len(models)} MTP model(s):")
|
||||||
|
for m in models:
|
||||||
|
print(f" • {m['name']} → {m['gguf']}")
|
||||||
|
|
||||||
|
print(f"\n🧰 Toolboxes: {', '.join(toolboxes.keys())}")
|
||||||
|
print(f"📊 Modes: {', '.join(MODES.keys())}")
|
||||||
|
total = len(models) * len(toolboxes) * len(MODES)
|
||||||
|
print(f"📋 Total runs: {total}\n")
|
||||||
|
|
||||||
|
# Capture system info
|
||||||
|
capture_system_info(results_dir)
|
||||||
|
|
||||||
|
# Run benchmarks
|
||||||
|
run_count = 0
|
||||||
|
for tb_name, tb_config in toolboxes.items():
|
||||||
|
for model in models:
|
||||||
|
for mode_name, spec_flags in MODES.items():
|
||||||
|
run_count += 1
|
||||||
|
out_file = results_dir / f"{model['name']}__{tb_name}__{mode_name}.json"
|
||||||
|
|
||||||
|
print(f"\n{'─' * 80}")
|
||||||
|
print(f"▶ [{run_count}/{total}] {model['name']} | {tb_name} | {mode_name}")
|
||||||
|
print(f" Output: {out_file.name}")
|
||||||
|
|
||||||
|
# Skip if results exist
|
||||||
|
if out_file.exists():
|
||||||
|
print(f" ⏩ Skipping — results already exist")
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Start server
|
||||||
|
if not start_server(tb_config, model["gguf"], models_dir,
|
||||||
|
spec_flags, args.port):
|
||||||
|
stop_container()
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Wait for health
|
||||||
|
if not wait_for_health(args.port):
|
||||||
|
stop_container()
|
||||||
|
continue
|
||||||
|
|
||||||
|
# Run benchmark
|
||||||
|
bench_data = run_benchmark(args.port, out_file)
|
||||||
|
|
||||||
|
# Stop server
|
||||||
|
print(f" 🛑 Stopping container...")
|
||||||
|
stop_container()
|
||||||
|
|
||||||
|
# Wrap result with metadata
|
||||||
|
if bench_data:
|
||||||
|
save_result(out_file, model, tb_name, mode_name,
|
||||||
|
spec_flags, bench_data)
|
||||||
|
print(f" ✅ Done — saved to {out_file.name}")
|
||||||
|
else:
|
||||||
|
print(f" ❌ No results collected")
|
||||||
|
|
||||||
|
# Cooldown
|
||||||
|
if run_count < total:
|
||||||
|
print(f" 💤 Cooldown ({COOLDOWN}s)...")
|
||||||
|
time.sleep(COOLDOWN)
|
||||||
|
|
||||||
|
# Print summary
|
||||||
|
print_summary(results_dir)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
Reference in New Issue
Block a user