diff --git a/cuda_core/cuda/core/__init__.py b/cuda_core/cuda/core/__init__.py index dc6fefdffea..c84351855f0 100644 --- a/cuda_core/cuda/core/__init__.py +++ b/cuda_core/cuda/core/__init__.py @@ -5,8 +5,11 @@ from cuda.core._version import __version__ +# TODO: remove this function altogether after wheel-variants become mainstream def _import_versioned_module() -> None: import importlib + import pathlib + import sys from cuda import bindings @@ -15,13 +18,13 @@ def _import_versioned_module() -> None: raise ImportError("cuda.bindings 12.x or 13.x must be installed") subdir = f"cu{cuda_major}" - try: - versioned_mod = importlib.import_module(f".{subdir}", __package__) - # Import all symbols from the module - globals().update(versioned_mod.__dict__) - except ImportError: - # This is not a wheel build, but a conda or local build, do nothing - pass + versioned_path = pathlib.Path(__file__).parent / subdir + # This is a wheel build with relevant modules in cuda/core/cu + # directory. Let's add it to module path so imports work as expected. + # cuda.core.cu is not meant to behave as a module itself, and + # does not belong in sys.modules + if versioned_path.is_dir(): + __path__.append(str(versioned_path)) _import_versioned_module() diff --git a/cuda_core/tests/test_duplicate_imports.py b/cuda_core/tests/test_duplicate_imports.py new file mode 100644 index 00000000000..91c9aa6331a --- /dev/null +++ b/cuda_core/tests/test_duplicate_imports.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Tests cuda.core.__init__.py does not import duplicate modules.""" + +import sys + +from cuda import bindings + +cuda_major = bindings.__version__.split(".")[0] + + +def test_typing_module_imports(): + """ + Importing cuda.core.system should not also import cuda.core.cuXX.system + """ + + from cuda.core import * # NOQA + + assert "cuda.core.system" in sys.modules + assert f"cuda.core.cu{cuda_major}" not in sys.modules