""" bench.py — The orchestrator (torch-free at import; the runner is injected). Run matrix = models × reasoning × category × mode × N, iterated model-outer so a heavy checkpoint loads once and is freed before the next. Ground truth is loaded once per category so every model sees identical inputs (fairness). Durability (the project's standing pattern): config.json + a stream-written results.jsonl (one row per scored sample, written immediately) + a rejects.jsonl sidecar + an append-only run.log. Resume skips already-completed (model, reasoning, category, mode, image_id) keys, so a Colab disconnect costs at most one row. """ from __future__ import annotations import json import time from dataclasses import asdict, dataclass, field from datetime import datetime, timezone from pathlib import Path from typing import Callable, Optional from .datasets import load_gt from .metrics import score_vision_run, score_vision_sample from .report import write_reports from .tasks_vision import get_task, pilot_categories @dataclass class BenchConfig: models: list[str] categories: list[str] = field(default_factory=pilot_categories) reasonings: list[str] = field(default_factory=lambda: ["instruct"]) modes: list[str] = field(default_factory=lambda: ["json_mode"]) n: int = 50 dataset: str = "smoke" # "smoke" | "full" runner: str = "stub" # "stub" | "vlm" precision: str = "bf16" stub_behavior: str = "perfect" # stub only output_root: str = "runs/vision" gpu_hourly_rate: float = 2.0 clear_cache_after_model: bool = False # rm each model's HF cache after use (full-array sweeps) def _free_model_cache(model_key: str) -> None: """Delete a model's HF Hub cache from disk (full-array sweeps on a tight SSD).""" import os import shutil try: from .model_registry import get_model spec = get_model(model_key) repos = [spec.repo_id] + list(spec.quant_repo_ids.values()) if spec.thinking_repo_id: repos.append(spec.thinking_repo_id) except Exception: repos = [model_key] base = os.path.expanduser("~/.cache/huggingface/hub") for repo in repos: p = os.path.join(base, "models--" + repo.replace("/", "--")) if os.path.isdir(p): shutil.rmtree(p, ignore_errors=True) def _utc_stamp() -> str: # microsecond precision so back-to-back runs never collide into one run dir # (which would let resume fold one run's metrics into another's report) return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S_%fZ") def _default_runner_factory(config: BenchConfig) -> Callable[[str, str], object]: if config.runner == "stub": from .stub_runner import StubVLMRunner return lambda mk, rsn: StubVLMRunner(model_id=mk, behavior=config.stub_behavior, reasoning=rsn) # real VLM — imports torch lazily inside get_runner from .model_registry import get_runner return lambda mk, rsn: get_runner(mk, precision=config.precision, reasoning=rsn) def _completed_keys(results_path: Path) -> set: done = set() if not results_path.exists(): return done for line in results_path.read_text(encoding="utf-8").splitlines(): if not line.strip(): continue try: r = json.loads(line) done.add((r["model"], r["reasoning"], r["category"], r["mode"], r["image_id"])) except (json.JSONDecodeError, KeyError): continue return done def run_bench(config: BenchConfig, runner_factory: Optional[Callable] = None, run_dir: Optional[Path] = None) -> dict: runner_factory = runner_factory or _default_runner_factory(config) root = Path(config.output_root) run_dir = run_dir or (root / _utc_stamp()) run_dir.mkdir(parents=True, exist_ok=True) results_path = run_dir / "results.jsonl" rejects_path = run_dir / "rejects.jsonl" metrics_path = run_dir / "metrics.jsonl" log_path = run_dir / "run.log" (run_dir / "config.json").write_text(json.dumps(asdict(config), indent=2), encoding="utf-8") done = _completed_keys(results_path) def log(msg: str) -> None: stamp = datetime.now(timezone.utc).strftime("%H:%M:%S") with log_path.open("a", encoding="utf-8") as fh: fh.write(f"[{stamp}] {msg}\n") log(f"start config={asdict(config)}") metric_rows: list[dict] = [] n_total = n_valid = n_reject = n_skip = 0 with results_path.open("a", encoding="utf-8") as res_fh, \ rejects_path.open("a", encoding="utf-8") as rej_fh, \ metrics_path.open("a", encoding="utf-8") as met_fh: for model_key in config.models: for reasoning in config.reasonings: t_model = time.perf_counter() runner = runner_factory(model_key, reasoning) log(f"loaded {model_key}/{reasoning}") try: for category in config.categories: spec = get_task(category) gt_key = category if config.dataset == "smoke" else spec.gt_dataset samples = load_gt(gt_key, n=config.n, split=spec.gt_split, dataset=config.dataset) for mode in config.modes: cell: list = [] for s in samples: key = (model_key, reasoning, category, mode, s.image_id) if key in done: n_skip += 1 continue up = s.prompt if spec.per_sample_prompt else None res = runner.generate(spec, s.image, mode, image_id=s.image_id, image_size=s.size, gt=s.gt, user_prompt=up) mr = score_vision_sample( spec, res.raw_text, s.gt, mode=mode, image_id=s.image_id, image_size=s.size, grammar_conformant=res.grammar_conformant, n_output_tokens=res.n_output_tokens, gen_seconds=res.gen_seconds) cell.append(mr) n_total += 1 row = {"model": model_key, "reasoning": reasoning, **mr.to_dict()} res_fh.write(json.dumps(row) + "\n") res_fh.flush() if mr.schema_valid: n_valid += 1 else: n_reject += 1 rej_fh.write(json.dumps({**row, "raw_text": res.raw_text}) + "\n") rej_fh.flush() if cell: rm = score_vision_run(cell, model=model_key, reasoning=reasoning, category=category, mode=mode) row = asdict(rm) metric_rows.append(row) met_fh.write(json.dumps(row) + "\n") met_fh.flush() log(str(rm)) finally: close = getattr(runner, "close", None) if callable(close): close() if config.clear_cache_after_model and config.runner == "vlm": _free_model_cache(model_key) log(f"freed {model_key}/{reasoning} in {time.perf_counter() - t_model:.1f}s") # If resuming, fold in prior metric rows from metrics.jsonl for a complete report. if metrics_path.exists(): seen = {(r["model"], r["reasoning"], r["category"], r["mode"]) for r in metric_rows} for line in metrics_path.read_text(encoding="utf-8").splitlines(): if not line.strip(): continue try: r = json.loads(line) except json.JSONDecodeError: continue if (r["model"], r["reasoning"], r["category"], r["mode"]) not in seen: metric_rows.append(r) seen.add((r["model"], r["reasoning"], r["category"], r["mode"])) write_reports(run_dir, metric_rows, asdict(config)) summary = { "run_dir": str(run_dir), "n_total": n_total, "n_schema_valid": n_valid, "n_rejected": n_reject, "n_skipped": n_skip, "models": config.models, "categories": config.categories, } (run_dir / "run_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8") log(f"done {summary}") return summary