Handle NVMLError_NotSupported exception (#983)

Closes #982 Authors: - Richard (Rick) Zamora (https://github.com/rjzamora) - Peter Andreas Entschev (https://github.com/pentschev) Approvers: - Peter Andreas Entschev (https://github.com/pentschev) URL: #983
rapidsai · Aug 24, 2023 · 72d2ec9 · 72d2ec9
1 parent be4cdd9
commit 72d2ec9
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/ucp/__init__.py b/ucp/__init__.py
@@ -64,9 +64,14 @@ def _is_mig_device(handle):
             if _is_mig_device(handle):
                 continue
 
-            total_memory = pynvml.nvmlDeviceGetMemoryInfo(handle).total
-            bar1_total = pynvml.nvmlDeviceGetBAR1MemoryInfo(handle).bar1Total
+            try:
+                bar1_total = pynvml.nvmlDeviceGetBAR1MemoryInfo(handle).bar1Total
+            except pynvml.nvml.NVMLError_NotSupported:
+                # Bar1 access not supported on this device, set it to
+                # zero (always lower than device memory).
+                bar1_total = 0
 
+            total_memory = pynvml.nvmlDeviceGetMemoryInfo(handle).total
             if total_memory <= bar1_total:
                 large_bar1[dev_idx] = True