Improve NVML_NVLINK_MAX_LINKS dynamic handling

Author: mdboom

.pre-commit-config.yaml

@@ -115,6 +115,7 @@ repos:
         args: [--config-file=cuda_core/pyproject.toml, cuda_core/cuda/core]
         additional_dependencies:
           - numpy
+          - types-Deprecated
 
   - repo: https://github.com/rhysd/actionlint
     rev: "914e7df21a07ef503a81201c76d2b11c789d3fca"  # frozen: v1.7.12

cuda_core/AGENTS.md

@@ -145,3 +145,95 @@ so that they are documented but don't appear in the main index.
 ### API stability
 
 Reviews should point out where existing public APIs are broken.
+
+### Deprecation and API lifecycle
+
+`cuda.core` follows SemVer (see `docs/source/:
+
+- **New APIs** may be added at any time (`x.Y.0`).
+- **Breaking removals** only happen in **major releases** (`X.0.0`).
+- Per the support policy, a deprecation notice must be present for **at least
+  one minor release** before the API is actually removed.
+- Changes should be notated in the code and also in the release notes in the
+  "Deprecated APIs" section.
+
+**Annotating a new API** — Use the `deprecated.sphinx.versionadded` decorator:
+
+```python
+
+from deprecated.sphinx import versionadded
+
+@versionadded("1.2.0")
+def new_feature(...):
+    """Short description.
+    """
+```
+
+Alternatively, if the vagaries of how we implement functions in Cython does not
+allow this, you can add the reST `versionadded` directive directly:
+
+```python
+def new_feature(...):
+    """Short description.
+
+    .. versionadded:: 1.2.0
+    """
+```
+
+**Annotating a changed API** — Use the `deprecated.sphinx.versionchanged` decorator:
+
+```python
+
+from deprecated.sphinx import versionchanged
+
+@versionchanged("1.2.0", reason="The old version was broken because...")
+def new_feature(...):
+    """Short description.
+    """
+```
+
+Alternatively, if the vagaries of how we implement functions in Cython does not
+allow this, you can add the reST `versionchanged` directive directly:
+
+```python
+def new_feature(...):
+    """Short description.
+
+    .. versionchanged:: 1.2.0
+        The old version was broken because...
+    """
+```
+
+**Deprecating an existing API** — use the `@deprecated` decorator from the
+`Deprecated` PyPI package (`deprecated.sphinx` import name) and add a
+`.. deprecated::` directive in the docstring.  The decorator emits a
+`DeprecationWarning` at call time; the docstring directive surfaces it in the
+generated docs.
+
+```python
+from deprecated.sphinx import deprecated
+
+@deprecated(version="1.2.0", reason="Use `new_feature` instead.")
+def old_feature(...):
+    """Short description.
+    """
+```
+
+Rules to follow when deprecating:
+
+- The `version=` argument must be the **first** in which the
+  deprecation appears, not the release in which removal is planned.
+- The `reason=` string must name the replacement (if one exists) so users
+  know what to migrate to.
+- Keep the old implementation fully functional — do not change its behavior,
+  only add the decorator.
+- The deprecated API must remain in the codebase for **at least one full minor
+  release cycle** before it can be removed in a subsequent major release.
+
+**Removing a deprecated API** — removals land in a **major release**.  Before
+removing, verify that the deprecation has been present since at least the
+previous minor release.  Remove the decorator, the implementation, and any
+`__all__` entry; update `api.rst` and the release notes accordingly.
+
+At some point in the future, we will provide automation for removal of
+deprecated APIs.

cuda_core/cuda/core/system/_device.pyi

@@ -14,6 +14,7 @@ from cuda.core.system.typing import (AddressingMode, AffinityScope, ClockId,
                                      GpuTopologyLevel, InforomObject,
                                      TemperatureThresholds, ThermalController,
                                      ThermalTarget)
+from deprecated.sphinx import deprecated, versionadded, versionchanged
 
 
 class ClockOffsets:
@@ -787,11 +788,23 @@ class MigInfo:
             A list of all MIG devices corresponding to this GPU.
         """
 
-class NvlinkInfo:
+class _NvlinkInfoMeta(type):
+
+    @property
+    @deprecated(version='1.1.0', reason='Use Device.get_num_nvlinks instead to get the actual number of Nvlinks available on a specific device.')
+    def max_links(cls):
+        """
+        The statically-defined maximum number of Nvlinks available.  Defined in
+        upstream NVML as ``NVML_NVLINK_MAX_LINKS``.
+
+        To find the actual number of Nvlinks available on a device, use
+        :py:attr:`Device.get_num_nvlinks`.
+        """
+
+class _NvlinkInfo:
     """
     Nvlink information for a device.
     """
-    max_links = nvml.NVLINK_MAX_LINKS
 
     def __init__(self, device: Device, link: int):
         ...
@@ -824,6 +837,9 @@ class NvlinkInfo:
             `True` if the Nvlink is active.
         """
 
+class NvlinkInfo(_NvlinkInfo, metaclass=_NvlinkInfoMeta):
+    ...
+
 class PciInfo:
     """
     PCI information about a GPU device.
@@ -1719,13 +1735,30 @@ class Device:
         :obj:`~_device.MemoryInfo` object with memory information.
         """
 
+    @versionchanged(version='1.1.0', reason='Any link number not supported by this specific device will raise a `ValueError`.')
     def get_nvlink(self, link: int) -> NvlinkInfo:
         """
         Get :obj:`~NvlinkInfo` about this device.
 
         For devices with NVLink support.
         """
 
+    @versionadded(version='1.1.0')
+    def get_nvlink_count(self) -> int:
+        """
+        Get the number of NVLink links on this device.
+
+        For devices with NVLink support.
+        """
+
+    @versionadded(version='1.1.0')
+    def get_nvlinks(self) -> Iterable[NvlinkInfo]:
+        """
+        Get :obj:`~NvlinkInfo` about all NVLink links on this device.
+
+        For devices with NVLink support.
+        """
+
     @property
     def pci_info(self) -> PciInfo:
         """

cuda_core/cuda/core/system/_device.pyx

@@ -9,6 +9,8 @@ from multiprocessing import cpu_count
 from typing import Iterable, TYPE_CHECKING
 import warnings
 
+from deprecated.sphinx import deprecated, versionadded, versionchanged
+
 from cuda.bindings import nvml
 
 from ._nvml_context cimport initialize
@@ -884,16 +886,40 @@ cdef class Device:
     # NVLINK
     # See external class definitions in _nvlink.pxi
 
+    @versionchanged(
+        version="1.1.0",
+        reason="Any link number not supported by this specific device will raise a `ValueError`."
+    )
     def get_nvlink(self, link: int) -> NvlinkInfo:
         """
         Get :obj:`~NvlinkInfo` about this device.
 
         For devices with NVLink support.
         """
-        if link < 0 or link >= NvlinkInfo.max_links:
-            raise ValueError(f"Link index {link} is out of range [0, {NvlinkInfo.max_links})")
+        if link < 0 or link >= self.get_num_nvlinks():
+            raise ValueError(f"Link index {link} is out of range [0, {self.get_num_nvlinks()})")
         return NvlinkInfo(self, link)
 
+    @versionadded(version="1.1.0")
+    def get_nvlink_count(self) -> int:
+        """
+        Get the number of NVLink links on this device.
+
+        For devices with NVLink support.
+        """
+        return self.get_field_values([FieldId.DEV_NVLINK_LINK_COUNT])[0].value
+
+    @versionadded(version="1.1.0")
+    def get_nvlinks(self) -> Iterable[NvlinkInfo]:
+        """
+        Get :obj:`~NvlinkInfo` about all NVLink links on this device.
+
+        For devices with NVLink support.
+        """
+        num_links = self.get_num_nvlinks()
+        for link in range(num_links):
+            yield self.get_nvlink(link)
+
     ##########################################################################
     # PCI INFO
     # See external class definitions in _pci_info.pxi

cuda_core/cuda/core/system/_nvlink.pxi

@@ -18,7 +18,24 @@ if _NVLINK_VERSION_6_0 is not None:
     _NVLINK_VERSION_MAPPING[_NVLINK_VERSION_6_0] = (6, 0)
 
 
-cdef class NvlinkInfo:
+class _NvlinkInfoMeta(type):
+    @property
+    @deprecated(
+        version="1.1.0",
+        reason="Use Device.get_num_nvlinks instead to get the actual number of Nvlinks available on a specific device."
+    )
+    def max_links(cls):
+        """
+        The statically-defined maximum number of Nvlinks available.  Defined in
+        upstream NVML as ``NVML_NVLINK_MAX_LINKS``.
+
+        To find the actual number of Nvlinks available on a device, use
+        :py:attr:`Device.get_num_nvlinks`.
+        """
+        return nvml.NVLINK_MAX_LINKS
+
+
+cdef class _NvlinkInfo:
     """
     Nvlink information for a device.
     """
@@ -67,4 +84,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

cuda_core/pyproject.toml

@@ -49,6 +49,7 @@ classifiers = [
 ]
 dependencies = [
     "cuda-pathfinder >=1.4.2",
+    "Deprecated",
     "numpy",
     "backports.strenum; python_version < '3.11'",
 ]

cuda_core/tests/system/test_system_device.py

@@ -763,27 +763,40 @@ def test_compute_running_processes():
 
 def test_nvlink():
     for device in system.Device.get_all_devices():
-        max_links = _device.NvlinkInfo.max_links
-        assert isinstance(max_links, int)
-        assert max_links > 0
-
-        for link in range(max_links):
-            with unsupported_before(device, None):
-                nvlink_info = device.get_nvlink(link)
-            assert isinstance(nvlink_info, _device.NvlinkInfo)
-
-            with unsupported_before(device, None):
-                state = nvlink_info.state
-            assert isinstance(state, bool)
-
-            if not state:
-                continue
-
-            with unsupported_before(device, None):
-                version = nvlink_info.version
-            assert isinstance(version, tuple)
-            assert len(version) == 2
-            assert all(isinstance(i, int) for i in version)
+        with unsupported_before(device, None):
+            for link in range(device.get_nvlink_count()):
+                with unsupported_before(device, None):
+                    nvlink_info = device.get_nvlink(link)
+                assert isinstance(nvlink_info, _device.NvlinkInfo)
+
+                with unsupported_before(device, None):
+                    state = nvlink_info.state
+                assert isinstance(state, bool)
+
+                if not state:
+                    continue
+
+                with unsupported_before(device, None):
+                    version = nvlink_info.version
+                assert isinstance(version, tuple)
+                assert len(version) == 2
+                assert all(isinstance(i, int) for i in version)
+
+            for nvlink_info in device.get_nvlinks():
+                assert isinstance(nvlink_info, _device.NvlinkInfo)
+
+                with unsupported_before(device, None):
+                    state = nvlink_info.state
+                assert isinstance(state, bool)
+
+                if not state:
+                    continue
+
+                with unsupported_before(device, None):
+                    version = nvlink_info.version
+                assert isinstance(version, tuple)
+                assert len(version) == 2
+                assert all(isinstance(i, int) for i in version)
 
 
 def test_utilization():