Use a __getattr__ approach

Author: mdboom

cuda_bindings/cuda/bindings/nvml.pyx

@@ -1593,15 +1593,6 @@ class FieldId(_FastEnum):
 
     MAX = 289
 
-# This value changed in CTK 13.3.  We need to build a binary that works across
-# all versions, so the only way to support this is to check the version at
-# runtime and set the value accordingly.
-
-if tuple(int(x) for x in system_get_nvml_version().split(".")) < (3, 13):
-    NVLINK_MAX_LINKS = 18
-else:
-    NVLINK_MAX_LINKS = 36
-
 
 class RUSD(_FastEnum):
     POLL_NONE = (0x0, "Disable RUSD polling on all metric groups")
@@ -28458,3 +28449,17 @@ cpdef str vgpu_type_get_name(unsigned int vgpu_type_id):
 device_get_virtualization_mode.__doc__ = device_get_virtualization_mode.__doc__.replace("NVML_GPU_VIRTUALIZATION_?", "``NVML_GPU_VIRTUALIZATION_?``")
 device_set_virtualization_mode.__doc__ = device_set_virtualization_mode.__doc__.replace("NVML_GPU_VIRTUALIZATION_?", "``NVML_GPU_VIRTUALIZATION_?``")
 GpmMetricId.GPM_METRIC_DRAM_BW_UTIL.__doc__ = "Percentage of DRAM bw used vs theoretical maximum. ``0.0 - 100.0 *\u200d/``."
+
+
+def __getattr__(name: str):
+    # This value changed in CTK 13.3.  We need to build a binary that works across
+    # all versions, so the only way to support this is to check the version at
+    # runtime and set the value accordingly.
+
+    if name == "NVLINK_MAX_LINKS":
+        if tuple(int(x) for x in system_get_nvml_version().split(".")) < (3, 13):
+            return 18
+        else:
+            return 36
+
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")

cuda_bindings/tests/nvml/test_nvlink.py

@@ -32,3 +32,12 @@ def test_nvlink_get_link_count(all_devices):
         # can't be more specific about how many links we should find.
         if value.nvml_return == nvml.Return.SUCCESS:
             assert value.value.ui_val[0] <= nvml.NVLINK_MAX_LINKS, f"Unexpected link count {value.value.ui_val[0]}"
+
+
+def test_nvlink_max_links():
+    nvml_version = tuple(int(x) for x in nvml.system_get_nvml_version().split("."))
+
+    if nvml_version < (13, 3):
+        assert nvml.NVLINK_MAX_LINKS == 18
+    else:
+        assert nvml.NVLINK_MAX_LINKS == 36

cuda_core/cuda/core/system/_nvlink.pxi

@@ -18,7 +18,13 @@ if _NVLINK_VERSION_6_0 is not None:
     _NVLINK_VERSION_MAPPING[_NVLINK_VERSION_6_0] = (6, 0)
 
 
-cdef class NvlinkInfo:
+cdef class _NvlinkInfoMeta(type):
+    @property
+    def max_links(cls):
+        return nvml.NVLINK_MAX_LINKS
+
+
+cdef class _NvlinkInfo:
     """
     Nvlink information for a device.
     """
@@ -67,4 +73,6 @@ cdef class NvlinkInfo:
             nvml.device_get_nvlink_state(self._device._handle, self._link) == nvml.EnableState.FEATURE_ENABLED
         )
 
-    max_links = nvml.NVLINK_MAX_LINKS
+
+class NvlinkInfo(_NvlinkInfo, metaclass=_NvlinkInfoMeta):
+    pass