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)
scripts/benchmark_tool_calling.py
+26 -9
View File
@@ -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
scripts/ollama_residency.py
+121
View File
@@ -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)