test(bench): restore GPU residency after a benchmark run
Benchmarking swaps models on the GPU production is serving from. Ollama evicts to make room, so the first run unpinned gemma4:e2b and left gemma4:e4b resident: the next voice turn would have paid a ~36s cold load, and only the monitoring noticing unexpected_models caught it. Snapshot residency and pinning before the run, then evict whatever the benchmark loaded and re-pin what was pinned before. The restore is wired to SIGTERM as well as the normal exit path. Python runs `finally` for SIGINT, which arrives as KeyboardInterrupt, but the default SIGTERM action terminates outright — so a `timeout`, a systemd stop or a plain `kill` skipped the guard entirely. That was not theoretical: the first SIGTERM after adding this bypassed it, and the pinned model survived only because the run had not reached the second model yet. Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -35,6 +35,7 @@ from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import signal
|
||||
import statistics
|
||||
import sys
|
||||
import time
|
||||
@@ -84,6 +85,65 @@ 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"])
|
||||
@@ -158,7 +218,21 @@ 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)
|
||||
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--models", default=",".join(DEFAULT_MODELS))
|
||||
parser.add_argument("--think", default=",".join(DEFAULT_THINK),
|
||||
@@ -180,21 +254,31 @@ def main() -> int:
|
||||
|
||||
cells: dict[str, Any] = {}
|
||||
with 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)")
|
||||
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)
|
||||
|
||||
RESULTS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
stamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
|
||||
|
||||
Reference in New Issue
Block a user