ywkang
049e3d8bb3
Move benches/ -> src/kernbench/benches/ and src/kernbench/cli/probe.py -> src/kernbench/probes/probe.py. Each bench self-registers via @bench(name=..., description=...); kernbench list enumerates benches with auto-assigned indices, --bench accepts kebab-case name or numeric index. Audit at package-import time fails if any non-underscore module forgets the decorator. ADR-0010 (EN + KO) updated to reflect the new resolver path, list subcommand, and probes package separation. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
107 lines
3.3 KiB
Python
107 lines
3.3 KiB
Python
"""Bench registry: @bench decorator + name/index resolution.
|
|
|
|
Each bench module under ``kernbench.benches`` MUST register its callable
|
|
via ``@bench(name=..., description=...)``. Indices are assigned
|
|
alphabetically by name after eager import; they are a CLI convenience,
|
|
not a stable API.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass
|
|
from importlib import import_module
|
|
from pkgutil import iter_modules
|
|
|
|
BenchFn = Callable[..., object]
|
|
|
|
_NAME_RE = re.compile(r"^[a-z][a-z0-9]*(-[a-z0-9]+)*$")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BenchSpec:
|
|
index: int
|
|
name: str
|
|
description: str
|
|
run: BenchFn
|
|
|
|
|
|
_PENDING: list[tuple[str, str, BenchFn]] = []
|
|
_REGISTERED_MODULES: set[str] = set()
|
|
_REGISTRY: dict[str, BenchSpec] = {}
|
|
|
|
|
|
def bench(*, name: str, description: str) -> Callable[[BenchFn], BenchFn]:
|
|
if not isinstance(name, str) or not _NAME_RE.match(name):
|
|
raise ValueError(
|
|
f"bench name {name!r} must be kebab-case (lowercase, digits, dashes; "
|
|
f"starts with a letter)."
|
|
)
|
|
if not isinstance(description, str) or not description.strip():
|
|
raise ValueError(f"bench {name!r}: description must be a non-empty string.")
|
|
|
|
def deco(fn: BenchFn) -> BenchFn:
|
|
_PENDING.append((name, description, fn))
|
|
_REGISTERED_MODULES.add(fn.__module__)
|
|
return fn
|
|
|
|
return deco
|
|
|
|
|
|
def _finalize() -> None:
|
|
if _REGISTRY:
|
|
return
|
|
seen: set[str] = set()
|
|
for n, _, _ in _PENDING:
|
|
if n in seen:
|
|
raise RuntimeError(f"duplicate bench name: {n!r}")
|
|
seen.add(n)
|
|
for i, (n, d, f) in enumerate(sorted(_PENDING, key=lambda t: t[0]), start=1):
|
|
_REGISTRY[n] = BenchSpec(index=i, name=n, description=d, run=f)
|
|
|
|
|
|
def list_all() -> list[BenchSpec]:
|
|
_finalize()
|
|
return sorted(_REGISTRY.values(), key=lambda s: s.index)
|
|
|
|
|
|
def resolve(identifier: str) -> BenchSpec:
|
|
_finalize()
|
|
if not isinstance(identifier, str) or not identifier.strip():
|
|
raise ValueError("bench identifier must be a non-empty string.")
|
|
ident = identifier.strip()
|
|
if ident.isdigit():
|
|
idx = int(ident)
|
|
for s in _REGISTRY.values():
|
|
if s.index == idx:
|
|
return s
|
|
raise ValueError(
|
|
f"No bench with index {idx}. Use 'kernbench list' to see options."
|
|
)
|
|
if ident in _REGISTRY:
|
|
return _REGISTRY[ident]
|
|
raise ValueError(
|
|
f"Unknown bench {ident!r}. Use 'kernbench list' to see options."
|
|
)
|
|
|
|
|
|
def _audit_modules(imported: list[str], registered: set[str]) -> None:
|
|
missing = sorted(m for m in imported if m not in registered)
|
|
if missing:
|
|
raise RuntimeError(
|
|
f"Bench module(s) missing @bench decorator: {missing}. "
|
|
f"Each file under kernbench.benches/ must register at least one bench "
|
|
f"via @bench(...), or be renamed with a leading underscore if it is a "
|
|
f"helper."
|
|
)
|
|
|
|
|
|
def _eager_import_and_audit(pkg_path: list[str], pkg_name: str) -> None:
|
|
imported: list[str] = []
|
|
for m in iter_modules(pkg_path):
|
|
if m.name == "registry" or m.name.startswith("_"):
|
|
continue
|
|
mod = import_module(f"{pkg_name}.{m.name}")
|
|
imported.append(mod.__name__)
|
|
_audit_modules(imported, _REGISTERED_MODULES)
|