184 lines
6.6 KiB
Python
184 lines
6.6 KiB
Python
"""Registry: deklarative Slots/Adapter (YAML) plus Laufzeitzustand (Store)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
from typing import Any, Literal
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
from .store import Store
|
|
from .util import load_yaml
|
|
|
|
SlotStatus = Literal["ok", "degraded", "unavailable"]
|
|
|
|
|
|
class SlotSpec(BaseModel):
|
|
name: str
|
|
tier: int
|
|
device: Literal["gpu", "cpu", "external"] = "gpu"
|
|
engine: str = "vllm"
|
|
node: str | None = None
|
|
litellm_model: str
|
|
base_model: str
|
|
max_context: int
|
|
effective_context: int | None = None
|
|
supports_lora: bool = False
|
|
fallbacks: list[str] = Field(default_factory=list)
|
|
kind: Literal["chat", "embedding", "rerank"] = "chat"
|
|
|
|
|
|
class AdapterSpec(BaseModel):
|
|
name: str
|
|
base_model: str
|
|
cluster: str
|
|
version: int = 1
|
|
status: Literal["active", "candidate", "retired"] = "active"
|
|
|
|
|
|
class Slot(BaseModel):
|
|
"""Effektive Sicht: Spezifikation + Laufzeitzustand."""
|
|
|
|
spec: SlotSpec
|
|
status: SlotStatus = "ok"
|
|
reason: str | None = None
|
|
set_by: str | None = None
|
|
calibrated_context: int | None = None
|
|
base_model_override: str | None = None
|
|
canary_model: str | None = None
|
|
canary_share: float = 0.0
|
|
canary_cluster: str | None = None
|
|
|
|
@property
|
|
def name(self) -> str:
|
|
return self.spec.name
|
|
|
|
@property
|
|
def base_model(self) -> str:
|
|
return self.base_model_override or self.spec.base_model
|
|
|
|
@property
|
|
def effective_context(self) -> int:
|
|
for value in (self.calibrated_context, self.spec.effective_context):
|
|
if value:
|
|
return min(value, self.spec.max_context)
|
|
return self.spec.max_context
|
|
|
|
@property
|
|
def available(self) -> bool:
|
|
return self.status == "ok"
|
|
|
|
|
|
class Registry:
|
|
def __init__(self, slots: dict[str, SlotSpec], adapters: list[AdapterSpec], store: Store):
|
|
self._slots = slots
|
|
self._adapters = {a.name: a for a in adapters}
|
|
self.store = store
|
|
|
|
@classmethod
|
|
def from_file(cls, path: str | Path, store: Store) -> Registry:
|
|
data = load_yaml(path)
|
|
return cls.from_dict(data, store)
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: dict[str, Any], store: Store) -> Registry:
|
|
slots = {name: SlotSpec(name=name, **spec) for name, spec in (data.get("slots") or {}).items()}
|
|
adapters = [AdapterSpec(**a) for a in data.get("adapters") or []]
|
|
for spec in slots.values():
|
|
for fb in spec.fallbacks:
|
|
if fb not in slots:
|
|
raise ValueError(f"Slot {spec.name}: unbekannter Fallback {fb}")
|
|
return cls(slots, adapters, store)
|
|
|
|
# --- Slots ------------------------------------------------------------------------------
|
|
def slot_names(self) -> list[str]:
|
|
return list(self._slots)
|
|
|
|
def has_slot(self, name: str) -> bool:
|
|
return name in self._slots
|
|
|
|
def slot(self, name: str) -> Slot:
|
|
spec = self._slots[name]
|
|
state = self.store.get_slot_state(name) or {}
|
|
return Slot(
|
|
spec=spec,
|
|
status=state.get("status") or "ok",
|
|
reason=state.get("reason"),
|
|
set_by=state.get("set_by"),
|
|
calibrated_context=state.get("effective_context"),
|
|
base_model_override=state.get("base_model"),
|
|
canary_model=state.get("canary_model"),
|
|
canary_share=state.get("canary_share") or 0.0,
|
|
canary_cluster=state.get("canary_cluster"),
|
|
)
|
|
|
|
def slots(self) -> list[Slot]:
|
|
return [self.slot(n) for n in self._slots]
|
|
|
|
def strong_slot(self) -> Slot:
|
|
chat = [s for s in self.slots() if s.spec.kind == "chat" and s.spec.tier == 0]
|
|
if not chat:
|
|
raise LookupError("Kein Slot mit tier 0 konfiguriert")
|
|
available = [s for s in chat if s.available]
|
|
return (available or chat)[0]
|
|
|
|
def set_status(self, name: str, status: SlotStatus, reason: str = "", set_by: str = "human") -> None:
|
|
if name not in self._slots:
|
|
raise KeyError(name)
|
|
self.store.set_slot_state(name, status=status, reason=reason, set_by=set_by)
|
|
|
|
def set_calibrated_context(self, name: str, tokens: int | None) -> None:
|
|
self.store.set_slot_state(name, effective_context=tokens)
|
|
|
|
def set_base_model(self, name: str, base_model: str | None) -> None:
|
|
self.store.set_slot_state(name, base_model=base_model)
|
|
|
|
def set_canary(self, name: str, model: str | None, share: float = 0.0, cluster: str | None = None) -> None:
|
|
"""Leitet `share` des Stufe-1-Verkehrs (optional nur eines Clusters) auf ein Kandidatenmodell."""
|
|
self.store.set_slot_state(name, canary_model=model, canary_share=share if model else 0.0,
|
|
canary_cluster=cluster if model else None)
|
|
|
|
# --- Adapter ----------------------------------------------------------------------------
|
|
def adapters(self) -> list[AdapterSpec]:
|
|
merged = {n: a.model_copy() for n, a in self._adapters.items()}
|
|
for row in self.store.all_adapter_states():
|
|
if row["name"] in merged:
|
|
if row.get("status"):
|
|
merged[row["name"]].status = row["status"]
|
|
else:
|
|
merged[row["name"]] = AdapterSpec(
|
|
name=row["name"],
|
|
base_model=row["base_model"],
|
|
cluster=row["cluster"],
|
|
version=row.get("version") or 1,
|
|
status=row.get("status") or "candidate",
|
|
)
|
|
return list(merged.values())
|
|
|
|
def active_adapter(self, base_model: str, cluster: str) -> AdapterSpec | None:
|
|
found = [
|
|
a for a in self.adapters()
|
|
if a.base_model == base_model and a.cluster == cluster and a.status == "active"
|
|
]
|
|
return max(found, key=lambda a: a.version) if found else None
|
|
|
|
def set_adapter(self, adapter: AdapterSpec) -> None:
|
|
self.store.set_adapter_state(
|
|
adapter.name,
|
|
base_model=adapter.base_model,
|
|
cluster=adapter.cluster,
|
|
version=adapter.version,
|
|
status=adapter.status,
|
|
)
|
|
|
|
def export(self) -> dict[str, Any]:
|
|
"""Effektiver Stand als YAML-fähiges Dict (für Rückführung nach Git)."""
|
|
slots = {}
|
|
for s in self.slots():
|
|
data = s.spec.model_dump(exclude={"name"})
|
|
data["base_model"] = s.base_model
|
|
if s.calibrated_context:
|
|
data["effective_context"] = s.calibrated_context
|
|
slots[s.name] = data
|
|
return {"slots": slots, "adapters": [a.model_dump() for a in self.adapters()]}
|