Skip to content

[NVBUG: 5472822] Skip memory monitoring if not available #374

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Sep 26, 2025
Merged

[NVBUG: 5472822] Skip memory monitoring if not available #374

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion modelopt/torch/utils/memory_monitor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,7 +131,7 @@ def stop(self):
nvmlShutdown()


def launch_memory_monitor(monitor_interval: float = 1.0) -> GPUMemoryMonitor:
def launch_memory_monitor(monitor_interval: float = 1.0) -> GPUMemoryMonitor | None:
"""Launch a GPU memory monitor in a separate thread.

Args:
Expand All @@ -140,6 +140,11 @@ def launch_memory_monitor(monitor_interval: float = 1.0) -> GPUMemoryMonitor:
Returns:
GPUMemoryMonitor: The monitor instance that was launched
"""
try:
nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(0))
except Exception as e:
print(f"Failed to get GPU memory info: {e}. Stopping GPU memory monitor.")
return None
Comment on lines +143 to +147
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

Pre-check calls NVML without initialization; narrow exception handling and balance init/shutdown.

Calling nvmlDeviceGetHandleByIndex / nvmlDeviceGetMemoryInfo before nvmlInit will raise Uninitialized and wrongly disable monitoring. Also, catching Exception is too broad. Initialize NVML for the probe, catch pynvml errors specifically, and ensure shutdown to balance the pre-check init.

Apply this diff to harden the pre-check:

-    try:
-        nvmlDeviceGetMemoryInfo(nvmlDeviceGetHandleByIndex(0))
-    except Exception as e:
-        print(f"Failed to get GPU memory info: {e}. Stopping GPU memory monitor.")
-        return None
+    try:
+        nvmlInit()
+        count = nvmlDeviceGetCount()
+        if count == 0:
+            print("No NVIDIA GPUs detected. Skipping GPU memory monitor.")
+            return None
+        handle = nvmlDeviceGetHandleByIndex(0)
+        nvmlDeviceGetMemoryInfo(handle)
+    except NVMLError_NotSupported as e:
+        print(f"NVML memory info not supported on this platform: {e}. Skipping GPU memory monitor.")
+        return None
+    except NVMLError as e:
+        print(f"NVML not available (init/count/handle/memoryinfo failed): {e}. Skipping GPU memory monitor.")
+        return None
+    finally:
+        try:
+            nvmlShutdown()
+        except NVMLError:
+            pass

Add the required exception imports near the existing NVML imports:

from pynvml import NVMLError, NVMLError_NotSupported

Alternatively, group with the existing import block.

🤖 Prompt for AI Agents 
In modelopt/torch/utils/memory_monitor.py around lines 143 to 147, the pre-check
calls nvmlDeviceGetHandleByIndex/nvmlDeviceGetMemoryInfo without calling
nvmlInit, catches Exception too broadly, and does not balance nvmlInit with
nvmlShutdown; initialize NVML before probing, wrap the probe in a try/except
that catches pynvml-specific errors (NVMLError and NVMLError_NotSupported)
instead of Exception, call nvmlShutdown in a finally block so NVML is always
shut down after the probe, and add the required imports from pynvml (NVMLError,
NVMLError_NotSupported) near the existing NVML imports.

monitor = GPUMemoryMonitor(monitor_interval)
monitor.start()
atexit.register(monitor.stop) # Ensure the monitor stops when the program exits
Expand Down
Loading