diff --git a/scripts/benchmark_routing.py b/scripts/benchmark_routing.py index 8b5876a..0d8559d 100644 --- a/scripts/benchmark_routing.py +++ b/scripts/benchmark_routing.py @@ -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")