commit 380badbd4ad2122c5ebd6abe6b0975569515b89f Author: forkless Date: Mon Jun 8 23:04:38 2026 +0200 test commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..1a5c693 --- /dev/null +++ b/.gitignore @@ -0,0 +1,6 @@ +__pycache__/ +*.pyc +.DS_Store +*.egg-info/ +dist/ +build/ diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..9efb593 --- /dev/null +++ b/__init__.py @@ -0,0 +1,51 @@ +""" +ComfyUI-XPUSYS-Monitor — GPU resource monitor plugin for ComfyUI. + +Entry point: auto-detects the GPU vendor, instantiates the appropriate +hardware provider, registers web routes, and starts the background +broadcast loop. +""" + +import asyncio +import logging + +logger = logging.getLogger("XPUSYSMonitor") + +WEB_DIRECTORY = "./web" +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} + +# --------------------------------------------------------------------------- +# Initialise backend +# --------------------------------------------------------------------------- + +try: + from .providers import auto_detect_provider + from .xpu_server import register_routes, set_provider, broadcast_loop + + _provider = auto_detect_provider(interval_ms=1000) + set_provider(_provider) + + # Attach HTTP routes to PromptServer + try: + from server import PromptServer + server = PromptServer.instance + if server is not None: + register_routes(server) + + # Schedule the WebSocket broadcast loop onto the server's event loop + loop = server.loop + if loop is not None: + asyncio.run_coroutine_threadsafe( + broadcast_loop(server, interval_s=1.0), loop + ) + except Exception as exc: + logger.warning(f"XPUSYSMonitor: could not attach to PromptServer — {exc}") + + logger.info("XPUSYSMonitor: plugin loaded successfully.") + +except Exception as exc: + logger.error(f"XPUSYSMonitor: failed to initialise — {exc}", exc_info=True) + NODE_CLASS_MAPPINGS = {} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/providers/__init__.py b/providers/__init__.py new file mode 100644 index 0000000..11268a6 --- /dev/null +++ b/providers/__init__.py @@ -0,0 +1,155 @@ +""" +providers/__init__.py — Provider registry and auto-detection factory. + +Usage: + from .providers import auto_detect_provider, BaseGPUProvider, GPUSnapshot + + provider = auto_detect_provider(interval_ms=1000) + snap = provider.get_snapshot() + +Detection strategy — align with ComfyUI's own device selection: + ComfyUI uses torch.cuda.is_available() for NVIDIA and torch.xpu.is_available() + for Intel. We follow the same signals so the monitor always tracks whichever + device ComfyUI is actually running on. + +Detection order: + 1. torch.cuda.is_available() + torch.version.roc → AMDProvider (ROCm) + 2. torch.cuda.is_available() → NvidiaProvider (NVIDIA) + 3. torch.xpu.is_available() → IntelProvider (Intel Arc) + 4. ze_loader.dll present → IntelProvider (fallback) + 5. pynvml available → NvidiaProvider (fallback) + 6. Last resort → IntelProvider (limited) + +Fallback tiers 4-5 cover edge cases where torch is not yet imported or the +user is running a non-standard environment without torch. +""" + +import logging +from .base import BaseGPUProvider, GPUSnapshot + +logger = logging.getLogger("XPUSYSMonitor") + + +def auto_detect_provider(interval_ms: int = 1000) -> BaseGPUProvider: + """ + Detect the available GPU hardware and return the appropriate provider. + + Primary strategy: mirror ComfyUI's own torch-based device selection so the + monitor always tracks the same device that ComfyUI is running inference on. + Fallback strategy: raw driver/library probing for non-standard environments. + """ + + # --- Primary: follow torch (mirrors ComfyUI model_management.py) --- + if _detect_nvidia_torch(): + # torch.cuda available — determine if NVIDIA or AMD via torch.version.roc + if _is_amd_rocme(): + # torch.version.roc is not None → AMD ROCm + logger.info("XPUSYSMonitor: torch.cuda + ROCm — using AMDProvider.") + from .amd import AMDProvider + return AMDProvider(interval_ms=interval_ms) + else: + # torch.version.roc is None → NVIDIA + logger.info("XPUSYSMonitor: torch.cuda (NVIDIA) — using NvidiaProvider.") + from .nvidia import NvidiaProvider + return NvidiaProvider(interval_ms=interval_ms) + + if _detect_intel_torch(): + logger.info("XPUSYSMonitor: torch.xpu available — using IntelProvider.") + from .intel import IntelProvider + return IntelProvider(interval_ms=interval_ms) + + # --- Fallback: raw driver probing (torch not imported yet / non-std env) --- + if _detect_intel_driver(): + logger.info( + "XPUSYSMonitor: ze_loader.dll found (torch unavailable) — " + "using IntelProvider." + ) + from .intel import IntelProvider + return IntelProvider(interval_ms=interval_ms) + + if _detect_nvidia_driver(): + logger.info( + "XPUSYSMonitor: NVIDIA driver found (torch unavailable) — " + "using NvidiaProvider." + ) + from .nvidia import NvidiaProvider + return NvidiaProvider(interval_ms=interval_ms) + + # --- Last resort --- + logger.warning( + "XPUSYSMonitor: no supported GPU detected — " + "falling back to IntelProvider (limited functionality)." + ) + from .intel import IntelProvider + return IntelProvider(interval_ms=interval_ms) + + +# --------------------------------------------------------------------------- +# Primary detectors — torch-based (align with ComfyUI) +# --------------------------------------------------------------------------- + +def _detect_nvidia_torch() -> bool: + """Return True if torch has a working CUDA backend (mirrors ComfyUI).""" + try: + import torch + return torch.cuda.is_available() + except Exception: + return False + + +def _detect_intel_torch() -> bool: + """Return True if torch has a working XPU backend (mirrors ComfyUI).""" + try: + import torch + return torch.xpu.is_available() + except Exception: + return False + + +# --------------------------------------------------------------------------- +# Fallback detectors — raw driver / library probing +# Used when torch is not yet available or in non-standard environments. +# --------------------------------------------------------------------------- + +def _detect_intel_driver() -> bool: + """Return True if Intel Level Zero runtime (ze_loader.dll) is present.""" + import ctypes + try: + ctypes.WinDLL("ze_loader.dll") + return True + except OSError: + return False + + +def _detect_nvidia_driver() -> bool: + """Return True if pynvml is installed and NVIDIA driver is reachable.""" + try: + import pynvml + pynvml.nvmlInit() + count = pynvml.nvmlDeviceGetCount() + pynvml.nvmlShutdown() + return count > 0 + except Exception: + return False + + +def _is_amd_rocme() -> bool: + """ + Check if torch.cuda is backed by AMD ROCm. + + Returns True if torch.version.roc is not None (AMD ROCm PyTorch). + Returns False if torch.version.roc is None (NVIDIA or standard CUDA). + """ + try: + import torch + # torch.version.roc: True/str = AMD ROCm, None = NVIDIA/other + return torch.version.roc is not None + except Exception: + return False + + +__all__ = [ + "BaseGPUProvider", + "GPUSnapshot", + "auto_detect_provider", +] diff --git a/providers/_utils.py b/providers/_utils.py new file mode 100644 index 0000000..df782f3 --- /dev/null +++ b/providers/_utils.py @@ -0,0 +1,283 @@ +""" +providers/_utils.py — Shared CPU/RAM utility functions. + +Used by all providers to avoid duplicating Windows system-level reads. +""" +from __future__ import annotations + +import ctypes +import logging +from ctypes import wintypes +from typing import Any, Dict, Tuple + +logger = logging.getLogger("XPUSYSMonitor") + + +# --------------------------------------------------------------------------- +# Admin detection +# --------------------------------------------------------------------------- + +def _is_admin() -> bool: + try: + return bool(ctypes.windll.shell32.IsUserAnAdmin()) + except Exception: + return False + + +# --------------------------------------------------------------------------- +# CPU info (model name + thread count) +# --------------------------------------------------------------------------- + +def _get_cpu_info() -> Tuple[str, int]: + """Return (model_name, logical_thread_count). Called once at startup.""" + model = "" + try: + import winreg + key = winreg.OpenKey( + winreg.HKEY_LOCAL_MACHINE, + r"HARDWARE\DESCRIPTION\System\CentralProcessor\0", + ) + model, _ = winreg.QueryValueEx(key, "ProcessorNameString") + winreg.CloseKey(key) + model = " ".join(model.strip().split()) # collapse extra spaces + except Exception: + try: + import platform + model = platform.processor() + except Exception: + model = "Unknown CPU" + threads = 0 + try: + import psutil + threads = psutil.cpu_count(logical=True) or 0 + except Exception: + import os + threads = os.cpu_count() or 0 + return model, threads + + +# --------------------------------------------------------------------------- +# CPU / RAM polling +# --------------------------------------------------------------------------- + +def _read_cpu_ram_stats(psutil_ok: bool) -> Dict[str, Any]: + """ + Poll CPU utilisation, frequency, and RAM stats. + Returns a dict consumed by every provider's _poll(). + """ + out: Dict[str, Any] = { + "cpu_pct": 0.0, + "cpu_freq_ghz": 0.0, + "ram_pct": 0.0, + "ram_total_gb": 0.0, + "ram_used_gb": 0.0, + "ram_free_gb": 0.0, + "commit_used_gb": 0.0, + "commit_limit_gb": 0.0, + } + + if psutil_ok: + try: + import psutil + out["cpu_pct"] = psutil.cpu_percent(interval=None) + freq = psutil.cpu_freq() + if freq is not None: + out["cpu_freq_ghz"] = round(freq.current / 1000.0, 2) + mem = psutil.virtual_memory() + out["ram_pct"] = mem.percent + gb = 1024 ** 3 + out["ram_total_gb"] = round(mem.total / gb, 1) + out["ram_used_gb"] = round(mem.used / gb, 1) + out["ram_free_gb"] = round(mem.available / gb, 1) + except Exception: + logger.debug("XPUSYSMonitor: psutil read failed.", exc_info=True) + else: + # Fallback: GlobalMemoryStatusEx + try: + class MEMORYSTATUSEX(ctypes.Structure): + _fields_ = [ + ("dwLength", wintypes.DWORD), + ("dwMemoryLoad", wintypes.DWORD), + ("ullTotalPhys", ctypes.c_ulonglong), + ("ullAvailPhys", ctypes.c_ulonglong), + ("ullTotalPageFile", ctypes.c_ulonglong), + ("ullAvailPageFile", ctypes.c_ulonglong), + ("ullTotalVirtual", ctypes.c_ulonglong), + ("ullAvailVirtual", ctypes.c_ulonglong), + ("ullAvailExtendedVirtual", ctypes.c_ulonglong), + ] + state = MEMORYSTATUSEX() + state.dwLength = ctypes.sizeof(MEMORYSTATUSEX) + ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(state)) + gb = 1024 ** 3 + out["ram_pct"] = float(state.dwMemoryLoad) + out["ram_total_gb"] = round(state.ullTotalPhys / gb, 1) + out["ram_free_gb"] = round(state.ullAvailPhys / gb, 1) + out["ram_used_gb"] = round((state.ullTotalPhys - state.ullAvailPhys) / gb, 1) + out["commit_used_gb"] = round( + (state.ullTotalPageFile - state.ullAvailPageFile) / gb, 1 + ) + out["commit_limit_gb"] = round(state.ullTotalPageFile / gb, 1) + except Exception: + pass + + return out + + +# --------------------------------------------------------------------------- +# Windows Commit Charge (for RAM capsule) +# --------------------------------------------------------------------------- + +def _read_commit_charge() -> Tuple[float, float]: + """Return (commit_used_gb, commit_limit_gb) via GlobalMemoryStatusEx.""" + try: + class MEMORYSTATUSEX(ctypes.Structure): + _fields_ = [ + ("dwLength", wintypes.DWORD), + ("dwMemoryLoad", wintypes.DWORD), + ("ullTotalPhys", ctypes.c_ulonglong), + ("ullAvailPhys", ctypes.c_ulonglong), + ("ullTotalPageFile", ctypes.c_ulonglong), + ("ullAvailPageFile", ctypes.c_ulonglong), + ("ullTotalVirtual", ctypes.c_ulonglong), + ("ullAvailVirtual", ctypes.c_ulonglong), + ("ullAvailExtendedVirtual", ctypes.c_ulonglong), + ] + state = MEMORYSTATUSEX() + state.dwLength = ctypes.sizeof(MEMORYSTATUSEX) + ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(state)) + gb = 1024 ** 3 + used = (state.ullTotalPageFile - state.ullAvailPageFile) / gb + limit = state.ullTotalPageFile / gb + return round(used, 1), round(limit, 1) + except Exception: + return 0.0, 0.0 + + +# --------------------------------------------------------------------------- +# Windows Performance Data Helper (PDH) — GPU engine utilisation +# +# Uses pdh.dll via ctypes to query: +# \GPU Engine(*)\Utilization Percentage +# +# This is the same source Task Manager uses — zero pip dependencies. +# --------------------------------------------------------------------------- + +class _PdhQuery: + """Thin ctypes wrapper around PDH API for GPU engine utilisation.""" + + def __init__(self): + self._pdh = None + self._query = None + self._counters: list = [] + self._ok = False + + def init(self) -> bool: + if self._ok: + return True + try: + self._pdh = ctypes.windll.pdh # pdh.dll + # PdhOpenQueryW + self._pdh.PdhOpenQueryW.argtypes = [wintypes.LPCWSTR, wintypes.DWORD, ctypes.POINTER(ctypes.c_void_p)] + self._pdh.PdhOpenQueryW.restype = wintypes.LONG + + # PdhAddEnglishCounterW + self._pdh.PdhAddEnglishCounterW.argtypes = [ + ctypes.c_void_p, wintypes.LPCWSTR, wintypes.DWORD, + ctypes.POINTER(ctypes.c_void_p), + ] + self._pdh.PdhAddEnglishCounterW.restype = wintypes.LONG + + # PdhCollectQueryData + self._pdh.PdhCollectQueryData.argtypes = [ctypes.c_void_p] + self._pdh.PdhCollectQueryData.restype = wintypes.LONG + + # PdhGetFormattedCounterValue + self._pdh.PdhGetFormattedCounterValue.argtypes = [ + ctypes.c_void_p, wintypes.DWORD, + ctypes.POINTER(wintypes.DWORD), + ctypes.c_void_p, + ] + self._pdh.PdhGetFormattedCounterValue.restype = wintypes.LONG + + # PdhRemoveCounter / PdhCloseQuery + self._pdh.PdhRemoveCounter.argtypes = [ctypes.c_void_p] + self._pdh.PdhRemoveCounter.restype = wintypes.LONG + self._pdh.PdhCloseQuery.argtypes = [ctypes.c_void_p] + self._pdh.PdhCloseQuery.restype = wintypes.LONG + + # Open a query + self._query = ctypes.c_void_p() + ret = self._pdh.PdhOpenQueryW(None, 0, ctypes.byref(self._query)) + if ret != 0: # ERROR_SUCCESS + logger.warning(f"XPUSYSMonitor: PdhOpenQueryW failed — 0x{ret:08x}") + return False + + # Enumerate GPU engine counters + # The \GPU Engine(*)\Utilization Percentage counter path covers all GPU engines + # (3D, Compute, Copy) on both NVIDIA and AMD GPUs. + counter_path = "\\GPU Engine(*)\\Utilization Percentage" + counter_buf = ctypes.c_void_p() + ret = self._pdh.PdhAddEnglishCounterW( + self._query, counter_path, 0, ctypes.byref(counter_buf) + ) + if ret != 0: + logger.warning(f"XPUSYSMonitor: PDH GPU counter not available — 0x{ret:08x}") + self._pdh.PdhCloseQuery(self._query) + return False + + self._counters = [counter_buf] + self._ok = True + logger.info("XPUSYSMonitor: PDH GPU utilisation counters OK.") + return True + + except Exception as exc: + logger.debug(f"XPUSYSMonitor: PDH init error — {exc}") + return False + + def read_gpu_utilization(self) -> float: + """Query total GPU utilisation % across all engines.""" + if not self._ok: + return 0.0 + try: + # Collect + self._pdh.PdhCollectQueryData(self._query) + + # Read the wildcard counter — it aggregates across all GPU engines + # We use PDH_FMT_DOUBLE (0x00000200) | PDH_FMT_NOCAP100 (0x00008000) + fmt = 0x00008200 # PDH_FMT_DOUBLE | PDH_FMT_NOCAP100 + + class PDH_FMT_COUNTERVALUE_DOUBLE(ctypes.Structure): + _fields_ = [ + ("CStatus", wintypes.DWORD), + ("doubleValue", ctypes.c_double), + ] + + for counter in self._counters: + dwType = wintypes.DWORD(0) + val = PDH_FMT_COUNTERVALUE_DOUBLE() + ret = self._pdh.PdhGetFormattedCounterValue( + counter, fmt, ctypes.byref(dwType), ctypes.byref(val) + ) + if ret == 0 and val.CStatus == 0: # PDH_CSTATUS_VALID_DATA + return min(val.doubleValue, 100.0) + + return 0.0 + except Exception: + return 0.0 + + def close(self) -> None: + if self._pdh and self._query: + for c in self._counters: + self._pdh.PdhRemoveCounter(c) + self._pdh.PdhCloseQuery(self._query) + self._ok = False + + +__all__ = [ + "_is_admin", + "_get_cpu_info", + "_read_cpu_ram_stats", + "_read_commit_charge", + "_PdhQuery", +] diff --git a/providers/amd.py b/providers/amd.py new file mode 100644 index 0000000..6f15d8c --- /dev/null +++ b/providers/amd.py @@ -0,0 +1,241 @@ +""" +providers/amd.py — AMD GPU hardware provider for Windows ROCm. + +VRAM free/total : torch.cuda.mem_get_info(device) (driver-level, ROCm 6+) +PyTorch stats : torch.cuda.memory_allocated / memory_reserved +GPU load : Windows PDH API (\\GPU Engine(*)\\Utilization Percentage) +GPU frequency : Unavailable without vendor API -> 0 +GPU temperature : Unavailable without vendor API -> -1 +Power : Unavailable without vendor API -> -1 / False + +No dependency on rocm_smi_lib — works with native Windows ROCm PyTorch. +GPU utilisation via PDH (pdh.dll, zero pip deps). Temperature/freq/power +return unavailable sentinels on Windows where no vendor driver API exists. +""" + +import logging +import sys +from typing import Tuple + +from .base import BaseGPUProvider, GPUSnapshot +from ._utils import _get_cpu_info, _read_cpu_ram_stats, _PdhQuery, _is_admin + +logger = logging.getLogger("XPUSYSMonitor") + + +# --------------------------------------------------------------------------- +# AMDProvider +# --------------------------------------------------------------------------- + +class AMDProvider(BaseGPUProvider): + """ + Hardware provider for AMD GPUs on Windows ROCm. + + Uses torch.cuda for VRAM and PyTorch allocator stats. + Uses Windows PDH API for GPU engine utilisation. + Does NOT require rocm_smi_lib. + + Temperature, core clock, and power return unavailable sentinels + since no standard Python-accessible driver API exists on Windows. + """ + + GPU_VENDOR = "amd" + + def __init__(self, interval_ms: int = 1000): + self._torch_ok = False + self._psutil_ok = False + self._device_index = 0 + self._is_admin = _is_admin() + self._cpu_model = "" + self._cpu_threads = 0 + + self._check_torch() + self._check_psutil() + + # Windows PDH — GPU engine utilisation (graceful if unavailable) + self._pdh = _PdhQuery() + self._pdh_ok = self._pdh.init() + + # BaseGPUProvider.__init__ starts the polling thread — call last + super().__init__(interval_ms=interval_ms) + + logger.info( + f"XPUSYSMonitor: AMDProvider started " + f"(torch={self._torch_ok}, pdh={self._pdh_ok})" + ) + + # ------------------------------------------------------------------ + # Initialisation + # ------------------------------------------------------------------ + + def _check_torch(self) -> None: + """Check if torch.cuda is available (ROCm PyTorch on Windows).""" + try: + import torch + if torch.cuda.is_available(): + self._torch_ok = True + logger.info( + f"XPUSYSMonitor: torch.cuda OK (AMD ROCm), " + f"device count={torch.cuda.device_count()}" + ) + else: + logger.warning("XPUSYSMonitor: torch.cuda not available.") + except Exception as exc: + logger.warning(f"XPUSYSMonitor: torch import error — {exc}") + + def _check_psutil(self) -> None: + try: + import psutil + psutil.cpu_percent(interval=None) + self._psutil_ok = True + self._cpu_model, self._cpu_threads = _get_cpu_info() + logger.info( + f"XPUSYSMonitor: psutil OK — CPU={self._cpu_model!r}, " + f"threads={self._cpu_threads}" + ) + except Exception as exc: + logger.warning(f"XPUSYSMonitor: psutil not available — {exc}") + + # ------------------------------------------------------------------ + # Hardware reads + # ------------------------------------------------------------------ + + def _read_device_name(self) -> str: + """Return the GPU model name via torch.cuda.""" + if self._torch_ok: + try: + import torch + return torch.cuda.get_device_name(self._device_index) + except Exception: + pass + return "AMD GPU (ROCm)" + + def _read_vram(self) -> Tuple[float, float, float]: + """ + Return (free_gb, total_gb, driver_used_gb) via torch.cuda.mem_get_info. + + mem_get_info() returns (free_bytes, total_bytes) from the driver, + which works on ROCm 6+ PyTorch on Windows. + """ + if self._torch_ok: + try: + import torch + free_bytes, total_bytes = torch.cuda.mem_get_info(self._device_index) + gb = 1024 ** 3 + free_gb = free_bytes / gb + total_gb = total_bytes / gb + used_gb = max(0.0, total_gb - free_gb) + return free_gb, total_gb, used_gb + except Exception: + # Fallback: total from device properties + try: + import torch + total_gb = torch.cuda.get_device_properties(self._device_index).total_memory / (1024 ** 3) + return 0.0, total_gb, 0.0 + except Exception: + pass + return 0.0, 0.0, 0.0 + + def _read_torch_stats(self) -> Tuple[float, float]: + """Return (allocated_gb, reserved_gb) from torch.cuda allocator.""" + if not self._torch_ok: + return 0.0, 0.0 + try: + import torch + idx = self._device_index + gb = 1024 ** 3 + return ( + torch.cuda.memory_allocated(idx) / gb, + torch.cuda.memory_reserved(idx) / gb, + ) + except Exception: + return 0.0, 0.0 + + def _read_gpu_load(self) -> float: + """ + Return GPU utilisation % via Windows PDH API. + + Falls back to 0 if PDH is unavailable (non-Windows, or + counters not installed by the AMD driver). + """ + if self._pdh_ok: + return self._pdh.read_gpu_utilization() + return 0.0 + + def _read_gpu_freq_mhz(self) -> float: + """ + GPU core frequency in MHz. + + Unavailable on Windows without vendor driver API. + """ + return 0.0 + + def _read_gpu_temp_c(self) -> float: + """ + GPU core temperature in C. + + Unavailable on Windows without vendor driver API. + """ + return -1.0 + + def _read_power(self) -> Tuple[float, float, bool]: + """ + Return (power_w, tgp_w, power_available). + + Unavailable on Windows without vendor driver API. + """ + return -1.0, 0.0, False + + # ------------------------------------------------------------------ + # Poll — called by BaseGPUProvider._loop() every interval + # ------------------------------------------------------------------ + + def _poll(self) -> None: + """Collect all hardware metrics and push a fresh GPUSnapshot.""" + snap = GPUSnapshot(gpu_vendor=self.GPU_VENDOR) + snap.is_admin = self._is_admin + + if not self._torch_ok: + snap.error = "AMD ROCm (torch.cuda) unavailable" + else: + try: + snap.device_name = self._read_device_name() + + # VRAM — driver level via torch.cuda.mem_get_info + free_gb, total_gb, driver_used_gb = self._read_vram() + snap.vram_total_gb = total_gb + snap.vram_free_gb = free_gb + snap.vram_driver_used_gb = driver_used_gb + + # PyTorch allocator stats + snap.vram_allocated_gb, snap.vram_reserved_gb = self._read_torch_stats() + + # GPU metrics + snap.gpu_load_pct = self._read_gpu_load() + snap.gpu_freq_mhz = self._read_gpu_freq_mhz() + snap.gpu_temp_c = self._read_gpu_temp_c() + + # Power + snap.power_w, snap.tgp_w, snap.power_available = self._read_power() + + except Exception as exc: + logger.debug(f"XPUSYSMonitor: AMDProvider poll error — {exc}") + snap.error = str(exc) + + # CPU / RAM — always collected regardless of GPU state + sys_stats = _read_cpu_ram_stats(self._psutil_ok) + snap.cpu_pct = sys_stats.get("cpu_pct", 0.0) + snap.cpu_freq_ghz = sys_stats.get("cpu_freq_ghz", 0.0) + snap.cpu_model = self._cpu_model + snap.cpu_threads = self._cpu_threads + snap.ram_pct = sys_stats.get("ram_pct", 0.0) + snap.ram_total_gb = sys_stats.get("ram_total_gb", 0.0) + snap.ram_used_gb = sys_stats.get("ram_used_gb", 0.0) + snap.ram_free_gb = sys_stats.get("ram_free_gb", 0.0) + snap.commit_used_gb = sys_stats.get("commit_used_gb", 0.0) + snap.commit_limit_gb = sys_stats.get("commit_limit_gb", 0.0) + + self._update_snapshot(snap) + + +__all__ = ["AMDProvider"] diff --git a/providers/base.py b/providers/base.py new file mode 100644 index 0000000..c968ef9 --- /dev/null +++ b/providers/base.py @@ -0,0 +1,134 @@ +""" +providers/base.py — Abstract base class and shared data contract. + +All GPU provider implementations must subclass BaseGPUProvider and +implement get_snapshot(). The GPUSnapshot dataclass is the single +contract between the hardware layer and every consumer (server, +predictor, frontend). +""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass, field +from typing import Optional + + +# --------------------------------------------------------------------------- +# Shared data snapshot — the contract between providers and consumers +# --------------------------------------------------------------------------- + +@dataclass +class GPUSnapshot: + # --- VRAM (all in GB) --- + vram_total_gb: float = 0.0 # driver total + vram_free_gb: float = 0.0 # driver free + vram_driver_used_gb: float = 0.0 # total - free (all consumers) + vram_allocated_gb: float = 0.0 # torch allocated (ComfyUI workload) + vram_reserved_gb: float = 0.0 # torch cached / reserved pool + + # --- GPU --- + gpu_load_pct: float = 0.0 + gpu_freq_mhz: float = 0.0 # current GPU clock in MHz (0 = unavailable) + gpu_temp_c: float = -1.0 # GPU core temp in °C (-1 = unavailable) + + # --- Power --- + power_w: float = -1.0 # -1 = unavailable + power_available: bool = False + tgp_w: float = 0.0 # sustained TGP limit in W (0 = unknown) + device_name: str = "" + + # --- CPU --- + cpu_pct: float = 0.0 + cpu_freq_ghz: float = 0.0 + cpu_model: str = "" + cpu_threads: int = 0 + + # --- RAM --- + ram_pct: float = 0.0 + ram_total_gb: float = 0.0 + ram_used_gb: float = 0.0 + ram_free_gb: float = 0.0 + # Windows Commit Charge (≈ Task Manager "已提交") + commit_used_gb: float = 0.0 # CommitTotal = ullTotalPageFile - ullAvailPageFile + commit_limit_gb: float = 0.0 # CommitLimit = ullTotalPageFile + + # --- Meta --- + is_admin: bool = False + gpu_vendor: str = "" # "intel" | "nvidia" | "amd" | "unknown" + error: Optional[str] = None + + +# --------------------------------------------------------------------------- +# Abstract base provider +# --------------------------------------------------------------------------- + +class BaseGPUProvider: + """ + Abstract base class for all GPU hardware providers. + + Subclasses implement _poll() to fill a GPUSnapshot and call + _update_snapshot(snap) when done. The base class handles + thread-safe snapshot storage and the polling loop lifecycle. + + Consumers always call get_snapshot() — they never see the + concrete provider type. + """ + + # Subclasses set this to identify their vendor in GPUSnapshot.gpu_vendor + GPU_VENDOR: str = "unknown" + + def __init__(self, interval_ms: int = 1000): + self._interval = max(100, interval_ms) / 1000.0 + self._snapshot = GPUSnapshot(gpu_vendor=self.GPU_VENDOR) + self._lock = threading.Lock() + self._stop = threading.Event() + + self._thread = threading.Thread( + target=self._loop, daemon=True, name=f"XPUSYSMonitor-{self.GPU_VENDOR}" + ) + self._thread.start() + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def get_snapshot(self) -> GPUSnapshot: + """Return the latest hardware snapshot. Thread-safe.""" + with self._lock: + return self._snapshot + + def set_interval(self, ms: int) -> None: + """Adjust polling interval at runtime.""" + self._interval = max(100, ms) / 1000.0 + + def stop(self) -> None: + """Stop the background polling thread.""" + self._stop.set() + self._thread.join(timeout=5) + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _update_snapshot(self, snap: GPUSnapshot) -> None: + """Atomically replace the stored snapshot. Called from _poll().""" + with self._lock: + self._snapshot = snap + + def _loop(self) -> None: + """Background polling loop — calls _poll() every interval.""" + while not self._stop.is_set(): + try: + self._poll() + except Exception as exc: + snap = GPUSnapshot(gpu_vendor=self.GPU_VENDOR, error=str(exc)) + self._update_snapshot(snap) + self._stop.wait(self._interval) + + def _poll(self) -> None: + """ + Override in subclass: collect hardware data and call + _update_snapshot(snap) with a freshly built GPUSnapshot. + """ + raise NotImplementedError diff --git a/providers/nvidia.py b/providers/nvidia.py new file mode 100644 index 0000000..dc4744c --- /dev/null +++ b/providers/nvidia.py @@ -0,0 +1,261 @@ +""" +providers/nvidia.py — NVIDIA GPU hardware provider. + +VRAM free/total : pynvml — nvmlDeviceGetMemoryInfo +PyTorch stats : torch.cuda.memory_allocated / memory_reserved +GPU load : pynvml — nvmlDeviceGetUtilizationRates +GPU frequency : pynvml — nvmlDeviceGetClockInfo (GRAPHICS clock) +GPU temperature : pynvml — nvmlDeviceGetTemperature (no admin required) +Power / TGP : pynvml — nvmlDeviceGetPowerUsage / GetEnforcedPowerLimit + (no admin required on most NVIDIA drivers) + +Dependency: pip install pynvml + Included in nvidia-ml-py, which ships with most NVIDIA CUDA toolkits. + If pynvml is not installed or no NVIDIA driver is present, this provider + will raise ImportError / NVMLError at init time — the factory catches it. +""" + +import logging +import os +from typing import Tuple + +from .base import BaseGPUProvider, GPUSnapshot +from ._utils import _get_cpu_info, _read_cpu_ram_stats, _read_commit_charge, _is_admin + +logger = logging.getLogger("XPUSYSMonitor") + + +# --------------------------------------------------------------------------- +# NvidiaProvider +# --------------------------------------------------------------------------- + +class NvidiaProvider(BaseGPUProvider): + """ + Hardware provider for NVIDIA GPUs. + + Uses pynvml (nvidia-ml-py) for all GPU metrics. Unlike the Intel + provider, temperature and power do NOT require administrator privileges + on NVIDIA drivers. + + Only the first GPU (index 0) is monitored — multi-GPU support can be + added later if needed. + """ + + GPU_VENDOR = "nvidia" + + def __init__(self, interval_ms: int = 1000): + self._nvml_ok = False + self._handle = None # nvml device handle for GPU 0 + self._torch_ok = False + self._psutil_ok = False + self._device_index = 0 + self._is_admin = _is_admin() + self._cpu_model = "" + self._cpu_threads = 0 + + self._init_nvml() + self._check_torch() + self._check_psutil() + + # BaseGPUProvider.__init__ starts the polling thread — call last + super().__init__(interval_ms=interval_ms) + + logger.info( + f"XPUSYSMonitor: NvidiaProvider started " + f"(nvml={self._nvml_ok}, torch={self._torch_ok})" + ) + + # ------------------------------------------------------------------ + # Initialisation + # ------------------------------------------------------------------ + + def _init_nvml(self) -> None: + """Initialise pynvml and grab the device handle for GPU 0.""" + try: + import pynvml + pynvml.nvmlInit() + count = pynvml.nvmlDeviceGetCount() + if count == 0: + logger.warning("XPUSYSMonitor: pynvml — no NVIDIA devices found.") + return + self._handle = pynvml.nvmlDeviceGetHandleByIndex(self._device_index) + self._nvml_ok = True + name = pynvml.nvmlDeviceGetName(self._handle) + # pynvml may return bytes on older versions + if isinstance(name, bytes): + name = name.decode("utf-8", errors="replace") + logger.info(f"XPUSYSMonitor: pynvml OK — device[0] = {name!r}") + except ImportError: + logger.warning( + "XPUSYSMonitor: pynvml not installed — " + "run `pip install pynvml` to enable NVIDIA support." + ) + except Exception as exc: + logger.warning(f"XPUSYSMonitor: pynvml init error — {exc}") + + def _check_torch(self) -> None: + try: + import torch + if torch.cuda.is_available(): + self._torch_ok = True + logger.info( + f"XPUSYSMonitor: torch.cuda OK, " + f"device count={torch.cuda.device_count()}" + ) + else: + logger.warning("XPUSYSMonitor: torch.cuda not available.") + except Exception as exc: + logger.warning(f"XPUSYSMonitor: torch import error — {exc}") + + def _check_psutil(self) -> None: + try: + import psutil + psutil.cpu_percent(interval=None) # baseline call + self._psutil_ok = True + self._cpu_model, self._cpu_threads = _get_cpu_info() + logger.info( + f"XPUSYSMonitor: psutil OK — CPU={self._cpu_model!r}, " + f"threads={self._cpu_threads}" + ) + except Exception as exc: + logger.warning(f"XPUSYSMonitor: psutil not available — {exc}") + + # ------------------------------------------------------------------ + # Hardware reads + # ------------------------------------------------------------------ + + def _read_device_name(self) -> str: + try: + import pynvml + name = pynvml.nvmlDeviceGetName(self._handle) + if isinstance(name, bytes): + name = name.decode("utf-8", errors="replace") + return name + except Exception: + return "NVIDIA GPU" + + def _read_vram(self) -> Tuple[float, float, float]: + """Return (free_gb, total_gb, driver_used_gb).""" + try: + import pynvml + info = pynvml.nvmlDeviceGetMemoryInfo(self._handle) + gb = 1024 ** 3 + total = info.total / gb + free = info.free / gb + used = info.used / gb + return free, total, used + except Exception: + return 0.0, 0.0, 0.0 + + def _read_torch_stats(self) -> Tuple[float, float]: + """Return (allocated_gb, reserved_gb) from torch.cuda allocator.""" + if not self._torch_ok: + return 0.0, 0.0 + try: + import torch + idx = self._device_index + gb = 1024 ** 3 + return ( + torch.cuda.memory_allocated(idx) / gb, + torch.cuda.memory_reserved(idx) / gb, + ) + except Exception: + return 0.0, 0.0 + + def _read_gpu_load(self) -> float: + """Return GPU utilisation % via nvmlDeviceGetUtilizationRates.""" + try: + import pynvml + rates = pynvml.nvmlDeviceGetUtilizationRates(self._handle) + return float(rates.gpu) + except Exception: + return 0.0 + + def _read_gpu_freq_mhz(self) -> float: + """Return current GPU graphics clock in MHz.""" + try: + import pynvml + # NVML_CLOCK_GRAPHICS = 0 + return float(pynvml.nvmlDeviceGetClockInfo(self._handle, 0)) + except Exception: + return 0.0 + + def _read_gpu_temp_c(self) -> float: + """Return GPU temperature in °C. No admin required on NVIDIA.""" + try: + import pynvml + # NVML_TEMPERATURE_GPU = 0 + return float(pynvml.nvmlDeviceGetTemperature(self._handle, 0)) + except Exception: + return -1.0 + + def _read_power(self) -> Tuple[float, float, bool]: + """Return (power_w, tgp_w, power_available).""" + try: + import pynvml + # nvmlDeviceGetPowerUsage returns milliwatts + power_mw = pynvml.nvmlDeviceGetPowerUsage(self._handle) + power_w = power_mw / 1000.0 + + # Enforced power limit (TGP) in milliwatts + try: + tgp_mw = pynvml.nvmlDeviceGetEnforcedPowerLimit(self._handle) + tgp_w = tgp_mw / 1000.0 + except Exception: + tgp_w = 0.0 + + return power_w, tgp_w, True + except Exception: + return -1.0, 0.0, False + + # ------------------------------------------------------------------ + # Poll — called by BaseGPUProvider._loop() every interval + # ------------------------------------------------------------------ + + def _poll(self) -> None: + """Collect all hardware metrics and push a fresh GPUSnapshot.""" + snap = GPUSnapshot(gpu_vendor=self.GPU_VENDOR) + snap.is_admin = self._is_admin + + if not self._nvml_ok or self._handle is None: + # nvml unavailable — still collect CPU/RAM + snap.error = "pynvml unavailable" + else: + try: + snap.device_name = self._read_device_name() + + # VRAM + free_gb, total_gb, driver_used_gb = self._read_vram() + snap.vram_total_gb = total_gb + snap.vram_free_gb = free_gb + snap.vram_driver_used_gb = driver_used_gb + + # torch allocator stats + snap.vram_allocated_gb, snap.vram_reserved_gb = self._read_torch_stats() + + # GPU metrics + snap.gpu_load_pct = self._read_gpu_load() + snap.gpu_freq_mhz = self._read_gpu_freq_mhz() + snap.gpu_temp_c = self._read_gpu_temp_c() + + # Power + snap.power_w, snap.tgp_w, snap.power_available = self._read_power() + + except Exception as exc: + logger.debug(f"XPUSYSMonitor: NvidiaProvider poll error — {exc}") + snap.error = str(exc) + + # CPU / RAM — always collected regardless of GPU state + sys = _read_cpu_ram_stats(self._psutil_ok) + snap.cpu_pct = sys.get("cpu_pct", 0.0) + snap.cpu_freq_ghz = sys.get("cpu_freq_ghz", 0.0) + snap.cpu_model = self._cpu_model + snap.cpu_threads = self._cpu_threads + snap.ram_pct = sys.get("ram_pct", 0.0) + snap.ram_total_gb = sys.get("ram_total_gb", 0.0) + snap.ram_used_gb = sys.get("ram_used_gb", 0.0) + snap.ram_free_gb = sys.get("ram_free_gb", 0.0) + snap.commit_used_gb = sys.get("commit_used_gb", 0.0) + snap.commit_limit_gb = sys.get("commit_limit_gb", 0.0) + + self._update_snapshot(snap) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..638551c --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,14 @@ +[project] +name = "ComfyUI-XPUSYS-Monitor-AMD" +version = "1.0.3" +description = "AMD/ROCm Windows-native hardware monitor for ComfyUI. Uses torch.cuda.mem_get_info() for VRAM, PDH for GPU load, no rocm_smi_lib dependency." +license = { text = "MIT" } +dependencies = ["psutil", "pynvml"] + +[project.urls] +Repository = "https://github.com/forkless/ComfyUI-XPUSYS-Monitor-AMD" + +[tool.comfy] +PublisherId = "forkless" +DisplayName = "ComfyUI-XPUSYS-Monitor-AMD" +Icon = "" diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..72c7b1c --- /dev/null +++ b/requirements.txt @@ -0,0 +1,26 @@ +# CPU / RAM monitoring — required by all providers +psutil + +# NVIDIA GPU monitoring — required by NvidiaProvider only +# Safe to install on Intel/AMD machines; unused if no NVIDIA driver is present +pynvml + +# AMD GPU monitoring (Linux) — commented out for Windows-native ROCm +# rocm_smi_lib is Linux-only; Windows ROCm uses torch.cuda APIs instead. +# Install only if you are on Linux with ROCm and want driver-level GPU load: +# pip install rocm_smi_lib +# rocm_smi_lib + +# --------------------------------------------------------------------------- +# The following are provided by ComfyUI itself — DO NOT add them here. +# Listed for reference only so manual environment setups know the minimums. +# --------------------------------------------------------------------------- +# +# torch >= 2.5 +# torch.xpu (Intel Arc) was officially introduced in PyTorch 2.5. +# torch.cuda (NVIDIA/AMD ROCm) has been stable since PyTorch 1.x. +# Install the ROCm-enabled build from: https://pytorch.org/get-started/locally/ +# +# aiohttp >= 3.8 +# Used for HTTP routes and WebSocket broadcast. +# ComfyUI already depends on aiohttp; no separate installation needed. diff --git a/xpu_server.py b/xpu_server.py new file mode 100644 index 0000000..72dce14 --- /dev/null +++ b/xpu_server.py @@ -0,0 +1,155 @@ +""" +xpu_server.py — WebSocket / HTTP endpoint for GPU metrics. + +Registers a /xpusys/stats route on ComfyUI's PromptServer that returns +the latest GPUSnapshot as JSON. A background task also broadcasts data +to all connected WebSocket clients at the configured interval. + +This module is provider-agnostic: it works with any BaseGPUProvider +implementation (Intel, NVIDIA, AMD, ...). +""" + +import asyncio +import json +import logging +import os +from aiohttp import web + +try: + import folder_paths as _fp +except ImportError: + _fp = None + +logger = logging.getLogger("XPUSYSMonitor") + +# Will be set by __init__.py after creating the provider +_provider = None + + +def set_provider(provider): + global _provider + _provider = provider + + +def _snapshot_to_dict(snap) -> dict: + return { + # VRAM breakdown + "vram_total_gb": round(snap.vram_total_gb, 2), + "vram_free_gb": round(snap.vram_free_gb, 2), + "vram_driver_used_gb": round(snap.vram_driver_used_gb, 2), + "vram_allocated_gb": round(snap.vram_allocated_gb, 2), + "vram_reserved_gb": round(snap.vram_reserved_gb, 2), + # GPU + "gpu_load_pct": round(snap.gpu_load_pct, 1), + "gpu_freq_mhz": round(snap.gpu_freq_mhz, 0), + "gpu_temp_c": round(snap.gpu_temp_c, 1), + # Power + "power_w": round(snap.power_w, 1), + "power_available": snap.power_available, + "tgp_w": round(snap.tgp_w, 1), + "device_name": snap.device_name, + # CPU + "cpu_pct": round(snap.cpu_pct, 1), + "cpu_freq_ghz": round(snap.cpu_freq_ghz, 2), + "cpu_model": snap.cpu_model, + "cpu_threads": snap.cpu_threads, + # RAM + "ram_pct": round(snap.ram_pct, 1), + "ram_total_gb": round(snap.ram_total_gb, 2), + "ram_used_gb": round(snap.ram_used_gb, 2), + "ram_free_gb": round(snap.ram_free_gb, 2), + "commit_used_gb": round(snap.commit_used_gb, 2), + "commit_limit_gb": round(snap.commit_limit_gb, 2), + # Meta + "is_admin": snap.is_admin, + "gpu_vendor": snap.gpu_vendor, + "error": snap.error, + } + + +def register_routes(server): + """Call this with the PromptServer instance to attach our HTTP route.""" + + @server.routes.get("/xpusys/stats") + async def get_stats(request): + if _provider is None: + return web.json_response({"error": "provider not ready"}, status=503) + snap = _provider.get_snapshot() + return web.json_response(_snapshot_to_dict(snap)) + + @server.routes.post("/xpusys/model_sizes") + async def get_model_sizes(request): + """ + Accept a list of {name, type} model descriptors from the frontend, + resolve each to a physical file, and return {name, size} in GB. + Route name is intentionally distinct from /vram_predict/* used by + the standalone ComfyUI-Vram-Predictor plugin. + """ + if _fp is None: + return web.json_response({"models": []}) + try: + data = await request.json() + items = data.get("models", []) + _ALLOWED = {".safetensors", ".gguf", ".ckpt", ".pt", ".pth", ".bin", ".onnx", ".pkl"} + _SEARCH = ["checkpoints", "vae", "loras", "controlnet", "clip", + "upscale_models", "unet", "diffusion_models", + "ultralytics", "annotator", "bbox", "onnx", + "mmaudio", "audio", "rife", "vfi"] + results = [] + for m in items: + name = m.get("name", "") + model_path = m.get("path", "") + if not name: + continue + if os.path.splitext(name)[1].lower() not in _ALLOWED: + continue + + # 在所有搜索目录中按优先级查找 + path = None + for folder in _SEARCH: + # 优先使用完整路径查找（支持子文件夹） + if model_path: + p = _fp.get_full_path(folder, model_path) + if p and os.path.isfile(p): + path = p + break + + # 兜底：用文件名在所有目录下递归搜索 + if not path: + for folder in _SEARCH: + base_folder = _fp.get_full_path(folder, "") + if base_folder and os.path.isdir(base_folder): + for root, _, files in os.walk(base_folder): + if name in files: + path = os.path.join(root, name) + break + if path: + break + + if path and os.path.isfile(path): + size_gb = os.path.getsize(path) / (1024 ** 3) + if size_gb > 0.001: + results.append({"name": name, "size": round(size_gb, 2)}) + return web.json_response({"models": results}) + except Exception as exc: + logger.debug(f"XPUSYSMonitor: model_sizes error — {exc}") + return web.json_response({"models": []}) + + logger.info("XPUSYSMonitor: HTTP routes /xpusys/stats and /xpusys/model_sizes registered.") + + +async def broadcast_loop(server, interval_s: float = 1.0): + """ + Continuously broadcast XPU stats to all WebSocket clients via + ComfyUI's built-in send_json helper. + """ + while True: + await asyncio.sleep(interval_s) + if _provider is None: + continue + try: + snap = _provider.get_snapshot() + data = _snapshot_to_dict(snap) + await server.send_json("xpusys_stats", data) + except Exception as exc: + logger.debug(f"XPUSYSMonitor: broadcast error — {exc}")