Lumbridge Bench

kbench/registry/__init__.py

@@ -0,0 +1,167 @@
+"""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"]