qwen_test_runner/vision/bench.py

"""
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