#!/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()))