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

"""
Hält das ganze WARM-SET (Hirn + Augen/vision + Gedächtnis/embed) dauerhaft warm.

Hintergrund: llama-swap ist EIN-Gruppen-resident — lädt ein on-demand-Modell außerhalb
der `brains`-Gruppe, wird die ganze Gruppe verdrängt. `persist: true` verhindert nur
Idle-Unload, NICHT die Gruppen-Verdrängung; auch ein `-watch-config`-Reload (jede Config-
Änderung/Deploy) verwirft das Set, ohne llama-swaps ExecStartPost-Warmup neu auszulösen.
Dieser Wächter schließt die Lücke: ist die Box idle (nichts geladen), lädt er das ganze
Set über deploy/warmup.sh nach — NICHT nur das Hirn (sonst blieben Augen+embed kalt, live
vom Review-Wächter beobachtet). Während aktiver Last (irgendetwas geladen) hält er sich
raus, verdrängt also nie ein gerade genutztes Modell.

warmup.sh deckt korrekt ab: fast+vision über /v1/chat/completions, embed über /v1/embeddings
(anderer Endpunkt!) und das Vorkauen des ~70-KB-Agent-Prompts. Es ist selbst-detachend.

Abschaltbar/justierbar via Env: MC_REWARM_ENABLED=0, MC_REWARM_INTERVAL. Welche Modelle
warmup.sh lädt: MC_WARMUP_MODELS (Default 'fast vision') + MC_WARMUP_EMBED (Default 'embed').
"""

import asyncio
import logging
import os
import subprocess

import httpx

from config import LLAMA_SWAP_URL

log = logging.getLogger(__name__)

ENABLED = os.environ.get("MC_REWARM_ENABLED", "1") != "0"
INTERVAL = int(os.environ.get("MC_REWARM_INTERVAL", "90"))     # Sekunden zwischen Checks
START_DELAY = int(os.environ.get("MC_REWARM_START_DELAY", "25"))
NUDGE_GRACE = int(os.environ.get("MC_REWARM_NUDGE_GRACE", "3"))  # llama-swap den Reload abschließen lassen

# Das geteilte Warm-Skript (auch llama-swaps ExecStartPost) — EINE Quelle für „was ist das
# Warm-Set und wie wärmt man es korrekt", statt hier eine zweite, ärmere Logik zu pflegen.
_WARMUP_SH = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "deploy", "warmup.sh"))

# Wecksignal für einen sofortigen Vorwärm-Check (statt bis zum nächsten INTERVAL-Tick zu warten).
# Wird von write_config() nach einer Config-Änderung gesetzt: llama-swap (-watch-config) lädt die
# neue Config und verwirft dabei ALLE Modelle inkl. Hirn — ohne Nudge bliebe es bis zu INTERVAL
# Sekunden kalt liegen, bis der nächste Tick oder eine Anfrage es wieder lädt.
_loop: asyncio.AbstractEventLoop | None = None
_wake: asyncio.Event | None = None


def nudge() -> None:
    """Threadsicher: bittet den Wächter, nach einem Config-Reload bald vorzuwärmen. No-op,
    solange der Wächter (noch) nicht läuft."""
    if _loop is not None and _wake is not None and not _loop.is_closed():
        try:
            _loop.call_soon_threadsafe(_wake.set)
        except RuntimeError:
            pass


def _run_warmup() -> bool:
    """Volles Warm-Set via deploy/warmup.sh nachladen (fast+vision über /v1/chat/completions,
    embed über /v1/embeddings, plus Agent-Prompt-Prefill). Selbst-detachend, blockiert nicht.
    False, wenn das Skript fehlt."""
    if not os.path.exists(_WARMUP_SH):
        log.warning("rewarm: warmup.sh nicht gefunden (%s) — kein Nachwärmen möglich", _WARMUP_SH)
        return False
    try:
        subprocess.Popen(["bash", _WARMUP_SH],
                         env={**os.environ, "MC_LLAMA_SWAP_URL": LLAMA_SWAP_URL},
                         stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        return True
    except Exception:
        log.debug("rewarm: warmup.sh-Start fehlgeschlagen", exc_info=True)
        return False


def _warmset_members() -> set[str]:
    """Soll-Warm-Set = Mitglieder aller ko-residenten Gruppen (swap:false bzw. persist/persistent)."""
    try:
        from services import llamaswap
        groups = llamaswap.list_groups() or {}
    except Exception:
        return set()
    want: set[str] = set()
    for g in groups.values():
        if isinstance(g, dict) and (g.get("swap") is False or g.get("persist") or g.get("persistent")):
            want.update(g.get("members") or [])
    return want


async def _warmset_missing() -> list[str] | None:
    """Warm-Set-Mitglieder, die NICHT 'ready' in /running sind. [] = alles warm, None = /running
    nicht lesbar. Erkennt auch TEIL-Kälte (nur Hirn warm, Augen/embed rausgefallen) — genau der
    Fall, der nach einem watch-config-Reload/Deploy auftritt."""
    try:
        async with httpx.AsyncClient(timeout=8.0) as c:
            r = await c.get(f"{LLAMA_SWAP_URL}/running")
            data = r.json() or {}
    except Exception:
        return None
    ready = {str(x.get("model")) for x in (data.get("running") or []) if x.get("state") == "ready"}
    want = _warmset_members()
    if not want:                                   # Set unbekannt → alte Heuristik: nur bei ganz leer
        return [] if ready else ["<leer>"]
    return [m for m in want if m not in ready]


async def rewarm_loop() -> None:
    """Endlos-Schleife (Hintergrund-Task): hält das ganze Warm-Set warm. Prüft periodisch, ob ein
    Mitglied fehlt (Teil-Kälte!), und sofort nach einem Config-Reload-Nudge — lädt via warmup.sh nach."""
    global _loop, _wake
    _loop = asyncio.get_running_loop()
    _wake = asyncio.Event()
    await asyncio.sleep(START_DELAY)   # Box/Engine nach MC-Start setzen lassen
    while True:
        try:
            missing = await _warmset_missing()
            if missing:
                log.info("rewarm: Warm-Set unvollständig (%s) → warmup.sh", ", ".join(missing))
                _run_warmup()
        except Exception:
            log.debug("rewarm: Tick fehlgeschlagen", exc_info=True)
        # Bis zum nächsten Tick warten ODER sofort auf einen Config-Reload-Nudge reagieren.
        try:
            await asyncio.wait_for(_wake.wait(), timeout=INTERVAL)
            _wake.clear()
            await asyncio.sleep(NUDGE_GRACE)   # llama-swap den Reload abschließen lassen
            log.info("rewarm: Config-Reload → warmup.sh (volles Warm-Set nachladen)")
            _run_warmup()
        except asyncio.TimeoutError:
            pass