nisse/nisse-dev/src/nisse_dev/bench/results.py

"""Stored benchmark results, and the report that compares them.

Every run is one JSON line. The report is rebuilt from all of them each
time, so a model benchmarked today is read alongside every model measured
before it, as long as the task is still the same version.
"""

import json
import math
import statistics
from collections import defaultdict
from collections.abc import Iterable
from dataclasses import asdict, dataclass
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from nisse_dev.bench.runner import RunRecord
from nisse_dev.config import Tuning


def results_file(directory: Path, model: str, when: datetime | None = None) -> Path:
    day = (when or datetime.now(UTC)).strftime("%Y-%m-%d")
    return directory / f"{day}-{model}.jsonl"


def append(path: Path, record: RunRecord) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a") as out:
        out.write(json.dumps(record.to_json(), sort_keys=True) + "\n")


def load(directory: Path) -> list[dict[str, Any]]:
    records = []
    for path in sorted(directory.glob("*.jsonl")):
        for line in path.read_text().splitlines():
            if line.strip():
                records.append(json.loads(line))
    return records


def wilson(successes: int, runs: int, z: float = 1.96) -> tuple[float, float]:
    """A 95% confidence interval for a success rate, sound at small counts.

    Three successes in five runs and five in five overlap here, which is the
    point: small samples of a noisy model do not separate much."""
    if runs == 0:
        return 0.0, 1.0
    p = successes / runs
    denominator = 1 + z * z / runs
    centre = (p + z * z / (2 * runs)) / denominator
    half = z * math.sqrt(p * (1 - p) / runs + z * z / (4 * runs * runs)) / denominator
    return max(0.0, centre - half), min(1.0, centre + half)


@dataclass
class Row:
    model: str
    note: str
    task: str
    runs: int
    successes: int
    median_turns: float
    mean_wall_s: float  # working time: rate-limit waits are left out
    cost_per_run: float
    stale: bool
    rate_limited: int = 0  # runs that waited out the proxy's rate limit
    tuning: str = ""  # how the runs' tuning differs from today's defaults

    @property
    def interval(self) -> tuple[float, float]:
        return wilson(self.successes, self.runs)


def tuning_label(tuning: dict[str, Any]) -> str:
    """The settings in a run's tuning that differ from today's defaults, as
    "executor.max_turns=24"; empty when there are none. Relative to today's
    defaults so that rows at the current settings stay unlabelled."""
    defaults = asdict(Tuning())
    return ", ".join(
        f"{section}.{key}={value}"
        for section, values in sorted(tuning.items())
        for key, value in sorted((values or {}).items())
        if defaults.get(section, {}).get(key) != value
    )


def summarise(
    records: Iterable[dict[str, Any]], current_hashes: dict[str, str]
) -> list[Row]:
    """One row per (model, how it was served, task, task version, tuning).

    Tuning is part of the key because a knob changes the result: runs capped
    at 24 turns and runs allowed 40 are not the same measurement."""
    groups: dict[tuple[str, str, str, str, str], list[dict[str, Any]]] = defaultdict(
        list
    )
    for r in records:
        key = (
            r["backing"],
            r.get("note", ""),
            r["task"],
            r["task_hash"],
            tuning_label(r.get("tuning", {})),
        )
        groups[key].append(r)
    rows = []
    for (backing, note, task, task_hash, tuning), runs in sorted(groups.items()):
        rows.append(
            Row(
                model=backing,
                note=note,
                task=task,
                runs=len(runs),
                successes=sum(1 for r in runs if r["success"]),
                median_turns=statistics.median(r["turns"] for r in runs),
                mean_wall_s=statistics.fmean(
                    r["wall_s"] - r.get("rate_limit_wait_s", 0.0) for r in runs
                ),
                cost_per_run=statistics.fmean(r["cost"] for r in runs),
                stale=current_hashes.get(task) != task_hash,
                rate_limited=sum(1 for r in runs if r.get("rate_limit_wait_s")),
                tuning=tuning,
            )
        )
    return rows


def render(rows: list[Row]) -> str:
    # A tuning column only when some rows need one, so a report at today's
    # settings looks as it always has.
    tuned = any(row.tuning for row in rows)
    header = (
        f"{'model':<34} {'note':<22} "
        + (f"{'tuning':<24} " if tuned else "")
        + f"{'task':<15} {'success':>8} {'95% CI':>11} "
        f"{'turns':>6} {'time':>7} {'$/run':>8} {'429':>4}"
    )
    lines = [header, "-" * len(header)]
    for row in rows:
        low, high = row.interval
        task = row.task + (" (old)" if row.stale else "")
        tuning = f"{row.tuning[:24]:<24} " if tuned else ""
        lines.append(
            f"{row.model[:34]:<34} {row.note[:22]:<22} {tuning}{task[:15]:<15} "
            f"{row.successes:>3}/{row.runs:<4} {low:>4.0%}-{high:<5.0%} "
            f"{row.median_turns:>6.0f} {row.mean_wall_s:>6.0f}s {row.cost_per_run:>8.4f} "
            f"{row.rate_limited or '':>4}"
        )
    if any(row.rate_limited for row in rows):
        lines.append("")
        lines.append(
            "429 = runs that waited out the proxy's rate limit; time leaves the "
            "waits out."
        )
    if tuned:
        lines.append("")
        lines.append("tuning = settings that differ from today's defaults.")
    if any(row.stale for row in rows):
        lines.append("")
        lines.append(
            "(old) = measured against an earlier version of the task; "
            "not comparable with current runs."
        )
    return "\n".join(lines)