diff --git a/scripts/benchmark_routing.py b/scripts/benchmark_routing.py index 0d8559d..d813290 100644 --- a/scripts/benchmark_routing.py +++ b/scripts/benchmark_routing.py @@ -35,7 +35,6 @@ from __future__ import annotations import argparse import json -import signal import statistics import sys import time @@ -49,6 +48,10 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(PROJECT_ROOT)) from scripts.fixtures.routing_fixtures import FIXTURES # noqa: E402 +from scripts.ollama_residency import ( # noqa: E402 + install_sigterm_handler, + residency_guard, +) from src.agents.steward.agent import build_steward_prompt # noqa: E402 from src.agents.steward.service import _DELEGATE_LINE_RE, _extract_capabilities # noqa: E402 from src.core.startup import register_household_members # noqa: E402 @@ -85,65 +88,6 @@ def call(client: httpx.Client, body: dict[str, Any]) -> dict[str, Any] | None: return None -PINNED_THRESHOLD_SECONDS = 86400 - - -def snapshot_residency(client: httpx.Client) -> dict[str, bool]: - """Which models are resident, and which of those are pinned. - - Benchmarking swaps models on a GPU that production is using. Ollama evicts to - make room, so a run silently unpins the model Tatlock serves from and leaves - its own behind: the next voice turn then pays a ~36s cold load. Snapshot - before, restore after. - """ - try: - data = client.get(f"{OLLAMA_URL}/api/ps", timeout=10).json() - except Exception: # noqa: BLE001 - return {} - resident: dict[str, bool] = {} - now = datetime.now(UTC) - for model in data.get("models", []): - expires = model.get("expires_at", "") - pinned = False - try: - delta = datetime.fromisoformat(expires.replace("Z", "+00:00")) - now - pinned = delta.total_seconds() > PINNED_THRESHOLD_SECONDS - except ValueError: - pass - resident[model["name"]] = pinned - return resident - - -def set_keep_alive(client: httpx.Client, model: str, keep_alive: Any) -> bool: - """Load/unload/pin a model. Embedding models reject /api/generate.""" - payload = {"model": model, "keep_alive": keep_alive} - for endpoint in ("generate", "embed"): - try: - response = client.post(f"{OLLAMA_URL}/api/{endpoint}", json=payload, timeout=180) - 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 - except Exception: # noqa: BLE001 - return False - return False - - -def restore_residency(client: httpx.Client, before: dict[str, bool], used: list[str]) -> None: - """Evict what the benchmark loaded, then re-pin what production had pinned.""" - base = {name.split(":")[0] for name in before} - for model in used: - if model not in before and model.split(":")[0] not in base: - print(f" restoring: unloading benchmark model {model}") - set_keep_alive(client, model, 0) - for name, pinned in before.items(): - if pinned: - ok = set_keep_alive(client, name, -1) - print(f" restoring: re-pinned {name}" if ok - else f" restoring: FAILED to re-pin {name} -- run warmup-ollama.sh") - - def score(fixture: dict, found: list[str]) -> dict[str, Any]: expected = set(fixture["expect"]) forbidden = set(fixture["forbid"]) @@ -218,20 +162,8 @@ def summarise(rows: list[dict[str, Any]]) -> dict[str, Any]: } -def _raise_on_term(signum, _frame): - """Turn SIGTERM into an exception so the restore guard actually runs. - - Python runs `finally` for SIGINT, which arrives as KeyboardInterrupt, but the - default SIGTERM action terminates the process outright — a `timeout`, a - systemd stop or a plain `kill` would skip the restore and leave production - unpinned. Verified the hard way: an earlier SIGTERM here bypassed the guard - entirely and only luck kept the pinned model resident. - """ - raise KeyboardInterrupt(f"signal {signum}") - - def main() -> int: - signal.signal(signal.SIGTERM, _raise_on_term) + install_sigterm_handler() parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--models", default=",".join(DEFAULT_MODELS)) @@ -253,32 +185,24 @@ def main() -> int: f"x {args.repeats} repeats = {len(FIXTURES) * len(models) * len(think_modes) * args.repeats} calls\n") cells: dict[str, Any] = {} - with httpx.Client(timeout=args.timeout) as client: - before = snapshot_residency(client) - pinned_before = [n for n, p in before.items() if p] - print(f"resident before: {sorted(before)}" - f"{f' (pinned: {pinned_before})' if pinned_before else ''}\n") - try: - for model in models: - # Absorb the cold load (~36s) outside the measurements. - print(f"warming {model} ...", flush=True) - call(client, build_body(model, "hi", "false")) - for think in think_modes: - key = f"{model}|think={think}" - print(f" {key} ...", end=" ", flush=True) - started = time.perf_counter() - rows = run_cell(client, model, think, args.repeats) - summary = summarise(rows) - cells[key] = {"summary": summary, "rows": rows} - print(f"exact={summary.get('exact_pct')}% " - f"over={summary.get('over_routed_pct')}% " - f"median={summary.get('latency_ms_median')}ms " - f"({time.perf_counter() - started:.0f}s)") - finally: - # finally, not a tidy exit path: a Ctrl-C or a failed cell must not - # leave production unpinned behind us. - print("\nrestoring GPU residency ...") - restore_residency(client, before, models) + # The guard restores production's pinned models however this exits — a + # finished run, a failed cell, Ctrl-C or SIGTERM. + with residency_guard(models_used=models), httpx.Client(timeout=args.timeout) as client: + for model in models: + # Absorb the cold load (~36s) outside the measurements. + print(f"warming {model} ...", flush=True) + call(client, build_body(model, "hi", "false")) + for think in think_modes: + key = f"{model}|think={think}" + print(f" {key} ...", end=" ", flush=True) + started = time.perf_counter() + rows = run_cell(client, model, think, args.repeats) + summary = summarise(rows) + cells[key] = {"summary": summary, "rows": rows} + print(f"exact={summary.get('exact_pct')}% " + f"over={summary.get('over_routed_pct')}% " + f"median={summary.get('latency_ms_median')}ms " + f"({time.perf_counter() - started:.0f}s)") RESULTS_DIR.mkdir(parents=True, exist_ok=True) stamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ") @@ -302,4 +226,10 @@ def main() -> int: if __name__ == "__main__": - sys.exit(main()) + try: + sys.exit(main()) + except KeyboardInterrupt: + # The residency guard has already run by the time this is caught; + # a traceback here would just bury its output. + print("\ninterrupted", file=sys.stderr) + sys.exit(130) diff --git a/scripts/benchmark_tool_calling.py b/scripts/benchmark_tool_calling.py index f42d62d..f5f6893 100644 --- a/scripts/benchmark_tool_calling.py +++ b/scripts/benchmark_tool_calling.py @@ -17,12 +17,16 @@ import asyncio import json import re import statistics +import sys import time from dataclasses import dataclass, field from pathlib import Path import httpx +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +from scripts.ollama_residency import install_sigterm_handler, residency_guard + # --------------------------------------------------------------------------- # Configuration # --------------------------------------------------------------------------- @@ -525,18 +529,31 @@ async def main(): original_env = ENV_PATH.read_text() all_stats = [] - async with httpx.AsyncClient() as client: - for model in models: - stats = await benchmark_model(client, model, args.iterations) - all_stats.append(stats) - - # Restore original .env - ENV_PATH.write_text(original_env) - print(f"\n .env restored to original") + # Both restores must survive a crash or an interrupt. The .env one especially: + # this script rewrites OLLAMA_DEFAULT_MODEL and lets uvicorn reload onto it, + # so bailing out mid-run used to leave the *running server* pointed at the + # benchmark model — and DEFAULT_MODELS starts at mistral-nemo-large, the 9.2G + # model implicated in the 2026-08-07 VRAM outage. + install_sigterm_handler() + try: + with residency_guard(models_used=models): + async with httpx.AsyncClient() as client: + for model in models: + stats = await benchmark_model(client, model, args.iterations) + all_stats.append(stats) + finally: + ENV_PATH.write_text(original_env) + print("\n .env restored to original") print_comparison(all_stats) save_results(all_stats, Path(args.output)) if __name__ == "__main__": - asyncio.run(main()) + try: + asyncio.run(main()) + except KeyboardInterrupt: + # .env and GPU residency are both restored by now; do not bury that + # output under a traceback. + print("\ninterrupted", file=sys.stderr) + raise SystemExit(130) from None diff --git a/scripts/ollama_residency.py b/scripts/ollama_residency.py new file mode 100644 index 0000000..c9f876a --- /dev/null +++ b/scripts/ollama_residency.py @@ -0,0 +1,121 @@ +""" +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)