diff --git a/server.py b/server.py index db98b9a..9b12adf 100644 --- a/server.py +++ b/server.py @@ -383,7 +383,18 @@ async def api_switch_model(req: SwitchRequest): await vram_arbitrator.arbitrator.request_vram_for_ollama() res = await vram_arbitrator.switch_ollama_model(req.model, keep_alive=req.keep_alive or "30m") if not res.get("success"): - raise HTTPException(status_code=500, detail=res.get("error")) + # Reflect what actually went wrong. Ollama returns 400 for an unusable request -- + # asking an embedding model to generate, say -- and reporting that as 500 blames + # this service for the caller's mistake. A model that will not fit is neither: + # the request is valid and the service is healthy, there is simply no room. + upstream = res.get("upstream_status") + if res.get("vram_oom"): + status = 507 # Insufficient Storage + elif isinstance(upstream, int) and 400 <= upstream < 500: + status = 400 + else: + status = 502 if upstream else 500 + raise HTTPException(status_code=status, detail=res.get("error")) return res @app.post("/api/free-vram", summary="Soft-Yield Ollama VRAM", tags=["Orchestration"]) diff --git a/verify_arbitration.py b/verify_arbitration.py new file mode 100755 index 0000000..933ee6f --- /dev/null +++ b/verify_arbitration.py @@ -0,0 +1,251 @@ +#!/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())) diff --git a/vram_arbitrator.py b/vram_arbitrator.py index 01ca217..79ae9e6 100644 --- a/vram_arbitrator.py +++ b/vram_arbitrator.py @@ -880,6 +880,7 @@ async def switch_ollama_model(target_model: str, keep_alive: str = "30m", return retry return {"success": False, "error": f"HTTP {resp.status_code}: {body}", "duration_ms": total_duration_ms, + "upstream_status": resp.status_code, "vram_oom": looks_like_vram_oom(body)} except Exception as e: return {"success": False, "error": str(e),