jpmschweitzer/tatlock
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 5 Wiki Activity

refactor(bench): share the GPU residency guard, and guard tool calling too

Browse Source 
Extracts the residency snapshot/restore into scripts/ollama_residency.py
so the two benchmarks cannot drift, and applies it to
benchmark_tool_calling.py, which had no protection at all.

That script was the more dangerous of the two. It rewrites
OLLAMA_DEFAULT_MODEL in .env and lets uvicorn reload onto it, restoring
the original only after the loop — so any crash or interrupt left the
*running server* pointed at the benchmark model. Its DEFAULT_MODELS
begins with mistral-nemo-large, the 9.2G model implicated in the
2026-08-07 VRAM outage. Both the .env restore and the residency restore
now run from `finally`.

SIGTERM is handled explicitly in the shared module. Python runs `finally`
for SIGINT, which arrives as KeyboardInterrupt, but the default SIGTERM
action terminates outright, so `timeout` or a plain `kill` skipped the
guard entirely.

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
jpmschweitzer
2026-08-08 17:03:53 +02:00
co-authored by Claude
parent 4f42bc047a
commit bf13f9f0de
3 changed files with 177 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/benchmark_routing.py
+30 -100
View File
@@ -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)