robomello
87b0c6fd9c
Addresses #1 (NVML unavailable on Windows: NVMLError_LibraryNotFound). - Detect the dual-install conflict where both 'pynvml' and 'nvidia-ml-py' ship a top-level pynvml module and shadow each other; surface a clear "uninstall pynvml" hint in nvml_error. - On Windows, register the standard NVIDIA DLL directories (System32 DCH path + legacy NVSMI) and retry nvmlInit once before failing. - Replace the raw nvml_error string with an actionable diagnosis that probes for nvml.dll and explains the likely cause. - Document the conflict and the Windows LibraryNotFound case in the README troubleshooting section; add 0.1.1 changelog entry. The Linux/Docker init path is unchanged: _try_nvml_init re-raises immediately on non-Windows, so a healthy Linux setup behaves byte-for-byte as before (verified: provider NVIDIA, 2 GPUs, nvml_error None).
254 lines
8.0 KiB
Python
254 lines
8.0 KiB
Python
"""NVIDIA + system telemetry collector for ComfyUI-NVML-Monitor."""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
import time
|
|
import threading
|
|
import warnings
|
|
from typing import Any
|
|
|
|
import psutil
|
|
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore", FutureWarning)
|
|
import pynvml
|
|
|
|
|
|
_LOCK = threading.Lock()
|
|
_CACHE: dict[str, Any] = {"ts": 0.0, "data": None}
|
|
_CACHE_TTL = 0.5
|
|
|
|
_NVML_READY = False
|
|
_NVML_ERROR: str | None = None
|
|
_GPU_HANDLES: list[Any] = []
|
|
|
|
|
|
def _dist_installed(dist: str) -> bool:
|
|
"""True if a PyPI distribution of this exact name is installed."""
|
|
try:
|
|
import importlib.metadata as _md
|
|
_md.version(dist)
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def _windows_nvml_dirs() -> list[str]:
|
|
"""Known locations of nvml.dll on Windows, most current first."""
|
|
windir = os.environ.get("WINDIR", r"C:\Windows")
|
|
candidates = [
|
|
os.path.join(windir, "System32"), # DCH drivers (current)
|
|
os.path.join(os.environ.get("ProgramW6432", r"C:\Program Files"),
|
|
"NVIDIA Corporation", "NVSMI"), # legacy
|
|
os.path.join(os.environ.get("ProgramFiles", r"C:\Program Files"),
|
|
"NVIDIA Corporation", "NVSMI"),
|
|
]
|
|
seen: set[str] = set()
|
|
ordered: list[str] = []
|
|
for d in candidates:
|
|
if d not in seen:
|
|
seen.add(d)
|
|
ordered.append(d)
|
|
return ordered
|
|
|
|
|
|
def _diagnose_nvml_failure(exc: Exception) -> str:
|
|
"""Turn a raw NVML exception into an actionable nvml_error string."""
|
|
base = f"{type(exc).__name__}: {exc}"
|
|
hints: list[str] = []
|
|
|
|
# The classic trap: both 'pynvml' and 'nvidia-ml-py' provide a top-level
|
|
# `pynvml` module. Installed together they shadow each other and NVML can
|
|
# fail to load. NVIDIA ships the bindings as 'nvidia-ml-py'; drop 'pynvml'.
|
|
if _dist_installed("pynvml") and _dist_installed("nvidia-ml-py"):
|
|
hints.append(
|
|
"Both 'pynvml' and 'nvidia-ml-py' are installed and conflict. "
|
|
"Run: pip uninstall -y pynvml (keep nvidia-ml-py), then restart ComfyUI."
|
|
)
|
|
|
|
signature = f"{type(exc).__name__} {exc}"
|
|
if sys.platform.startswith("win") and "LibraryNotFound" in signature:
|
|
found = next(
|
|
(os.path.join(d, "nvml.dll") for d in _windows_nvml_dirs()
|
|
if os.path.isfile(os.path.join(d, "nvml.dll"))),
|
|
None,
|
|
)
|
|
if found:
|
|
hints.append(f"nvml.dll is present at {found}, so this is most likely the "
|
|
"package conflict rather than a missing driver.")
|
|
else:
|
|
hints.append("nvml.dll was not found in System32 or the NVSMI folder. "
|
|
"Install or repair a recent NVIDIA driver.")
|
|
|
|
return base + (" | " + " ".join(hints) if hints else "")
|
|
|
|
|
|
def _try_nvml_init() -> None:
|
|
"""Call nvmlInit. On Windows, retry once after registering NVIDIA DLL dirs.
|
|
|
|
On Linux the original error is re-raised immediately, so the working path is
|
|
byte-for-byte unchanged; only a Windows LibraryNotFound gets a second attempt.
|
|
"""
|
|
try:
|
|
pynvml.nvmlInit()
|
|
return
|
|
except Exception:
|
|
if not (sys.platform.startswith("win") and hasattr(os, "add_dll_directory")):
|
|
raise
|
|
for d in _windows_nvml_dirs():
|
|
if os.path.isdir(d):
|
|
try:
|
|
os.add_dll_directory(d)
|
|
except OSError:
|
|
pass
|
|
pynvml.nvmlInit()
|
|
|
|
|
|
def _init_nvml() -> None:
|
|
global _NVML_READY, _NVML_ERROR, _GPU_HANDLES
|
|
if _NVML_READY:
|
|
return
|
|
try:
|
|
_try_nvml_init()
|
|
count = pynvml.nvmlDeviceGetCount()
|
|
_GPU_HANDLES = [pynvml.nvmlDeviceGetHandleByIndex(i) for i in range(count)]
|
|
_NVML_READY = True
|
|
_NVML_ERROR = None
|
|
except Exception as exc:
|
|
_NVML_READY = False
|
|
_NVML_ERROR = _diagnose_nvml_failure(exc)
|
|
_GPU_HANDLES = []
|
|
|
|
|
|
def _as_str(value: Any) -> str:
|
|
return value.decode() if isinstance(value, bytes) else str(value)
|
|
|
|
|
|
def _safe(fn, default=None):
|
|
try:
|
|
return fn()
|
|
except Exception:
|
|
return default
|
|
|
|
|
|
def _gpu_processes(handle) -> list[dict]:
|
|
try:
|
|
procs = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
|
|
except Exception:
|
|
return []
|
|
out: list[dict] = []
|
|
for p in procs[:10]:
|
|
name = "unknown"
|
|
try:
|
|
name = psutil.Process(p.pid).name()
|
|
except Exception:
|
|
pass
|
|
mem_mb = (p.usedGpuMemory or 0) / (1024 * 1024) if p.usedGpuMemory else 0
|
|
out.append({"pid": p.pid, "name": name, "mem_mb": round(mem_mb, 1)})
|
|
return out
|
|
|
|
|
|
def _gpu_stats(index: int, handle) -> dict:
|
|
name = _as_str(_safe(lambda: pynvml.nvmlDeviceGetName(handle), "Unknown GPU"))
|
|
mem = _safe(lambda: pynvml.nvmlDeviceGetMemoryInfo(handle))
|
|
util = _safe(lambda: pynvml.nvmlDeviceGetUtilizationRates(handle))
|
|
temp = _safe(lambda: pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU))
|
|
power_draw = _safe(lambda: pynvml.nvmlDeviceGetPowerUsage(handle) / 1000.0)
|
|
power_limit = _safe(lambda: pynvml.nvmlDeviceGetPowerManagementLimit(handle) / 1000.0)
|
|
clk_graphics = _safe(lambda: pynvml.nvmlDeviceGetClockInfo(handle, pynvml.NVML_CLOCK_GRAPHICS))
|
|
clk_memory = _safe(lambda: pynvml.nvmlDeviceGetClockInfo(handle, pynvml.NVML_CLOCK_MEM))
|
|
fan = _safe(lambda: pynvml.nvmlDeviceGetFanSpeed(handle))
|
|
|
|
GB = 1024 ** 3
|
|
vram = {
|
|
"used_gb": round(mem.used / GB, 2) if mem else 0,
|
|
"total_gb": round(mem.total / GB, 2) if mem else 0,
|
|
"free_gb": round(mem.free / GB, 2) if mem else 0,
|
|
"percent": round((mem.used / mem.total) * 100, 1) if mem and mem.total else 0,
|
|
}
|
|
|
|
procs = _gpu_processes(handle)
|
|
visible_mb = sum(p["mem_mb"] for p in procs)
|
|
total_used_mb = (mem.used / (1024 * 1024)) if mem else 0
|
|
external_mb = max(0, total_used_mb - visible_mb)
|
|
|
|
return {
|
|
"index": index,
|
|
"name": name,
|
|
"vram": vram,
|
|
"util": {
|
|
"gpu": util.gpu if util else 0,
|
|
"memory": util.memory if util else 0,
|
|
},
|
|
"temp_c": temp,
|
|
"power": {
|
|
"draw_w": round(power_draw, 1) if power_draw is not None else None,
|
|
"limit_w": round(power_limit, 1) if power_limit is not None else None,
|
|
},
|
|
"clocks": {
|
|
"graphics_mhz": clk_graphics,
|
|
"memory_mhz": clk_memory,
|
|
},
|
|
"fan_percent": fan,
|
|
"processes": procs,
|
|
"external_mb": round(external_mb, 1),
|
|
}
|
|
|
|
|
|
def _system_stats() -> dict:
|
|
vm = psutil.virtual_memory()
|
|
GB = 1024 ** 3
|
|
return {
|
|
"cpu": {
|
|
"percent": round(psutil.cpu_percent(interval=None), 1),
|
|
"cores": [round(x, 1) for x in psutil.cpu_percent(interval=None, percpu=True)],
|
|
"count": psutil.cpu_count(logical=True),
|
|
},
|
|
"ram": {
|
|
"used_gb": round((vm.total - vm.available) / GB, 2),
|
|
"total_gb": round(vm.total / GB, 2),
|
|
"percent": round(vm.percent, 1),
|
|
},
|
|
}
|
|
|
|
|
|
def collect() -> dict:
|
|
now = time.time()
|
|
with _LOCK:
|
|
if _CACHE["data"] is not None and (now - _CACHE["ts"]) < _CACHE_TTL:
|
|
return _CACHE["data"]
|
|
|
|
_init_nvml()
|
|
driver = ""
|
|
gpus: list[dict] = []
|
|
if _NVML_READY:
|
|
try:
|
|
driver = _as_str(pynvml.nvmlSystemGetDriverVersion())
|
|
except Exception:
|
|
driver = ""
|
|
for i, h in enumerate(_GPU_HANDLES):
|
|
gpus.append(_gpu_stats(i, h))
|
|
|
|
data = {
|
|
"ts": now,
|
|
"provider": "NVIDIA" if _NVML_READY else "unavailable",
|
|
"driver": driver,
|
|
"nvml_error": _NVML_ERROR,
|
|
**_system_stats(),
|
|
"gpus": gpus,
|
|
}
|
|
_CACHE["ts"] = now
|
|
_CACHE["data"] = data
|
|
return data
|
|
|
|
|
|
def shutdown() -> None:
|
|
global _NVML_READY
|
|
if _NVML_READY:
|
|
try:
|
|
pynvml.nvmlShutdown()
|
|
except Exception:
|
|
pass
|
|
_NVML_READY = False
|