"""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 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 try: fan_pct = pynvml.nvmlDeviceGetFanSpeed(handle) except Exception: fan_pct = 0 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, "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)