alvis
b5aaceb65a
Per DESIGN-a2a-agents.md v2.1 §2-3b: models are the scarce queued resource, version-controlled here rather than hardcoded in callers. - model-registry.yaml: kimi (main reasoning, quota-gated), local-small (ollama/gemma3:4b, always-on cheap tier), bge-m3 (embedder + routing classifier, never-evict), tei-reranker (never-evict, interactive- critical), paid-fallback (metered, opt-in only, unreachable by default via empty routing.metered_opt_in). GPU residency policy carries the never-evict set, co-residency groups, and measured baseline (bge-m3+gemma3:4b+tei-reranker ~6.2/8GB on the GTX 1070). - model_registry.py: resolve(tier) picks an available model without the caller naming one, gated so a metered model is only reachable with both allow_metered=True and an opted-in virtual key; to_probe_config() bridges registry quota data into kb_worker.py's existing Probe classes (no duplicated probe logic); preload_check() expresses the §3b pre-load VRAM check purely from registry data. Gap noted for follow-up: bge-m3 has no litellm-config.yaml model_list entry yet (embedder there still points at ollama/nomic-embed-text on a different port) — out of scope here, registry documents it as-is.
302 lines
13 KiB
Python
Executable File
302 lines
13 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""model_registry — reads model-registry.yaml (kb#133, A2A-1).
|
|
|
|
Per DESIGN-a2a-agents.md v2.1 §2-3b: the model registry is data, not logic.
|
|
a(t) probe *mechanics* (QuotaProbe, GPUResidencyProbe, ...) already live in
|
|
kanboard/bin/kb_worker.py — this module does not reimplement them. It gives
|
|
callers two things instead:
|
|
|
|
* resolve(tier) — "an available model for tier X" without the
|
|
caller naming a model. Structural availability
|
|
(lifecycle, metered opt-in) is decided here from
|
|
registry data; live a(t) truth (is the quota
|
|
window open right now, is the GPU actually free)
|
|
is decided by an optional `probe_check` callback
|
|
the caller supplies (e.g. wired to kb_worker's
|
|
Probe classes via to_probe_config()).
|
|
* preload_check(...) — the §3b GPU pre-load check, expressed purely from
|
|
registry numbers (never-evict reservations +
|
|
candidate footprint) plus a headroom figure the
|
|
caller supplies. It does not shell nvidia-smi
|
|
itself — kb_worker.GPUResidencyProbe (or
|
|
`nvidia-smi` directly) is the live-read path;
|
|
this stays pure/testable.
|
|
|
|
Usage (library):
|
|
from model_registry import load_registry, resolve, to_probe_config, preload_check
|
|
reg = load_registry()
|
|
model = resolve(reg, tier="large") # -> the "kimi" entry
|
|
cfg = to_probe_config(reg, "kimi") # -> kb_worker probe config dict
|
|
ok, reason = preload_check(reg, "local-small", headroom_mb=1900)
|
|
|
|
Usage (CLI, for manual verification):
|
|
./model_registry.py resolve --tier large
|
|
./model_registry.py resolve --tier large --allow-metered --opted-in agent:torgash
|
|
./model_registry.py probe-config --id kimi
|
|
./model_registry.py preload-check --id local-small --headroom-mb 1900
|
|
./model_registry.py preload-check --id local-small --live # shells nvidia-smi
|
|
./model_registry.py list
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
|
|
import yaml
|
|
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
DEFAULT_REGISTRY_PATH = os.path.join(HERE, "model-registry.yaml")
|
|
|
|
|
|
class RegistryError(Exception):
|
|
pass
|
|
|
|
|
|
def load_registry(path=None):
|
|
"""Load and lightly validate model-registry.yaml."""
|
|
path = path or DEFAULT_REGISTRY_PATH
|
|
with open(path) as f:
|
|
reg = yaml.safe_load(f)
|
|
if not reg or "models" not in reg:
|
|
raise RegistryError(f"{path}: missing top-level 'models' list")
|
|
ids = [m["id"] for m in reg["models"]]
|
|
if len(ids) != len(set(ids)):
|
|
raise RegistryError(f"{path}: duplicate model ids in {ids}")
|
|
return reg
|
|
|
|
|
|
def get_model(registry, model_id):
|
|
for m in registry["models"]:
|
|
if m["id"] == model_id:
|
|
return m
|
|
raise RegistryError(f"unknown model id: {model_id!r}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# resolve — "an available model for tier X" without the caller naming one.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def resolve(registry, tier, allow_metered=False, opted_in_key=None, probe_check=None):
|
|
"""Return the first model in `tier`'s pool that is structurally usable,
|
|
and (if probe_check is given) currently available.
|
|
|
|
Structural filter (from registry data alone):
|
|
- candidate must be listed under routing.tiers[tier]
|
|
- a metered model is only a candidate at all when the CALLER passes
|
|
allow_metered=True AND opted_in_key appears in routing.metered_opt_in
|
|
(§3a: "no metered API by default" — an empty metered_opt_in list, the
|
|
shipped default, makes every metered model structurally unreachable
|
|
regardless of allow_metered).
|
|
|
|
Live filter (optional): probe_check(model_dict) -> bool. Wire this to
|
|
kb_worker's Probe.available() (via to_probe_config below) when the
|
|
caller wants real a(t) truth instead of just structural eligibility.
|
|
"""
|
|
pools = registry.get("routing", {}).get("tiers", {})
|
|
if tier not in pools:
|
|
raise RegistryError(f"unknown tier: {tier!r} (have: {sorted(pools)})")
|
|
opted_in = set(registry.get("routing", {}).get("metered_opt_in", []) or [])
|
|
|
|
candidates = []
|
|
for model_id in pools[tier]:
|
|
m = get_model(registry, model_id)
|
|
if m.get("metered"):
|
|
if not allow_metered:
|
|
continue
|
|
if not m.get("opt_in_required", True):
|
|
# Registry says this metered model doesn't need opt-in — treat
|
|
# as a data error rather than silently routing to it.
|
|
raise RegistryError(
|
|
f"model {model_id!r} is metered but opt_in_required=false; "
|
|
"fix the registry entry, this helper will not assume implicit access"
|
|
)
|
|
if opted_in_key is None or opted_in_key not in opted_in:
|
|
continue
|
|
candidates.append(m)
|
|
|
|
for m in candidates:
|
|
if probe_check is None or probe_check(m):
|
|
return m
|
|
|
|
raise RegistryError(
|
|
f"no available model for tier={tier!r} "
|
|
f"(allow_metered={allow_metered}, opted_in_key={opted_in_key!r}); "
|
|
f"checked candidates: {[m['id'] for m in candidates] or pools[tier]}"
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# to_probe_config — bridges registry quota data into kb_worker's probe cfg
|
|
# shape (kanboard/bin/kb_worker.py PROBE_BUILDERS), so probes read registry
|
|
# numbers instead of the registry re-implementing probe logic.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def to_probe_config(registry, model_id):
|
|
"""Return a dict matching kb_worker.py's `build_probe(cfg)` input shape
|
|
for `model_id`'s lifecycle. Raises if the model has no probe-relevant
|
|
lifecycle (e.g. cost-gated with no probe_command wired yet)."""
|
|
m = get_model(registry, model_id)
|
|
lifecycle = m["lifecycle"]
|
|
|
|
if lifecycle == "always-on":
|
|
return {"type": "always_on"}
|
|
|
|
if lifecycle == "quota-gated":
|
|
q = m.get("quota") or {}
|
|
windows = q.get("windows") or []
|
|
if not windows:
|
|
raise RegistryError(f"{model_id}: quota-gated but no quota.windows configured")
|
|
# kb_worker's QuotaProbe checks one field; the tightest (first-to-hit)
|
|
# window in practice is the short one — default to the first entry,
|
|
# callers needing multi-window gating build one probe per window.
|
|
window = windows[0]
|
|
return {
|
|
"type": "quota",
|
|
"command": q["probe_command"],
|
|
"field": window["field"],
|
|
"threshold_pct": q.get("threshold_pct", 95),
|
|
}
|
|
|
|
if lifecycle == "cost-gated":
|
|
q = m.get("quota") or {}
|
|
if not q.get("probe_command"):
|
|
raise RegistryError(
|
|
f"{model_id}: cost-gated but no budget probe wired yet "
|
|
"(opt-in path incomplete — see registry comment)"
|
|
)
|
|
return {
|
|
"type": "budget",
|
|
"command": q["probe_command"],
|
|
"field": q["field"],
|
|
"limit": q["limit"],
|
|
}
|
|
|
|
if lifecycle == "on-demand":
|
|
ep = (m.get("endpoints") or [{}])[0]
|
|
if not ep.get("health_url"):
|
|
raise RegistryError(f"{model_id}: on-demand but no endpoint.health_url configured")
|
|
return {"type": "on_demand", "url": ep["health_url"]}
|
|
|
|
raise RegistryError(f"{model_id}: unknown lifecycle {lifecycle!r}")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# preload_check — §3b GPU pre-load check, pure registry-data math. The live
|
|
# VRAM headroom READ is the caller's job (kb_worker.GPUResidencyProbe or
|
|
# nvidia-smi directly) — see live_headroom_mb() below for a thin convenience
|
|
# wrapper used only by this module's own CLI, not by the check itself.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def preload_check(registry, candidate_id, headroom_mb, resident_ids=None):
|
|
"""Would loading `candidate_id` fit, given `headroom_mb` free VRAM right
|
|
now (as reported by a live probe) and `resident_ids` already loaded?
|
|
|
|
Never-evict models are never subtracted from headroom by a caller in the
|
|
first place (their footprint is a standing reservation baked into any
|
|
correct live headroom read) — this function just re-asserts that policy
|
|
from registry data: if `candidate_id` itself is never-evict, it always
|
|
passes (it's not something a worker "loads speculatively" and might be
|
|
told to skip); anything else must fit in the reported headroom.
|
|
"""
|
|
policy = registry.get("gpu_residency_policy") or {}
|
|
m = get_model(registry, candidate_id)
|
|
gr = m.get("gpu_residency")
|
|
if not gr:
|
|
return True, f"{candidate_id} has no gpu_residency entry (not a GPU-resident model)"
|
|
|
|
if gr.get("never_evict"):
|
|
return True, f"{candidate_id} is in the never-evict set — always resident by policy"
|
|
|
|
required_mb = gr["vram_mb"]
|
|
never_evict_ids = set(policy.get("never_evict_ids", []))
|
|
resident_ids = set(resident_ids or [])
|
|
# Sanity: if the live headroom read already accounts for never-evict
|
|
# reservations (the expected contract — see kb_worker.GPUResidencyProbe's
|
|
# never_evict_reserved_mb param), this is just a straight comparison.
|
|
# If a caller passes raw total-minus-used instead, warn via the reason
|
|
# string rather than silently under/over-reserving.
|
|
reserved_hint = sum(
|
|
get_model(registry, mid)["gpu_residency"]["vram_mb"]
|
|
for mid in never_evict_ids
|
|
if mid not in resident_ids # already counted as "used" if resident_ids says so
|
|
)
|
|
ok = headroom_mb >= required_mb
|
|
reason = (
|
|
f"headroom={headroom_mb}MB required={required_mb}MB "
|
|
f"(never-evict reserve expected already netted out by the caller's probe: "
|
|
f"~{reserved_hint}MB across {sorted(never_evict_ids)})"
|
|
)
|
|
return ok, reason
|
|
|
|
|
|
def live_headroom_mb(never_evict_reserved_mb=0):
|
|
"""Convenience for manual CLI checks only — NOT used by preload_check()
|
|
itself. Shells nvidia-smi the same way kb_worker.GPUResidencyProbe does."""
|
|
out = subprocess.run(
|
|
["nvidia-smi", "--query-gpu=memory.used,memory.total", "--format=csv,noheader,nounits"],
|
|
capture_output=True, text=True, timeout=5, check=True,
|
|
).stdout.strip().splitlines()[0]
|
|
used, total = (int(x) for x in out.split(","))
|
|
return total - used - never_evict_reserved_mb
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# CLI — manual verification only, not part of the library contract.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
ap.add_argument("--registry", default=None, help="path to model-registry.yaml (default: sibling file)")
|
|
sub = ap.add_subparsers(dest="cmd", required=True)
|
|
|
|
p = sub.add_parser("resolve")
|
|
p.add_argument("--tier", required=True)
|
|
p.add_argument("--allow-metered", action="store_true")
|
|
p.add_argument("--opted-in", default=None, help="virtual key claimed to be opted in")
|
|
|
|
p = sub.add_parser("probe-config")
|
|
p.add_argument("--id", required=True)
|
|
|
|
p = sub.add_parser("preload-check")
|
|
p.add_argument("--id", required=True)
|
|
p.add_argument("--headroom-mb", type=float, default=None)
|
|
p.add_argument("--live", action="store_true", help="read live headroom via nvidia-smi instead of --headroom-mb")
|
|
p.add_argument("--resident", action="append", default=[], help="repeatable: id already resident")
|
|
|
|
sub.add_parser("list")
|
|
|
|
args = ap.parse_args()
|
|
reg = load_registry(args.registry)
|
|
|
|
try:
|
|
if args.cmd == "resolve":
|
|
m = resolve(reg, args.tier, allow_metered=args.allow_metered, opted_in_key=args.opted_in)
|
|
print(json.dumps(m, indent=2))
|
|
elif args.cmd == "probe-config":
|
|
print(json.dumps(to_probe_config(reg, args.id), indent=2))
|
|
elif args.cmd == "preload-check":
|
|
headroom = args.headroom_mb
|
|
if args.live:
|
|
policy = reg.get("gpu_residency_policy") or {}
|
|
reserved = sum(get_model(reg, mid)["gpu_residency"]["vram_mb"]
|
|
for mid in policy.get("never_evict_ids", []))
|
|
headroom = live_headroom_mb(never_evict_reserved_mb=reserved)
|
|
if headroom is None:
|
|
raise RegistryError("preload-check needs --headroom-mb or --live")
|
|
ok, reason = preload_check(reg, args.id, headroom, resident_ids=args.resident)
|
|
print(json.dumps({"ok": ok, "reason": reason}))
|
|
sys.exit(0 if ok else 1)
|
|
elif args.cmd == "list":
|
|
for m in reg["models"]:
|
|
print(f"{m['id']:16} tier={m['tier']:6} lifecycle={m['lifecycle']:13} "
|
|
f"metered={m['metered']} role={m['role']}")
|
|
except RegistryError as exc:
|
|
print(f"error: {exc}", file=sys.stderr)
|
|
sys.exit(2)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|