Redo of Make cuda.bindings.nvml public (NVIDIA#1544) (NVIDIA#1548)

mdboom · rwgk · kkraus14 · web-flow · commit f39bdc93526b · 2026-01-30T21:10:29.000Z
* Make cuda.bindings.nvml public (NVIDIA#1544) * cuda.bindings.nvml: Small follow-on to NVIDIA#1524 * cuda.bindings.nvml: Make public * Fix NVML is compatible * Update cuda_core/cuda/core/system/_system.pyx Co-authored-by: Ralf W. Grosse-Kunstleve <rwgkio@gmail.com> --------- Co-authored-by: Ralf W. Grosse-Kunstleve <rwgkio@gmail.com> * Update cuda_core/cuda/core/system/_system.pyx Co-authored-by: Keith Kraus <keith.j.kraus@gmail.com> --------- Co-authored-by: Ralf W. Grosse-Kunstleve <rwgkio@gmail.com> Co-authored-by: Keith Kraus <keith.j.kraus@gmail.com>
diff --git a/cuda_bindings/cuda/bindings/_internal/nvml.pxd b/cuda_bindings/cuda/bindings/_internal/nvml.pxd
@@ -4,7 +4,7 @@
 #
 # This code was automatically generated across versions from 12.9.1 to 13.1.1. Do not modify it directly.
 
-from ..cy_nvml cimport *
+from ..cynvml cimport *
 
 
 ###############################################################################
diff --git a/cuda_bindings/cuda/bindings/_internal/nvml_linux.pyx b/cuda_bindings/cuda/bindings/_internal/nvml_linux.pyx
diff --git a/cuda_bindings/cuda/bindings/_internal/nvml_windows.pyx b/cuda_bindings/cuda/bindings/_internal/nvml_windows.pyx
diff --git a/cuda_bindings/cuda/bindings/_test_helpers/arch_check.py b/cuda_bindings/cuda/bindings/_test_helpers/arch_check.py
@@ -5,7 +5,7 @@
 from contextlib import contextmanager
 
 import pytest
-from cuda.bindings import _nvml as nvml
+from cuda.bindings import nvml
 
 
 @contextmanager
diff --git a/cuda_bindings/cuda/bindings/cynvml.pxd b/cuda_bindings/cuda/bindings/cynvml.pxd
diff --git a/cuda_bindings/cuda/bindings/cynvml.pyx b/cuda_bindings/cuda/bindings/cynvml.pyx
@@ -4,7 +4,7 @@
 #
 # This code was automatically generated across versions from 12.9.1 to 13.1.1. Do not modify it directly.
 
-from ._internal cimport _nvml as _nvml
+from ._internal cimport nvml as _nvml
 
 
 ###############################################################################
diff --git a/cuda_bindings/cuda/bindings/nvml.pxd b/cuda_bindings/cuda/bindings/nvml.pxd
@@ -6,7 +6,7 @@
 
 from libc.stdint cimport intptr_t
 
-from .cy_nvml cimport *
+from .cynvml cimport *
 
 
 ###############################################################################
diff --git a/cuda_bindings/cuda/bindings/nvml.pyx b/cuda_bindings/cuda/bindings/nvml.pyx
diff --git a/cuda_bindings/tests/nvml/conftest.py b/cuda_bindings/tests/nvml/conftest.py
@@ -1,10 +1,10 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
 from collections import namedtuple
 
 import pytest
-from cuda.bindings import _nvml as nvml
+from cuda.bindings import nvml
 from cuda.bindings._test_helpers.arch_check import unsupported_before  # noqa: F401
 
 
diff --git a/cuda_bindings/tests/nvml/test_compute_mode.py b/cuda_bindings/tests/nvml/test_compute_mode.py
@@ -5,7 +5,7 @@
 import sys
 
 import pytest
-from cuda.bindings import _nvml as nvml
+from cuda.bindings import nvml
 
 from .conftest import unsupported_before
 
diff --git a/cuda_bindings/tests/nvml/test_cuda.py b/cuda_bindings/tests/nvml/test_cuda.py
@@ -1,8 +1,8 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
 import cuda.bindings.driver as cuda
-from cuda.bindings import _nvml as nvml
+from cuda.bindings import nvml
 
 from .conftest import NVMLInitializer
 
diff --git a/cuda_bindings/tests/nvml/test_device.py b/cuda_bindings/tests/nvml/test_device.py
@@ -5,7 +5,7 @@
 from functools import cache
 
 import pytest
-from cuda.bindings import _nvml as nvml
+from cuda.bindings import nvml
 
 from .conftest import unsupported_before
 
diff --git a/cuda_bindings/tests/nvml/test_gpu.py b/cuda_bindings/tests/nvml/test_gpu.py
@@ -3,7 +3,7 @@
 
 import numpy as np
 import pytest
-from cuda.bindings import _nvml as nvml
+from cuda.bindings import nvml
 
 from . import util
 from .conftest import unsupported_before
diff --git a/cuda_bindings/tests/nvml/test_init.py b/cuda_bindings/tests/nvml/test_init.py
@@ -1,11 +1,11 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
 import sys
 import warnings
 
 import pytest
-from cuda.bindings import _nvml as nvml
+from cuda.bindings import nvml
 
 
 def assert_nvml_is_initialized():
diff --git a/cuda_bindings/tests/nvml/test_nvlink.py b/cuda_bindings/tests/nvml/test_nvlink.py
@@ -1,8 +1,8 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
 
-from cuda.bindings import _nvml as nvml
+from cuda.bindings import nvml
 
 
 def test_nvlink_get_link_count(all_devices):
diff --git a/cuda_bindings/tests/nvml/test_page_retirement.py b/cuda_bindings/tests/nvml/test_page_retirement.py
@@ -1,8 +1,8 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
 import pytest
-from cuda.bindings import _nvml as nvml
+from cuda.bindings import nvml
 
 from . import util
 
diff --git a/cuda_bindings/tests/nvml/test_pci.py b/cuda_bindings/tests/nvml/test_pci.py
@@ -4,7 +4,7 @@
 
 import contextlib
 
-from cuda.bindings import _nvml as nvml
+from cuda.bindings import nvml
 
 from .conftest import unsupported_before
 
diff --git a/cuda_bindings/tests/nvml/test_pynvml.py b/cuda_bindings/tests/nvml/test_pynvml.py
@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
 # A set of tests ported from https://github.com/gpuopenanalytics/pynvml/blob/11.5.3/pynvml/tests/test_nvml.py
@@ -7,7 +7,7 @@
 import time
 
 import pytest
-from cuda.bindings import _nvml as nvml
+from cuda.bindings import nvml
 
 from . import util
 from .conftest import unsupported_before
diff --git a/cuda_bindings/tests/nvml/util.py b/cuda_bindings/tests/nvml/util.py
@@ -1,12 +1,12 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: LicenseRef-NVIDIA-SOFTWARE-LICENSE
 
 
 import functools
 import platform
 from pathlib import Path
 
-from cuda.bindings import _nvml as nvml
+from cuda.bindings import nvml
 
 current_os = platform.system()
 if current_os == "VMkernel":
diff --git a/cuda_core/cuda/core/system/_device.pyx b/cuda_core/cuda/core/system/_device.pyx
@@ -8,7 +8,7 @@ from libc.math cimport ceil
 from multiprocessing import cpu_count
 from typing import Iterable
 
-from cuda.bindings import _nvml as nvml
+from cuda.bindings import nvml
 
 from ._nvml_context cimport initialize
 
diff --git a/cuda_core/cuda/core/system/_nvml_context.pyx b/cuda_core/cuda/core/system/_nvml_context.pyx
@@ -1,10 +1,10 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
 
 import threading
 
-from cuda.bindings import _nvml as nvml
+from cuda.bindings import nvml
 
 from . import exceptions
 
diff --git a/cuda_core/cuda/core/system/_system.pyx b/cuda_core/cuda/core/system/_system.pyx
@@ -1,9 +1,9 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
 
 
-# This file needs to either use NVML exclusively, or when `cuda.bindings._nvml`
+# This file needs to either use NVML exclusively, or when `cuda.bindings.nvml`
 # isn't available, fall back to non-NVML-based methods for backward
 # compatibility.
 
@@ -19,13 +19,17 @@ else:
 
 
 if CUDA_BINDINGS_NVML_IS_COMPATIBLE:
-    from cuda.bindings import _nvml as nvml
-    # TODO: We need to be even more specific than version numbers for development.
-    # This can be removed once we have a release including everything we need.
-    for member in ["FieldId", "ClocksEventReasons"]:
-        if not hasattr(nvml, member):
-            CUDA_BINDINGS_NVML_IS_COMPATIBLE = False
-            break
+    try:
+        from cuda.bindings import nvml
+    except ImportError:
+        CUDA_BINDINGS_NVML_IS_COMPATIBLE = False
+    else:
+        # TODO: We need to be even more specific than version numbers for development.
+        # This can be removed once we have a release including everything we need.
+        for member in ["FieldId", "ClocksEventReasons"]:
+            if not hasattr(nvml, member):
+                CUDA_BINDINGS_NVML_IS_COMPATIBLE = False
+                break
 
 if CUDA_BINDINGS_NVML_IS_COMPATIBLE:
     from ._nvml_context import initialize
diff --git a/cuda_core/cuda/core/system/_system_events.pyx b/cuda_core/cuda/core/system/_system_events.pyx
@@ -1,11 +1,11 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
 
 
 from libc.stdint cimport intptr_t
 
-from cuda.bindings import _nvml as nvml
+from cuda.bindings import nvml
 
 from ._nvml_context cimport initialize
 
diff --git a/cuda_core/cuda/core/system/exceptions.py b/cuda_core/cuda/core/system/exceptions.py
@@ -1,9 +1,9 @@
-# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # SPDX-License-Identifier: Apache-2.0
 
 
-from cuda.bindings import _nvml as nvml
+from cuda.bindings import nvml
 
 NvmlError = nvml.NvmlError
 UninitializedError = nvml.UninitializedError
diff --git a/cuda_core/tests/system/test_system_device.py b/cuda_core/tests/system/test_system_device.py
@@ -19,7 +19,7 @@
 from cuda.core import system
 
 if system.CUDA_BINDINGS_NVML_IS_COMPATIBLE:
-    from cuda.bindings import _nvml as nvml
+    from cuda.bindings import nvml
     from cuda.core.system import DeviceArch, _device
 
 

Original file line number	Diff line number	Diff line change
`@@ -4,7 +4,7 @@`
`4`	`4`	`#`
`5`	`5`	`# This code was automatically generated across versions from 12.9.1 to 13.1.1. Do not modify it directly.`
`6`	`6`
`7`		`-from ..cy_nvml cimport *`
	`7`	`+from ..cynvml cimport *`
`8`	`8`
`9`	`9`
`10`	`10`	`###############################################################################`
Original file line number	Diff line number	Diff line change
`@@ -4,7 +4,7 @@`
`4`	`4`	`#`
`5`	`5`	`# This code was automatically generated across versions from 12.9.1 to 13.1.1. Do not modify it directly.`
`6`	`6`
`7`		`-from ._internal cimport _nvml as _nvml`
	`7`	`+from ._internal cimport nvml as _nvml`
`8`	`8`
`9`	`9`
`10`	`10`	`###############################################################################`