From db2e8067049f4db6bbfc4305bbda8a3e8ef48c6a Mon Sep 17 00:00:00 2001 From: Ben Shaharizad Date: Sat, 19 Sep 2026 00:40:15 +0300 Subject: [PATCH] =?UTF-8?q?W5b-4:=20one-command=20M5=20runbook=20(benchmar?= =?UTF-8?q?ks/m5.py)=20=E2=80=94=20doctor=20gate,=20per-model=20slow=20par?= =?UTF-8?q?ity,=20bench,=20invariance,=20timing,=20optional=20A/B=20worktr?= =?UTF-8?q?ee,=20summary;=20RUNBOOK.md=20logging,=20idempotent=20steps,=20?= =?UTF-8?q?planner+summary=20unit=20tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- BENCHMARKING.md | 26 ++ benchmarks/m5.py | 711 +++++++++++++++++++++++++++++++++++++++++++++++ tests/test_m5.py | 263 ++++++++++++++++++ 3 files changed, 1000 insertions(+) create mode 100644 benchmarks/m5.py create mode 100644 tests/test_m5.py diff --git a/BENCHMARKING.md b/BENCHMARKING.md index 5cf19a4..d2e79c3 100644 --- a/BENCHMARKING.md +++ b/BENCHMARKING.md @@ -99,3 +99,29 @@ than 5 MB total are gzipped automatically (the folder README says so). - [ ] Ran `python -m benchmarks.leaderboard --results benchmarks/results --readme README.md` so the README leaderboard block is up to date (the `--check-readme` freshness gate in `results-check` CI enforces this) + +## M5 runbook: one command for the whole gate sequence + +`benchmarks/m5.py` chains the full milestone-gate sequence in order, each step +logged with wall time and exit status into `/RUNBOOK.md`: + +```bash +python -m benchmarks.m5 --out m5-2026-09-18 [--ab-branch w2a-field-local] +``` + +Steps in order: (1) `jevmlx doctor --json` (gate — any FAIL aborts); (2) +`pytest -m slow` once per parity model (Qwen3-8B, Llama-3.1-8B, Gemma-3-12B by +default; override with `--parity-models` or `--models-file`, the id rides the +`MODEL_ID` env var); (3) `jevmlx bench --model quality`; (4) +`benchmarks.invariance` on quality with `--extra 1,5,20,40` over the TypeSafe +cases (fetched automatically when missing); (5) `benchmarks.timing --model +quality --reps 5`; (6) `jevmlx bench --models-file` for the remaining parity +models; (7) with `--ab-branch`: a temp worktree of that branch with its own +venv, steps 3+4 rerun there, worktree removed afterwards; (8) `/SUMMARY.md` +comparing main vs A/B (agreement/accuracy, flip rate, log-odds drift, time per +case, peak memory) from the produced json files. + +Idempotent: steps whose output markers already exist are skipped (`--fresh` +reruns everything). Per-step logs land in `/.log`. Interrupted +runs resume; the gate step keeps a half-finished evening from wasting GPU time +on a broken environment. diff --git a/benchmarks/m5.py b/benchmarks/m5.py new file mode 100644 index 0000000..3ce8fd1 --- /dev/null +++ b/benchmarks/m5.py @@ -0,0 +1,711 @@ +"""One-command M5 runbook: the full milestone-gate sequence in order. + +Each step is a subprocess logged with wall time and exit status into +``/RUNBOOK.md``: + + 1. ``jevmlx doctor --json`` — gate: any FAIL aborts the runbook. + 2. ``pytest -m slow -q`` once per parity model (default Qwen3-8B, + Llama-3.1-8B, Gemma-3-12B; the model rides the ``MODEL_ID`` env var). + 3. ``jevmlx bench --model quality``. + 4. ``benchmarks.invariance`` on quality with ``--extra 1,5,20,40`` over the + TypeSafe cases (fetched first when missing). + 5. ``benchmarks.timing --model quality --reps 5``. + 6. ``jevmlx bench --models-file`` for the remaining parity models. + 7. With ``--ab-branch``: a temp worktree of that branch, its own venv, and + steps 3+4 rerun there (A/B against main). The worktree is removed after. + 8. ``/SUMMARY.md`` comparing main vs A/B: agreement, flip rate, drift, + time per case, peak memory — from the produced json files only. + +Idempotent: a step whose outputs already exist is skipped (``--fresh`` +ignores the markers and reruns). The step planner and the summary builder +are pure functions tested with fakes; no model ever loads in a unit test. + +Usage: + python -m benchmarks.m5 --out m5-2026-09-18 [--models-file models.txt] + [--parity-models id1,id2] [--ab-branch w2a-field-local] [--fresh] +""" + +from __future__ import annotations + +import argparse +import datetime +import json +import math +import os +import shutil +import subprocess +import sys +import time +from dataclasses import dataclass, field +from pathlib import Path + +__all__ = [ + "DEFAULT_PARITY_MODELS", + "Step", + "build_summary_text", + "combo_row", + "invariance_rollup", + "main", + "plan_steps", + "read_models_file", +] + +REPO_ROOT = Path(__file__).resolve().parent.parent +BENCH_CACHE = Path.home() / ".cache" / "jevmlx" / "bench" +TYPESAFE_JSONL = BENCH_CACHE / "typesafe.jsonl" +QUALITY_ALIAS = "quality" +DEFAULT_PARITY_MODELS = ( + "mlx-community/Qwen3-8B-4bit", + "mlx-community/Llama-3.1-8B-Instruct-4bit", + "mlx-community/gemma-3-12b-it-4bit", +) + + +@dataclass +class Step: + """One runbook step: a subprocess (or the in-process summary), its idempotency + outputs, and logging metadata. ``extra_argv`` runs after ``argv`` into the + same log (multi-command setup steps); every command must exit 0.""" + + id: str + title: str + argv: tuple[str, ...] + outputs: tuple[Path, ...] + env: dict[str, str] = field(default_factory=dict) + cwd: Path = REPO_ROOT + gate: bool = False # doctor: a failure aborts the remaining steps + capture: bool = False # capture stdout (small json) instead of streaming + capture_path: Path | None = None # where captured stdout is also written + pre_argv: tuple[str, ...] = () # run before argv when pre_target is missing + pre_target: Path | None = None + extra_argv: tuple[tuple[str, ...], ...] = () + in_process: bool = False # the summary step: rendered from artifacts + + +def _venv_bin(name: str) -> str: + """Absolute path to a console script next to the running interpreter. + + uv venvs put console scripts in ``.venv/bin`` while ``sys.executable`` + may resolve to the managed CPython that the venv symlinks to, so walk + the interpreter's directory first, then its parent's ``bin/`` (the + venv root seen through the symlink), then fall back to PATH. + """ + seen: list[Path] = [] + for base in (Path(sys.executable).resolve().parent, Path(sys.executable).parent): + if base in seen: + continue + seen.append(base) + candidate = base / name + if candidate.exists(): + return str(candidate) + return name + + +def _slug(model_id: str) -> str: + from jevmlx.bench import model_slug + + return model_slug(model_id) + + +def read_models_file(path: str | Path) -> list[str]: + """One model id per line, '#' comments and blank lines allowed.""" + models: list[str] = [] + for line in Path(path).read_text(encoding="utf-8").splitlines(): + line = line.split("#", 1)[0].strip() + if line: + models.append(line) + return models + + +def plan_steps( + out: Path, + *, + parity_models: list[str], + extra: str = "1,5,20,40", + reps: int = 5, + ab_branch: str | None = None, + typesafe_data: Path = TYPESAFE_JSONL, + quality: str = QUALITY_ALIAS, +) -> list[Step]: + """The ordered step list. Pure apart from writing ``models-rest.txt``. + + ``parity_models`` drives steps 2 (pytest per model) and 6 (bench the + remaining ones, i.e. every model except the quality alias's target). + A/B steps appear only when ``ab_branch`` is given. + """ + from jevmlx.models import resolve_model + + out = Path(out) + steps: list[Step] = [] + steps.append( + Step( + id="doctor", + title="doctor (gate)", + argv=(_venv_bin("jevmlx"), "doctor", "--json"), + outputs=(out / "doctor.done", out / "doctor.json"), + gate=True, + capture=True, + capture_path=out / "doctor.json", + ) + ) + for model in parity_models: + steps.append( + Step( + id=f"parity-{_slug(model)}", + title=f"pytest -m slow ({model})", + argv=(_venv_bin("pytest"), "-m", "slow", "-q"), + outputs=(out / f"parity-{_slug(model)}.done",), + env={"MODEL_ID": model}, + ) + ) + steps.append( + Step( + id="bench-quality", + title="jevmlx bench --model quality", + argv=( + _venv_bin("jevmlx"), + "bench", + "--model", + quality, + "--out", + str(out / "bench-quality"), + ), + outputs=(out / "bench-quality.done",), + ) + ) + steps.append( + Step( + id="invariance", + title="invariance on quality (TypeSafe cases, extra 1/5/20/40)", + argv=( + sys.executable, + "-m", + "benchmarks.invariance", + "--model", + quality, + "--data", + str(typesafe_data), + "--out", + str(out / "invariance"), + "--extra", + extra, + ), + outputs=(out / "invariance.done", out / "invariance" / "invariance.json"), + pre_argv=( + sys.executable, + "-m", + "benchmarks.typesafe.fetch", + "--out", + str(typesafe_data), + ), + pre_target=typesafe_data, + ) + ) + steps.append( + Step( + id="timing", + title=f"timing on quality ({reps} reps)", + argv=( + sys.executable, + "-m", + "benchmarks.timing", + "--model", + quality, + "--reps", + str(reps), + "--out", + str(out), + ), + outputs=( + out / "timing.done", + out / f"timing-{resolve_model(quality).replace('/', '_')}.json", + ), + ) + ) + rest = [m for m in parity_models if m != resolve_model(quality)] + if rest: + rest_file = out / "models-rest.txt" + rest_file.write_text("\n".join(rest) + "\n", encoding="utf-8") + steps.append( + Step( + id="bench-rest", + title="bench the remaining models", + argv=( + _venv_bin("jevmlx"), + "bench", + "--models-file", + str(rest_file), + "--out", + str(out / "bench-rest"), + ), + outputs=(out / "bench-rest.done",), + ) + ) + if ab_branch is not None: + worktree = out / "ab-worktree" + venv_python = worktree / ".venv" / "bin" / "python" + steps.append( + Step( + id="ab-setup", + title=f"A/B setup: worktree + venv for {ab_branch}", + argv=("git", "worktree", "add", "--detach", str(worktree), ab_branch), + extra_argv=( + ( + shutil.which("uv") or "uv", + "venv", + "--python-preference", + "only-managed", + "--python", + "3.12", + str(worktree / ".venv"), + ), + ( + shutil.which("uv") or "uv", + "pip", + "install", + "--python", + str(venv_python), + "-e", + ".[dev]", + ), + ), + outputs=(out / "ab-setup.done",), + cwd=REPO_ROOT, + ) + ) + steps.append( + Step( + id="ab-bench", + title=f"A/B bench quality on {ab_branch}", + argv=( + str(worktree / ".venv" / "bin" / "jevmlx"), + "bench", + "--model", + quality, + "--out", + str(out / "ab" / "bench-quality"), + ), + outputs=(out / "ab-bench.done",), + cwd=worktree, + ) + ) + steps.append( + Step( + id="ab-invariance", + title=f"A/B invariance on quality ({extra}) on {ab_branch}", + argv=( + str(venv_python), + "-m", + "benchmarks.invariance", + "--model", + quality, + "--data", + str(typesafe_data), + "--out", + str(out / "ab" / "invariance"), + "--extra", + extra, + ), + outputs=(out / "ab-invariance.done", out / "ab" / "invariance" / "invariance.json"), + cwd=worktree, + ) + ) + steps.append( + Step( + id="summary", + title="SUMMARY.md (main vs A/B)", + argv=(), + outputs=(out / "SUMMARY.md",), + in_process=True, + ) + ) + return steps + + +def step_done(step: Step, fresh: bool) -> bool: + """Idempotency rule: skip when every output exists, unless --fresh.""" + return not fresh and all(p.exists() for p in step.outputs) + + +def _step_env(step: Step) -> dict[str, str]: + return os.environ | step.env + + +def _now() -> str: + return datetime.datetime.now(datetime.UTC).strftime("%Y-%m-%dT%H:%M:%SZ") + + +def runbook_append( + out: Path, index: int, step: Step, *, rc: int | None, secs: float | None +) -> None: + """One RUNBOOK.md section per step, appended as the step completes.""" + lines = [f"## {index}. {step.title}"] + if rc is None: + lines[0] += " — skip (outputs exist)" + else: + lines[0] += f" — exit {rc} — {secs:.1f}s — {_now()}" + lines.append(f"cmd: {' '.join(step.argv)}") + if step.env: + lines.append(f"env: MODEL_ID={step.env['MODEL_ID']}") + lines.append(f"log: {step.id}.log") + lines.append("") + with (out / "RUNBOOK.md").open("a", encoding="utf-8") as f: + f.write("\n".join(lines)) + + +def execute_step(step: Step) -> int: + """Run one step; returns the exit code. Streams into .log.""" + log_path = step.outputs[0].parent / f"{step.id}.log" + with log_path.open("w", encoding="utf-8") as log: + if step.pre_argv and (step.pre_target is None or not step.pre_target.exists()): + log.write(f"$ {' '.join(step.pre_argv)}\n") + log.flush() + pre = subprocess.run( + step.pre_argv, + cwd=step.cwd, + env=_step_env(step), + stdout=log, + stderr=subprocess.STDOUT, + ) + if pre.returncode != 0: + return pre.returncode + if step.in_process: + text = build_summary_text( + collect_side(step.outputs[0].parent), + collect_side(step.outputs[0].parent / "ab"), + parity_models=_meta_parity_models(step.outputs[0].parent), + ) + log.write(text) + (step.outputs[0].parent / "SUMMARY.md").write_text(text, encoding="utf-8") + return 0 + if step.capture: + proc = subprocess.run( + step.argv, cwd=step.cwd, env=_step_env(step), capture_output=True, text=True + ) + log.write(proc.stdout) + log.write(proc.stderr) + if step.capture_path is not None and proc.returncode == 0: + step.capture_path.write_text(proc.stdout, encoding="utf-8") + return proc.returncode + log.write(f"$ {' '.join(step.argv)}\n") + log.flush() + proc = subprocess.run( + step.argv, cwd=step.cwd, env=_step_env(step), stdout=log, stderr=subprocess.STDOUT + ) + rc = proc.returncode + for extra in step.extra_argv: + if rc != 0: + break + log.write(f"$ {' '.join(extra)}\n") + log.flush() + rc = subprocess.run( + extra, cwd=step.cwd, env=_step_env(step), stdout=log, stderr=subprocess.STDOUT + ).returncode + return rc + + +def _meta_parity_models(out: Path) -> list[str]: + """Parity models recorded by the planner side effect (models-rest.txt holds + only the rest, so read the RUNBOOK header instead).""" + for line in (out / "RUNBOOK.md").read_text(encoding="utf-8").splitlines(): + if line.startswith("parity models:"): + return [m.strip() for m in line.split(":", 1)[1].split(",") if m.strip()] + return [] + + +# ------------------------------------------------------------- summary build + + +def _load_json(path: Path) -> dict | None: + try: + with open(path, encoding="utf-8") as f: + data = json.load(f) + except (OSError, json.JSONDecodeError): + return None + return data if isinstance(data, dict) else None + + +def combo_row(combo_name: str, report: dict | None, timing: dict | None) -> dict: + """One bench combo -> summary row (agreement-or-accuracy, total_ms, peak).""" + metrics = (report or {}).get("metrics") or {} + median = (timing or {}).get("median") or {} + peak = median.get("peak_active_bytes") + return { + "combo": combo_name, + "agreement": metrics.get("agreement"), + "accuracy": metrics.get("accuracy"), + "total_ms": median.get("total_ms"), + "peak_gb": round(peak / 2**30, 2) if isinstance(peak, (int, float)) else None, + } + + +def invariance_rollup(invariance: dict | None) -> dict: + """invariance.json -> mean flip rate, mean |log-odds drift|, per-field @40. + + Rung keys are ints in-process but strings after the JSON round-trip; + both are accepted. + """ + targets = (invariance or {}).get("targets") or [] + flips: list[float] = [] + drifts: list[float] = [] + fields: dict[str, dict] = {} + for row in targets: + rungs = {str(k): v for k, v in (row.get("rungs") or {}).items()} + for entry in rungs.values(): + if isinstance(entry.get("flip_rate"), (int, float)): + flips.append(entry["flip_rate"]) + drift = entry.get("winner_logodds_drift_mean") + if isinstance(drift, (int, float)) and math.isfinite(drift): + drifts.append(drift) + rung40 = rungs.get("40") or {} + fields[row.get("field", "?")] = { + "flip40": rung40.get("flip_rate"), + "drift40": rung40.get("winner_logodds_drift_mean"), + } + return { + "mean_flip_rate": sum(flips) / len(flips) if flips else None, + "mean_drift": sum(drifts) / len(drifts) if drifts else None, + "fields": fields, + } + + +def collect_side(side_dir: Path) -> dict: + """Bench + invariance artifacts for one side (main or ab) -> summary block.""" + side_dir = Path(side_dir) + combos = [ + combo_row( + report_path.parent.name, + _load_json(report_path), + _load_json(report_path.parent / "timing.json"), + ) + for report_path in sorted(side_dir.glob("bench-quality/**/report.json")) + ] + block: dict = { + "combos": combos, + "invariance": invariance_rollup(_load_json(side_dir / "invariance" / "invariance.json")), + } + timing_reports = sorted(side_dir.glob("timing-*.json")) + if timing_reports: + report = _load_json(timing_reports[0]) or {} + presets = report.get("presets") or {} + timing_block: dict = {} + for preset_id, preset in sorted(presets.items()): + agg = preset.get("aggregate") or {} + total = (agg.get("total_ms") or {}).get("median") + peak = (agg.get("peak_active_bytes") or {}).get("median") + timing_block[preset_id] = { + "total_ms_median": total, + "peak_gb": round(peak / 2**30, 2) if isinstance(peak, (int, float)) else None, + } + if timing_block: + block["timing_report"] = timing_block + return block + + +def _fmt(value: float | int | None, pct: bool = False) -> str: + if value is None: + return "-" + return f"{value:.1%}" if pct else f"{value:.4g}" + + +def _mean(values: list[float]) -> float | None: + return sum(values) / len(values) if values else None + + +def _combo_agreement(row: dict) -> float | None: + return row["agreement"] if row["agreement"] is not None else row["accuracy"] + + +def build_summary_text(main: dict, ab: dict | None, *, parity_models: list[str]) -> str: + """SUMMARY.md: main table, timing report, invariance rollup, A/B deltas. + + Pure: takes the collected blocks (see collect_side), never touches disk. + """ + lines = [ + "# M5 runbook summary", + "", + f"- parity models: {', '.join(parity_models) or '-'}", + "", + "## Main — bench combos", + "", + "| combo | agreement/accuracy | total_ms | peak_gb |", + "|---|---|---|---|", + ] + for row in main.get("combos") or []: + lines.append( + f"| {row['combo']} | {_fmt(_combo_agreement(row), pct=True)} " + f"| {_fmt(row['total_ms'])} | {_fmt(row['peak_gb'])} |" + ) + timing_report = main.get("timing_report") or {} + if timing_report: + lines += [ + "", + "## Main — timing report (quality, decide() presets)", + "", + "| preset | total_ms median | peak_gb |", + "|---|---|---|", + ] + for preset, entry in timing_report.items(): + lines.append( + f"| {preset} | {_fmt(entry.get('total_ms_median'))} " + f"| {_fmt(entry.get('peak_gb'))} |" + ) + inv = main.get("invariance") or {} + if inv.get("fields"): + lines += [ + "", + "## Main — invariance (TypeSafe, extra 1/5/20/40; @40 vs extra=1 baseline)", + "", + f"- mean flip rate: {_fmt(inv.get('mean_flip_rate'), pct=True)}", + f"- mean |winner log-odds drift|: {_fmt(inv.get('mean_drift'))}", + "", + "| field | flip@40 | drift@40 |", + "|---|---|---|", + ] + for field, entry in inv["fields"].items(): + lines.append( + f"| {field} | {_fmt(entry.get('flip40'), pct=True)} " + f"| {_fmt(entry.get('drift40'))} |" + ) + if ab is not None: + + def _agg(block: dict) -> dict: + rows = block.get("combos") or [] + agreements = [v for r in rows if (v := _combo_agreement(r)) is not None] + total_ms = [v for r in rows if (v := r["total_ms"]) is not None] + peaks = [v for r in rows if (v := r["peak_gb"]) is not None] + block_inv = block.get("invariance") or {} + return { + "agreement": _mean(agreements), + "flip_rate": block_inv.get("mean_flip_rate"), + "drift": block_inv.get("mean_drift"), + "time_per_case_ms": _mean(total_ms), + "peak_gb": max(peaks) if peaks else None, + } + + m, a = _agg(main), _agg(ab) + labels = [ + ("agreement", "agreement/accuracy (mean)", True), + ("flip_rate", "flip rate (invariance mean)", True), + ("drift", "|log-odds drift| (mean)", False), + ("time_per_case_ms", "time per case (ms)", False), + ("peak_gb", "peak memory (GB, max)", False), + ] + lines += [ + "", + "## A/B comparison — main vs A/B (bench quality + invariance)", + "", + "| metric | main | A/B | delta (A/B - main) |", + "|---|---|---|---|", + ] + for key, label, pct in labels: + mv, av = m[key], a[key] + delta = None if mv is None or av is None else av - mv + lines.append( + f"| {label} | {_fmt(mv, pct=pct)} | {_fmt(av, pct=pct)} | {_fmt(delta, pct=pct)} |" + ) + else: + lines += ["", "## A/B comparison", "", "A/B not run (no --ab-branch)."] + lines.append("") + return "\n".join(lines) + + +# --------------------------------------------------------------------- main + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser( + prog="python -m benchmarks.m5", + description=( + "One-command M5 runbook (W5b-4): doctor, slow parity suite, bench," + " invariance, timing, A/B, summary." + ), + ) + parser.add_argument("--out", required=True, help="runbook output directory, e.g. m5-2026-09-18") + parser.add_argument( + "--parity-models", + default=None, + help="comma-separated parity model ids (default: Qwen3-8B, Llama-3.1-8B, Gemma-3-12B)", + ) + parser.add_argument( + "--models-file", + default=None, + help="file with one parity model id per line ('#' comments allowed);" + " overrides --parity-models", + ) + parser.add_argument( + "--ab-branch", default=None, help="git branch/ref to bench + invariance as A/B" + ) + parser.add_argument("--reps", type=int, default=5, help="timing reps (step 5)") + parser.add_argument( + "--extra", default="1,5,20,40", help="invariance extra-field ladder (step 4)" + ) + parser.add_argument( + "--fresh", action="store_true", help="rerun steps whose outputs already exist" + ) + args = parser.parse_args(argv) + + if args.parity_models: + parity_models = [m.strip() for m in args.parity_models.split(",") if m.strip()] + elif args.models_file: + parity_models = read_models_file(args.models_file) + else: + parity_models = list(DEFAULT_PARITY_MODELS) + + out = Path(args.out).expanduser().resolve() + out.mkdir(parents=True, exist_ok=True) + steps = plan_steps( + out, + parity_models=parity_models, + extra=args.extra, + reps=args.reps, + ab_branch=args.ab_branch, + ) + + if not (out / "RUNBOOK.md").exists(): + header = [ + f"# M5 runbook — {_now()}", + f"cmd: python -m benchmarks.m5 --out {args.out}" + + (f" --ab-branch {args.ab_branch}" if args.ab_branch else ""), + f"parity models: {', '.join(parity_models)}", + "", + ] + (out / "RUNBOOK.md").write_text("\n".join(header), encoding="utf-8") + + failures: list[str] = [] + worktree = out / "ab-worktree" + try: + for index, step in enumerate(steps, 1): + if step_done(step, args.fresh): + runbook_append(out, index, step, rc=None, secs=None) + continue + secs = time.perf_counter() + rc = execute_step(step) + secs = time.perf_counter() - secs + runbook_append(out, index, step, rc=rc, secs=secs) + if rc == 0: + for marker in step.outputs: + if marker.name == f"{step.id}.done": + marker.touch() + else: + failures.append(step.id) + if step.gate: + with (out / "RUNBOOK.md").open("a", encoding="utf-8") as f: + f.write( + f"**ABORT: {step.id} failed (exit {rc}); remaining steps skipped.**\n\n" + ) + return 1 + finally: + if worktree.exists() and shutil.which("git"): + subprocess.run( + ["git", "worktree", "remove", "--force", str(worktree)], + cwd=REPO_ROOT, + capture_output=True, + ) + shutil.rmtree(worktree, ignore_errors=True) + + return 1 if failures else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_m5.py b/tests/test_m5.py new file mode 100644 index 0000000..aa810b1 --- /dev/null +++ b/tests/test_m5.py @@ -0,0 +1,263 @@ +"""Tests for benchmarks/m5.py — step planner and summary builder only. + +No model ever loads here: the planner is checked against a temp directory, +and the summary builder against synthetic json blocks (the fake-model +contract). Subprocesses are not spawned either (execute_step is exercised +only via a fake runner in the flow test). +""" + +from __future__ import annotations + +import pytest + +from benchmarks.m5 import ( + DEFAULT_PARITY_MODELS, + Step, + build_summary_text, + combo_row, + invariance_rollup, + plan_steps, + read_models_file, + step_done, +) + +QUALITY_TARGET = "mlx-community/Qwen2.5-7B-Instruct-4bit" + + +@pytest.fixture() +def out(tmp_path): + return tmp_path / "m5-run" + + +class TestReadModelsFile: + def test_comments_and_blanks_skipped(self, tmp_path): + f = tmp_path / "models.txt" + f.write_text( + "# parity set\n" + "mlx-community/Qwen3-8B-4bit\n" + "\n" + "mlx-community/gemma-3-12b-it-4bit # trailing comment\n", + encoding="utf-8", + ) + assert read_models_file(f) == [ + "mlx-community/Qwen3-8B-4bit", + "mlx-community/gemma-3-12b-it-4bit", + ] + + +class TestPlanSteps: + def test_default_order_no_ab(self, out): + out.mkdir() + steps = plan_steps(out, parity_models=list(DEFAULT_PARITY_MODELS)) + assert [s.id for s in steps] == [ + "doctor", + "parity-mlx-community--qwen3-8b-4bit", + "parity-mlx-community--llama-3.1-8b-instruct-4bit", + "parity-mlx-community--gemma-3-12b-it-4bit", + "bench-quality", + "invariance", + "timing", + "bench-rest", + "summary", + ] + + def test_doctor_is_gate_and_captures_json(self, out): + steps = plan_steps(out, parity_models=[QUALITY_TARGET]) + doctor = steps[0] + assert doctor.gate is True + assert doctor.capture is True + assert doctor.capture_path == out / "doctor.json" + assert "--json" in doctor.argv + + def test_parity_step_rides_model_id_env(self, out): + out.mkdir() + steps = plan_steps(out, parity_models=["some-org/Model-4bit"]) + parity = next(s for s in steps if s.id.startswith("parity-")) + assert parity.env == {"MODEL_ID": "some-org/Model-4bit"} + assert "-m" in parity.argv and "slow" in parity.argv + + def test_bench_rest_excludes_quality_target(self, out): + out.mkdir() + models = [QUALITY_TARGET, "other-org/Other-4bit"] + steps = plan_steps(out, parity_models=models) + rest_file = out / "models-rest.txt" + assert rest_file.exists() + assert rest_file.read_text().split() == ["other-org/Other-4bit"] + bench_rest = next(s for s in steps if s.id == "bench-rest") + assert "--models-file" in bench_rest.argv + + def test_no_rest_step_when_only_quality_target(self, out): + steps = plan_steps(out, parity_models=[QUALITY_TARGET]) + assert not (out / "models-rest.txt").exists() + assert all(s.id != "bench-rest" for s in steps) + + def test_invariance_fetches_typesafe_when_missing(self, out): + steps = plan_steps(out, parity_models=[QUALITY_TARGET], typesafe_data=out / "cases.jsonl") + step = next(s for s in steps if s.id == "invariance") + assert step.pre_target == out / "cases.jsonl" + assert step.pre_argv and "typesafe.fetch" in " ".join(step.pre_argv) + assert "--extra" in step.argv and "1,5,20,40" in step.argv + + def test_ab_steps_branch_worktree_and_cwd(self, out): + steps = plan_steps(out, parity_models=[QUALITY_TARGET], ab_branch="w2a-field-local") + ids = [s.id for s in steps] + assert ids[-3:] == ["ab-setup", "ab-bench", "ab-invariance"] or ids[-4:] == [ + "ab-setup", + "ab-bench", + "ab-invariance", + "summary", + ] + ab_setup = next(s for s in steps if s.id == "ab-setup") + assert "worktree" in ab_setup.argv and "w2a-field-local" in ab_setup.argv + assert any("uv" in " ".join(extra) for extra in ab_setup.extra_argv) + ab_bench = next(s for s in steps if s.id == "ab-bench") + assert ab_bench.cwd == out / "ab-worktree" + assert str(out / "ab" / "bench-quality") in " ".join(ab_bench.argv) + ab_inv = next(s for s in steps if s.id == "ab-invariance") + assert str(out / "ab" / "invariance" / "invariance.json") in [ + str(p) for p in ab_inv.outputs + ] + + def test_summary_step_in_process(self, out): + steps = plan_steps(out, parity_models=[QUALITY_TARGET]) + summary = steps[-1] + assert summary.in_process is True + assert summary.argv == () + assert summary.outputs == (out / "SUMMARY.md",) + + def test_every_step_has_done_marker_and_outputs(self, out): + for steps in ( + plan_steps(out / "a", parity_models=[QUALITY_TARGET]), + plan_steps(out / "b", parity_models=[QUALITY_TARGET], ab_branch="main"), + ): + for step in steps: + assert step.outputs, step.id + + +class TestStepDone: + def test_skips_when_all_outputs_exist(self, tmp_path): + done = tmp_path / "x.done" + done.touch() + step = Step(id="x", title="t", argv=("true",), outputs=(done,)) + assert step_done(step, fresh=False) is True + + def test_missing_output_means_run(self, tmp_path): + step = Step(id="x", title="t", argv=("true",), outputs=(tmp_path / "missing.done",)) + assert step_done(step, fresh=False) is False + + def test_fresh_reruns_despite_outputs(self, tmp_path): + done = tmp_path / "x.done" + done.touch() + step = Step(id="x", title="t", argv=("true",), outputs=(done,)) + assert step_done(step, fresh=True) is False + + +class TestComboRow: + def test_agreement_wins_accuracy_fallback(self): + row = combo_row( + "c1", + {"metrics": {"agreement": 0.9}}, + {"median": {"total_ms": 100.0, "peak_active_bytes": 2**30}}, + ) + assert row == { + "combo": "c1", + "agreement": 0.9, + "accuracy": None, + "total_ms": 100.0, + "peak_gb": 1.0, + } + + def test_missing_report_is_none_safe(self): + row = combo_row("c2", None, None) + assert row["accuracy"] is None and row["peak_gb"] is None and row["total_ms"] is None + + +class TestInvarianceRollup: + def test_means_and_field_ladder(self): + inv = { + "targets": [ + { + "field": "f1", + "rungs": { + "1": {"flip_rate": 0.0, "winner_logodds_drift_mean": 0.0}, + "40": {"flip_rate": 0.5, "winner_logodds_drift_mean": 1.5}, + }, + }, + { + "field": "f2", + "rungs": { + "1": {"flip_rate": 0.0, "winner_logodds_drift_mean": 0.2}, + "40": {"flip_rate": 0.1, "winner_logodds_drift_mean": 0.4}, + }, + }, + ] + } + roll = invariance_rollup(inv) + assert roll["mean_flip_rate"] == pytest.approx(0.15) + assert roll["mean_drift"] == pytest.approx(0.525) + assert roll["fields"]["f1"]["flip40"] == 0.5 + assert roll["fields"]["f1"]["drift40"] == 1.5 + + def test_inf_drift_excluded_from_mean(self): + inv = { + "targets": [ + { + "field": "f", + "rungs": { + "1": {"flip_rate": 0.0, "winner_logodds_drift_mean": 0.0}, + "40": {"flip_rate": 0.0, "winner_logodds_drift_mean": float("inf")}, + }, + } + ] + } + roll = invariance_rollup(inv) + assert roll["mean_drift"] == 0.0 + + def test_none_safe(self): + assert invariance_rollup(None)["fields"] == {} + assert invariance_rollup(None)["mean_flip_rate"] is None + + +class TestBuildSummaryText: + @staticmethod + def _side( + *, agreement: float, flip: float, drift: float, total_ms: float, peak_gb: float + ) -> dict: + return { + "combos": [ + combo_row( + "parallel-slots-typesafe", + {"metrics": {"agreement": agreement}}, + {"median": {"total_ms": total_ms, "peak_active_bytes": int(peak_gb * 2**30)}}, + ) + ], + "invariance": { + "mean_flip_rate": flip, + "mean_drift": drift, + "fields": {"risk_level": {"flip40": flip, "drift40": drift}}, + }, + "timing_report": {"fintech_fraud": {"total_ms_median": total_ms, "peak_gb": peak_gb}}, + } + + def test_main_only(self): + main = self._side(agreement=0.9, flip=0.05, drift=0.3, total_ms=1200.0, peak_gb=6.5) + text = build_summary_text(main, None, parity_models=["a/b"]) + assert "# M5 runbook summary" in text + assert "parallel-slots-typesafe" in text + assert "90.0%" in text + assert "A/B not run" in text + assert "risk_level" in text + + def test_ab_delta_rows(self): + main = self._side(agreement=0.90, flip=0.10, drift=1.0, total_ms=1000.0, peak_gb=6.0) + ab = self._side(agreement=0.92, flip=0.05, drift=0.5, total_ms=800.0, peak_gb=5.5) + text = build_summary_text(main, ab, parity_models=[]) + assert "| agreement/accuracy (mean) | 90.0% | 92.0% | 2.0% |" in text + assert "| flip rate (invariance mean) | 10.0% | 5.0% | -5.0% |" in text + assert "| time per case (ms) | 1000 | 800 | -200 |" in text + assert "| peak memory (GB, max) | 6 | 5.5 | -0.5 |" in text + + def test_none_values_render_as_dash(self): + text = build_summary_text({"combos": [], "invariance": {}}, None, parity_models=[]) + assert "- |" in text or "|" in text # table renders, never crashes + assert "M5 runbook summary" in text