Download scripts/bench_accept.py from reyden009/speculative-decoding-lab: direct link, hf CLI and curl.
- Browser
- Download file 35 kB
-
https://huggingface.co/reyden009/speculative-decoding-lab/resolve/main/scripts/bench_accept.py
- Command line
-
hf download hf://reyden009/speculative-decoding-lab/scripts/bench_accept.py
-
curl -L -o bench_accept.py https://huggingface.co/reyden009/speculative-decoding-lab/resolve/main/scripts/bench_accept.py
35 kB
| #!/usr/bin/env python3 | |
| """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" | |
| # llama-cli default n_max in this build (common.h) — parity with F1/F2/F3. | |
| DEFAULT_N_MAX = 3 | |
| # Sampling of the final run: true greedy (T=0 = pure argmax, ignores top_k/top_p), | |
| # fixed seed 42. F1/F2/F3 ran at T=0.8 (llama-cli defaults) — not greedy. | |
| DEFAULT_SAMPLING = {"temperature": 0.0, "top_k": 40, "top_p": 0.95} | |
| # Fixed results.csv header (no text column; None → empty cell). | |
| CSV_HEADER = [ | |
| "id", | |
| "domain", | |
| "config", | |
| "tok_per_s", | |
| "alpha", | |
| "tau", | |
| "draft_n", | |
| "prompt_ms", | |
| "predicted_ms", | |
| "elapsed_s", | |
| "attempts", | |
| "ts", | |
| ] | |
| # Shared sha256 cache of GGUFs (keyed path|size_bytes|mtime_ns) — avoids | |
| # re-hashing ~12 GB per config (30-60 s) across the 25-config chain. | |
| 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( # noqa: S603 | |
| [ | |
| "/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) # noqa: S310 — localhost | |
| else: | |
| req = urllib.request.Request( # noqa: S310 — localhost | |
| url, | |
| data=json.dumps(payload).encode("utf-8"), | |
| headers={"Content-Type": "application/json"}, | |
| ) | |
| with urllib.request.urlopen(req, timeout=timeout) as resp: # noqa: S310 — localhost | |
| 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( # noqa: S603 — internally controlled command | |
| 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: # noqa: BLE001, S110 — still loading | |
| 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( # noqa: S603 — fixed internal repo | |
| ["git", "-C", str(repo), "rev-parse", "HEAD"], # noqa: S607 — git from the env PATH | |
| 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( # noqa: S603, S607 — internally controlled binary | |
| [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: # noqa: BLE001, S110 — server down | |
| 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) # noqa: S311 — deterministic shuffle with fixed 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" | |
| # Reproducibility: commit + version + sha256/size of GGUFs with the shared | |
| # cache (experiments/runs/model-hashes.json); per-out fallback. | |
| 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: | |
| # R3-003: with no pending prompts we do not start a server — protects | |
| # server.log from truncation ("acc per pos" curves of already-completed | |
| # configs) and avoids loading model/VRAM during resume walks. | |
| print( | |
| "[bench] resume: no pending prompts — not starting the server " | |
| "(protects server.log, avoids model/VRAM load)", | |
| file=sys.stderr, | |
| ) | |
| # Power loss: if the config ended up without metrics/csv (crash before | |
| # the final aggregation), they are regenerated from results.jsonl — no server/GPU. | |
| 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) | |
| # Anti-rerun lock BEFORE starting the server (R4-001): if another runner | |
| # writes this --out, we abort without spawning processes or using the GPU. | |
| 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: # noqa: BLE001 | |
| 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, # per prompt | |
| "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: # noqa: BLE001 — overnight hardening | |
| 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() | |
| # Final metrics: re-read the full results.jsonl (resume-safe) → CSV + JSON. | |
| 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()) | |