aeba1b47fd
The health check now reports what unmanaged VRAM actually costs rather than just how much of it there is: "0.82 GB held by python (842 MB)" becomes "5 model(s) fit within 15.42 GB but not the 14.60 GB actually available", naming them. Getting that arithmetic right took a correction. The first version subtracted only the desktop and the unmanaged process, and so reported a 14.93 GB model as fitting against a real ceiling of 14.60 GB -- the same model the service had just refused with 507. ComfyUI keeps a few hundred MB of CUDA context for as long as the process lives, which a purge does not free, so it is not available either. The floor is taken from the minimum ComfyUI VRAM in recent telemetry rather than its current value, which could be a 7 GB checkpoint mid-generation. The verifier's reclaim stage now re-runs a graph immediately beforehand to reset the 30 s idle window, since a large model takes longer than that to load and the purge was freeing ComfyUI mid-load, so the reclaim path was never reached. Tests: 199 (was 192). The new ones pin the ceiling arithmetic, including that a model too large to fit on the card at all is not blamed on the third-party process. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
339 lines
15 KiB
Python
Executable File
339 lines
15 KiB
Python
Executable File
#!/usr/bin/env python
|
|
"""End-to-end verification of the arbitration cycle, against real hardware.
|
|
|
|
The unit suite covers logic in isolation; this exercises the promise the whole service
|
|
exists to make -- that an LLM and a diffusion pipeline can share one 16 GB card without
|
|
either failing -- and reports what actually happened at each stage.
|
|
|
|
It is deliberately not part of `pytest tests/`: it loads real models, runs a real
|
|
diffusion graph and moves real VRAM, taking a few minutes. Run it when you want proof
|
|
the system works on this machine:
|
|
|
|
python verify_arbitration.py # full cycle
|
|
python verify_arbitration.py --quick # skip the diffusion stages
|
|
|
|
Every stage restores what it changed, and the script refuses to start if ComfyUI is
|
|
already busy.
|
|
"""
|
|
import argparse
|
|
import asyncio
|
|
import sys
|
|
import time
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import httpx
|
|
|
|
BASE = "http://localhost:9090"
|
|
|
|
PASS, FAIL, SKIP, WARN = "PASS", "FAIL", "SKIP", "WARN"
|
|
_results: List[Dict[str, Any]] = []
|
|
|
|
|
|
def record(stage: str, status: str, detail: str, evidence: str = "") -> None:
|
|
_results.append({"stage": stage, "status": status, "detail": detail,
|
|
"evidence": evidence})
|
|
colour = {"PASS": "\033[32m", "FAIL": "\033[31m",
|
|
"SKIP": "\033[33m", "WARN": "\033[33m"}[status]
|
|
print(f" {colour}{status:<4}\033[0m {stage:<38} {detail}")
|
|
if evidence:
|
|
print(f" {evidence}")
|
|
|
|
|
|
async def api(client: httpx.AsyncClient, method: str, path: str,
|
|
allow_error: bool = False, **kw) -> Any:
|
|
r = await client.request(method, f"{BASE}{path}", **kw)
|
|
if allow_error:
|
|
# Some stages deliberately provoke a failure and need to read it.
|
|
body = r.json() if r.headers.get("content-type", "").startswith("application/json") else {}
|
|
return {"_status": r.status_code, **(body if isinstance(body, dict) else {})}
|
|
r.raise_for_status()
|
|
return r.json()
|
|
|
|
|
|
async def stage_preflight(c: httpx.AsyncClient) -> Optional[Dict[str, Any]]:
|
|
health = await api(c, "GET", "/api/health")
|
|
failed = health.get("failed") or []
|
|
if failed:
|
|
record("preflight: dependencies", FAIL,
|
|
f"{len(failed)} dependency check(s) failing", ", ".join(failed))
|
|
return None
|
|
record("preflight: dependencies", PASS, health["summary"])
|
|
|
|
stats = await api(c, "GET", "/api/stats")
|
|
if stats["comfyui"].get("executing") or stats["comfyui"].get("queue_remaining"):
|
|
record("preflight: ComfyUI idle", FAIL, "ComfyUI is busy; refusing to interfere")
|
|
return None
|
|
record("preflight: ComfyUI idle", PASS, "queue empty")
|
|
return stats
|
|
|
|
|
|
async def stage_llm_load(c: httpx.AsyncClient, model: str) -> bool:
|
|
res = await api(c, "POST", "/api/switch-model",
|
|
json={"model": model, "keep_alive": "5m"}, timeout=300)
|
|
if not res.get("success"):
|
|
record("LLM loads", FAIL, res.get("error", "")[:90])
|
|
return False
|
|
gbps, status = res.get("load_gbps"), res.get("cache_status")
|
|
record("LLM loads", PASS, f"{res['model_size_gb']} GB in {res['load_duration_ms']:.0f} ms",
|
|
f"{gbps} GB/s -> {status}; {res['tokens_per_sec']} tok/s")
|
|
|
|
# The classification must follow from the measured bandwidth, not a fixed duration.
|
|
if gbps is not None:
|
|
expected = ("RAM Cache Hit" if gbps >= 2.0
|
|
else "Partial Cache" if gbps >= 0.8 else "Cold Disk Load")
|
|
ok = expected.split()[0] in (status or "")
|
|
record("load classified by bandwidth", PASS if ok else FAIL,
|
|
f"{gbps} GB/s reported as '{status}'",
|
|
"" if ok else f"expected something matching '{expected}'")
|
|
return True
|
|
|
|
|
|
async def stage_yield(c: httpx.AsyncClient) -> bool:
|
|
before = await api(c, "GET", "/api/gpu")
|
|
held = before["breakdown"]["ollama_gb"]
|
|
res = await api(c, "POST", "/api/free-vram", timeout=60)
|
|
outcome = res.get("outcome")
|
|
|
|
if outcome == "released":
|
|
after = await api(c, "GET", "/api/gpu")
|
|
freed = held - after["breakdown"]["ollama_gb"]
|
|
# The barrier's promise: on return, the VRAM is genuinely gone.
|
|
ok = after["breakdown"]["ollama_gb"] < 0.3
|
|
record("VRAM yield is confirmed", PASS if ok else FAIL,
|
|
f"released in {res['confirm_ms']} ms",
|
|
f"{freed:.2f} GB freed; NVML now reports "
|
|
f"{after['breakdown']['ollama_gb']} GB held by Ollama")
|
|
return ok
|
|
if outcome == "busy":
|
|
record("VRAM yield is confirmed", WARN,
|
|
"model was mid-generation, so the unload was deferred",
|
|
res.get("error", ""))
|
|
return True
|
|
record("VRAM yield is confirmed", FAIL, f"outcome={outcome}", res.get("error", ""))
|
|
return False
|
|
|
|
|
|
async def stage_diffusion(c: httpx.AsyncClient) -> bool:
|
|
sys.path.insert(0, "/home/drjones/unified-model-manager")
|
|
import autotune, vram_arbitrator # noqa: E402 (imported late; needs the service's deps)
|
|
|
|
# The first run loads the checkpoint from disk. Timing that and calling the result
|
|
# "it/s" understates throughput by roughly 10x -- 0.67 it/s against a steady-state
|
|
# 6.7 -- so the load is measured separately and reported as what it is.
|
|
first = await autotune._diffusion_benchmark()
|
|
if not first.get("ok"):
|
|
record("diffusion runs", FAIL, first.get("error", "")[:90])
|
|
return False
|
|
record("diffusion runs (cold, includes checkpoint load)", PASS,
|
|
f"{first['exec_ms']} ms", f"{first['it_per_sec']} it/s including load")
|
|
|
|
res = await autotune._diffusion_benchmark()
|
|
if not res.get("ok"):
|
|
record("diffusion runs (warm)", FAIL, res.get("error", "")[:90])
|
|
return False
|
|
record("diffusion throughput (warm)", PASS,
|
|
f"SDXL 1024/20 steps in {res['exec_ms']} ms", f"{res['it_per_sec']} it/s")
|
|
|
|
snap = vram_arbitrator.get_process_vram_bytes()
|
|
comfy_gb = snap["comfyui_bytes"] / (1024 ** 3)
|
|
record("ComfyUI holds its checkpoint", PASS if comfy_gb > 0.5 else WARN,
|
|
f"{comfy_gb:.2f} GB retained",
|
|
"held for the idle window rather than purged between iterations")
|
|
return True
|
|
|
|
|
|
async def stage_idle_purge(c: httpx.AsyncClient) -> bool:
|
|
# The completion event arrives over the ComfyUI websocket, so the flag is set a
|
|
# moment after the graph returns. Checking instantly raced it.
|
|
arb = {}
|
|
for _ in range(12):
|
|
arb = (await api(c, "GET", "/api/stats"))["arbitrator"]
|
|
if arb.get("pending_purge"):
|
|
break
|
|
await asyncio.sleep(0.5)
|
|
if not arb.get("pending_purge"):
|
|
record("purge is deferred, not immediate", WARN,
|
|
"no purge pending after 6 s (ComfyUI may already be clean)")
|
|
return True
|
|
idle_s = arb.get("comfy_idle_s")
|
|
record("purge is deferred, not immediate", PASS,
|
|
f"holding checkpoints for {arb.get('idle_purge_after_s')} s",
|
|
f"idle {idle_s} s so far" if idle_s is not None
|
|
else "idle timer just started")
|
|
return True
|
|
|
|
|
|
async def stage_reclaim(c: httpx.AsyncClient, model: str) -> bool:
|
|
"""The direction that used to fail outright: an LLM that will not fit.
|
|
|
|
This only proves anything if the chosen model genuinely cannot fit in what ComfyUI
|
|
has left free. A small model fits alongside the checkpoint and the stage passes
|
|
without exercising the reclaim path at all, so pick the largest model that will not
|
|
fit and say plainly when no such model exists.
|
|
"""
|
|
gpu = await api(c, "GET", "/api/gpu")
|
|
comfy_gb = gpu["breakdown"]["comfyui_gb"]
|
|
free_gb = gpu["vram_free_gb"]
|
|
if comfy_gb < 0.5:
|
|
record("reclaims VRAM for the LLM", SKIP,
|
|
f"ComfyUI only holds {comfy_gb} GB; nothing to reclaim")
|
|
return True
|
|
|
|
models = (await api(c, "GET", "/api/models"))["ollama_models"]
|
|
EMBED = {"bert", "nomic-bert", "gte", "jina-bert"}
|
|
usable = [m for m in models
|
|
if (m.get("details", {}).get("family") or "").lower() not in EMBED
|
|
and "embed" not in m["name"].lower()]
|
|
# On-disk weight size is not the VRAM footprint: measured on this box, a 12.87 GB
|
|
# blob occupies 14.9 GB once context and KV cache are allocated. Sizing the test off
|
|
# disk size picks a model that cannot fit even after a successful reclaim.
|
|
VRAM_OVERHEAD = 1.18
|
|
HEADROOM_GB = 0.4
|
|
|
|
def vram_need(m):
|
|
return m.get("size", 0) / (1024 ** 3) * VRAM_OVERHEAD
|
|
|
|
# Unmanaged VRAM never comes back, so it is not part of what a reclaim can offer.
|
|
# Ignoring it picked a model that failed even after a correct reclaim -- on this box
|
|
# an 842 MB third-party process is the difference between a 14.9 GB model fitting
|
|
# and not.
|
|
reclaimable_gb = free_gb + comfy_gb - HEADROOM_GB
|
|
too_big = [m for m in usable
|
|
if vram_need(m) > free_gb and vram_need(m) < reclaimable_gb]
|
|
if too_big:
|
|
target = max(too_big, key=lambda m: m.get("size", 0))
|
|
model = target["name"]
|
|
print(f" using {model} ({target['size'] / (1024**3):.1f} GB on disk, "
|
|
f"~{vram_need(target):.1f} GB in VRAM) — will not fit in "
|
|
f"{free_gb:.1f} GB free, should fit after reclaiming {comfy_gb:.1f} GB")
|
|
else:
|
|
unmanaged = gpu["breakdown"].get("unmanaged_gb", 0)
|
|
record("reclaims VRAM for the LLM", SKIP,
|
|
f"no installed model needs between {free_gb:.1f} and "
|
|
f"{reclaimable_gb:.1f} GB of VRAM",
|
|
f"reclaimable ceiling excludes {unmanaged} GB held by processes "
|
|
f"HyperSwap cannot free")
|
|
return True
|
|
|
|
# Re-run a graph first. The idle purge fires 30 s after ComfyUI goes quiet, and a
|
|
# large model takes longer than that to load -- so without resetting the timer the
|
|
# purge frees ComfyUI mid-load and the reclaim path is never reached.
|
|
sys.path.insert(0, "/home/drjones/unified-model-manager")
|
|
import autotune # noqa: E402
|
|
await autotune._diffusion_benchmark()
|
|
gpu = await api(c, "GET", "/api/gpu")
|
|
print(f" reset the idle window; ComfyUI holds "
|
|
f"{gpu['breakdown']['comfyui_gb']} GB, {gpu['vram_free_gb']} GB free")
|
|
|
|
res = await api(c, "POST", "/api/switch-model", allow_error=True,
|
|
json={"model": model, "keep_alive": "2m"}, timeout=600)
|
|
if res.get("_status") == 507:
|
|
record("reclaims VRAM for the LLM", FAIL,
|
|
"reclaim ran but the model still did not fit",
|
|
str(res.get("detail", ""))[:150])
|
|
return False
|
|
if res.get("_status", 200) >= 400 or not res.get("success"):
|
|
record("reclaims VRAM for the LLM", FAIL,
|
|
f"HTTP {res.get('_status')}", str(res.get("detail", ""))[:120])
|
|
return False
|
|
if res.get("_status") and res.get("_status") != 200:
|
|
pass
|
|
if res.get("reclaimed_from_comfyui_gb"):
|
|
record("reclaims VRAM for the LLM", PASS,
|
|
f"reclaimed {res['reclaimed_from_comfyui_gb']} GB and retried",
|
|
f"'{model}' then loaded at {res.get('load_gbps')} GB/s")
|
|
else:
|
|
# It fit anyway, so nothing was proven; do not report that as a pass. The usual
|
|
# cause is the idle purge firing during the load and freeing ComfyUI first.
|
|
record("reclaims VRAM for the LLM", WARN,
|
|
"model fit without a reclaim, so the path was not exercised",
|
|
"the idle purge most likely freed ComfyUI during the load; "
|
|
f"loaded at {res.get('load_gbps')} GB/s")
|
|
return True
|
|
|
|
|
|
async def stage_accounting(c: httpx.AsyncClient) -> bool:
|
|
"""Reported VRAM must add up, and reported settings must match the hardware."""
|
|
gpu = await api(c, "GET", "/api/gpu")
|
|
bd = gpu["breakdown"]
|
|
parts = bd["ollama_gb"] + bd["comfyui_gb"] + bd["system_gb"]
|
|
used = gpu["vram_used_gb"]
|
|
# Driver overhead means the parts never sum exactly; a large gap means mis-attribution.
|
|
ok = abs(parts - used) < 1.5
|
|
record("VRAM attribution adds up", PASS if ok else FAIL,
|
|
f"parts {parts:.2f} GB vs NVML used {used:.2f} GB",
|
|
f"ollama {bd['ollama_gb']} + comfy {bd['comfyui_gb']} + system {bd['system_gb']} "
|
|
f"(desktop {bd.get('desktop_gb')}, unmanaged {bd.get('unmanaged_gb')})")
|
|
|
|
oc = await api(c, "GET", "/api/overclock")
|
|
drift = oc["drift"]
|
|
ok2 = not drift["drifted"]
|
|
record("reported GPU state matches hardware", PASS if ok2 else FAIL,
|
|
f"profile '{drift['profile']}' asks {drift['power_limit_intended_w']} W, "
|
|
f"card reports {drift['power_limit_actual_w']} W",
|
|
drift.get("reason") or "")
|
|
return ok and ok2
|
|
|
|
|
|
async def main() -> int:
|
|
ap = argparse.ArgumentParser(description=__doc__,
|
|
formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
ap.add_argument("--quick", action="store_true",
|
|
help="skip the diffusion and reclaim stages")
|
|
ap.add_argument("--model", default=None,
|
|
help="Ollama model to test with (default: smallest installed)")
|
|
args = ap.parse_args()
|
|
|
|
print("\nHyperSwap arbitration verification\n" + "=" * 62)
|
|
async with httpx.AsyncClient(timeout=60.0) as c:
|
|
baseline = await stage_preflight(c)
|
|
if baseline is None:
|
|
print("\nPreflight failed; not continuing.\n")
|
|
return 2
|
|
|
|
model = args.model
|
|
if not model:
|
|
models = (await api(c, "GET", "/api/models"))["ollama_models"]
|
|
# Embedding models have no /api/generate endpoint, and the smallest model
|
|
# installed is often one of them.
|
|
EMBED_FAMILIES = {"bert", "nomic-bert", "gte", "jina-bert"}
|
|
usable = [m for m in models
|
|
if (m.get("details", {}).get("family") or "").lower()
|
|
not in EMBED_FAMILIES and "embed" not in m["name"].lower()]
|
|
if not usable:
|
|
record("choose a test model", FAIL,
|
|
"no generative Ollama models installed "
|
|
f"({len(models)} found, all embedding-only)")
|
|
return 2
|
|
model = min(usable, key=lambda m: m.get("size", 0))["name"]
|
|
print(f"\n using model: {model}\n")
|
|
|
|
await stage_llm_load(c, model)
|
|
await stage_yield(c)
|
|
|
|
if not args.quick:
|
|
if await stage_diffusion(c):
|
|
await stage_idle_purge(c)
|
|
await stage_reclaim(c, model)
|
|
else:
|
|
record("diffusion stages", SKIP, "--quick")
|
|
|
|
await stage_accounting(c)
|
|
|
|
# Leave the box as we found it.
|
|
await api(c, "POST", "/api/free-vram", timeout=60)
|
|
|
|
failed = [r for r in _results if r["status"] == FAIL]
|
|
warned = [r for r in _results if r["status"] == WARN]
|
|
print("=" * 62)
|
|
print(f" {len(_results) - len(failed) - len(warned)} passed, "
|
|
f"{len(warned)} warned, {len(failed)} failed")
|
|
for r in failed:
|
|
print(f" FAILED: {r['stage']} — {r['detail']}")
|
|
print()
|
|
return 1 if failed else 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(asyncio.run(main()))
|