diff --git a/nvmath/__init__.py b/nvmath/__init__.py index b0f008d..a9dff5a 100644 --- a/nvmath/__init__.py +++ b/nvmath/__init__.py @@ -2,31 +2,50 @@ # # SPDX-License-Identifier: Apache-2.0 -import importlib.metadata - - -# Attempt to preload libraries. Fail silently if preload fails. -def _force_lib_load(): - from nvmath._utils import module_init_force_cupy_lib_load - - module_init_force_cupy_lib_load() - +"""NVIDIA Math Python libraries.""" -_force_lib_load() +import importlib +import importlib.metadata +import os -from nvmath import ( # noqa: E402 - bindings, # noqa: E402 - fft, - linalg, - sparse, - tensor, -) -from nvmath._utils import ( # noqa: E402 - ComputeType, - CudaDataType, - LibraryPropertyType, -) -from nvmath.memory import BaseCUDAMemoryManager, BaseCUDAMemoryManagerAsync, MemoryPointer # noqa: E402 +# High-level subpackages (fft, linalg, sparse, tensor, bindings) and public +# names from nvmath._utils / nvmath.memory are loaded on first access. Eager +# imports here would pull in every library wrapper whenever a nested module +# such as nvmath.bindings.cycurand is imported (the CuPy bindings-build path). +# Native-library preload (module_init_force_cupy_lib_load) is likewise deferred +# until a high-level subpackage is used; it must not run for cycurand-only +# imports. +# +# We deliberately hand-roll this instead of using scientific-python's +# lazy_loader (SPEC 1), because lazy_loader cannot express two behaviors we +# must preserve: +# +# 1. Optional bindings wrappers (cufftMp, cublasMp, cusolverMp, nvshmem) +# resolve to None when their library is unavailable. lazy_loader.attach +# only ever returns a module or raises ImportError. +# +# 2. A host-API subpackage (fft/linalg/sparse/tensor), on first access, must +# first run module_init_force_cupy_lib_load() -- which dlopens the NVIDIA +# libraries (libcublas, libcufft, ...) so CuPy can find them (cupy#9127). +# This preload must NOT run for a bindings-only import such as CuPy's +# `from nvmath.bindings.cycurand cimport ...`, which needs a single +# wrapper and none of the host APIs. lazy_loader.attach offers no hook to +# run a side effect on access of some names but not others. +# +# Both would require wrapping lazy_loader anyway, so a small module-level +# __getattr__ (PEP 562) is the simpler fit. We still adopt lazy_loader's +# EAGER_IMPORT convention so broken imports can be surfaced. +_LAZY_MODULES = frozenset({"bindings", "fft", "linalg", "sparse", "tensor"}) +_LAZY_ATTRS = { + "BaseCUDAMemoryManager": "nvmath.memory", + "BaseCUDAMemoryManagerAsync": "nvmath.memory", + "ComputeType": "nvmath._utils", + "CudaDataType": "nvmath._utils", + "LibraryPropertyType": "nvmath._utils", + "MemoryPointer": "nvmath.memory", +} +_PRELOAD_ON = frozenset({"fft", "linalg", "sparse", "tensor"}) +_libs_preloaded = False __all__ = [ "BaseCUDAMemoryManager", @@ -42,4 +61,46 @@ def _force_lib_load(): "tensor", ] + +def _ensure_libs_preloaded() -> None: + global _libs_preloaded + if _libs_preloaded: + return + _libs_preloaded = True + from nvmath._utils import module_init_force_cupy_lib_load + + module_init_force_cupy_lib_load() + + +def _eager_import_enabled() -> bool: + # See lazy_loader's EAGER_IMPORT convention: a truthy value forces the lazy + # imports to run now, so missing/broken submodules fail at import time. + return os.environ.get("EAGER_IMPORT", "").lower() not in ("", "0", "false") + + +def __getattr__(name: str): + if name in _LAZY_MODULES: + if name in _PRELOAD_ON: + _ensure_libs_preloaded() + mod = importlib.import_module(f"{__name__}.{name}") + globals()[name] = mod + return mod + origin = _LAZY_ATTRS.get(name) + if origin is not None: + value = getattr(importlib.import_module(origin), name) + globals()[name] = value + return value + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__() -> list[str]: + return sorted(set(globals()) | _LAZY_MODULES | _LAZY_ATTRS.keys()) + + +if _eager_import_enabled(): + for _name in sorted(_LAZY_MODULES | _LAZY_ATTRS.keys()): + __getattr__(_name) + del _name + + __version__ = importlib.metadata.version("nvmath-python") diff --git a/nvmath/bindings/__init__.py b/nvmath/bindings/__init__.py index 08da134..e69445e 100644 --- a/nvmath/bindings/__init__.py +++ b/nvmath/bindings/__init__.py @@ -4,41 +4,71 @@ # type: ignore -from nvmath.bindings import cublas -from nvmath.bindings import cublasLt -from nvmath.bindings import cudss -from nvmath.bindings import cufft -from nvmath.bindings import curand -from nvmath.bindings import cusolver -from nvmath.bindings import cusolverDn -from nvmath.bindings import cusolverSp -from nvmath.bindings import cusparse -from nvmath.bindings import cusparseLt -from nvmath.bindings import cutensor - -try: - # cufftMp is Linux-only. - from nvmath.bindings import cufftMp -except ImportError: - cufftMp = None - -try: - # nvshmem is Linux-only. - from nvmath.bindings import nvshmem -except ImportError: - nvshmem = None - -try: - # cublasMp is Linux-only. - from nvmath.bindings import cublasMp -except ImportError: - cublasMp = None - -try: - # cusolverMp is Linux-only. - from nvmath.bindings import cusolverMp -except ImportError: - cusolverMp = None +import importlib +import os + +# Load library wrappers on first access (PEP 562). Eager imports here would +# map every wrapper .so (including optional *Mp / NVSHMEM modules) as soon as +# any nvmath.bindings Cython submodule is loaded — for example CuPy's +# `from nvmath.bindings.cycurand cimport ...`. +_REQUIRED = frozenset( + { + "cublas", + "cublasLt", + "cudss", + "cufft", + "curand", + "cusolver", + "cusolverDn", + "cusolverSp", + "cusparse", + "cusparseLt", + "cutensor", + "nvpl", + } +) +_OPTIONAL = frozenset({"cufftMp", "nvshmem", "cublasMp", "cusolverMp"}) +_SUBMODULES = _REQUIRED | _OPTIONAL + + +def _eager_import_enabled() -> bool: + # Same opt-in convention as scientific-python's lazy_loader: EAGER_IMPORT + # set to a truthy value forces the lazy imports to run now, so a broken or + # missing wrapper fails at import time instead of on first attribute access. + return os.environ.get("EAGER_IMPORT", "").lower() not in ("", "0", "false") + + +def __getattr__(name): + if name in _SUBMODULES: + try: + mod = importlib.import_module(f"{__name__}.{name}") + except ImportError: + if name in _OPTIONAL: + globals()[name] = None + return None + raise + globals()[name] = mod + return mod + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__(): + return sorted(set(globals()) | _SUBMODULES) + + +if _eager_import_enabled(): + # Force required wrappers to load now (hard failure if one is broken). + # Optional wrappers stay best-effort: force-loading them would raise on + # platforms where their library is absent, so they keep None-on-ImportError. + for _name in sorted(_REQUIRED): + __getattr__(_name) + for _name in sorted(_OPTIONAL): + try: + __getattr__(_name) + except ImportError: + globals()[_name] = None + del _name + __all__ = [ "cublas", diff --git a/nvmath/bindings/_internal/__init__.py b/nvmath/bindings/_internal/__init__.py index e69de29..c4f3137 100644 --- a/nvmath/bindings/_internal/__init__.py +++ b/nvmath/bindings/_internal/__init__.py @@ -0,0 +1,19 @@ +# Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# +# SPDX-License-Identifier: Apache-2.0 + +# type: ignore + +import importlib + + +def __getattr__(name): + # Submodules were previously imported as a side effect of eagerly loading + # every nvmath.bindings wrapper. Keep getattr() working now that wrappers + # are loaded on demand (see nvmath._utils.module_init_force_cupy_lib_load). + try: + mod = importlib.import_module(f"{__name__}.{name}") + except ImportError as exc: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc + globals()[name] = mod + return mod diff --git a/nvmath/fft/__init__.py b/nvmath/fft/__init__.py index 539a883..394de5d 100644 --- a/nvmath/fft/__init__.py +++ b/nvmath/fft/__init__.py @@ -2,6 +2,10 @@ # # SPDX-License-Identifier: Apache-2.0 -from ._configuration import * # noqa: F403 -from ._helpers import * # noqa: F403 -from .fft import * # noqa: F403 +from nvmath._utils import module_init_force_cupy_lib_load + +module_init_force_cupy_lib_load() + +from ._configuration import * # noqa: E402, F403 +from ._helpers import * # noqa: E402, F403 +from .fft import * # noqa: E402, F403 diff --git a/nvmath/linalg/__init__.py b/nvmath/linalg/__init__.py index 78bf49d..cfaf89f 100644 --- a/nvmath/linalg/__init__.py +++ b/nvmath/linalg/__init__.py @@ -2,10 +2,13 @@ # # SPDX-License-Identifier: Apache-2.0 +from nvmath._utils import module_init_force_cupy_lib_load from nvmath.bindings.cublas import ComputeType # type: ignore -from . import advanced -from .generic import ( +module_init_force_cupy_lib_load() + +from . import advanced # noqa: E402 +from .generic import ( # noqa: E402 DiagonalMatrixQualifier, DiagType, ExecutionCPU, @@ -23,7 +26,7 @@ matmul, matrix_qualifiers_dtype, ) -from .generic.solvermod import ( +from .generic.solvermod import ( # noqa: E402 DirectSolver, DirectSolverOptions, InvalidDirectSolverState, diff --git a/nvmath/sparse/__init__.py b/nvmath/sparse/__init__.py index 32a3d38..a1c8d5f 100644 --- a/nvmath/sparse/__init__.py +++ b/nvmath/sparse/__init__.py @@ -2,24 +2,28 @@ # # SPDX-License-Identifier: Apache-2.0 +from nvmath._utils import module_init_force_cupy_lib_load + +module_init_force_cupy_lib_load() + from . import advanced, ust # noqa: E402 -from .generic import ComputeType, ExecutionCUDA, Matmul, MatmulOptions, matmul -from .generic import ( +from .generic import ComputeType, ExecutionCUDA, Matmul, MatmulOptions, matmul # noqa: E402 +from .generic import ( # noqa: E402 compile_add as compile_matmul_add, ) -from .generic import ( +from .generic import ( # noqa: E402 compile_atomic_add as compile_matmul_atomic_add, ) -from .generic import ( +from .generic import ( # noqa: E402 compile_epilog as compile_matmul_epilog, ) -from .generic import ( +from .generic import ( # noqa: E402 compile_mul as compile_matmul_mul, ) -from .generic import ( +from .generic import ( # noqa: E402 compile_prolog as compile_matmul_prolog, ) -from .generic import matrix_qualifiers_dtype as matmul_matrix_qualifiers_dtype +from .generic import matrix_qualifiers_dtype as matmul_matrix_qualifiers_dtype # noqa: E402 __all__ = [ "advanced", diff --git a/nvmath/tensor/__init__.py b/nvmath/tensor/__init__.py index 4ae47a9..97e6db3 100644 --- a/nvmath/tensor/__init__.py +++ b/nvmath/tensor/__init__.py @@ -2,5 +2,9 @@ # # SPDX-License-Identifier: Apache-2.0 -from ._configuration import * # noqa: F403 -from .contract import * # noqa: F403 +from nvmath._utils import module_init_force_cupy_lib_load + +module_init_force_cupy_lib_load() + +from ._configuration import * # noqa: E402, F403 +from .contract import * # noqa: E402, F403 diff --git a/releasenotes/notes/lazy-bindings-init-7a3c9e2b1d4f8c60.yaml b/releasenotes/notes/lazy-bindings-init-7a3c9e2b1d4f8c60.yaml new file mode 100644 index 0000000..1f4fc51 --- /dev/null +++ b/releasenotes/notes/lazy-bindings-init-7a3c9e2b1d4f8c60.yaml @@ -0,0 +1,19 @@ +--- +# Uncomment a section header and replace with your own note. +#release_summary: +#api: +features: + - | + ``nvmath`` and ``nvmath.bindings`` now import high-level subpackages and + library wrappers on first access instead of loading every wrapper + (including optional ``*Mp`` and NVSHMEM modules) during package import. + Importing ``nvmath.bindings.cycurand`` therefore no longer pulls in + ``fft`` / ``linalg`` / ``sparse`` / ``tensor``. Public names, ``__all__``, + and the behavior that optional binding modules are ``None`` when + unavailable are unchanged. +#dependencies: +#deprecations: +#fixes: +#issues: +#security: +#doc: diff --git a/tests/nvmath_tests/test_bindings_init.py b/tests/nvmath_tests/test_bindings_init.py new file mode 100644 index 0000000..20e1db6 --- /dev/null +++ b/tests/nvmath_tests/test_bindings_init.py @@ -0,0 +1,233 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED. +# +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for lazy loading of nvmath and nvmath.bindings submodules.""" + +import os +import subprocess +import sys +import textwrap + +import nvmath +import nvmath.bindings as bindings + +# Required bindings wrappers that EAGER_IMPORT must force-load at import time. +REQUIRED_WRAPPERS = [ + "cublas", + "cublasLt", + "cudss", + "cufft", + "curand", + "cusolver", + "cusolverDn", + "cusolverSp", + "cusparse", + "cusparseLt", + "cutensor", + "nvpl", +] + +EXPECTED_BINDINGS_ALL = [ + "cublas", + "cublasLt", + "cublasMp", + "cudss", + "cufft", + "cufftMp", + "curand", + "cusolver", + "cusolverDn", + "cusolverSp", + "cusolverMp", + "cusparse", + "cusparseLt", + "cutensor", + "nvpl", + "nvshmem", +] + +EXPECTED_NVMATH_ALL = [ + "BaseCUDAMemoryManager", + "BaseCUDAMemoryManagerAsync", + "bindings", + "ComputeType", + "CudaDataType", + "fft", + "LibraryPropertyType", + "linalg", + "MemoryPointer", + "sparse", + "tensor", +] + + +def _run(code: str, env: dict | None = None) -> str: + # A fresh subprocess is required: the eager/lazy behavior is decided at + # import time of nvmath[.bindings], and those modules are already cached in + # the pytest session's sys.modules. It also keeps EAGER_IMPORT's whole-tree + # load from polluting sys.modules for the rest of the session. + run_env = None + if env is not None: + run_env = os.environ.copy() + run_env.update(env) + result = subprocess.run( + [sys.executable, "-I", "-c", textwrap.dedent(code)], + check=True, + capture_output=True, + text=True, + env=run_env, + ) + return result.stdout.strip() + + +def test_bindings_all_unchanged(): + assert bindings.__all__ == EXPECTED_BINDINGS_ALL + + +def test_nvmath_all_unchanged(): + assert nvmath.__all__ == EXPECTED_NVMATH_ALL + + +def test_dir_lists_public_wrappers(): + names = _run( + """ + import nvmath.bindings as bindings + print(','.join(n for n in bindings.__all__ if n in dir(bindings))) + """ + ) + listed = names.split(",") if names else [] + assert listed == EXPECTED_BINDINGS_ALL + + +def test_dir_lists_nvmath_public_names(): + names = _run( + """ + import nvmath + print(','.join(n for n in nvmath.__all__ if n in dir(nvmath))) + """ + ) + listed = names.split(",") if names else [] + assert listed == EXPECTED_NVMATH_ALL + + +def test_cycurand_does_not_import_host_or_optional_wrappers(): + loaded = _run( + """ + import sys + import nvmath.bindings.cycurand # noqa: F401 + forbidden = [ + "nvmath.fft", + "nvmath.linalg", + "nvmath.sparse", + "nvmath.tensor", + "nvmath.bindings.cublas", + "nvmath.bindings.cufft", + "nvmath.bindings.cublasMp", + "nvmath.bindings.cufftMp", + "nvmath.bindings.cusolverMp", + "nvmath.bindings.nvshmem", + ] + loaded = [n for n in forbidden if n in sys.modules] + print(','.join(loaded)) + """ + ) + assert loaded == "" + + +def test_attribute_access_returns_wrapper_or_none(): + status = _run( + """ + import nvmath.bindings as bindings + import types + + curand = bindings.curand + cublas = bindings.cublas + assert isinstance(curand, types.ModuleType) + assert isinstance(cublas, types.ModuleType) + assert curand is bindings.curand + assert cublas is bindings.cublas + + for name in ("cublasMp", "cufftMp", "cusolverMp", "nvshmem"): + value = getattr(bindings, name) + assert value is None or isinstance(value, types.ModuleType), name + print("ok") + """ + ) + assert status == "ok" + + +def test_nvmath_lazy_public_attributes(): + status = _run( + """ + import sys + import types + import nvmath + + assert "nvmath.fft" not in sys.modules + assert "nvmath.linalg" not in sys.modules + assert isinstance(nvmath.ComputeType, type) + assert isinstance(nvmath.fft, types.ModuleType) + assert "nvmath.fft" in sys.modules + assert nvmath.fft is nvmath.fft + print("ok") + """ + ) + assert status == "ok" + + +def test_star_import_and_unknown_attribute(): + status = _run( + """ + from nvmath.bindings import * # noqa: F403 + import nvmath.bindings as bindings + + assert bindings.cublas is not None + try: + bindings.not_a_real_binding + except AttributeError: + print("ok") + else: + raise SystemExit("expected AttributeError") + """ + ) + assert status == "ok" + + +def test_eager_import_loads_required_wrappers_at_import(): + # With EAGER_IMPORT set, merely importing nvmath.bindings must pull in every + # required wrapper (and the host-API subpackages) before any attribute + # access, so a broken/missing import fails loudly at import time. + required = ",".join(REQUIRED_WRAPPERS) + missing = _run( + f""" + import sys + import nvmath.bindings # noqa: F401 + + required = "{required}".split(",") + missing = [n for n in required if f"nvmath.bindings.{{n}}" not in sys.modules] + # Host-API subpackages should also be eagerly loaded via nvmath.__init__. + for sub in ("nvmath.fft", "nvmath.linalg", "nvmath.sparse", "nvmath.tensor"): + if sub not in sys.modules: + missing.append(sub) + print(",".join(missing)) + """, + env={"EAGER_IMPORT": "1"}, + ) + assert missing == "" + + +def test_default_import_stays_lazy(): + # The complement of the eager test: with EAGER_IMPORT unset, importing + # nvmath.bindings must not eagerly load the wrappers. + loaded = _run( + """ + import sys + import nvmath.bindings # noqa: F401 + + watch = ("nvmath.bindings.cublas", "nvmath.bindings.cufft", "nvmath.fft") + print(",".join(n for n in watch if n in sys.modules)) + """, + env={"EAGER_IMPORT": ""}, + ) + assert loaded == ""