Wrap NVML calls inside try/except.

Author: lohedges

src/somd2/runner/_repex.py

@@ -516,18 +516,23 @@ def _check_device_memory(index):
         index: int
             The index of the CUDA device.
         """
-        from pynvml import (
-            nvmlInit,
-            nvmlShutdown,
-            nvmlDeviceGetHandleByIndex,
-            nvmlDeviceGetMemoryInfo,
-        )
+        try:
+            from pynvml import (
+                nvmlInit,
+                nvmlShutdown,
+                nvmlDeviceGetHandleByIndex,
+                nvmlDeviceGetMemoryInfo,
+            )
+
+            nvmlInit()
+            handle = nvmlDeviceGetHandleByIndex(index)
+            info = nvmlDeviceGetMemoryInfo(handle)
+            result = (info.used, info.free, info.total)
+            nvmlShutdown()
+        except Exception as e:
+            msg = f"Could not determine memory usage for device {index}: {e}"
+            _logger.error(msg)
 
-        nvmlInit()
-        handle = nvmlDeviceGetHandleByIndex(index)
-        info = nvmlDeviceGetMemoryInfo(handle)
-        result = (info.used, info.free, info.total)
-        nvmlShutdown()
         return result