fix: fall back to CUDA runtime mem_get_info when NVML memory info is unsupported (e.g. NVIDIA GB10) - #288
Draft
zbrad wants to merge 1 commit into
Draft
fix: fall back to CUDA runtime mem_get_info when NVML memory info is unsupported (e.g. NVIDIA GB10)#288zbrad wants to merge 1 commit into
zbrad wants to merge 1 commit into
Conversation
nvmlDeviceGetMemoryInfo (and nvmlDeviceGetBAR1MemoryInfo) return NVML_ERROR_NOT_SUPPORTED on NVIDIA GB10 (a unified coherent-memory SoC, confirmed independently via nvidia-smi -q -d MEMORY showing all fields N/A) even though NVML itself, utilization, and temperature queries all work fine. Previously this made deviceGetMemoryInfo() raise on first call, permanently disabling the VRAM panel for the whole session. Add probeMemoryInfoSource(), run once per device at startup: try nvmlDeviceGetMemoryInfo, and on NVMLError_NotSupported specifically, fall back to testing torch.cuda.mem_get_info() (the CUDA runtime's own query, already used by ComfyUI's own model_management.py for VRAM accounting) against the same device index. Cache the winning source per device index and dispatch off it on every subsequent call. This is a runtime capability probe, not an arch/name check, so it keeps working correctly on discrete-VRAM GPUs (nvml path, unchanged behavior), on GB10-class unified-memory SoCs (torch fallback, new), and on any future architecture/driver combination without needing a code change -- including a mixed-GPU machine, since the probe and the cached source are both per-device-index. Assisted-by: Claude Sonnet 5
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
On NVIDIA GB10 (a unified coherent-memory SoC — Grace Blackwell "DGX Spark"-class hardware, not a discrete VRAM card),
pynvml.nvmlDeviceGetMemoryInfo()raisesNVMLError_NotSupported. This happens even though NVML itself, GPU utilization, and temperature queries all work fine — it's specifically the memory counters that are unpopulated for this architecture.Confirmed independently at the driver level:
nvidia-smi -q -d MEMORYshowsFB Memory UsageandBAR1 Memory UsageasN/Aon this hardware (driver 580.173.02). Not a pynvml version gap either — there's nonvmlDeviceGetMemoryInfo_v2in this pynvml build to fall back to.Today,
deviceGetMemoryInfo()just raises on the first call, which trips theexceptingetStatus()and setsswitchVRAM = False— permanently disabling the VRAM panel for the whole session, with only a genericCould not get GPU memory info. Not Supportedlog line.Fix
torch.cuda.mem_get_info()(the CUDA runtime API — a different driver subsystem from NVML) works fine on GB10 and reports real numbers against the same coherent memory pool. It's the same primitive ComfyUI's ownmodel_management.pyalready relies on for VRAM accounting, so it's a reasonable fallback source here too.Rather than hardcoding this by GPU name or architecture string, this adds a one-time runtime probe (
probeMemoryInfoSource) per device at startup: trynvmlDeviceGetMemoryInfo, and only onNVMLError_NotSupportedspecifically, testtorch.cuda.mem_get_info()against the same device index. Whichever source actually works gets cached per device index inself.gpusVRAMSource, anddeviceGetMemoryInfo()dispatches off that cache on every subsequent poll.This keeps existing behavior fully unchanged on discrete-VRAM GPUs (nvml path untouched), fixes it for GB10-class unified-memory SoCs (new torch fallback), and should keep working correctly for any future architecture/driver combination — including a machine with mixed GPU types, since the probe and its result are both per-device-index rather than global.
Testing
Verified live against a GB10 box running ComfyUI + this node as a systemd service. Startup log went from:
to:
with the VRAM panel now reporting real numbers (cross-checked against
/proc/meminfo) instead of staying disabled for the session. GPU utilization and temperature reporting are untouched by this change.Opening as draft since I haven't exercised the actual
crystools.monitorwebsocket message end-to-end in automated form — verified via the startup log's source-selection line and a direct call to the sametorch.cuda.mem_get_info()path instead. Happy to add a lighter-weight test or tighten anything if useful.