nisse/scripts/benchmark.py

"""Benchmark driver models on canned milestones, and compare the results.

    uv run python scripts/benchmark.py run --models A,B [--runs 5] [--tasks T,...]
        [--note "vLLM 0.29, AWQ, fp8 KV"] [--config PATH]
    uv run python scripts/benchmark.py report [--models A,B] [--tasks T,...]

`run` appends one JSON line per run to bench/results/, which is committed:
the report compares every run ever stored, so benchmarking a new model means
running only that model. See docs/Tuning-Guide.md, "Benchmarking models".
"""

import argparse
import os
import sys
from pathlib import Path

from nisse_core.config import ConfigError
from nisse_core.envfile import load_env_file
from nisse_core.models.catalog import ModelCatalog
from nisse_dev.bench.results import append, load, render, results_file, summarise
from nisse_dev.bench.runner import agent_commit, run_task
from nisse_dev.bench.tasks import load_tasks
from nisse_dev.config import AgentConfig
from nisse_dev.loop.build import build_host, make_llm_clients

REPO = Path(__file__).resolve().parents[1]
RESULTS = REPO / "bench" / "results"


def _names(raw: str | None) -> list[str] | None:
    return [n.strip() for n in raw.split(",") if n.strip()] if raw else None


def run(args: argparse.Namespace) -> int:
    load_env_file(Path(os.environ.get("NISSE_ENV", ".env")).expanduser())
    try:
        config = (
            AgentConfig.from_file(Path(args.config))
            if args.config
            else AgentConfig.load()
        )
    except ConfigError as exc:
        print(exc, file=sys.stderr)
        return 2
    models = _names(args.models) or []
    tasks = load_tasks(_names(args.tasks))
    catalog = ModelCatalog.discover(config.llm_base_url, config.llm_token)
    catalog.require_limits({m: m for m in models})
    normalized, _passthrough = make_llm_clients(config, catalog)
    # Runs happen on the first project's build VM, as the agent's do.
    host = build_host(config, config.projects[0])
    commit = agent_commit(REPO)

    for model in models:
        facts = catalog.facts(model)
        out = results_file(args.results, model)
        for task in tasks:
            for n in range(args.runs):
                # The driver always takes the normalized route in the agent,
                # so it does here too.
                record = run_task(
                    task,
                    facts,
                    normalized,
                    config.tuning,
                    host=host,
                    note=args.note,
                    commit=commit,
                    keep=args.keep_runs,
                )
                append(out, record)
                verdict = "pass" if record.success else f"FAIL ({record.outcome})"
                extra = (
                    f" fallback to {record.fallbacks[0]}!" if record.fallbacks else ""
                )
                print(
                    f"{model} {task.name} run {n + 1}/{args.runs}: {verdict}, "
                    f"{record.turns} turns, {record.wall_s:.0f}s, "
                    f"${record.cost:.4f}{extra}",
                    flush=True,
                )
    print()
    return report(args)


def report(args: argparse.Namespace) -> int:
    records = load(args.results)
    models = _names(getattr(args, "models", None))
    tasks = _names(getattr(args, "tasks", None))
    if models:
        records = [r for r in records if r["model"] in models or r["backing"] in models]
    if tasks:
        records = [r for r in records if r["task"] in tasks]
    if not records:
        print("no stored results match", file=sys.stderr)
        return 1
    current = {task.name: task.content_hash for task in load_tasks()}
    print(render(summarise(records, current)))
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    parser.add_argument("--results", type=Path, default=RESULTS, help=argparse.SUPPRESS)
    sub = parser.add_subparsers(dest="command", required=True)

    run_p = sub.add_parser("run", help="benchmark models, then report")
    run_p.add_argument(
        "--models", required=True, help="LiteLLM aliases, comma separated"
    )
    run_p.add_argument("--runs", type=int, default=5, help="runs per task (default 5)")
    run_p.add_argument("--tasks", help="task names, comma separated (default: all)")
    run_p.add_argument(
        "--note", default="", help='how the model was served, e.g. "vLLM 0.29, AWQ"'
    )
    run_p.add_argument("--config", help="agent config (default: NISSE_CONFIG)")
    run_p.add_argument(
        "--keep-runs",
        type=Path,
        help="keep each run's state store and call record here, for the status page",
    )
    run_p.set_defaults(handler=run)

    report_p = sub.add_parser("report", help="compare stored results")
    report_p.add_argument("--models", help="aliases or backing models to include")
    report_p.add_argument("--tasks", help="tasks to include")
    report_p.set_defaults(handler=report)

    args = parser.parse_args()
    return args.handler(args)


if __name__ == "__main__":
    raise SystemExit(main())