Add GPU fan control, live telemetry, and per-profile fan curves in HyperSwap dashboard

vram_arbitrator.py

@@ -6,6 +6,12 @@ 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()
@@ -21,6 +27,7 @@ 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:
@@ -41,10 +48,27 @@ def get_gpu_hardware_stats() -> Dict[str, Any]:
         except Exception:
             power_w = 0.0
             
+        fan_pct = 0
+        fans = []
         try:
-            fan_pct = pynvml.nvmlDeviceGetFanSpeed(handle)
+            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:
-            fan_pct = 0
+            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)
@@ -119,6 +143,8 @@ def get_gpu_hardware_stats() -> Dict[str, Any]:
             "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": {
@@ -320,3 +346,159 @@ async def switch_ollama_model(target_model: str, keep_alive: str = "30m") -> Dic
 
 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()
+
+