drjones/gpu-program-swapper
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add barrier-confirmed yielding, measured residency, persistence and closed-loop tuning

Browse Source 
Nine changes, in rough order of how much they affect real behaviour:

1. VRAM yield is now a barrier. Posting keep_alive:0 only asks Ollama to unload;
   measured here, the HTTP call returns in 63ms while the driver takes a further
   77ms to release 14.9GB. Returning inside that window is how ComfyUI ends up
   allocating into VRAM that is still occupied. instant_free_ollama_vram() polls
   NVML until the allocation is actually gone and reports request/confirm split.

2. ComfyUI VRAM is no longer purged 1.5s after every prompt, which forced a full
   checkpoint reload on each workflow iteration. It is held for 30s of genuinely
   empty queue, with an immediate purge when Ollama actually asks for the memory.

3. Cache-hit classification uses achieved bandwidth (size / load duration) rather
   than a fixed `load_duration < 2500ms`. That constant called a 12.9GB model read
   at 2.9GB/s a cold load, and a 0.5GB model read from NVMe a cache hit.

4. Page-cache residency is measured, not assumed. mincore(2) reported 128GB
   resident on a box with 46GB of page cache: the kernel only permits page-cache
   introspection on files you own, and the Ollama blobs are owned by uid ollama,
   for which mincore answers "all resident" instead of failing. Uses cachestat(2)
   where permitted and a randomised read-rate probe elsewhere, labelling which was
   used. Fixed-offset probing was self-fulfilling, so windows are random and cold
   ones are returned with FADV_DONTNEED.

5. Warming is budgeted and ranked by recency/frequency instead of reading every
   file top-to-bottom, which on 64GB of RAM just evicts whatever was warmed first.

6. Telemetry and events persist to SQLite (~0.38 MB/hour) instead of living in a
   50-entry in-memory deque, so /api/analytics/profiles can finally answer whether
   an overclock profile actually delivers more tok/s.

7. Thermal governor walks the overclock back on sustained heat or hardware
   throttling, with hysteresis, fed from the existing sampler.

8. Autotune sweeps a clock offset, benchmarks decode at each step, watches for Xid
   errors and degenerate output, and restores the profile in a finally block.

9. Stock clocks/power/fans are restored on shutdown and via systemd ExecStopPost.
   Nothing previously undid a locked clock or a manually pinned fan.

Also: one shared 1Hz telemetry sampler fanned out to SSE subscribers rather than
every client re-running the whole snapshot; wall-clock timestamps in place of the
event loop's monotonic clock; cached nvidia-smi shell-outs; quieter httpx logging.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
drjones
2026-08-28 08:57:35 -07:00
parent 6acf235bd7
commit 5431144b2e
12 changed files with 2779 additions and 362 deletions
Download Patch File Download Diff File Expand all files Collapse all files

721
vram_arbitrator.py
View File

@@ -11,6 +11,8 @@ import json
import websockets
import overclock_manager
import ram_optimizer
import telemetry_store
try:
import pynvml
@@ -24,9 +26,116 @@ logger = logging.getLogger("vram_arbitrator")
OLLAMA_API_BASE = "http://localhost:11434"
COMFY_API_BASE = "http://127.0.0.1:8188"
# Circular buffer for transition events
# Circular buffer for transition events (the durable log lives in telemetry_store)
SWITCH_HISTORY = deque(maxlen=50)
# Bandwidth thresholds used to classify how a model actually got into VRAM.
# PCIe 4.0 x16 tops out near 31.5 GB/s; this NVMe sustains well under 2 GB/s.
RAM_HIT_GBPS = 5.0
PARTIAL_HIT_GBPS = 1.5
# How long Ollama's VRAM may take to actually drain before we stop waiting.
YIELD_CONFIRM_TIMEOUT_S = 3.0
YIELD_CONFIRM_POLL_S = 0.02
YIELD_RESIDUAL_BYTES = 256 * 1024 ** 2 # treat <256 MB as "released"
# Connection-pooled clients. Re-creating an AsyncClient per call meant a fresh TCP
# handshake on every one of the watchdog's polls.
_clients: Dict[str, httpx.AsyncClient] = {}
def _client(base_url: str, timeout: float) -> httpx.AsyncClient:
key = f"{base_url}|{timeout}"
c = _clients.get(key)
if c is None or c.is_closed:
c = httpx.AsyncClient(
base_url=base_url,
timeout=timeout,
limits=httpx.Limits(max_keepalive_connections=4, max_connections=8),
)
_clients[key] = c
return c
async def close_clients() -> None:
for c in list(_clients.values()):
try:
await c.aclose()
except Exception:
pass
_clients.clear()
# Bit flags from nvmlDeviceGetCurrentClocksThrottleReasons, decoded for the governor.
THROTTLE_REASONS = {
0x0000000000000001: "gpu_idle",
0x0000000000000002: "applications_clocks_setting",
0x0000000000000004: "sw_power_cap",
0x0000000000000008: "hw_slowdown",
0x0000000000000010: "sync_boost",
0x0000000000000020: "sw_thermal_slowdown",
0x0000000000000040: "hw_thermal_slowdown",
0x0000000000000080: "hw_power_brake_slowdown",
0x0000000000000100: "display_clock_setting",
}
def decode_throttle_reasons(bits: int) -> List[str]:
return [name for mask, name in THROTTLE_REASONS.items() if bits & mask]
def get_process_vram_bytes() -> Dict[str, int]:
"""Fast NVML-only VRAM attribution, used by the yield barrier's tight poll loop.
Deliberately avoids psutil lookups: this runs every 20 ms while we wait for VRAM
to actually drain.
"""
out = {"ollama_bytes": 0, "comfyui_bytes": 0, "other_bytes": 0, "free_bytes": 0}
if not NVML_AVAILABLE:
return out
try:
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
out["free_bytes"] = pynvml.nvmlDeviceGetMemoryInfo(handle).free
procs = list(pynvml.nvmlDeviceGetComputeRunningProcesses(handle))
try:
procs += list(pynvml.nvmlDeviceGetGraphicsRunningProcesses(handle))
except Exception:
pass
merged: Dict[int, int] = {}
for p in procs:
merged[p.pid] = max(merged.get(p.pid, 0), p.usedGpuMemory or 0)
for pid, used in merged.items():
kind = _PID_KIND_CACHE.get(pid)
if kind is None:
kind = _classify_pid(pid)
_PID_KIND_CACHE[pid] = kind
if kind == "ollama":
out["ollama_bytes"] += used
elif kind == "comfy":
out["comfyui_bytes"] += used
else:
out["other_bytes"] += used
except Exception as e:
logger.debug(f"get_process_vram_bytes failed: {e}")
return out
_PID_KIND_CACHE: Dict[int, str] = {}
def _classify_pid(pid: int) -> str:
try:
proc = psutil.Process(pid)
pname = proc.name().lower()
cmdline = " ".join(proc.cmdline()).lower()
except Exception:
return "other"
if "ollama" in pname or "llama-server" in cmdline:
return "ollama"
if "comfy" in cmdline or "main.py" in cmdline:
return "comfy"
return "other"
def get_gpu_hardware_stats() -> Dict[str, Any]:
"""Retrieve comprehensive GPU hardware and process metrics via NVML."""
@@ -77,6 +186,49 @@ def get_gpu_hardware_stats() -> Dict[str, Any]:
clock_graphics = 0
clock_mem = 0
# PCIe throughput (KB/s) — TX + RX. Key metric for the RAM-cache
# PCIe-speed swap thesis (assimilated from pmady/gpu-mcp-server).
pcie_tx_kbps = 0
pcie_rx_kbps = 0
try:
pcie_tx_kbps = pynvml.nvmlDeviceGetPcieThroughput(handle, pynvml.NVML_PCIE_UTIL_TX_BYTES)
except Exception:
pass
try:
pcie_rx_kbps = pynvml.nvmlDeviceGetPcieThroughput(handle, pynvml.NVML_PCIE_UTIL_RX_BYTES)
except Exception:
pass
# Why the GPU is not running at full clocks — the thermal governor reads this.
throttle_bits = 0
throttle_reasons: List[str] = []
try:
throttle_bits = pynvml.nvmlDeviceGetCurrentClocksThrottleReasons(handle)
throttle_reasons = decode_throttle_reasons(throttle_bits)
except Exception:
pass
# Power management limit (watts) — the OC ceiling.
power_limit_w = 0.0
try:
power_limit_w = round(pynvml.nvmlDeviceGetPowerManagementLimit(handle) / 1000.0, 1)
except Exception:
pass
# Driver + CUDA version (completeness).
driver_version = ""
cuda_version = ""
try:
dv = pynvml.nvmlSystemGetDriverVersion()
driver_version = dv.decode("utf-8") if isinstance(dv, bytes) else str(dv)
except Exception:
pass
try:
cv = pynvml.nvmlSystemGetCudaDriverVersion()
cuda_version = cv # int like 12030 == CUDA 12.3
except Exception:
pass
# Discover processes on GPU
proc_breakdown = {
"ollama_bytes": 0,
@@ -142,6 +294,13 @@ def get_gpu_hardware_stats() -> Dict[str, Any]:
"mem_util_pct": util_rates.memory,
"temperature_c": temp_c,
"power_w": power_w,
"power_limit_w": power_limit_w,
"pcie_tx_kbps": pcie_tx_kbps,
"pcie_rx_kbps": pcie_rx_kbps,
"throttle_bits": throttle_bits,
"throttle_reasons": throttle_reasons,
"driver_version": driver_version,
"cuda_version": cuda_version,
"fan_pct": fan_pct,
"fans": fans,
"num_fans": len(fans),
@@ -174,25 +333,25 @@ async def get_ollama_live_state() -> Dict[str, Any]:
"installed_models": []
}
try:
async with httpx.AsyncClient(timeout=3.0) as client:
# Check running models (ps)
ps_resp = await client.get(f"{OLLAMA_API_BASE}/api/ps")
if ps_resp.status_code == 200:
state["online"] = True
models = ps_resp.json().get("models", [])
state["loaded_models"] = models
if models:
first = models[0]
state["active_model_name"] = first.get("name")
vram_bytes = first.get("size_vram", first.get("size", 0))
state["active_model_vram_gb"] = round(vram_bytes / (1024**3), 2)
state["active_context"] = first.get("context_length", 0)
state["expires_at"] = first.get("expires_at")
# Check all tags
tags_resp = await client.get(f"{OLLAMA_API_BASE}/api/tags")
if tags_resp.status_code == 200:
state["installed_models"] = tags_resp.json().get("models", [])
client = _client(OLLAMA_API_BASE, 3.0)
# Check running models (ps)
ps_resp = await client.get("/api/ps")
if ps_resp.status_code == 200:
state["online"] = True
models = ps_resp.json().get("models", [])
state["loaded_models"] = models
if models:
first = models[0]
state["active_model_name"] = first.get("name")
vram_bytes = first.get("size_vram", first.get("size", 0))
state["active_model_vram_gb"] = round(vram_bytes / (1024**3), 2)
state["active_context"] = first.get("context_length", 0)
state["expires_at"] = first.get("expires_at")
# Check all tags
tags_resp = await client.get("/api/tags")
if tags_resp.status_code == 200:
state["installed_models"] = tags_resp.json().get("models", [])
except Exception as e:
logger.debug(f"Ollama check error: {e}")
@@ -211,156 +370,286 @@ async def get_comfyui_live_state() -> Dict[str, Any]:
"vram_total_mb": 0,
}
try:
async with httpx.AsyncClient(timeout=3.0) as client:
# Check system stats
stats_resp = await client.get(f"{COMFY_API_BASE}/system_stats")
if stats_resp.status_code == 200:
state["online"] = True
data = stats_resp.json()
devices = data.get("devices", [])
if devices:
dev = devices[0]
state["vram_free_mb"] = round(dev.get("vram_free", 0) / (1024**2), 1)
state["vram_total_mb"] = round(dev.get("vram_total", 0) / (1024**2), 1)
# Check queue
queue_resp = await client.get(f"{COMFY_API_BASE}/queue")
if queue_resp.status_code == 200:
qdata = queue_resp.json()
running = qdata.get("queue_running", [])
pending = qdata.get("queue_pending", [])
state["queue_running"] = len(running)
state["queue_remaining"] = len(pending)
state["executing"] = len(running) > 0
if running:
state["current_prompt_id"] = running[0][1] if len(running[0]) > 1 else str(running[0])
client = _client(COMFY_API_BASE, 3.0)
# Check system stats
stats_resp = await client.get("/system_stats")
if stats_resp.status_code == 200:
state["online"] = True
data = stats_resp.json()
devices = data.get("devices", [])
if devices:
dev = devices[0]
state["vram_free_mb"] = round(dev.get("vram_free", 0) / (1024**2), 1)
state["vram_total_mb"] = round(dev.get("vram_total", 0) / (1024**2), 1)
# Check queue
queue_resp = await client.get("/queue")
if queue_resp.status_code == 200:
qdata = queue_resp.json()
running = qdata.get("queue_running", [])
pending = qdata.get("queue_pending", [])
state["queue_running"] = len(running)
state["queue_remaining"] = len(pending)
state["executing"] = len(running) > 0
if running:
state["current_prompt_id"] = running[0][1] if len(running[0]) > 1 else str(running[0])
except Exception as e:
logger.debug(f"ComfyUI check error: {e}")
return state
async def instant_free_ollama_vram(model_name: Optional[str] = None) -> Dict[str, Any]:
"""Tell Ollama to instantly yield VRAM without evicting from OS page cache."""
async def _await_vram_release(baseline_bytes: int,
timeout_s: float = YIELD_CONFIRM_TIMEOUT_S) -> Dict[str, Any]:
"""Block until Ollama's VRAM has actually drained, or we give up.
Posting keep_alive:0 only *asks* Ollama to unload; the driver frees the allocation
some milliseconds later. Returning before that happens is how ComfyUI ends up
allocating into VRAM that is still occupied, which shows up as a CUDA OOM mid-graph.
"""
t0 = time.perf_counter()
last = baseline_bytes
while True:
snap = get_process_vram_bytes()
last = snap["ollama_bytes"]
if last <= YIELD_RESIDUAL_BYTES:
return {
"confirmed": True,
"confirm_ms": round((time.perf_counter() - t0) * 1000, 2),
"residual_bytes": last,
"free_bytes": snap["free_bytes"],
}
if (time.perf_counter() - t0) >= timeout_s:
return {
"confirmed": False,
"confirm_ms": round((time.perf_counter() - t0) * 1000, 2),
"residual_bytes": last,
"free_bytes": snap["free_bytes"],
"error": f"Ollama still holding {round(last / (1024**3), 2)} GB after {timeout_s}s",
}
await asyncio.sleep(YIELD_CONFIRM_POLL_S)
def _record(event: Dict[str, Any]) -> None:
"""Push an event to both the in-memory ring and the durable store."""
event.setdefault("ts", time.time())
event.setdefault("timestamp", time.strftime("%H:%M:%S"))
SWITCH_HISTORY.appendleft(event)
telemetry_store.record_event(event, profile=overclock_manager.ACTIVE_PROFILE)
async def instant_free_ollama_vram(model_name: Optional[str] = None,
confirm: bool = True) -> Dict[str, Any]:
"""Yield Ollama's VRAM and wait for the driver to actually release it.
The returned duration_ms is now the real end-to-end release time, not just how long
the HTTP POST took.
"""
t0 = time.perf_counter()
if not model_name:
ollama_state = await get_ollama_live_state()
model_name = ollama_state.get("active_model_name")
if not model_name:
return {"success": True, "message": "No active Ollama model in VRAM", "duration_ms": 0}
return {"success": True, "message": "No active Ollama model in VRAM",
"duration_ms": 0, "confirmed": True}
baseline = get_process_vram_bytes()["ollama_bytes"]
try:
async with httpx.AsyncClient(timeout=5.0) as client:
resp = await client.post(
f"{OLLAMA_API_BASE}/api/generate",
json={"model": model_name, "keep_alive": 0},
)
duration_ms = round((time.perf_counter() - t0) * 1000, 2)
event = {
"timestamp": time.strftime("%H:%M:%S"),
"event_type": "Ollama VRAM Yield",
"source": model_name,
"target": "VRAM 0MB (Kept in RAM)",
"duration_ms": duration_ms,
"cache_status": "RAM-Cached",
}
SWITCH_HISTORY.appendleft(event)
return {"success": True, "model": model_name, "duration_ms": duration_ms}
client = _client(OLLAMA_API_BASE, 5.0)
await client.post("/api/generate", json={"model": model_name, "keep_alive": 0})
request_ms = round((time.perf_counter() - t0) * 1000, 2)
barrier = {"confirmed": None, "confirm_ms": 0.0, "residual_bytes": baseline}
if confirm:
barrier = await _await_vram_release(baseline)
duration_ms = round((time.perf_counter() - t0) * 1000, 2)
freed_gb = round(max(baseline - barrier.get("residual_bytes", 0), 0) / (1024**3), 2)
_record({
"event_type": "Ollama VRAM Yield",
"source": model_name,
"target": "VRAM 0MB (Kept in RAM)",
"duration_ms": duration_ms,
"yield_confirm_ms": barrier.get("confirm_ms"),
"cache_status": "RAM-Cached" if barrier.get("confirmed") else "Yield Timeout",
"detail": barrier.get("error"),
})
return {
"success": True,
"model": model_name,
"duration_ms": duration_ms,
"request_ms": request_ms,
"confirm_ms": barrier.get("confirm_ms"),
"confirmed": barrier.get("confirmed"),
"freed_gb": freed_gb,
"residual_gb": round(barrier.get("residual_bytes", 0) / (1024**3), 2),
"free_vram_gb": round(barrier.get("free_bytes", 0) / (1024**3), 2),
"error": barrier.get("error"),
}
except Exception as e:
return {"success": False, "error": str(e), "duration_ms": round((time.perf_counter() - t0) * 1000, 2)}
return {"success": False, "error": str(e),
"duration_ms": round((time.perf_counter() - t0) * 1000, 2)}
async def instant_free_comfyui_vram() -> Dict[str, Any]:
"""Tell ComfyUI to purge loaded diffusion models from VRAM."""
t0 = time.perf_counter()
try:
async with httpx.AsyncClient(timeout=5.0) as client:
resp = await client.post(
f"{COMFY_API_BASE}/free",
json={"unload_models": True, "free_memory": True},
)
duration_ms = round((time.perf_counter() - t0) * 1000, 2)
event = {
"timestamp": time.strftime("%H:%M:%S"),
"event_type": "ComfyUI VRAM Purge",
"source": "ComfyUI Pipeline",
"target": "VRAM Free",
"duration_ms": duration_ms,
"cache_status": "Cleaned",
}
SWITCH_HISTORY.appendleft(event)
return {"success": True, "duration_ms": duration_ms}
client = _client(COMFY_API_BASE, 5.0)
await client.post("/free", json={"unload_models": True, "free_memory": True})
duration_ms = round((time.perf_counter() - t0) * 1000, 2)
snap = get_process_vram_bytes()
_record({
"event_type": "ComfyUI VRAM Purge",
"source": "ComfyUI Pipeline",
"target": "VRAM Free",
"duration_ms": duration_ms,
"cache_status": "Cleaned",
})
return {"success": True, "duration_ms": duration_ms,
"free_vram_gb": round(snap["free_bytes"] / (1024**3), 2)}
except Exception as e:
return {"success": False, "error": str(e), "duration_ms": round((time.perf_counter() - t0) * 1000, 2)}
return {"success": False, "error": str(e),
"duration_ms": round((time.perf_counter() - t0) * 1000, 2)}
_MODEL_SIZE_CACHE: Dict[str, int] = {}
def _model_size_bytes(model_name: str) -> int:
"""On-disk weight size for an Ollama model, used to turn load time into bandwidth."""
if model_name in _MODEL_SIZE_CACHE:
return _MODEL_SIZE_CACHE[model_name]
try:
for f in ram_optimizer.find_ollama_model_files():
_MODEL_SIZE_CACHE[f["model"]] = f["size_bytes"]
except Exception as e:
logger.debug(f"model size lookup failed: {e}")
return _MODEL_SIZE_CACHE.get(model_name, 0)
def classify_load(size_bytes: int, load_duration_ms: float) -> Dict[str, Any]:
"""Classify how a model reached VRAM, from achieved bandwidth rather than a constant.
The old rule was `load_duration_ms < 2500`, which called a 27B Q2_K read from NVMe a
cache hit and a small model read from RAM a cold load. Bandwidth separates them
cleanly: page cache feeds PCIe at many GB/s, this NVMe does not.
"""
if load_duration_ms <= 1.0:
return {"cache_status": "Already in VRAM", "load_gbps": None, "is_ram_hit": True}
if not size_bytes:
# No size on record — fall back to the old heuristic, but say so.
return {
"cache_status": "RAM Cache Hit ⚡" if load_duration_ms < 2500 else "Cold Disk Load 💾",
"load_gbps": None,
"is_ram_hit": load_duration_ms < 2500,
"detail": "size unknown, fell back to duration heuristic",
}
gbps = (size_bytes / (1024**3)) / (load_duration_ms / 1000.0)
if gbps >= RAM_HIT_GBPS:
status = "RAM Cache Hit ⚡"
elif gbps >= PARTIAL_HIT_GBPS:
status = "Partial Cache 🌤"
else:
status = "Cold Disk Load 💾"
return {"cache_status": status, "load_gbps": round(gbps, 2), "is_ram_hit": gbps >= RAM_HIT_GBPS}
async def switch_ollama_model(target_model: str, keep_alive: str = "30m") -> Dict[str, Any]:
"""High-speed hot-swap to target Ollama model, tracking swap metrics."""
t0 = time.perf_counter()
cur_state = await get_ollama_live_state()
prev_model = cur_state.get("active_model_name") or "None"
try:
async with httpx.AsyncClient(timeout=180.0) as client:
resp = await client.post(
f"{OLLAMA_API_BASE}/api/generate",
json={"model": target_model, "prompt": "Ready check", "stream": False, "keep_alive": keep_alive},
)
total_duration = time.perf_counter() - t0
total_duration_ms = round(total_duration * 1000, 2)
if resp.status_code == 200:
data = resp.json()
load_dur_ms = round(data.get("load_duration", 0) / 1e6, 2)
eval_dur_ms = round(data.get("eval_duration", 0) / 1e6, 2)
eval_count = data.get("eval_count", 0)
tokens_per_sec = round((eval_count / (eval_dur_ms / 1000)) if eval_dur_ms > 0 else 0, 1)
# Check if it was a RAM cache hit (load duration < 1500ms for large model indicates RAM hit)
is_ram_hit = load_dur_ms < 2500
event = {
"timestamp": time.strftime("%H:%M:%S"),
"event_type": "LLM Model Switch",
"source": prev_model,
"target": target_model,
"duration_ms": total_duration_ms,
"load_duration_ms": load_dur_ms,
"tokens_per_sec": tokens_per_sec,
"cache_status": "RAM Cache Hit ⚡" if is_ram_hit else "Cold Disk Load 💾",
}
SWITCH_HISTORY.appendleft(event)
return {
"success": True,
"prev_model": prev_model,
"target_model": target_model,
"total_duration_ms": total_duration_ms,
"load_duration_ms": load_dur_ms,
"tokens_per_sec": tokens_per_sec,
"is_ram_hit": is_ram_hit,
"response": data.get("response", ""),
}
else:
return {"success": False, "error": f"HTTP {resp.status_code}: {resp.text}", "duration_ms": total_duration_ms}
client = _client(OLLAMA_API_BASE, 180.0)
resp = await client.post(
"/api/generate",
json={"model": target_model, "prompt": "Ready check", "stream": False,
"keep_alive": keep_alive},
)
total_duration_ms = round((time.perf_counter() - t0) * 1000, 2)
if resp.status_code == 200:
data = resp.json()
load_dur_ms = round(data.get("load_duration", 0) / 1e6, 2)
eval_dur_ms = round(data.get("eval_duration", 0) / 1e6, 2)
eval_count = data.get("eval_count", 0)
tokens_per_sec = round((eval_count / (eval_dur_ms / 1000)) if eval_dur_ms > 0 else 0, 1)
size_bytes = _model_size_bytes(target_model)
cls = classify_load(size_bytes, load_dur_ms)
_record({
"event_type": "LLM Model Switch",
"source": prev_model,
"target": target_model,
"duration_ms": total_duration_ms,
"load_duration_ms": load_dur_ms,
"tokens_per_sec": tokens_per_sec,
"bytes_loaded": size_bytes,
"load_gbps": cls["load_gbps"],
"cache_status": cls["cache_status"],
"detail": cls.get("detail"),
})
return {
"success": True,
"prev_model": prev_model,
"target_model": target_model,
"total_duration_ms": total_duration_ms,
"load_duration_ms": load_dur_ms,
"tokens_per_sec": tokens_per_sec,
"model_size_gb": round(size_bytes / (1024**3), 2) if size_bytes else None,
"load_gbps": cls["load_gbps"],
"cache_status": cls["cache_status"],
"is_ram_hit": cls["is_ram_hit"],
"response": data.get("response", ""),
}
return {"success": False, "error": f"HTTP {resp.status_code}: {resp.text}",
"duration_ms": total_duration_ms}
except Exception as e:
return {"success": False, "error": str(e), "duration_ms": round((time.perf_counter() - t0) * 1000, 2)}
return {"success": False, "error": str(e),
"duration_ms": round((time.perf_counter() - t0) * 1000, 2)}
def get_switch_history() -> List[Dict[str, Any]]:
return list(SWITCH_HISTORY)
class AutoArbitrator:
"""Real-time bidirectional background arbitrator for seamless Ollama <-> ComfyUI hot-swapping."""
"""Real-time bidirectional background arbitrator for seamless Ollama <-> ComfyUI hot-swapping.
Two behavioural changes worth knowing about:
* ComfyUI's VRAM is no longer purged 1.5 s after every finished prompt. Iterating on
a workflow is the common case, and purging between runs forced a full checkpoint
reload each time. The purge now waits for COMFY_IDLE_PURGE_S of genuinely empty
queue, and happens immediately only when Ollama actually needs the VRAM.
* The watchdog no longer polls two ComfyUI endpoints every 300 ms. The WebSocket is
the primary signal; polling is a fallback that runs at 1 Hz and only hits /queue,
backing off further while the socket is healthy.
"""
COMFY_IDLE_PURGE_S = 30.0
WATCHDOG_INTERVAL_S = 1.0
WATCHDOG_INTERVAL_WS_OK_S = 3.0
def __init__(self):
self.running = False
self.ws_task: Optional[asyncio.Task] = None
self.poll_task: Optional[asyncio.Task] = None
self.idle_task: Optional[asyncio.Task] = None
self.last_yield_time = 0.0
self.last_comfy_free_time = 0.0
self.connected_ws = False
self.last_action = "Idle"
self.comfy_was_active = False
self.comfy_idle_since: Optional[float] = None
self.oc_profile = None
self.pending_purge = False
self.stats = {"yields": 0, "purges": 0, "yield_timeouts": 0, "deferred_purges": 0}
async def start(self):
if self.running:
@@ -368,24 +657,27 @@ class AutoArbitrator:
self.running = True
self.ws_task = asyncio.create_task(self._ws_listener())
self.poll_task = asyncio.create_task(self._poll_watchdog())
self.idle_task = asyncio.create_task(self._idle_purge_loop())
logger.info("AutoArbitrator background engine started (Bidirectional).")
# Apply the default (balanced) overclock profile on startup.
try:
await asyncio.get_event_loop().run_in_executor(None, overclock_manager.apply_profile, "balanced")
await asyncio.get_running_loop().run_in_executor(
None, overclock_manager.apply_profile, "balanced")
self.oc_profile = "balanced"
except Exception as e:
logger.warning(f"Startup overclock apply failed: {e}")
async def stop(self):
self.running = False
if self.ws_task:
self.ws_task.cancel()
if self.poll_task:
self.poll_task.cancel()
for task in (self.ws_task, self.poll_task, self.idle_task):
if task:
task.cancel()
await close_clients()
logger.info("AutoArbitrator background engine stopped.")
async def trigger_comfy_priority(self, reason: str = "ComfyUI prompt detected"):
"""Instantly yield Ollama VRAM to 0MB when ComfyUI needs to run diffusion models."""
"""Yield Ollama's VRAM — and confirm it is gone — before diffusion allocates."""
self.comfy_was_active = True
self.comfy_idle_since = None
self._apply_oc_profile("comfy")
now = time.time()
if now - self.last_yield_time < 1.0:
@@ -394,89 +686,143 @@ class AutoArbitrator:
ollama_state = await get_ollama_live_state()
if ollama_state.get("active_model_name"):
model = ollama_state["active_model_name"]
logger.info(f"⚡ ComfyUI active ({reason}) -> Auto-yielding Ollama model '{model}' from VRAM...")
logger.info(f"⚡ ComfyUI active ({reason}) -> Auto-yielding Ollama model '{model}'...")
self.last_yield_time = time.time()
res = await instant_free_ollama_vram(model)
dur = res.get("duration_ms", 0)
self.last_action = f"Auto-yielded '{model}' for ComfyUI ({dur}ms)"
res = await instant_free_ollama_vram(model, confirm=True)
self.stats["yields"] += 1
if res.get("confirmed"):
self.last_action = (f"Yielded '{model}' for ComfyUI in "
f"{res.get('confirm_ms')}ms (confirmed {res.get('freed_gb')}GB free)")
else:
self.stats["yield_timeouts"] += 1
self.last_action = (f"⚠ Yield of '{model}' NOT confirmed: "
f"{res.get('residual_gb')}GB still held")
logger.warning(f"VRAM yield barrier timed out: {res.get('error')}")
logger.info(f"Ollama auto-yield completed: {res}")
async def trigger_comfy_completed(self):
"""Purge ComfyUI VRAM cache when generation finishes, keeping VRAM 100% free for Ollama."""
async def trigger_comfy_completed(self, immediate: bool = False):
"""Mark the end of a generation. The actual purge is deferred unless forced."""
self.comfy_was_active = False
if self.comfy_idle_since is None:
self.comfy_idle_since = time.time()
self._apply_oc_profile("ollama")
if immediate:
await self._purge_comfy_now("Ollama needs VRAM")
else:
self.pending_purge = True
self.stats["deferred_purges"] += 1
self.last_action = (f"ComfyUI idle — holding its checkpoints for "
f"{int(self.COMFY_IDLE_PURGE_S)}s in case you iterate")
async def _purge_comfy_now(self, reason: str):
now = time.time()
if now - self.last_comfy_free_time < 3.0:
return
self.last_comfy_free_time = now
self.comfy_was_active = False
self._apply_oc_profile("ollama")
logger.info("⚡ ComfyUI finished generation -> Auto-purging ComfyUI VRAM cache for Ollama...")
self.pending_purge = False
logger.info(f"⚡ Purging ComfyUI VRAM cache ({reason})...")
res = await instant_free_comfyui_vram()
dur = res.get("duration_ms", 0)
self.last_action = f"Auto-purged ComfyUI VRAM ({dur}ms) - Ready for Ollama"
logger.info(f"ComfyUI auto-purge completed: {res}")
self.stats["purges"] += 1
self.last_action = f"Purged ComfyUI VRAM ({res.get('duration_ms')}ms) — {reason}"
logger.info(f"ComfyUI purge completed: {res}")
async def _idle_purge_loop(self):
"""Purge ComfyUI's VRAM only after a real idle gap, not between iterations."""
while self.running:
try:
if self.pending_purge and self.comfy_idle_since and not self.comfy_was_active:
idle_for = time.time() - self.comfy_idle_since
if idle_for >= self.COMFY_IDLE_PURGE_S:
await self._purge_comfy_now(
f"idle {int(idle_for)}s")
except Exception as e:
logger.debug(f"idle purge loop error: {e}")
await asyncio.sleep(2.0)
async def request_vram_for_ollama(self, needed_gb: float = 0.0) -> Dict[str, Any]:
"""Called when Ollama needs VRAM now: purge ComfyUI immediately rather than waiting."""
snap = get_process_vram_bytes()
free_gb = snap["free_bytes"] / (1024**3)
if needed_gb and free_gb >= needed_gb:
return {"purged": False, "free_gb": round(free_gb, 2), "reason": "enough free VRAM"}
if snap["comfyui_bytes"] > YIELD_RESIDUAL_BYTES:
await self._purge_comfy_now(f"Ollama requested {needed_gb or '?'}GB")
snap = get_process_vram_bytes()
return {"purged": True, "free_gb": round(snap["free_bytes"] / (1024**3), 2)}
return {"purged": False, "free_gb": round(free_gb, 2), "reason": "ComfyUI holds no VRAM"}
async def _ws_listener(self):
client_id = "hyperswap-arbitrator"
ws_url = f"ws://127.0.0.1:8188/ws?clientId={client_id}"
backoff = 2.0
while self.running:
try:
async with websockets.connect(ws_url, ping_interval=10, ping_timeout=10) as ws:
self.connected_ws = True
backoff = 2.0
logger.info("AutoArbitrator connected to ComfyUI WebSocket.")
while self.running:
msg = await ws.recv()
if isinstance(msg, str):
try:
data = json.loads(msg)
msg_type = data.get("type")
msg_data = data.get("data", {})
if not isinstance(msg, str):
continue
try:
data = json.loads(msg)
msg_type = data.get("type")
msg_data = data.get("data", {})
if msg_type == "status":
queue_rem = msg_data.get("status", {}).get("exec_info", {}).get("queue_remaining", 0)
if queue_rem > 0:
await self.trigger_comfy_priority(f"Queue remaining: {queue_rem}")
elif queue_rem == 0 and self.comfy_was_active:
# Prompt queue finished
await asyncio.sleep(1.5)
await self.trigger_comfy_completed()
elif msg_type in ("execution_start", "execution_cached"):
await self.trigger_comfy_priority(f"Event: {msg_type}")
elif msg_type == "executing":
node = msg_data.get("node")
if node is not None:
await self.trigger_comfy_priority(f"Executing node: {node}")
elif node is None and self.comfy_was_active:
# Finished executing graph
await asyncio.sleep(1.5)
await self.trigger_comfy_completed()
elif msg_type == "execution_success":
await asyncio.sleep(1.5)
if msg_type == "status":
queue_rem = (msg_data.get("status", {})
.get("exec_info", {}).get("queue_remaining", 0))
if queue_rem > 0:
await self.trigger_comfy_priority(f"Queue remaining: {queue_rem}")
elif queue_rem == 0 and self.comfy_was_active:
await self.trigger_comfy_completed()
except Exception as e:
logger.debug(f"WS parse error: {e}")
elif msg_type in ("execution_start", "execution_cached"):
await self.trigger_comfy_priority(f"Event: {msg_type}")
elif msg_type == "executing":
node = msg_data.get("node")
if node is not None:
await self.trigger_comfy_priority(f"Executing node: {node}")
elif self.comfy_was_active:
await self.trigger_comfy_completed()
elif msg_type == "execution_success":
await self.trigger_comfy_completed()
elif msg_type == "execution_error":
logger.warning(f"ComfyUI execution error: {msg_data}")
await self.trigger_comfy_completed()
except Exception as e:
logger.debug(f"WS parse error: {e}")
except (websockets.exceptions.ConnectionClosed, OSError, asyncio.CancelledError):
self.connected_ws = False
except Exception as e:
self.connected_ws = False
logger.debug(f"WS connection error: {e}")
await asyncio.sleep(2.0)
await asyncio.sleep(backoff)
backoff = min(backoff * 1.5, 15.0)
async def _poll_watchdog(self):
"""Watchdog polling /queue every 300ms for robust bidirectional arbitration."""
"""Fallback for when the WebSocket is down. One cheap /queue call, 1 Hz.
The previous version hit /system_stats and /queue every 300 ms on fresh TCP
connections — roughly 6.6 requests/second against ComfyUI, forever.
"""
while self.running:
interval = self.WATCHDOG_INTERVAL_WS_OK_S if self.connected_ws else self.WATCHDOG_INTERVAL_S
try:
comfy_state = await get_comfyui_live_state()
is_executing = comfy_state.get("queue_running", 0) > 0 or comfy_state.get("queue_remaining", 0) > 0 or comfy_state.get("executing", False)
if is_executing:
await self.trigger_comfy_priority("Polling detected active queue/execution")
elif self.comfy_was_active and not is_executing:
await asyncio.sleep(1.5)
await self.trigger_comfy_completed()
client = _client(COMFY_API_BASE, 3.0)
resp = await client.get("/queue")
if resp.status_code == 200:
q = resp.json()
busy = len(q.get("queue_running", [])) > 0 or len(q.get("queue_pending", [])) > 0
if busy:
await self.trigger_comfy_priority("Watchdog saw an active queue")
elif self.comfy_was_active:
await self.trigger_comfy_completed()
except Exception:
pass
await asyncio.sleep(0.3)
await asyncio.sleep(interval)
def _apply_oc_profile(self, profile: str):
"""Apply an overclock profile in a background thread; only fire on transition."""
@@ -484,18 +830,25 @@ class AutoArbitrator:
return
self.oc_profile = profile
try:
loop = asyncio.get_event_loop()
loop = asyncio.get_running_loop()
loop.run_in_executor(None, overclock_manager.apply_profile, profile)
logger.info(f"🎛️ Overclock profile switched -> '{profile}'")
except Exception as e:
logger.warning(f"Overclock profile switch failed ({profile}): {e}")
def get_status(self) -> Dict[str, Any]:
idle_for = (time.time() - self.comfy_idle_since) if self.comfy_idle_since else None
return {
"running": self.running,
"connected_ws": self.connected_ws,
"last_action": self.last_action,
"mode": "Bidirectional Hot-Swap (ComfyUI <-> Ollama)",
"comfy_active": self.comfy_was_active,
"pending_purge": self.pending_purge,
"comfy_idle_s": round(idle_for, 1) if idle_for is not None else None,
"idle_purge_after_s": self.COMFY_IDLE_PURGE_S,
"oc_profile": self.oc_profile,
"counters": dict(self.stats),
}