Teknium
388 lines
17 KiB
Python
388 lines
17 KiB
Python
"""Curated starter catalog for the managed local runtime.
|
|
|
|
Small and honest: every entry carries the estimator inputs (measured on real GGUFs) so the picker
|
|
can price a model BEFORE the user downloads gigabytes. Once a file is on disk, profile_from_gguf()
|
|
is the authority and the catalog numbers are only used for the download decision.
|
|
|
|
Validation lifecycle: builds proven end-to-end on real hardware are marked validated. Day-0 entries
|
|
ship before that proof (they simply lack the validated flag) — ensure_model_ready's touch generation
|
|
still gates every first load at runtime.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import re
|
|
import threading
|
|
import time
|
|
import urllib.request
|
|
from dataclasses import dataclass, field
|
|
from pathlib import PurePosixPath
|
|
|
|
from hermes_cli.local_runtime.context_policy import (
|
|
FLOOR,
|
|
RUNTIME_OVERHEAD_BYTES,
|
|
TARGET_WINDOW,
|
|
ub_logits_bytes,
|
|
)
|
|
from hermes_cli.local_runtime.estimator import (
|
|
HardwareBudget,
|
|
LayerKind,
|
|
ModelProfile,
|
|
ctx_bytes,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_PART_SUFFIX = re.compile(r"-\d{5}-of-\d{5}$")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AssetFile:
|
|
"""One downloadable file: repo-relative path and exact bytes (the size feeds the estimator and the
|
|
download progress bar; there is no download-time integrity check by design — a corrupt file
|
|
surfaces as a llama.cpp load error). ``local`` overrides the on-disk name (repos reuse generic
|
|
names like mmproj-BF16.gguf across models). Non-model extras live under the models dir's assets/
|
|
subdirectory so the router never lists them.
|
|
"""
|
|
|
|
path: str # repo-relative (may include a subdir)
|
|
size_bytes: int
|
|
local: str | None = None
|
|
|
|
@property
|
|
def local_name(self) -> str:
|
|
return self.local or PurePosixPath(self.path).name
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class QuantVariant:
|
|
"""One downloadable build of a model. Split GGUFs list every part in
|
|
files; the model loads from the first part."""
|
|
|
|
quant: str # e.g. "UD-Q4_K_M"
|
|
files: tuple # AssetFile, first = the load target
|
|
validated: bool = False # proven end-to-end on real hardware
|
|
|
|
@property
|
|
def model_id(self) -> str:
|
|
stem = PurePosixPath(self.files[0].path).name.removesuffix(".gguf")
|
|
return _PART_SUFFIX.sub("", stem)
|
|
|
|
@property
|
|
def size_bytes(self) -> int:
|
|
return sum(f.size_bytes for f in self.files)
|
|
|
|
@property
|
|
def weights_bytes(self) -> int:
|
|
"""Pre-download weights estimate: GGUF bytes ≈ tensor bytes + a small header (<2%) — a safe,
|
|
slightly conservative stand-in until profile_from_gguf reads the real table.
|
|
"""
|
|
return self.size_bytes
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class CatalogEntry:
|
|
id: str # stable family id (variant-independent)
|
|
display_name: str
|
|
description: str # one line, plain language
|
|
repo: str # HF repo
|
|
variants: tuple # QuantVariant (exactly one, Q4-class)
|
|
# Estimator inputs (measured or config-derived; quant changes weights,
|
|
# never KV). Entries with gated upstream configs carry a conservative
|
|
# same-family prior — the GGUF header is the authority after download.
|
|
n_ctx_train: int
|
|
full_layers: int
|
|
recurrent_layers: int
|
|
per_layer_f16: int # KV bytes/token per full-attention layer
|
|
swa_layers: int = 0
|
|
swa_window: int = 0
|
|
moe: bool = False
|
|
mtp: bool = False # ships MTP heads (spec decode when loaded)
|
|
# Speculative draft depth for MTP models. Per-model and measured:
|
|
# deeper drafting pays only while draft acceptance holds, and the
|
|
# break-even depth differs by model.
|
|
mtp_draft_depth: int = 3
|
|
# Vocab size prices the GPU logits buffers (ubatch x vocab x fp32,
|
|
# doubled under MTP backend sampling) — a multi-GiB term at large
|
|
# vocab sizes that a weights-only fit would miss.
|
|
n_vocab: int = 0
|
|
mmproj: "AssetFile | None" = None # vision projector, downloads with model
|
|
draft: "AssetFile | None" = None # spec-decode draft model (e.g. DSpark)
|
|
sampling: dict = field(default_factory=dict) # INI long-form launch defaults
|
|
# Oldest llama.cpp release tag that can load this model (day-0
|
|
# architectures need the release where their support landed). Empty
|
|
# means any installed engine. The pane gates download/activate on it.
|
|
min_engine: str = ""
|
|
# Editorial quality ordering (higher = smarter), authored once,
|
|
# globally, at catalog-authoring time — Artificial Analysis-informed
|
|
# where they cover the model (scripts/aa_quality_sync.py proposes,
|
|
# the commit decides), editorial elsewhere. Ranks entries for the
|
|
# per-machine recommendation; never displayed as a score (it grades
|
|
# the full-precision model, not our Q4 build).
|
|
quality: int = 0
|
|
# Fraction of the build's bytes read per decoded token: 1.0 for dense
|
|
# models (every weight streams every token), the active slice for MoE
|
|
# (attention + shared + routed experts over total). With memory
|
|
# bandwidth this predicts decode speed — the physics half of the
|
|
# recommendation.
|
|
decode_fraction: float = 1.0
|
|
|
|
def profile(self, variant: QuantVariant) -> ModelProfile:
|
|
layers = ([(LayerKind.FULL, self.per_layer_f16)] * self.full_layers
|
|
+ [(LayerKind.SWA, self.per_layer_f16)] * self.swa_layers
|
|
+ [(LayerKind.RECURRENT, 0)] * self.recurrent_layers)
|
|
return ModelProfile(
|
|
name=variant.model_id, weights_bytes=variant.weights_bytes,
|
|
embd_table_bytes=0, n_ctx_train=self.n_ctx_train,
|
|
layers=layers, swa_window=self.swa_window, moe=self.moe,
|
|
n_vocab=self.n_vocab,
|
|
kv_scale=1.2 if self.mtp else 1.0)
|
|
|
|
def download_files(self, variant: QuantVariant) -> tuple:
|
|
"""Everything a download job fetches for this variant, in order."""
|
|
extras = tuple(a for a in (self.mmproj, self.draft) if a is not None)
|
|
return tuple(variant.files) + extras
|
|
|
|
def download_bytes(self, variant: QuantVariant) -> int:
|
|
return sum(f.size_bytes for f in self.download_files(variant))
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class VariantChoice:
|
|
"""Selection result: which build this machine should download and why.
|
|
reason_key is a UI-copy discriminator, not display text."""
|
|
|
|
variant: QuantVariant
|
|
zero_spill: bool
|
|
reason_key: str # "best-large-window" | "best-fits" | "smallest-fits-spilled"
|
|
|
|
|
|
def select_variant(entry: CatalogEntry, budget: HardwareBudget) -> VariantChoice | None:
|
|
"""Fit the entry's one build (Q4-class) to this machine.
|
|
|
|
Every entry ships exactly one variant (see the module docstring for why there is no quant
|
|
ladder); headroom buys a bigger window, never a bigger quant. The fit shapes:
|
|
|
|
- "best-large-window": zero-spills at TARGET_WINDOW - "best-fits": zero-spills at the 64K floor
|
|
- "smallest-fits-spilled": weights spill to host RAM, priced honestly - None: even spilled,
|
|
physics refuses (the machine can't run it)
|
|
"""
|
|
overhead = (RUNTIME_OVERHEAD_BYTES
|
|
+ (entry.mmproj.size_bytes if entry.mmproj else 0)
|
|
+ ub_logits_bytes(entry.n_vocab, mtp_capable=entry.mtp))
|
|
native = entry.n_ctx_train or FLOOR
|
|
variant = entry.variants[-1]
|
|
profile = entry.profile(variant)
|
|
need = variant.weights_bytes + overhead
|
|
vram = budget.usable_vram_bytes
|
|
if need + ctx_bytes(profile, min(TARGET_WINDOW, native)) <= vram:
|
|
return VariantChoice(variant, zero_spill=True, reason_key="best-large-window")
|
|
floor_kv = ctx_bytes(profile, min(FLOOR, native))
|
|
if need + floor_kv <= vram:
|
|
return VariantChoice(variant, zero_spill=True, reason_key="best-fits")
|
|
if need + floor_kv <= vram + budget.ram_available_bytes:
|
|
return VariantChoice(variant, zero_spill=False, reason_key="smallest-fits-spilled")
|
|
return None
|
|
|
|
|
|
# ── recommendation: best quality that fits and isn't miserably slow ──
|
|
#
|
|
# Two axes, each living where it belongs. QUALITY is a judgment made once,
|
|
# globally, at authoring time (entry.quality — AA-informed, editorially
|
|
# owned). SPEED is physics computed per machine: decode is memory-bound,
|
|
# so predicted tok/s ≈ bandwidth / bytes-read-per-token, and the bytes per
|
|
# token are the build's size scaled by its decode fraction (dense reads
|
|
# everything; MoE reads the active slice). The pick: highest quality among
|
|
# entries that run resident and clear a pleasant speed floor; else the
|
|
# fastest resident entry; else the least-painful spilled one.
|
|
#
|
|
# The bandwidth axis is the `uma` flag for now: every discrete card that
|
|
# matters is 900+ GB/s GDDR while the unified-memory class measures ~1/5th
|
|
# of that, so the flag IS the high/low split. A measured per-machine
|
|
# bandwidth (one cached memcpy probe) can replace these class constants
|
|
# without touching the rule; predictions order candidates and gate the
|
|
# floor — they are not display values.
|
|
|
|
_DISCRETE_BANDWIDTH_GB_S = 1000.0 # representative GDDR6X/GDDR7 class
|
|
_UMA_BANDWIDTH_GB_S = 210.0 # measured on unified-memory NVIDIA
|
|
_HOST_BANDWIDTH_GB_S = 80.0 # spilled weights stream over host DRAM
|
|
|
|
# The one editorial constant in the tree: below this predicted decode
|
|
# speed a model stops feeling pleasant for agentic use (roughly reading
|
|
# speed with headroom for tool-call bursts). Distinct from the growth
|
|
# policy's 6 tok/s compress floor, which marks unusable, not unpleasant.
|
|
PLEASANT_FLOOR_TOK_S = 20.0
|
|
|
|
|
|
def predicted_decode_tok_s(entry: CatalogEntry, variant: QuantVariant,
|
|
budget: HardwareBudget, *,
|
|
spilled: bool = False) -> float:
|
|
"""Memory-bound decode prediction for ordering and floor-gating."""
|
|
bandwidth = (_HOST_BANDWIDTH_GB_S if spilled
|
|
else _UMA_BANDWIDTH_GB_S if budget.uma
|
|
else _DISCRETE_BANDWIDTH_GB_S)
|
|
bytes_per_token = max(1.0, variant.size_bytes * entry.decode_fraction)
|
|
return bandwidth * 1e9 / bytes_per_token
|
|
|
|
|
|
def recommended_entry(budget: HardwareBudget,
|
|
entries: "tuple[CatalogEntry, ...] | None" = None
|
|
) -> "tuple[CatalogEntry, str] | None":
|
|
"""The catalog's default pick for THIS machine, with its reason.
|
|
|
|
Callers pass pre-filtered entries when some are ineligible for reasons the catalog can't know
|
|
(engine too old); default is the full catalog.
|
|
|
|
best-quality-resident quality won among resident entries that clear the pleasant floor speed-
|
|
gated-quality same, but the floor eliminated a HIGHER quality candidate — the exact 'why not the
|
|
big model?' a unified-memory owner asks fastest-resident nothing resident clears the floor; the
|
|
quickest resident entry wins least-painful-spilled nothing runs resident; fastest from host
|
|
memory (MoE by construction)
|
|
"""
|
|
pool = CATALOG if entries is None else entries
|
|
fitting = [(e, c) for e in pool if (c := select_variant(e, budget)) is not None]
|
|
if not fitting:
|
|
return None
|
|
|
|
def speed(t, spilled=False):
|
|
return predicted_decode_tok_s(t[0], t[1].variant, budget, spilled=spilled)
|
|
|
|
resident = [(e, c) for e, c in fitting if c.zero_spill]
|
|
pleasant = [t for t in resident if speed(t) >= PLEASANT_FLOOR_TOK_S]
|
|
if pleasant:
|
|
pick = max(pleasant, key=lambda t: (t[0].quality, -t[1].variant.size_bytes))[0]
|
|
floor_gated = any(e.quality > pick.quality for e, _ in resident)
|
|
return (pick, "speed-gated-quality" if floor_gated else "best-quality-resident")
|
|
if resident:
|
|
return (max(resident, key=speed)[0], "fastest-resident")
|
|
# Everything spills: take the least painful — fastest predicted decode
|
|
# from host memory (MoE wins here by construction; a dense spill
|
|
# streams every weight over the host bus).
|
|
return (max(fitting, key=lambda t: speed(t, spilled=True))[0], "least-painful-spilled")
|
|
|
|
|
|
# ── catalog data: packaged JSON, refreshed from GitHub in memory ─
|
|
#
|
|
# The catalog DATA lives in catalog.json (checked in beside this module
|
|
# and shipped as package data); this module keeps all policy. At import
|
|
# we load the packaged copy — no network on the import path. A TTL-gated
|
|
# background refresh fetches the same file from the repo's main branch
|
|
# and swaps it in memory only: nothing on disk changes, so a git
|
|
# checkout never sees a dirty tracked file and the packaged copy remains
|
|
# the offline truth. A reverted commit on main heals every install on
|
|
# its next fetch, and day-0 entries reach users without an app release.
|
|
|
|
_CATALOG_URL = ("https://raw.githubusercontent.com/NousResearch/hermes-agent"
|
|
"/main/hermes_cli/local_runtime/catalog.json")
|
|
_SCHEMA_VERSION = 1
|
|
_REFRESH_TTL_S = 6 * 3600
|
|
_refresh_lock = threading.Lock()
|
|
_last_refresh_attempt = 0.0
|
|
|
|
|
|
def _asset_from(d: "dict | None") -> "AssetFile | None":
|
|
if not d:
|
|
return None
|
|
return AssetFile(path=d["path"], size_bytes=int(d["size_bytes"]),
|
|
local=d.get("local"))
|
|
|
|
|
|
# Scalar CatalogEntry fields parsed from JSON: key -> (coerce, default);
|
|
# a None default means the key is required.
|
|
_SCALAR_FIELDS = {
|
|
"n_ctx_train": (int, None), "full_layers": (int, None),
|
|
"recurrent_layers": (int, None), "per_layer_f16": (int, None),
|
|
"swa_layers": (int, 0), "swa_window": (int, 0),
|
|
"moe": (bool, False), "mtp": (bool, False), "mtp_draft_depth": (int, 3),
|
|
"n_vocab": (int, 0), "sampling": (dict, {}), "min_engine": (str, ""),
|
|
"quality": (int, 0), "decode_fraction": (float, 1.0),
|
|
}
|
|
|
|
|
|
def _load_catalog(doc: dict) -> "tuple[CatalogEntry, ...]":
|
|
"""Parse a catalog document into entries. Unknown fields are ignored
|
|
(newer catalogs stay readable by older apps); a major schema bump is
|
|
the signal that they wouldn't be, and the caller skips the document."""
|
|
if int(doc.get("schema_version", 0)) != _SCHEMA_VERSION:
|
|
raise ValueError(f"catalog schema {doc.get('schema_version')!r} "
|
|
f"(this build reads {_SCHEMA_VERSION})")
|
|
entries = []
|
|
for m in doc["models"]:
|
|
variants = tuple(
|
|
QuantVariant(quant=v["quant"],
|
|
files=tuple(_asset_from(f) for f in v["files"]),
|
|
validated=bool(v.get("validated")))
|
|
for v in m["variants"])
|
|
scalars = {k: coerce(m[k] if default is None else m.get(k, default))
|
|
for k, (coerce, default) in _SCALAR_FIELDS.items()}
|
|
entries.append(CatalogEntry(
|
|
id=m["id"], display_name=m["display_name"],
|
|
description=m["description"], repo=m["repo"], variants=variants,
|
|
mmproj=_asset_from(m.get("mmproj")), draft=_asset_from(m.get("draft")),
|
|
**scalars))
|
|
return tuple(entries)
|
|
|
|
|
|
def _packaged_catalog() -> "tuple[CatalogEntry, ...]":
|
|
from importlib.resources import files
|
|
|
|
raw = files("hermes_cli.local_runtime").joinpath("catalog.json").read_text(
|
|
encoding="utf-8")
|
|
return _load_catalog(json.loads(raw))
|
|
|
|
|
|
CATALOG: "tuple[CatalogEntry, ...]" = _packaged_catalog()
|
|
|
|
|
|
def refresh_catalog(force: bool = False) -> bool:
|
|
"""Fetch the current catalog from the repo and swap it in memory.
|
|
|
|
Best-effort by design: any failure (offline, GitHub down, unreadable schema) leaves the running
|
|
catalog untouched and retries after the TTL. Returns True when a fetched document replaced the
|
|
catalog.
|
|
"""
|
|
global CATALOG, _last_refresh_attempt
|
|
|
|
now = time.monotonic()
|
|
with _refresh_lock:
|
|
if not force and now - _last_refresh_attempt < _REFRESH_TTL_S:
|
|
return False
|
|
_last_refresh_attempt = now
|
|
try:
|
|
req = urllib.request.Request(
|
|
_CATALOG_URL, headers={"User-Agent": "hermes-local-runtime"})
|
|
with urllib.request.urlopen(req, timeout=10) as r:
|
|
fetched = _load_catalog(json.load(r))
|
|
except Exception as exc: # noqa: BLE001
|
|
logger.debug("catalog refresh skipped: %s", exc)
|
|
return False
|
|
if fetched != CATALOG:
|
|
logger.info("catalog refreshed from repo (%d models)", len(fetched))
|
|
CATALOG = fetched
|
|
return True
|
|
|
|
|
|
def refresh_catalog_soon() -> None:
|
|
"""TTL-gated background refresh; returns immediately. The caller's current request serves the
|
|
catalog it already has — the refresh lands for the next one.
|
|
"""
|
|
if time.monotonic() - _last_refresh_attempt < _REFRESH_TTL_S:
|
|
return
|
|
threading.Thread(target=refresh_catalog, daemon=True,
|
|
name="catalog-refresh").start()
|
|
|
|
|
|
def catalog_by_id() -> dict[str, CatalogEntry]:
|
|
return {entry.id: entry for entry in CATALOG}
|
|
|
|
|
|
def find_entry_for_model(model_id: str) -> "tuple[CatalogEntry, QuantVariant] | None":
|
|
"""Locate the entry + variant that owns a staged model id."""
|
|
for entry in CATALOG:
|
|
for variant in entry.variants:
|
|
if variant.model_id == model_id:
|
|
return entry, variant
|
|
return None
|