tatlock/scripts/ollama_residency.py

"""
Guard production's GPU residency across a benchmark run.

Benchmarks swap models on the card production is serving from. Ollama evicts to
make room, so a run leaves its own models resident and the production one gone:
the next voice turn pays a ~36s cold load, and the pin that prevented it is
silently lost. That happened on 2026-08-08 — a routing benchmark evicted
gemma4:e2b and left gemma4:e4b behind, and only the monitoring noticing
`unexpected_models` caught it.

Snapshot before, restore after, and wire the restore to SIGTERM as well as the
normal path. Python runs `finally` for SIGINT, which arrives as
KeyboardInterrupt, but the default SIGTERM action terminates outright — so
`timeout`, a systemd stop or a plain `kill` would skip the guard entirely.

    from scripts.ollama_residency import residency_guard, install_sigterm_handler

    install_sigterm_handler()
    with residency_guard(models_used=["gemma4:e4b"]):
        ...
"""
from __future__ import annotations

import signal
from collections.abc import Iterator
from contextlib import contextmanager
from datetime import UTC, datetime
from typing import Any

import httpx

OLLAMA_URL = "http://localhost:11434"

# keep_alive:-1 yields a year-2318 expiry, so "pinned" is simply "expires more
# than a day out". Matches check-ai-pipeline.sh in system-admin-toj.
PINNED_THRESHOLD_SECONDS = 86400


def install_sigterm_handler() -> None:
    """Make SIGTERM raise, so `finally` blocks and context managers still run."""
    def _raise(signum, _frame):
        raise KeyboardInterrupt(f"signal {signum}")

    signal.signal(signal.SIGTERM, _raise)


def snapshot_residency(client: httpx.Client | None = None) -> dict[str, bool]:
    """Resident models mapped to whether each is pinned."""
    owns = client is None
    client = client or httpx.Client(timeout=30)
    try:
        data = client.get(f"{OLLAMA_URL}/api/ps", timeout=10).json()
    except Exception:  # noqa: BLE001 - a missing snapshot must not abort the run
        return {}
    finally:
        if owns:
            client.close()

    resident: dict[str, bool] = {}
    now = datetime.now(UTC)
    for model in data.get("models", []):
        pinned = False
        try:
            expires = datetime.fromisoformat(model.get("expires_at", "").replace("Z", "+00:00"))
            pinned = (expires - now).total_seconds() > PINNED_THRESHOLD_SECONDS
        except ValueError:
            pass
        resident[model["name"]] = pinned
    return resident


def set_keep_alive(model: str, keep_alive: Any, client: httpx.Client | None = None) -> bool:
    """Load, unload or pin a model. Embedding models reject /api/generate."""
    owns = client is None
    client = client or httpx.Client(timeout=180)
    payload = {"model": model, "keep_alive": keep_alive}
    try:
        for endpoint in ("generate", "embed"):
            try:
                response = client.post(f"{OLLAMA_URL}/api/{endpoint}", json=payload, timeout=180)
            except Exception:  # noqa: BLE001
                return False
            if response.status_code == 200:
                return True
            if response.status_code == 400 and "does not support generate" in response.text:
                continue  # embedding-only model; try /api/embed
            return False
        return False
    finally:
        if owns:
            client.close()


def restore_residency(before: dict[str, bool], used: list[str]) -> None:
    """Evict what the benchmark loaded, then re-pin what was pinned before."""
    base = {name.split(":")[0] for name in before}
    with httpx.Client(timeout=180) as client:
        for model in used:
            if model not in before and model.split(":")[0] not in base:
                print(f"  residency: unloading benchmark model {model}")
                set_keep_alive(model, 0, client)
        for name, pinned in before.items():
            if not pinned:
                continue
            ok = set_keep_alive(name, -1, client)
            print(f"  residency: re-pinned {name}" if ok
                  else f"  residency: FAILED to re-pin {name} -- run warmup-ollama.sh")


@contextmanager
def residency_guard(models_used: list[str]) -> Iterator[dict[str, bool]]:
    """Snapshot residency on entry, restore it on exit however that happens."""
    before = snapshot_residency()
    pinned = [n for n, p in before.items() if p]
    print(f"  residency: resident before {sorted(before)}"
          f"{f' (pinned: {pinned})' if pinned else ''}")
    try:
        yield before
    finally:
        print("  residency: restoring ...")
        restore_residency(before, models_used)