505 lines
21 KiB
Python
505 lines
21 KiB
Python
"""VRAM Arbitrator and High-Speed Switch Manager for Ollama and ComfyUI."""
|
|
import time
|
|
import httpx
|
|
import psutil
|
|
import logging
|
|
from typing import Dict, List, Any, Optional
|
|
from collections import deque
|
|
|
|
import asyncio
|
|
import json
|
|
import websockets
|
|
|
|
import overclock_manager
|
|
|
|
try:
|
|
import pynvml
|
|
pynvml.nvmlInit()
|
|
NVML_AVAILABLE = True
|
|
except Exception as e:
|
|
NVML_AVAILABLE = False
|
|
|
|
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
|
|
SWITCH_HISTORY = deque(maxlen=50)
|
|
|
|
|
|
def get_gpu_hardware_stats() -> Dict[str, Any]:
|
|
"""Retrieve comprehensive GPU hardware and process metrics via NVML."""
|
|
if not NVML_AVAILABLE:
|
|
return {"available": False, "error": "NVML not initialized"}
|
|
|
|
try:
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
|
|
name = pynvml.nvmlDeviceGetName(handle)
|
|
if isinstance(name, bytes):
|
|
name = name.decode("utf-8")
|
|
|
|
mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
|
util_rates = pynvml.nvmlDeviceGetUtilizationRates(handle)
|
|
temp_c = pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU)
|
|
try:
|
|
power_mw = pynvml.nvmlDeviceGetPowerUsage(handle)
|
|
power_w = round(power_mw / 1000.0, 1)
|
|
except Exception:
|
|
power_w = 0.0
|
|
|
|
fan_pct = 0
|
|
fans = []
|
|
try:
|
|
num_fans = pynvml.nvmlDeviceGetNumFans(handle)
|
|
for i in range(num_fans):
|
|
try:
|
|
fans.append(pynvml.nvmlDeviceGetFanSpeed_v2(handle, i))
|
|
except Exception:
|
|
pass
|
|
if fans:
|
|
fan_pct = max(fans)
|
|
else:
|
|
fan_pct = pynvml.nvmlDeviceGetFanSpeed(handle)
|
|
fans = [fan_pct]
|
|
except Exception:
|
|
try:
|
|
fan_pct = pynvml.nvmlDeviceGetFanSpeed(handle)
|
|
fans = [fan_pct]
|
|
except Exception:
|
|
fan_pct = 0
|
|
fans = []
|
|
|
|
try:
|
|
clock_graphics = pynvml.nvmlDeviceGetClockInfo(handle, pynvml.NVML_CLOCK_GRAPHICS)
|
|
clock_mem = pynvml.nvmlDeviceGetClockInfo(handle, pynvml.NVML_CLOCK_MEM)
|
|
except Exception:
|
|
clock_graphics = 0
|
|
clock_mem = 0
|
|
|
|
# Discover processes on GPU
|
|
proc_breakdown = {
|
|
"ollama_bytes": 0,
|
|
"comfyui_bytes": 0,
|
|
"system_bytes": 0,
|
|
"processes": []
|
|
}
|
|
|
|
try:
|
|
procs = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
|
|
graphics_procs = pynvml.nvmlDeviceGetGraphicsRunningProcesses(handle)
|
|
all_procs = {p.pid: p.usedGpuMemory for p in procs}
|
|
for p in graphics_procs:
|
|
all_procs[p.pid] = max(all_procs.get(p.pid, 0), p.usedGpuMemory or 0)
|
|
|
|
for pid, used_mem in all_procs.items():
|
|
pname = "Unknown"
|
|
cmdline = ""
|
|
try:
|
|
proc = psutil.Process(pid)
|
|
pname = proc.name()
|
|
cmdline = " ".join(proc.cmdline())
|
|
except Exception:
|
|
pass
|
|
|
|
is_ollama = "ollama" in pname.lower() or "llama-server" in cmdline.lower()
|
|
is_comfy = "comfy" in cmdline.lower() or "main.py" in cmdline.lower()
|
|
|
|
if is_ollama:
|
|
proc_breakdown["ollama_bytes"] += used_mem
|
|
elif is_comfy:
|
|
proc_breakdown["comfyui_bytes"] += used_mem
|
|
else:
|
|
proc_breakdown["system_bytes"] += used_mem
|
|
|
|
proc_breakdown["processes"].append({
|
|
"pid": pid,
|
|
"name": pname,
|
|
"cmdline": cmdline[:60],
|
|
"vram_bytes": used_mem,
|
|
"vram_mb": round(used_mem / (1024**2), 1),
|
|
"is_ollama": is_ollama,
|
|
"is_comfy": is_comfy,
|
|
})
|
|
except Exception as e:
|
|
logger.error(f"Error enumerating GPU processes: {e}")
|
|
|
|
total_vram = mem_info.total
|
|
used_vram = mem_info.used
|
|
free_vram = mem_info.free
|
|
|
|
return {
|
|
"available": True,
|
|
"device_name": name,
|
|
"vram_total_bytes": total_vram,
|
|
"vram_total_gb": round(total_vram / (1024**3), 2),
|
|
"vram_used_bytes": used_vram,
|
|
"vram_used_gb": round(used_vram / (1024**3), 2),
|
|
"vram_free_bytes": free_vram,
|
|
"vram_free_gb": round(free_vram / (1024**3), 2),
|
|
"vram_used_pct": round((used_vram / total_vram * 100) if total_vram > 0 else 0, 1),
|
|
"gpu_util_pct": util_rates.gpu,
|
|
"mem_util_pct": util_rates.memory,
|
|
"temperature_c": temp_c,
|
|
"power_w": power_w,
|
|
"fan_pct": fan_pct,
|
|
"fans": fans,
|
|
"num_fans": len(fans),
|
|
"clock_graphics_mhz": clock_graphics,
|
|
"clock_mem_mhz": clock_mem,
|
|
"breakdown": {
|
|
"ollama_mb": round(proc_breakdown["ollama_bytes"] / (1024**2), 1),
|
|
"ollama_gb": round(proc_breakdown["ollama_bytes"] / (1024**3), 2),
|
|
"comfyui_mb": round(proc_breakdown["comfyui_bytes"] / (1024**2), 1),
|
|
"comfyui_gb": round(proc_breakdown["comfyui_bytes"] / (1024**3), 2),
|
|
"system_mb": round(proc_breakdown["system_bytes"] / (1024**2), 1),
|
|
"system_gb": round(proc_breakdown["system_bytes"] / (1024**3), 2),
|
|
"free_mb": round(free_vram / (1024**2), 1),
|
|
"free_gb": round(free_vram / (1024**3), 2),
|
|
"processes": proc_breakdown["processes"],
|
|
}
|
|
}
|
|
except Exception as e:
|
|
return {"available": False, "error": str(e)}
|
|
|
|
async def get_ollama_live_state() -> Dict[str, Any]:
|
|
"""Get active models, running status, and VRAM expiration from Ollama."""
|
|
state = {
|
|
"online": False,
|
|
"loaded_models": [],
|
|
"active_model_name": None,
|
|
"active_model_vram_gb": 0.0,
|
|
"active_context": 0,
|
|
"expires_at": None,
|
|
"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", [])
|
|
except Exception as e:
|
|
logger.debug(f"Ollama check error: {e}")
|
|
|
|
return state
|
|
|
|
async def get_comfyui_live_state() -> Dict[str, Any]:
|
|
"""Get prompt queue, device status, and active execution from ComfyUI."""
|
|
state = {
|
|
"online": False,
|
|
"executing": False,
|
|
"queue_remaining": 0,
|
|
"queue_running": 0,
|
|
"current_node": None,
|
|
"current_prompt_id": None,
|
|
"vram_free_mb": 0,
|
|
"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])
|
|
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."""
|
|
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}
|
|
|
|
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}
|
|
except Exception as e:
|
|
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}
|
|
except Exception as e:
|
|
return {"success": False, "error": str(e), "duration_ms": round((time.perf_counter() - t0) * 1000, 2)}
|
|
|
|
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}
|
|
except Exception as e:
|
|
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."""
|
|
|
|
def __init__(self):
|
|
self.running = False
|
|
self.ws_task: Optional[asyncio.Task] = None
|
|
self.poll_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.oc_profile = None
|
|
|
|
async def start(self):
|
|
if self.running:
|
|
return
|
|
self.running = True
|
|
self.ws_task = asyncio.create_task(self._ws_listener())
|
|
self.poll_task = asyncio.create_task(self._poll_watchdog())
|
|
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")
|
|
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()
|
|
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."""
|
|
self.comfy_was_active = True
|
|
self._apply_oc_profile("comfy")
|
|
now = time.time()
|
|
if now - self.last_yield_time < 1.0:
|
|
return
|
|
|
|
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...")
|
|
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)"
|
|
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."""
|
|
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...")
|
|
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}")
|
|
|
|
async def _ws_listener(self):
|
|
client_id = "hyperswap-arbitrator"
|
|
ws_url = f"ws://127.0.0.1:8188/ws?clientId={client_id}"
|
|
|
|
while self.running:
|
|
try:
|
|
async with websockets.connect(ws_url, ping_interval=10, ping_timeout=10) as ws:
|
|
self.connected_ws = True
|
|
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 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)
|
|
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)
|
|
|
|
async def _poll_watchdog(self):
|
|
"""Watchdog polling /queue every 300ms for robust bidirectional arbitration."""
|
|
while self.running:
|
|
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()
|
|
except Exception:
|
|
pass
|
|
await asyncio.sleep(0.3)
|
|
|
|
def _apply_oc_profile(self, profile: str):
|
|
"""Apply an overclock profile in a background thread; only fire on transition."""
|
|
if self.oc_profile == profile:
|
|
return
|
|
self.oc_profile = profile
|
|
try:
|
|
loop = asyncio.get_event_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]:
|
|
return {
|
|
"running": self.running,
|
|
"connected_ws": self.connected_ws,
|
|
"last_action": self.last_action,
|
|
"mode": "Bidirectional Hot-Swap (ComfyUI <-> Ollama)",
|
|
}
|
|
|
|
|
|
arbitrator = AutoArbitrator()
|
|
|
|
|