nvbug6084457: Fix device architecture handling and NVLink link count query

Author: mdboom

cuda_bindings/tests/nvml/test_init.py

@@ -25,11 +25,18 @@ def test_devices_are_the_same_architecture(all_devices):
     # they won't be tested properly.  This tests for the (hopefully rare) case
     # where a system has devices of different architectures and produces a warning.
 
-    all_arches = {nvml.DeviceArch(nvml.device_get_architecture(device)) for device in all_devices}
+    def get_architecture_name(arch):
+        try:
+            arch = nvml.DeviceArch(arch)
+            return arch.name
+        except ValueError:
+            return f"UNKNOWN({arch})"
+
+    all_arches = {nvml.device_get_architecture(device) for device in all_devices}
 
     if len(all_arches) > 1:
         warnings.warn(
-            f"System has devices of multiple architectures ({', '.join(x.name for x in all_arches)}). "
+            f"System has devices of multiple architectures ({', '.join(get_architecture_name(x) for x in all_arches)}). "
             f" Some tests may be skipped unexpectedly",
             UserWarning,
         )

cuda_bindings/tests/nvml/test_nvlink.py

@@ -26,4 +26,4 @@ def test_nvlink_get_link_count(all_devices):
         # The feature_nvlink_supported detection is not robust, so we
         # can't be more specific about how many links we should find.
         if value.nvml_return == nvml.Return.SUCCESS:
-            assert value.value.ui_val <= nvml.NVLINK_MAX_LINKS, f"Unexpected link count {value.value.ui_val}"
+            assert value.value.ui_val[0] <= nvml.NVLINK_MAX_LINKS, f"Unexpected link count {value.value.ui_val[0]}"

cuda_core/cuda/core/system/_device.pyx

@@ -165,7 +165,12 @@ cdef class Device:
         "VOLTA"``, and RTX A6000 will report ``DeviceArchitecture.name ==
         "AMPERE"``.
         """
-        return DeviceArch(nvml.device_get_architecture(self._handle))
+        arch = nvml.device_get_architecture(self._handle)
+        try:
+            arch = DeviceArch(arch)
+            return arch
+        except ValueError:
+            return nvml.DeviceArch.UNKNOWN
 
     @property
     def name(self) -> str: