| |
| """Acceptance (alpha/tau) benchmark via llama-server (requires gaps G1/G3/G4). |
| |
| Why a server and not llama-cli: |
| - llama-cli does NOT expose speculative-decoding statistics; llama-server DOES: |
| every /completion response returns timings.draft_n and timings.draft_n_accepted |
| (server-task.cpp result_timings::to_json), plus predicted_per_second. |
| - Additionally, with --verbose it writes per-request log lines: |
| "draft acceptance = 0.xxxxx (N accepted / M generated), mean len = X.XX" |
| and at TRC level "acc per pos = (r1, r2, ...)" (per-position curve, optional: |
| parse <out>/server.log after the run). |
| |
| Protocol parity with F1/F2/F3 (llama-cli, default n_max = 3 in common.h): |
| - --spec-draft-n-max 3, -c 2048, -t 8, seed 42, n_predict 256. |
| - Explicit sampling in the body: temperature 0.0 (pure greedy argmax; ignores |
| top_k/top_p), top_k 40, top_p 0.95 — the final run corrects F1/F2/F3 which |
| ran at T=0.8 (llama-cli defaults). Non-thinking: --reasoning off in the |
| server command (if raw mode rejects it, the documented escape is in --help). |
| - The numbers from this script are ONLY comparable among themselves |
| (server-consistent methodology: each config starts its own server). DO NOT |
| mix with the F1/F2/F3 tok/s (fresh process per prompt vs persistent server). |
| |
| Anti-rerun design (identical to bench_spec.py): |
| - <out>/results.jsonl (one line per OK prompt), <out>/errors.jsonl (failures, |
| retried with --resume), exclusive lock (exit 3 if another runner writes the |
| same --out), vram.json (max_gpu_mib + max_power_w sampled during the run), |
| config.json (args + sampling + reproducibility: llama.cpp commit, version, |
| sha256/size of the GGUFs), metrics.json (total and per-domain aggregates) |
| and results.csv (without text column) at the end. |
| - exit 0 = clean; 2 = failures; 3 = lock busy; 1 = server did not start or |
| died mid-run (abort after 2 consecutive connection failures). |
| |
| Usage (example): |
| # 1) create a per-domain stratified subset (seed 42): |
| python scripts/bench_accept.py --make-subset-from experiments/prompts/f1-sample.jsonl \ |
| --prompts experiments/prompts/acc-sample.jsonl --subset-per-domain 60 |
| |
| # 2) run one config (target + drafter): |
| python scripts/bench_accept.py --model models/Qwen3-8B-Q4_K_M.gguf \ |
| --draft models/drafts/dflash_qwen3_8b_block7.gguf --spec-type draft-dflash \ |
| --config-name q4-dflash-q4 --prompts experiments/prompts/acc-sample.jsonl \ |
| --out experiments/runs/acc-q4-dflash-q4 --resume |
| |
| # 3) target-solo (no drafter): omit --draft. |
| # 4) smoke/mini-runs: --max-prompts N (first N pending prompts). |
| RUN ONLY WITH THE GPU FREE (never while the F3 chain is measuring). |
| |
| Per-record output: {id, domain, text, config, tok_per_s, alpha, tau, draft_n, |
| prompt_ms, predicted_ms, elapsed_s, attempts, ts} with alpha = draft_n_accepted / |
| draft_n (None if no draft), tau = draft_n_accepted (accepted tokens), |
| prompt_ms = TTFT (prefill), predicted_ms = total generation time. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import fcntl |
| import hashlib |
| import json |
| import math |
| import os |
| import random |
| import re |
| import shlex |
| import socket |
| import subprocess |
| import sys |
| import threading |
| import time |
| import urllib.request |
| from pathlib import Path |
|
|
| LLAMA_BIN = Path(os.environ.get("LLAMA_CPP_BIN", os.path.expanduser("~/llama.cpp/build/bin"))) |
| SERVER_BIN = LLAMA_BIN / "llama-server" |
|
|
| |
| DEFAULT_N_MAX = 3 |
|
|
| |
| |
| DEFAULT_SAMPLING = {"temperature": 0.0, "top_k": 40, "top_p": 0.95} |
|
|
| |
| CSV_HEADER = [ |
| "id", |
| "domain", |
| "config", |
| "tok_per_s", |
| "alpha", |
| "tau", |
| "draft_n", |
| "prompt_ms", |
| "predicted_ms", |
| "elapsed_s", |
| "attempts", |
| "ts", |
| ] |
|
|
| |
| |
| SHARED_HASH_CACHE = Path("experiments/runs/model-hashes.json") |
|
|
|
|
| class VramSampler: |
| """Samples VRAM (memory.used) and power draw (power.draw) every 3 s in a daemon thread.""" |
|
|
| def __init__(self, out_dir: Path) -> None: |
| self.out_dir = out_dir |
| self._max_mib = 0 |
| self._max_power = 0.0 |
| self._lock = threading.Lock() |
| self._stop = threading.Event() |
| self._thread = threading.Thread(target=self._run, daemon=True) |
|
|
| def start(self) -> None: |
| self._thread.start() |
|
|
| def stop(self) -> None: |
| self._stop.set() |
| self._thread.join(timeout=5) |
| with self._lock: |
| max_mib = self._max_mib |
| max_power = self._max_power |
| (self.out_dir / "vram.json").write_text( |
| json.dumps( |
| { |
| "max_gpu_mib": max_mib, |
| "max_power_w": max_power, |
| "sample_interval_s": 3, |
| "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"), |
| } |
| ) |
| ) |
|
|
| def _run(self) -> None: |
| while not self._stop.is_set(): |
| try: |
| out = subprocess.run( |
| [ |
| "/usr/bin/nvidia-smi", |
| "--query-gpu=memory.used,power.draw", |
| "--format=csv,noheader,nounits", |
| ], |
| capture_output=True, |
| text=True, |
| timeout=10, |
| check=True, |
| ) |
| parts = [p.strip() for p in out.stdout.split(",")] |
| used = int(parts[0]) |
| pwr_raw = parts[1] if len(parts) > 1 else "" |
| pwr = float(pwr_raw) if pwr_raw not in ("", "[N/A]") else 0.0 |
| with self._lock: |
| self._max_mib = max(self._max_mib, used) |
| self._max_power = max(self._max_power, pwr) |
| except (subprocess.SubprocessError, ValueError, IndexError): |
| pass |
| self._stop.wait(3) |
|
|
|
|
| def _free_port() -> int: |
| with socket.socket() as s: |
| s.bind(("127.0.0.1", 0)) |
| return s.getsockname()[1] |
|
|
|
|
| def _http_json(url: str, payload: dict | None = None, timeout: float = 30.0) -> tuple[int, dict]: |
| """GET (payload=None) or POST JSON; returns (status, json).""" |
| if payload is None: |
| req = urllib.request.Request(url) |
| else: |
| req = urllib.request.Request( |
| url, |
| data=json.dumps(payload).encode("utf-8"), |
| headers={"Content-Type": "application/json"}, |
| ) |
| with urllib.request.urlopen(req, timeout=timeout) as resp: |
| return resp.status, json.loads(resp.read().decode("utf-8")) |
|
|
|
|
| class Server: |
| """llama-server lifecycle: start → wait_health → stop (always).""" |
|
|
| def __init__(self, cmd: list[str]) -> None: |
| self.cmd = cmd |
| self.proc: subprocess.Popen | None = None |
| self.url = "" |
|
|
| def start(self, port: int) -> None: |
| self.url = f"http://127.0.0.1:{port}" |
| self.proc = subprocess.Popen( |
| self.cmd, |
| stdout=subprocess.DEVNULL, |
| stderr=subprocess.DEVNULL, |
| ) |
|
|
| def wait_health(self, timeout: float = 300.0) -> str: |
| """Wait for /health == 200. Returns "" if OK, otherwise the reason.""" |
| deadline = time.time() + timeout |
| while time.time() < deadline: |
| if self.proc is not None and self.proc.poll() is not None: |
| return f"server died on startup (rc={self.proc.returncode})" |
| try: |
| status, _ = _http_json(self.url + "/health", timeout=5) |
| if status == 200: |
| return "" |
| except Exception: |
| pass |
| time.sleep(2) |
| return f"timeout waiting for /health ({int(timeout)}s)" |
|
|
| def stop(self) -> None: |
| if self.proc is not None and self.proc.poll() is None: |
| self.proc.terminate() |
| try: |
| self.proc.wait(timeout=10) |
| except subprocess.TimeoutExpired: |
| self.proc.kill() |
| self.proc = None |
|
|
|
|
| def _tail(text: str, n: int = 300) -> str: |
| return text[-n:] if text else "" |
|
|
|
|
| def _parse_completion( |
| data: dict, |
| ) -> tuple[float | None, float | None, int | None, int | None, float | None, float | None]: |
| """(tok_per_s, alpha, tau, draft_n, prompt_ms, predicted_ms) from /completion. |
| |
| alpha = draft_n_accepted / draft_n; tau = draft_n_accepted. With draft_n 0 |
| or missing (target-solo configs) alpha/tau/draft_n are None (spec). prompt_ms |
| (TTFT) and predicted_ms are draft-independent and always present. |
| """ |
| tim = data.get("timings", {}) |
| tok_s = tim.get("predicted_per_second") |
| draft_n = tim.get("draft_n") |
| draft_acc = tim.get("draft_n_accepted") |
| alpha: float | None = None |
| tau: int | None = None |
| if draft_n: |
| if draft_acc is not None: |
| alpha = round(draft_acc / draft_n, 4) |
| tau = draft_acc |
| else: |
| draft_n = None |
| prompt_ms = tim.get("prompt_ms") |
| predicted_ms = tim.get("predicted_ms") |
| return tok_s, alpha, tau, draft_n, prompt_ms, predicted_ms |
|
|
|
|
| def read_git_commit(repo: str) -> str | None: |
| """HEAD commit of the pinned repo (fixed -C); None non-fatal if it fails. |
| |
| The run proceeds without commit reproducibility (threat matrix: git). |
| """ |
| try: |
| out = subprocess.run( |
| ["git", "-C", str(repo), "rev-parse", "HEAD"], |
| capture_output=True, |
| text=True, |
| timeout=10, |
| check=True, |
| ) |
| except (subprocess.SubprocessError, FileNotFoundError): |
| return None |
| return out.stdout.strip() or None |
|
|
|
|
| def parse_llama_version(bin_path: Path) -> str | None: |
| """Binary version: 'version: 22 (0713275)' → '22'; 'build: 10249' → '10249'. |
| |
| None if the binary does not answer or the format is not recognized. |
| """ |
| try: |
| out = subprocess.run( |
| [str(bin_path), "--version"], |
| capture_output=True, |
| text=True, |
| timeout=15, |
| check=True, |
| ) |
| except (subprocess.SubprocessError, FileNotFoundError): |
| return None |
| combined = (out.stdout or "") + "\n" + (out.stderr or "") |
| m = re.search(r"version:\s*(\S+)", combined) |
| if m: |
| return m.group(1) |
| m = re.search(r"build:\s*(\d+)", combined) |
| return m.group(1) if m else None |
|
|
|
|
| def file_sha256(path: Path, cache: dict[tuple[str, int, int], str] | None = None) -> str | None: |
| """SHA-256 of a file; cache keyed (path, size_bytes, mtime_ns) (hit → no re-hash). |
| |
| None if the file does not exist or is not readable. |
| """ |
| try: |
| st = path.stat() |
| except OSError: |
| return None |
| key = (str(path), st.st_size, st.st_mtime_ns) |
| if cache is not None and key in cache: |
| return cache[key] |
| h = hashlib.sha256() |
| try: |
| with path.open("rb") as f: |
| for chunk in iter(lambda: f.read(1 << 20), b""): |
| h.update(chunk) |
| except OSError: |
| return None |
| digest = h.hexdigest() |
| if cache is not None: |
| cache[key] = digest |
| return digest |
|
|
|
|
| def _cache_key_str(key: tuple[str, int, int]) -> str: |
| return f"{key[0]}|{key[1]}|{key[2]}" |
|
|
|
|
| def load_hash_cache(path: Path) -> dict[tuple[str, int, int], str]: |
| """Load the shared cache (JSON with 'path|size|mtime_ns' keys) → dict of tuples.""" |
| cache: dict[tuple[str, int, int], str] = {} |
| try: |
| if path.exists(): |
| raw = json.loads(path.read_text()) |
| for k, v in raw.items(): |
| p, s, m = k.rsplit("|", 2) |
| cache[(p, int(s), int(m))] = v |
| except (OSError, ValueError): |
| pass |
| return cache |
|
|
|
|
| def save_hash_cache(path: Path, cache: dict[tuple[str, int, int], str]) -> bool: |
| """Save the shared cache; True if OK, False if it failed (per-out fallback).""" |
| try: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| path.write_text(json.dumps({_cache_key_str(k): v for k, v in cache.items()}, indent=2)) |
| return True |
| except OSError: |
| return False |
|
|
|
|
| def _gguf_meta(path: str, cache: dict[tuple[str, int, int], str]) -> dict: |
| p = Path(path) |
| st = p.stat() if p.exists() else None |
| return {"path": path, "size_bytes": st.st_size if st else None, "sha256": file_sha256(p, cache)} |
|
|
|
|
| def read_reproducibility( |
| model: str, draft: str | None, cache: dict[tuple[str, int, int], str] |
| ) -> dict: |
| """Reproducibility metadata: llama.cpp commit + version + sha256/size of GGUFs. |
| |
| Failed git → commit None + warn (the run continues; null reproducibility). |
| """ |
| repo = os.path.expanduser("~/llama.cpp") |
| commit = read_git_commit(repo) |
| if commit is None: |
| print( |
| f"[bench] WARN: could not read the commit of {repo} → incomplete reproducibility", |
| file=sys.stderr, |
| ) |
| return { |
| "llama_cpp_commit": commit, |
| "llama_cpp_version": parse_llama_version(SERVER_BIN), |
| "host": socket.gethostname(), |
| "ts": time.strftime("%Y-%m-%dT%H:%M:%S"), |
| "model": _gguf_meta(model, cache), |
| "draft": _gguf_meta(draft, cache) if draft else None, |
| } |
|
|
|
|
| def build_server_cmd( |
| args: argparse.Namespace, port: int, out_dir: Path, log_name: str = "server.log" |
| ) -> list[str]: |
| """llama-server command: target (± drafter), non-thinking (--reasoning off). |
| |
| --reasoning off goes after --verbose and before the extras; no auto-retry |
| (if raw mode rejects it, the escape is --extra reasoning_effort:"none", |
| see --help of --extra). |
| """ |
| cmd = [ |
| str(SERVER_BIN), |
| "-m", |
| args.model, |
| "-ngl", |
| str(args.n_gpu_layers), |
| "-c", |
| str(args.ctx), |
| "-t", |
| str(args.threads), |
| "-np", |
| "1", |
| "--host", |
| "127.0.0.1", |
| "--port", |
| str(port), |
| "--log-file", |
| str(out_dir / log_name), |
| "--verbose", |
| ] |
| if args.draft: |
| cmd += [ |
| "-md", |
| args.draft, |
| "--spec-type", |
| args.spec_type, |
| "--spec-draft-n-max", |
| str(args.spec_draft_n_max), |
| "-ngld", |
| str(args.draft_ngl), |
| ] |
| if args.spec_draft_p_min is not None: |
| cmd += ["--spec-draft-p-min", str(args.spec_draft_p_min)] |
| cmd += ["--reasoning", "off"] |
| cmd += [a for pair in args.extra for a in shlex.split(pair)] |
| return cmd |
|
|
|
|
| def resume_done(results_path: Path, errors_path: Path) -> set[str]: |
| """Ids already processed (results.jsonl) or known failures (errors.jsonl) for resume. |
| |
| Includes failures so prompts that always fail are not retried (e.g. the 8 |
| arena-hard-v2 ones > ctx 2048) and errors.jsonl entries are not duplicated. |
| Tolerates corrupt lines (e.g. power loss) and missing files. |
| """ |
| done: set[str] = set() |
| for path in (results_path, errors_path): |
| if not path.exists(): |
| continue |
| for line in path.read_text(encoding="utf-8", errors="replace").splitlines(): |
| try: |
| done.add(json.loads(line)["id"]) |
| except (json.JSONDecodeError, KeyError): |
| continue |
| return done |
|
|
|
|
| def pending_prompts(records: list[dict], done: set[str], max_prompts: int | None) -> list[dict]: |
| """Pending = records not completed (resume), truncated to max_prompts (N>0).""" |
| pending = [p for p in records if p["id"] not in done] |
| if max_prompts is not None and max_prompts > 0: |
| pending = pending[:max_prompts] |
| return pending |
|
|
|
|
| def read_results(results_path: Path) -> list[dict]: |
| """Read results.jsonl tolerating corrupt lines (power loss) and a missing file.""" |
| records: list[dict] = [] |
| if results_path.exists(): |
| for line in results_path.read_text(encoding="utf-8", errors="replace").splitlines(): |
| try: |
| records.append(json.loads(line)) |
| except json.JSONDecodeError: |
| continue |
| return records |
|
|
|
|
| def server_log_path(out_dir: Path) -> Path: |
| """llama-server session log: server.log the 1st time, server-2.log/3.log on |
| resumes. llama-server TRUNCATES --log-file (fopen "w", log.cpp:322) → a resume |
| of a partial config must not erase the "acc per pos" curves of its previous session.""" |
| first = out_dir / "server.log" |
| if not first.exists(): |
| return first |
| n = 2 |
| while (out_dir / f"server-{n}.log").exists(): |
| n += 1 |
| return out_dir / f"server-{n}.log" |
|
|
|
|
| def _server_alive(url: str) -> bool: |
| """Probe /health: True if the server answers 200 (not down).""" |
| try: |
| status, _ = _http_json(url + "/health", timeout=5) |
| return status == 200 |
| except Exception: |
| return False |
|
|
|
|
| def _mean(xs: list[float]) -> float | None: |
| return round(sum(xs) / len(xs), 4) if xs else None |
|
|
|
|
| def _median(xs: list[float]) -> float | None: |
| if not xs: |
| return None |
| s = sorted(xs) |
| n = len(s) |
| mid = n // 2 |
| if n % 2 == 1: |
| return round(s[mid], 4) |
| return round((s[mid - 1] + s[mid]) / 2, 4) |
|
|
|
|
| def _pct(xs: list[float], p: float) -> float | None: |
| if not xs: |
| return None |
| s = sorted(xs) |
| idx = min(len(s) - 1, max(0, math.ceil(p / 100 * len(s)) - 1)) |
| return round(s[idx], 4) |
|
|
|
|
| def write_metrics(out_dir: Path, records: list[dict], failed: int, duration_s: float) -> Path: |
| """Write metrics.json with total and per-domain aggregates (math/code/chat). |
| |
| Resume-safe: the caller passes ALL records (this run + those re-read from |
| results.jsonl, which include the previous ones). Aggregates ignore None |
| (alpha/tau/draft_n of target-solo configs). sampling/spec/reproducibility |
| are copied from config.json; vram/power from vram.json. |
| """ |
| cfg: dict = {} |
| vram: dict = {} |
| try: |
| cfg = json.loads((out_dir / "config.json").read_text()) |
| except (OSError, ValueError): |
| pass |
| try: |
| vram = json.loads((out_dir / "vram.json").read_text()) |
| except (OSError, ValueError): |
| pass |
|
|
| toks = [r["tok_per_s"] for r in records if r.get("tok_per_s") is not None] |
| alphas = [r["alpha"] for r in records if r.get("alpha") is not None] |
| taus = [r["tau"] for r in records if r.get("tau") is not None] |
| draft_ns = [r["draft_n"] for r in records if r.get("draft_n") is not None] |
| ttfts = [r["prompt_ms"] for r in records if r.get("prompt_ms") is not None] |
|
|
| per_domain: dict[str, dict] = {} |
| for dom in ("math", "code", "chat"): |
| dr = [r for r in records if r.get("domain") == dom] |
| dtoks = [r["tok_per_s"] for r in dr if r.get("tok_per_s") is not None] |
| dalphas = [r["alpha"] for r in dr if r.get("alpha") is not None] |
| dtaus = [r["tau"] for r in dr if r.get("tau") is not None] |
| dttfts = [r["prompt_ms"] for r in dr if r.get("prompt_ms") is not None] |
| per_domain[dom] = { |
| "n": len(dr), |
| "tok_per_s.mean": _mean(dtoks), |
| "alpha.mean": _mean(dalphas), |
| "tau.mean": _mean(dtaus), |
| "ttft.mean": _mean(dttfts), |
| } |
|
|
| metrics = { |
| "prompts": {"total": len(records) + failed, "ok": len(records), "failed": failed}, |
| "tok_per_s": { |
| "mean": _mean(toks), |
| "median": _median(toks), |
| "p50": _median(toks), |
| "p95": _pct(toks, 95), |
| "min": round(min(toks), 4) if toks else None, |
| "max": round(max(toks), 4) if toks else None, |
| }, |
| "alpha": {"mean": _mean(alphas), "median": _median(alphas)}, |
| "tau": {"mean": _mean(taus), "median": _median(taus)}, |
| "draft_n": {"total": sum(draft_ns), "mean": _mean(draft_ns), "median": _median(draft_ns)}, |
| "ttft": { |
| "prompt_ms.mean": _mean(ttfts), |
| "prompt_ms.median": _median(ttfts), |
| "prompt_ms.p95": _pct(ttfts, 95), |
| }, |
| "vram": {"max_gpu_mib": vram.get("max_gpu_mib")}, |
| "power": {"max_power_w": vram.get("max_power_w")}, |
| "duration_s": round(duration_s, 2), |
| "errors": failed, |
| "per_domain": per_domain, |
| "sampling": cfg.get("sampling", {}), |
| "spec": cfg.get("spec", {}), |
| "reproducibility": cfg.get("reproducibility", {}), |
| } |
| path = out_dir / "metrics.json" |
| path.write_text(json.dumps(metrics, indent=2, ensure_ascii=False) + "\n") |
| return path |
|
|
|
|
| def export_csv(results_path: Path) -> Path: |
| """Export results.jsonl → results.csv (fixed header, no text column).""" |
| csv_path = results_path.with_suffix(".csv") |
| with csv_path.open("w", newline="") as f: |
| w = csv.writer(f) |
| w.writerow(CSV_HEADER) |
| for line in results_path.read_text().splitlines(): |
| if not line.strip(): |
| continue |
| try: |
| r = json.loads(line) |
| except ValueError: |
| continue |
| w.writerow([r.get(h) for h in CSV_HEADER]) |
| return csv_path |
|
|
|
|
| def make_subset(source: Path, dest: Path, per_domain: int, seed: int) -> int: |
| """Stratified subset: up to per_domain prompts per domain (shuffle seed).""" |
| by_domain: dict[str, list[dict]] = {} |
| for line in source.read_text().splitlines(): |
| if not line.strip(): |
| continue |
| p = json.loads(line) |
| by_domain.setdefault(p.get("domain", "unknown"), []).append(p) |
| rng = random.Random(seed) |
| total = 0 |
| with dest.open("w") as f: |
| for dom in sorted(by_domain): |
| chosen = by_domain[dom][:] |
| rng.shuffle(chosen) |
| chosen = chosen[:per_domain] |
| for p in chosen: |
| f.write(json.dumps(p, ensure_ascii=False) + "\n") |
| total += len(chosen) |
| print(f"[subset] {dom}: {len(chosen)}/{len(by_domain[dom])}", file=sys.stderr) |
| print(f"[subset] {total} prompts → {dest}", file=sys.stderr) |
| return 0 |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser( |
| description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter |
| ) |
| ap.add_argument("--model", default=None, help="Target GGUF") |
| ap.add_argument("--draft", default=None, help="Drafter GGUF (omit = target-solo)") |
| ap.add_argument( |
| "--spec-type", |
| default="none", |
| help="draft-simple|draft-eagle3|draft-mtp|draft-dflash|draft-dspark", |
| ) |
| ap.add_argument( |
| "--spec-draft-n-max", type=int, default=DEFAULT_N_MAX, help="llama-cli parity n_max" |
| ) |
| ap.add_argument("--spec-draft-p-min", type=float, default=None, help="DSpark confidence cutoff") |
| ap.add_argument("--config-name", default=None, help="Label in every record (default: stems)") |
| ap.add_argument("--prompts", required=True, type=Path) |
| ap.add_argument("--out", default=None, type=Path) |
| ap.add_argument("--n-tokens", type=int, default=256) |
| ap.add_argument("--seed", type=int, default=42) |
| ap.add_argument("--temperature", type=float, default=DEFAULT_SAMPLING["temperature"]) |
| ap.add_argument("--top-k", type=int, default=DEFAULT_SAMPLING["top_k"]) |
| ap.add_argument("--top-p", type=float, default=DEFAULT_SAMPLING["top_p"]) |
| ap.add_argument("--n-gpu-layers", type=int, default=99) |
| ap.add_argument("--draft-ngl", type=int, default=99, help="GPU layers of the drafter (-ngld)") |
| ap.add_argument("--ctx", type=int, default=2048) |
| ap.add_argument("--threads", type=int, default=8) |
| ap.add_argument("--retries", type=int, default=2) |
| ap.add_argument("--server-timeout", type=float, default=300.0, help="Wait for /health (s)") |
| ap.add_argument( |
| "--max-prompts", |
| type=int, |
| default=None, |
| help="Only the first N pending prompts (smoke/mini-runs/OOM-CHECK)", |
| ) |
| ap.add_argument("--resume", action="store_true") |
| ap.add_argument( |
| "--extra", |
| action="append", |
| default=[], |
| help=( |
| "Extra flags for llama-server (shlex). If raw mode rejects --reasoning off, " |
| 'documented escape: --extra reasoning_effort:"none"' |
| ), |
| ) |
| ap.add_argument( |
| "--make-subset-from", type=Path, default=None, help="Subset mode: source f1-sample.jsonl" |
| ) |
| ap.add_argument("--subset-per-domain", type=int, default=60) |
| args = ap.parse_args() |
|
|
| if args.make_subset_from is not None: |
| return make_subset(args.make_subset_from, args.prompts, args.subset_per_domain, args.seed) |
|
|
| if args.model is None or args.out is None: |
| ap.error("--model and --out are required outside --make-subset-from mode") |
|
|
| if not SERVER_BIN.exists(): |
| print(f"[bench] ERROR: {SERVER_BIN} does not exist (did llama.cpp build?)", file=sys.stderr) |
| return 1 |
|
|
| out_dir = args.out |
| out_dir.mkdir(parents=True, exist_ok=True) |
| results_path = out_dir / "results.jsonl" |
| errors_path = out_dir / "errors.jsonl" |
|
|
| |
| |
| hash_cache = load_hash_cache(SHARED_HASH_CACHE) |
| reproducibility = read_reproducibility(args.model, args.draft, hash_cache) |
| if not save_hash_cache(SHARED_HASH_CACHE, hash_cache): |
| save_hash_cache(out_dir / "model-hashes.json", hash_cache) |
|
|
| cfg = {k: (str(v) if isinstance(v, Path) else v) for k, v in vars(args).items()} |
| cfg["sampling"] = { |
| "temperature": args.temperature, |
| "top_k": args.top_k, |
| "top_p": args.top_p, |
| "seed": args.seed, |
| } |
| cfg["spec"] = { |
| "type": args.spec_type, |
| "draft_n_max": args.spec_draft_n_max, |
| "p_min": args.spec_draft_p_min, |
| } |
| cfg["reproducibility"] = reproducibility |
| (out_dir / "config.json").write_text(json.dumps(cfg, indent=2, ensure_ascii=False)) |
|
|
| done: set[str] = set() |
| if args.resume: |
| done = resume_done(results_path, errors_path) |
| print( |
| f"[bench] resume: {len(done)} prompts already processed (results + known errors)", |
| file=sys.stderr, |
| ) |
|
|
| records_all: list[dict] = [] |
| for line in args.prompts.read_text().splitlines(): |
| if not line.strip(): |
| continue |
| try: |
| records_all.append(json.loads(line)) |
| except json.JSONDecodeError: |
| continue |
| pending = pending_prompts(records_all, done, args.max_prompts) |
| if args.max_prompts is not None: |
| print( |
| f"[bench] max-prompts={args.max_prompts}: {len(pending)} prompts in this run", |
| file=sys.stderr, |
| ) |
|
|
| if not pending: |
| |
| |
| |
| print( |
| "[bench] resume: no pending prompts — not starting the server " |
| "(protects server.log, avoids model/VRAM load)", |
| file=sys.stderr, |
| ) |
| |
| |
| if not (out_dir / "metrics.json").exists() or not (out_dir / "results.csv").exists(): |
| recs = read_results(results_path) |
| export_csv(results_path) |
| write_metrics(out_dir, recs, 0, 0.0) |
| print( |
| "[bench] early-exit: metrics.json/results.csv regenerated (were missing)", |
| file=sys.stderr, |
| ) |
| return 0 |
|
|
| port = _free_port() |
| cmd = build_server_cmd(args, port, out_dir, log_name=server_log_path(out_dir).name) |
|
|
| config_name = args.config_name or ( |
| Path(args.model).stem |
| if not args.draft |
| else f"{Path(args.model).stem}+{Path(args.draft).stem}" |
| ) |
|
|
| server = Server(cmd) |
| |
| |
| guard = results_path.open("a") |
| try: |
| fcntl.flock(guard.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) |
| except OSError: |
| guard.close() |
| print( |
| "[bench] ERROR: another runner is writing this --out (lock busy). " |
| "Wait for it to finish or check it with ps.", |
| file=sys.stderr, |
| ) |
| return 3 |
|
|
| t_start = time.time() |
| print(f"[bench] starting llama-server at {server.url} (pid launched)", file=sys.stderr) |
| server.start(port) |
| try: |
| reason = server.wait_health(args.server_timeout) |
| if reason: |
| server.stop() |
| guard.close() |
| print(f"[bench] ERROR: {reason} → see {out_dir / 'server.log'}", file=sys.stderr) |
| return 1 |
| print("[bench] server OK", file=sys.stderr) |
| except Exception: |
| server.stop() |
| guard.close() |
| raise |
|
|
| n_new = 0 |
| n_failed = 0 |
| aborted = False |
| vram = VramSampler(out_dir) |
| vram.start() |
| try: |
| with results_path.open("a") as fout, errors_path.open("a") as ferr: |
| body = { |
| "prompt": None, |
| "n_predict": args.n_tokens, |
| "seed": args.seed, |
| "temperature": args.temperature, |
| "top_k": args.top_k, |
| "top_p": args.top_p, |
| "cache_prompt": False, |
| "stream": False, |
| } |
| consec_conn = 0 |
| for p in pending: |
| if p["id"] in done: |
| continue |
| n_new += 1 |
| record: dict | None = None |
| last_err = "" |
| detail_text = "" |
| for attempt in range(1, args.retries + 1): |
| t0 = time.time() |
| try: |
| body["prompt"] = p["text"] |
| status, data = _http_json(server.url + "/completion", body, timeout=900) |
| tok_s, alpha, tau, draft_n, prompt_ms, predicted_ms = _parse_completion( |
| data |
| ) |
| if status != 200: |
| detail_text = _tail(str(data), 300) |
| except Exception as e: |
| status = 0 |
| tok_s, alpha, tau, draft_n, prompt_ms, predicted_ms = ( |
| None, |
| None, |
| None, |
| None, |
| None, |
| None, |
| ) |
| last_err = repr(e) |
| elapsed = time.time() - t0 |
| if status == 200 and tok_s is not None: |
| consec_conn = 0 |
| record = { |
| "id": p["id"], |
| "domain": p.get("domain"), |
| "text": p["text"], |
| "config": config_name, |
| "tok_per_s": round(tok_s, 3), |
| "alpha": alpha, |
| "tau": tau, |
| "draft_n": draft_n, |
| "prompt_ms": round(prompt_ms, 1) if prompt_ms is not None else None, |
| "predicted_ms": round(predicted_ms, 1) |
| if predicted_ms is not None |
| else None, |
| "elapsed_s": round(elapsed, 2), |
| "attempts": attempt, |
| "ts": time.strftime("%Y-%m-%dT%H:%M:%S"), |
| } |
| break |
| if status == 0: |
| consec_conn += 1 |
| else: |
| consec_conn = 0 |
| if consec_conn >= 2 and not _server_alive(server.url): |
| print( |
| f"[bench] ERROR: server not responding (2 consecutive " |
| f"connection failures: {last_err}) → abort", |
| file=sys.stderr, |
| ) |
| aborted = True |
| break |
| last_err = f"status={status} {last_err}" |
| print( |
| f"[bench] {p['id']} attempt {attempt} failed ({last_err})", file=sys.stderr |
| ) |
| if record is not None: |
| fout.write(json.dumps(record, ensure_ascii=False) + "\n") |
| fout.flush() |
| print( |
| f"[bench] {p['id']}: {record['tok_per_s']} tok/s " |
| f"α={record['alpha']} τ={record['tau']} ({record['elapsed_s']}s)", |
| file=sys.stderr, |
| ) |
| else: |
| n_failed += 1 |
| ferr.write( |
| json.dumps( |
| { |
| "id": p["id"], |
| "domain": p.get("domain"), |
| "error": last_err, |
| "detail": detail_text or last_err, |
| "attempts": args.retries, |
| }, |
| ensure_ascii=False, |
| ) |
| + "\n" |
| ) |
| ferr.flush() |
| print( |
| f"[bench] {p['id']} FAILED after {args.retries} attempts → errors.jsonl", |
| file=sys.stderr, |
| ) |
| if aborted: |
| break |
| finally: |
| server.stop() |
| vram.stop() |
| guard.close() |
|
|
| |
| records = read_results(results_path) |
| export_csv(results_path) |
| write_metrics(out_dir, records, n_failed, time.time() - t_start) |
|
|
| print( |
| f"[bench] done: {n_new - n_failed} new valid, {n_failed} failures, " |
| f"{len(done)} previous → {results_path}" |
| ) |
| if aborted: |
| return 1 |
| return 0 if n_failed == 0 else 2 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|