mission-control-v2/backend/services/budget.py

"""
Speicher-Budget & SETUP-BEWUSSTE ctx-Vergabe — EINE Quelle der Wahrheit.

Modelliert die auf der Box VERIFIZIERTE Residenz-Realität (llama-swap, Ein-Gruppen-
Residenz, GTT ~124 GB):
  • Das Agent-Hirn (Rolle `hermes`) ist IMMER resident.
  • Weitere persist-Mitglieder (fast/vision) dürfen verdrängt werden, wenn ein großes
    on-demand-Modell lädt.
Daraus folgt, wie viel Speicher NEBEN einem Zielmodell reserviert bleiben muss —
und damit der größte Kontext, der wirklich passt (nicht nur für das Modell allein).

Vorher rechnete nur der Hirn-Wechsel (agent.py) setup-bewusst; die allgemeine
ctx-Vergabe nahm den Gesamt-RAM in Isolation. Dieses Modul vereint beides.
"""

import re

import psutil

from services.fit import (
    QUANT_BYTES_PER_PARAM,
    estimate_memory_gb,
    extract_params_b,
    max_ctx_in_budget,
)

HEADROOM_GB = 4.0   # OS/Treiber/Fragmentierung


def gtt_budget_gb() -> float:
    """GPU-adressierbarer Speicher (GTT) in GB — die harte Obergrenze. Liest
    amdgpu.gttsize aus /proc/cmdline, sonst RAM minus OS-Reserve."""
    try:
        with open("/proc/cmdline") as f:
            m = re.search(r"amdgpu\.gttsize=(\d+)", f.read())
            if m:
                return round(int(m.group(1)) / 1024.0, 1)
    except Exception:
        pass
    return round(psutil.virtual_memory().total / (1024 ** 3) - 6.0, 1)


def params_of_model(model: dict) -> float:
    """Robuste Params (Mrd.) eines INSTALLIERTEN Modells: MAXIMUM aus Caps-Schätzung und
    Dateigröße. Deckt 'Coder-Next' ohne Größe im Namen (→ aus Datei) und Split-GGUFs
    (size_bytes = nur erster Teil → ignoriert) ab."""
    caps = model.get("capabilities") or {}
    quant = model.get("quant") or "Q4_K_M"
    bpp = QUANT_BYTES_PER_PARAM.get(quant.upper(), 0.55)
    size_gb = (model.get("size_bytes") or 0) / (1024 ** 3)
    pb_size = (size_gb / bpp) if size_gb > 1.0 else 0.0
    return max(float(caps.get("params_b") or 0), pb_size, 7.0)


def footprint_gb(model: dict) -> float:
    """Loaded-Footprint eines Modells = Gewichte + kalibrierter KV-Anteil (bei seinem
    aktuellen ctx)."""
    quant = model.get("quant") or "Q4_K_M"
    ctx = int(model.get("ctx") or 32768)
    bpp = QUANT_BYTES_PER_PARAM.get(quant.upper(), 0.55)
    size_gb = (model.get("size_bytes") or 0) / (1024 ** 3)
    pb = params_of_model(model)
    weights = max(pb * bpp, size_gb)
    kv = estimate_memory_gb(pb, quant, ctx) - pb * bpp
    return weights + max(kv, 0.0)


def params_b_for(name: str) -> float:
    """Parameter (Mrd.) für einen Modell-/Repo-Namen: KATALOG (echte Metadaten) zuerst,
    sonst Namens-Schätzung. Gemeinsam für Fit-Vorschau und ctx-Vergabe."""
    from services import catalog
    meta = catalog.meta_for_name(name) if name else None
    if meta and meta.get("total_params_b"):
        return float(meta["total_params_b"])
    return extract_params_b(name)


def _coresident_members(groups: dict) -> set:
    """Modelle, die GLEICHZEITIG warm sind: Mitglieder aller `swap:false`-Gruppen
    (Ko-Residenz, z.B. brains = Hirn+embed+vision). Ein Modell AUSSERHALB dieser Gruppen
    ist on-demand und verdrängt beim Laden die GANZE Gruppe — llama-swap swappt Gruppen
    (live verifiziert: heavy laden → brains-Gruppe komplett raus, heavy läuft allein)."""
    out: set = set()
    for g in (groups or {}).values():
        if isinstance(g, dict) and g.get("swap") is False:
            out.update(g.get("members") or [])
    return out


def reserved_gb(role: str | None) -> dict:
    """Speicher, der NEBEN einem Zielmodell der gegebenen Rolle resident bleibt — gemäß der
    VERIFIZIERTEN llama-swap-Gruppen-Swap-Semantik (NICHT der früheren Annahme „Hirn bleibt
    immer"). Nur die ko-residente `swap:false`-Gruppe läuft gemeinsam; ein on-demand-Modell
    verdrängt die Gruppe und läuft ALLEIN mit dem vollen GTT.

    - Modell IN der Ko-Residenz-Gruppe (Hirn/embed/vision): koexistiert mit den ÜBRIGEN
      Gruppen-Mitgliedern → reserviert deren Summe.
    - Modell AUSSERHALB (heavy/coder/coder-lite/fast/scout): läuft allein (Gruppe wird beim
      Laden rausgeswappt) → reserviert NICHTS, darf den vollen GTT für Kontext nutzen.
    """
    from services import llamaswap
    models = llamaswap.list_models()
    groups = llamaswap.list_groups()
    cores = _coresident_members(groups)
    brain = next((m for m in models if (m.get("role") == "hermes")), None)
    brain_gb = footprint_gb(brain) if brain else 0.0
    role = (role or "").strip().lower()

    holder = next((m for m in models if (m.get("role") == role)), None) if role else None
    holder_name = holder["name"] if holder else None
    # Hirn (hermes) ist per Definition Teil der Ko-Residenz-Gruppe; sonst Gruppen-Mitgliedschaft prüfen.
    in_group = role == "hermes" or bool(holder_name and holder_name in cores)

    if in_group:
        others = sum(footprint_gb(m) for m in models
                     if m["name"] in cores and m["name"] != holder_name)
        return {"reserved_gb": others, "mode": "co-resident", "brain_gb": brain_gb}
    # on-demand: verdrängt die Ko-Residenz-Gruppe → läuft allein, voller GTT für Kontext.
    return {"reserved_gb": 0.0, "mode": "ondemand-alone", "brain_gb": brain_gb}


def setup_aware_ctx(params_b: float, quant: str, role: str | None = None) -> dict:
    """Größter Kontext, der für ein Modell (params_b/quant) der gegebenen Rolle NEBEN dem
    bestehenden Setup passt. Gibt ctx + die Budget-Herleitung zurück (für UI/Transparenz)."""
    gtt = gtt_budget_gb()
    r = reserved_gb(role)
    budget = max(gtt - r["reserved_gb"] - HEADROOM_GB, 0.0)
    ctx = max_ctx_in_budget(params_b, quant, budget)
    return {
        "ctx": ctx,
        "gtt_gb": gtt,
        "reserved_gb": round(r["reserved_gb"], 1),
        "budget_gb": round(budget, 1),
        "mode": r["mode"],
    }


def setup_aware_ctx_for_model(model: dict) -> dict:
    """Setup-bewusster Optimal-ctx für ein INSTALLIERTES Modell (aus seiner Rolle,
    Params & Quant). Für den 'Auto'-Button an der Modellkarte."""
    return setup_aware_ctx(params_of_model(model),
                           model.get("quant") or "Q4_K_M",
                           role=model.get("role"))