bench/kbench/registry/__init__.py

"""Model registry.

A *target* is not a model. It is a specific way of serving a specific model:

    (model × host × quantization × serving config × checkpoint)

This is the central design decision of the bench and it cannot be retrofitted
later without invalidating every result already recorded. On unified-memory
boxes like the DGX Spark, serving flags move throughput more than a model swap
does -- `--enforce-eager` alone can cost 30%+, and `--gpu-memory-utilization`
determines how much KV cache exists, which sets the concurrency ceiling. A
registry keyed on model name alone produces numbers nobody can reproduce.

`checkpoint` is first-class for the same reason: once we fine-tune our own
models, we will compare `karti-7b@step2000` against `karti-7b@step4000` far
more often than we compare Qwen against Llama.
"""

from __future__ import annotations

import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import yaml

REGISTRY_PATH = Path(__file__).parent / "models.yaml"


@dataclass(frozen=True)
class Host:
    """A machine that can serve models."""

    id: str
    ssh: str | None = None
    hardware: str | None = None
    memory_gb: int | None = None
    memory_bandwidth_gbs: int | None = None
    notes: str | None = None


@dataclass(frozen=True)
class Serving:
    """How a target is served. Part of the identity of a result."""

    engine: str  # vllm | sglang | llama.cpp | anthropic | openai | ...
    model_name: str  # value sent as "model" in the API request
    base_url: str | None = None
    service: str | None = None  # inspect openai-api service prefix
    max_model_len: int | None = None
    flags: dict[str, Any] = field(default_factory=dict)
    # How to sample this model during quality runs. Part of the identity of a result: the
    # same model at temperature 1 and at temperature 0.6 is two different measurements.
    #
    # Not a global constant, because the right value is a property of the model. Greedy
    # decoding looks like the obvious choice for reproducibility and is actively wrong for
    # reasoning models -- Qwen3 in thinking mode degenerates into repetition at
    # temperature 0, and a 2-sample probe that took two minutes ran past twelve without
    # terminating. Reproducibility comes from pinning the seed, not from killing entropy.
    sampling: dict[str, Any] = field(default_factory=dict)

    @property
    def is_local(self) -> bool:
        return self.engine in {"vllm", "sglang", "llama.cpp", "ollama"}


@dataclass(frozen=True)
class Target:
    """One benchmarkable configuration."""

    id: str
    display_name: str
    serving: Serving
    host: str
    family: str | None = None
    params_b: float | None = None
    active_params_b: float | None = None  # MoE: params active per token
    quant: str | None = None
    checkpoint: str | None = None
    tier: str = "candidate"  # production | candidate | reference | baseline
    aliases: list[str] = field(default_factory=list)
    notes: str | None = None

    @property
    def inspect_model(self) -> str:
        """The model string Inspect uses to address this target."""
        s = self.serving
        if s.engine in {"anthropic", "openai", "google", "openrouter", "grok"}:
            return f"{s.engine}/{s.model_name}"
        if s.is_local:
            if not s.service:
                raise ValueError(f"target {self.id}: local serving requires a 'service'")
            # Inspect convention: openai-api/<service>/<model>, with base_url and
            # api key read from <SERVICE>_BASE_URL / <SERVICE>_API_KEY.
            return f"openai-api/{s.service}/{s.model_name}"
        raise ValueError(f"target {self.id}: unsupported engine {s.engine!r}")

    @property
    def env(self) -> dict[str, str]:
        """Environment variables Inspect needs to reach this target."""
        s = self.serving
        if not s.is_local or not s.service:
            return {}
        prefix = s.service.upper().replace("-", "_")
        env = {f"{prefix}_API_KEY": os.environ.get(f"{prefix}_API_KEY", "no-key-required")}
        if s.base_url:
            env[f"{prefix}_BASE_URL"] = s.base_url
        return env

    def apply_env(self) -> None:
        """Export this target's env vars into the current process."""
        os.environ.update(self.env)

    @property
    def slug(self) -> str:
        """Filesystem-safe identifier used in results filenames."""
        return self.id.replace("/", "_").replace("@", "__")


@dataclass(frozen=True)
class Registry:
    hosts: dict[str, Host]
    targets: dict[str, Target]

    def target(self, ident: str) -> Target:
        """Look up a target by id or alias."""
        if ident in self.targets:
            return self.targets[ident]
        for t in self.targets.values():
            if ident in t.aliases:
                return t
        known = ", ".join(sorted(self.targets))
        raise KeyError(f"unknown target {ident!r}. known targets: {known}")

    def host(self, ident: str) -> Host:
        if ident not in self.hosts:
            raise KeyError(f"unknown host {ident!r}")
        return self.hosts[ident]

    def by_tier(self, tier: str) -> list[Target]:
        return [t for t in self.targets.values() if t.tier == tier]


def load_registry(path: Path | None = None) -> Registry:
    raw = yaml.safe_load((path or REGISTRY_PATH).read_text())

    hosts = {
        hid: Host(id=hid, **(spec or {}))
        for hid, spec in (raw.get("hosts") or {}).items()
    }

    targets: dict[str, Target] = {}
    for spec in raw.get("targets") or []:
        spec = dict(spec)
        serving = Serving(**spec.pop("serving"))
        target = Target(serving=serving, **spec)
        if target.host not in hosts:
            raise ValueError(f"target {target.id}: unknown host {target.host!r}")
        if target.id in targets:
            raise ValueError(f"duplicate target id {target.id!r}")
        targets[target.id] = target

    return Registry(hosts=hosts, targets=targets)


__all__ = ["Host", "Registry", "Serving", "Target", "load_registry"]