bench/kbench/results.py

"""Result schema.

One run produces one JSON file in results/, committed to git. The commit
history IS the longitudinal record -- there is no database to migrate or rot.

Two properties of this schema are load-bearing and expensive to retrofit:

1. PER-SAMPLE OUTCOMES, not just aggregates. An aggregate tells you a
   checkpoint got worse. Per-sample diffs tell you *which capability broke*.
   That difference is what separates a leaderboard from a debugging tool, and
   it is the entire reason this bench is useful for fine-tuning.

2. THE FULL TARGET IS SNAPSHOTTED INTO THE RESULT. Registry entries change.
   A result that merely references a target id becomes unreproducible the
   first time someone edits models.yaml. Every run carries its own copy of
   the serving flags it actually ran against.

Output storage policy: we store the score plus a short excerpt per sample.
Full model outputs stay in the raw Inspect .eval log, which is gitignored --
they are large, and for signal-tier tasks they can echo sensitive content
from the prompts they were derived from. Deep debugging uses the local log;
git carries the shape of the failure, not the whole transcript.
"""

from __future__ import annotations

import json
import platform
from dataclasses import asdict, dataclass, field
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

SCHEMA_VERSION = 1
RESULTS_DIR = Path(__file__).resolve().parent.parent / "results"


def _json_safe(value: Any) -> Any:
    """Map non-finite floats to None, recursively.

    Python's json emits bare `NaN`/`Infinity`, which its own loader accepts and almost
    nothing else does — JSON.parse, Go, and serde all reject them. A score card is meant
    to be a portable, durable record, so a card that only Python can read is a broken one.
    An ungraded sample surfaced as NaN here and made the first committed card unparseable
    by the website that has to render it.

    None means "no score", which is exactly what a non-finite score meant anyway.
    """
    if isinstance(value, float):
        # NaN != NaN, and inf fails the range check; both are unrepresentable in JSON.
        return None if value != value or value in (float("inf"), float("-inf")) else value
    if isinstance(value, dict):
        return {k: _json_safe(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [_json_safe(v) for v in value]
    return value

# How much of each model output to keep in the committed record.
EXCERPT_CHARS = 400


def utc_now() -> str:
    return datetime.now(UTC).isoformat(timespec="seconds")


@dataclass
class SampleOutcome:
    """One sample's result. The unit of regression debugging."""

    sample_id: str
    score: float
    passed: bool
    excerpt: str | None = None  # truncated model output; full text in .eval log
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class QualityResult:
    """One task's scores against one target."""

    task: str
    tier: str  # "reference" (public, contaminated, calibration only) | "signal"
    dataset_version: str | None = None
    # Content hash of the samples actually run. dataset_version is inspect's static
    # task_version and does not move when the data changes; this does. Two cards are
    # only comparable if their fingerprints match.
    dataset_fingerprint: str | None = None
    n_samples: int = 0
    metrics: dict[str, float] = field(default_factory=dict)
    samples: list[SampleOutcome] = field(default_factory=list)
    duration_s: float | None = None
    error: str | None = None


@dataclass
class PerfPoint:
    """Serving performance at one concurrency level."""

    concurrency: int
    input_tokens: int
    output_tokens: int
    n_requests: int
    completed: int
    failed: int
    duration_s: float
    # Aggregate throughput across all streams -- the number that matters for
    # multi-user serving, and the one single-stream reviews miss entirely.
    output_tps_total: float
    # Per-stream decode rate -- what a single user actually feels.
    output_tps_per_stream: float
    ttft_p50_ms: float | None = None
    ttft_p95_ms: float | None = None
    tpot_p50_ms: float | None = None  # time per output token, inter-token latency
    prefill_tps: float | None = None
    # Speculative decoding, when the engine reports it. Throughput alone cannot say whether
    # drafting is working: a draft rejected almost every time still emits tokens, having paid
    # for both passes. `None` means the engine exposed no counters, not that the rate was zero.
    spec_acceptance_rate: float | None = None
    spec_draft_tokens: int | None = None
    # True when output length was NOT pinned with ignore_eos. Comparability across points is
    # weaker in that mode, and a reader has to know which they are looking at.
    natural_stop: bool = False
    output_token_stdev: float | None = None
    error: str | None = None


@dataclass
class PerfResult:
    engine: str
    points: list[PerfPoint] = field(default_factory=list)
    peak_memory_gb: float | None = None
    idle_memory_gb: float | None = None
    notes: str | None = None


@dataclass
class Run:
    """A complete score card: one target, one point in time."""

    run_id: str
    timestamp: str
    target_id: str
    target: dict[str, Any]  # full snapshot -- see module docstring
    host: dict[str, Any]
    quality: list[QualityResult] = field(default_factory=list)
    perf: PerfResult | None = None
    verdict: dict[str, Any] = field(default_factory=dict)
    runner: dict[str, Any] = field(default_factory=dict)
    # How the model was sampled. Two cards are only comparable if this matches: a score
    # taken at temperature 1 is a draw from a distribution, and differs run to run by more
    # than most real regressions do. Recorded rather than assumed, because the quality half
    # sampled unpinned for its entire life and nothing in the card said so.
    sampling: dict[str, Any] = field(default_factory=dict)
    schema_version: int = SCHEMA_VERSION

    # ---------------------------------------------------------------- io --

    @property
    def filename(self) -> str:
        date = self.timestamp[:10]
        slug = self.target.get("slug") or self.target_id.replace("@", "__")
        return f"{date}__{slug}__{self.run_id[:8]}.json"

    def to_dict(self) -> dict[str, Any]:
        return _json_safe(asdict(self))

    def save(self, results_dir: Path | None = None) -> Path:
        d = results_dir or RESULTS_DIR
        d.mkdir(parents=True, exist_ok=True)
        path = d / self.filename
        # allow_nan=False so this raises instead of emitting a token that is not JSON.
        # to_dict has already mapped the non-finite values to null; if a new one appears,
        # the write should fail loudly rather than commit an unparseable score card.
        path.write_text(
            json.dumps(self.to_dict(), indent=2, sort_keys=False, allow_nan=False) + "\n"
        )
        return path

    @staticmethod
    def load(path: Path) -> dict[str, Any]:
        return json.loads(Path(path).read_text())


def runner_info() -> dict[str, Any]:
    """Provenance of the machine that drove the benchmark (not the one serving)."""
    return {
        "host": platform.node(),
        "python": platform.python_version(),
        "platform": platform.platform(),
    }


def truncate(text: str | None, limit: int = EXCERPT_CHARS) -> str | None:
    if text is None:
        return None
    text = text.strip()
    return text if len(text) <= limit else text[:limit] + f"... [+{len(text) - limit} chars]"


def compute_verdict(quality: list[QualityResult], perf: PerfResult | None) -> dict[str, Any]:
    """Derive a self-hosting verdict from the raw numbers.

    Deliberately simple and explicit rather than clever. The thresholds encode
    what makes a model usable in our own stack, and they live in code (and
    therefore in git history) so that a changed verdict is always traceable to
    a changed rule rather than to drift.
    """
    verdict: dict[str, Any] = {}

    signal = [q for q in quality if q.tier == "signal" and q.error is None]
    reference = [q for q in quality if q.tier == "reference" and q.error is None]

    def mean_accuracy(results: list[QualityResult]) -> float | None:
        vals = [q.metrics.get("accuracy") for q in results]
        vals = [v for v in vals if v is not None]
        return round(sum(vals) / len(vals), 4) if vals else None

    verdict["signal_score"] = mean_accuracy(signal)
    verdict["reference_score"] = mean_accuracy(reference)

    # Contamination gate. A probe nobody reads is the same as no probe, so this does not just
    # record a number — when a probe fires it takes `signal_score` away.
    #
    # Nulling it is deliberate. Anything downstream that reads `verdict["signal_score"]` —
    # `kbench results`, a comparison, a published card — would otherwise happily render a
    # memorised score next to a small warning, and the number is what people quote. The raw
    # value is preserved under a name nobody reads by accident, so nothing is lost for
    # forensics and nothing is quotable by mistake.
    canary = [q for q in quality if q.tier == "canary" and q.error is None]
    canary_score = mean_accuracy(canary)
    verdict["canary_score"] = canary_score
    if canary_score is None:
        # No probe ran. The signal score stands, but it is unverified, and saying so is the
        # difference between "clean" and "nobody checked".
        verdict["contamination"] = "unverified"
    elif canary_score > 0:
        verdict["contamination"] = "detected"
        verdict["signal_score_unverified"] = verdict["signal_score"]
        verdict["signal_score"] = None
    else:
        verdict["contamination"] = "clean"

    if perf and perf.points:
        single = next((p for p in perf.points if p.concurrency == 1), None)
        best = max(perf.points, key=lambda p: p.output_tps_total)
        if single:
            verdict["single_stream_tps"] = round(single.output_tps_per_stream, 1)
            # Below ~15 tok/s a model reads slower than a person, which rules
            # out anything interactive regardless of how good its answers are.
            verdict["interactive_viable"] = single.output_tps_per_stream >= 15.0
        verdict["peak_throughput_tps"] = round(best.output_tps_total, 1)
        verdict["peak_throughput_concurrency"] = best.concurrency

    return verdict