#!/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, **kw) -> Any: r = await client.request(method, f"{BASE}{path}", **kw) 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) t0 = time.perf_counter() res = await autotune._diffusion_benchmark() if not res.get("ok"): record("diffusion runs", FAIL, res.get("error", "")[:90]) return False record("diffusion runs", 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: stats = await api(c, "GET", "/api/stats") arb = stats["arbitrator"] if not arb.get("pending_purge"): record("purge is deferred, not immediate", WARN, "no purge pending (ComfyUI may already be clean)") return True record("purge is deferred, not immediate", PASS, f"holding checkpoints for {arb.get('idle_purge_after_s')} s", f"idle {arb.get('comfy_idle_s')} s so far") 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.""" gpu = await api(c, "GET", "/api/gpu") comfy_gb = gpu["breakdown"]["comfyui_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 res = await api(c, "POST", "/api/switch-model", json={"model": model, "keep_alive": "2m"}, timeout=300) if not res.get("success"): record("reclaims VRAM for the LLM", FAIL, res.get("error", "")[:90], str(res.get("unmanaged_blockers", ""))[:120]) return False 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"loaded at {res.get('load_gbps')} GB/s") else: record("reclaims VRAM for the LLM", PASS, "fit without needing a reclaim", f"{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()))