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 _persist_members(groups: dict) -> set:
    out: set = set()
    for g in (groups or {}).values():
        if isinstance(g, dict) and g.get("persist"):
            out.update(g.get("members") or [])
    return out


def reserved_gb(role: str | None) -> dict:
    """Speicher, der NEBEN einem Zielmodell der gegebenen Rolle resident bleiben muss —
    gemäß verifizierter Box-Residenz. Liest die aktuelle llama-swap-Config, passt sich
    also der echten Gruppen-/persist-Konfiguration an (nicht hartkodiert).

    - Rolle `hermes` (das Hirn): muss mit dem GRÖSSTEN on-demand-Modell koexistieren.
    - Rolle, deren aktueller Träger im warmen (persist) Set liegt (z.B. fast/vision):
      koexistiert mit Hirn + den ÜBRIGEN warmen Mitgliedern.
    - sonst (heavy/coder/scout/keine): on-demand → verdrängt fast/vision, nur das Hirn bleibt.
    """
    from services import llamaswap
    models = llamaswap.list_models()
    groups = llamaswap.list_groups()
    persist = _persist_members(groups)
    brain = next((m for m in models if (m.get("role") == "hermes")), None)
    brain_name = brain["name"] if brain else None
    brain_gb = footprint_gb(brain) if brain else 0.0
    role = (role or "").strip().lower()

    if role == "hermes":
        reserved = max((footprint_gb(m) for m in models if m["name"] not in persist),
                       default=0.0)
        return {"reserved_gb": reserved, "mode": "brain", "brain_gb": brain_gb}

    holder = next((m for m in models if (m.get("role") == role)), None) if role else None
    warm = bool(holder and holder["name"] in persist)
    if warm:
        others = sum(footprint_gb(m) for m in models
                     if m["name"] in persist and m["name"] not in {brain_name, holder["name"]})
        return {"reserved_gb": brain_gb + others, "mode": "warm", "brain_gb": brain_gb}

    return {"reserved_gb": brain_gb, "mode": "ondemand", "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"))