[Common] Fix CUDA library loading for pip-installed frameworks

docs/troubleshooting.rst

@@ -69,6 +69,10 @@ UV and virtual environments
        export CUDNN_HOME=$CUDNN_PATH
        export LD_LIBRARY_PATH=$CUDNN_PATH/lib:$LD_LIBRARY_PATH
 
+   At runtime, Transformer Engine searches for each CUDA library in ``<LIBRARY>_HOME``,
+   ``<LIBRARY>_PATH``, ``CUDA_HOME``, ``CUDA_PATH``, ``/usr/local/cuda``, and ``LD_LIBRARY_PATH``.
+   It then searches the Python environment, and finally the dynamic linker's default locations.
+
 3. **Building wheels**
 
    Use ``uv build --wheel --no-build-isolation -v`` when building the wheel and

tests/pytorch/test_sanity.py

@@ -8,8 +8,13 @@
 import pytest
 import os
 import copy
+import ctypes
+import subprocess
+import sys
+import textwrap
 
 import transformer_engine
+import transformer_engine.common as te_common
 import transformer_engine.pytorch as te
 from transformer_engine.pytorch.quantization import FP8GlobalStateManager
 from transformer_engine.pytorch.utils import (
@@ -1651,3 +1656,136 @@ def check_weights():
         with autocast(enabled=with_quantization, recipe=quantization_recipe):
             y = module(x, **kwargs)
     check_weights()
+
+
+_SEARCH_LIB = "tesearchordertest"
+
+
+def _run_cuda_library_search(monkeypatch, tmp_path, ld_library_path, linker_finds_library=True):
+    """Run the CUDA library search with a fake loader. Return its result and the loaded names."""
+    loaded = []
+
+    def fake_cdll(name, mode=0):
+        missing = not os.path.exists(name) if "/" in name else not linker_finds_library
+        if missing:
+            raise OSError(f"{name}: cannot open shared object file")
+        loaded.append(name)
+        return object()
+
+    monkeypatch.setattr(ctypes, "CDLL", fake_cdll)
+    monkeypatch.setattr(te_common.sysconfig, "get_path", lambda _: str(tmp_path / "python"))
+    for variable in ("CUDA_HOME", "CUDA_PATH", f"{_SEARCH_LIB.upper()}_HOME"):
+        monkeypatch.delenv(variable, raising=False)
+    monkeypatch.setenv("LD_LIBRARY_PATH", ld_library_path)
+    monkeypatch.setenv(f"{_SEARCH_LIB.upper()}_PATH", str(tmp_path / "lib_path"))
+
+    te_common._load_cuda_library_from_system.cache_clear()
+    te_common._load_cuda_library_from_python.cache_clear()
+    te_common._load_cuda_library_from_linker_default.cache_clear()
+    te_common._load_cuda_library.cache_clear()
+
+    found_in_system, _ = te_common._load_cuda_library(_SEARCH_LIB)
+    return found_in_system, loaded
+
+
+def _create_libraries(directory, *suffixes):
+    directory.mkdir(parents=True)
+    for suffix in suffixes:
+        (directory / f"lib{_SEARCH_LIB}.so{suffix}").touch()
+
+
+@pytest.mark.parametrize(
+    "available, expected",
+    [
+        (("linker",), "linker"),
+        (("linker", "python"), "python"),
+        (("linker", "python", "ld_library_path"), "ld_library_path"),
+        (("linker", "python", "ld_library_path", "lib_path"), "lib_path"),
+    ],
+)
+def test_cuda_library_search_order(monkeypatch, tmp_path, available, expected):
+    """Selected locations outrank the Python environment, which outranks a host-wide library."""
+    locations = {
+        "python": tmp_path / "python" / "nvidia" / _SEARCH_LIB / "lib",
+        "ld_library_path": tmp_path / "ld_library_path",
+        "lib_path": tmp_path / "lib_path" / "lib",
+    }
+    for name, location in locations.items():
+        if name in available:
+            _create_libraries(location, ".1")
+
+    found_in_system, loaded = _run_cuda_library_search(
+        monkeypatch,
+        tmp_path,
+        str(locations["ld_library_path"]),
+        linker_finds_library="linker" in available,
+    )
+
+    assert found_in_system == (expected != "python")
+    if expected == "linker":
+        assert loaded == [f"lib{_SEARCH_LIB}.so"]
+    else:
+        assert loaded == [str(locations[expected] / f"lib{_SEARCH_LIB}.so.1")]
+
+
+@pytest.mark.parametrize(
+    "suffixes, expected_suffix",
+    [(("", ".1", ".2"), ""), ((".1", ".2"), ".2")],
+)
+def test_cuda_library_ld_library_path_name_choice(monkeypatch, tmp_path, suffixes, expected_suffix):
+    """An LD_LIBRARY_PATH entry follows the unversioned name, then the newest versioned name."""
+    directory = tmp_path / "ld_library_path"
+    _create_libraries(directory, *suffixes)
+
+    _, loaded = _run_cuda_library_search(monkeypatch, tmp_path, str(directory))
+
+    assert loaded == [str(directory / f"lib{_SEARCH_LIB}.so{expected_suffix}")]
+
+
+@pytest.mark.parametrize("entry", ["$ORIGIN/lib", "${ORIGIN}/lib"])
+def test_cuda_library_ld_library_path_expands_origin(monkeypatch, tmp_path, entry):
+    """A $ORIGIN entry that has the library outranks the Python environment."""
+    _create_libraries(tmp_path / "python" / "nvidia" / _SEARCH_LIB / "lib", ".1")
+    _create_libraries(tmp_path / "bin" / "lib", "")
+    monkeypatch.setattr(sys, "executable", str(tmp_path / "bin" / "python"))
+
+    found_in_system, loaded = _run_cuda_library_search(monkeypatch, tmp_path, entry)
+
+    assert found_in_system
+    assert loaded == [str(tmp_path / "bin" / "lib" / f"lib{_SEARCH_LIB}.so")]
+
+
+def test_cuda_library_token_entry_without_library_keeps_python_priority(monkeypatch, tmp_path):
+    """A token entry that lacks the library must not let a host-wide library win."""
+    location = tmp_path / "python" / "nvidia" / _SEARCH_LIB / "lib"
+    _create_libraries(location, ".1")
+
+    found_in_system, loaded = _run_cuda_library_search(
+        monkeypatch, tmp_path, "$ORIGIN/missing:$LIB/missing"
+    )
+
+    assert not found_in_system
+    assert loaded == [str(location / f"lib{_SEARCH_LIB}.so.1")]
+
+
+def test_import_does_not_require_curand_library():
+    """TE only uses cuRAND's device headers, and torch no longer always installs its library."""
+    probe = textwrap.dedent(
+        """
+        import ctypes
+
+        cdll = ctypes.CDLL
+
+
+        def cdll_without_curand(name, *args, **kwargs):
+            if name is not None and "curand" in str(name):
+                raise OSError(f"{name}: cannot open shared object file")
+            return cdll(name, *args, **kwargs)
+
+
+        ctypes.CDLL = cdll_without_curand
+
+        import transformer_engine.pytorch
+        """
+    )
+    subprocess.run([sys.executable, "-c", probe], check=True)

transformer_engine/common/__init__.py

@@ -345,12 +345,11 @@ def _nvidia_cudart_include_dir(cuda_major_version: int) -> str:
 
 
 @functools.lru_cache(maxsize=None)
-def _load_cuda_library_from_python(lib_name: str, strict: bool = False):
+def _load_cuda_library_from_python(lib_name: str):
     """
     Attempts to load shared object file installed via python packages.
 
     `lib_name` : Name of package as found in the `nvidia` dir in python environment.
-    `strict` : If set to `True`, throw an error if lib is not found.
     """
 
     ext = _get_sys_extension()
@@ -382,20 +381,25 @@ def _load_cuda_library_from_python(lib_name: str, strict: bool = False):
         for so_path in so_paths:
             ctypes_handles.append(ctypes.CDLL(so_path, mode=ctypes.RTLD_GLOBAL))
 
-    if strict and not path_found:
-        raise RuntimeError(f"{lib_name} shared object not found.")
-
     return path_found, ctypes_handles
 
 
+def _newest_cuda_library(pattern: str) -> Optional[str]:
+    """Return the highest-sorted non-stub library that matches a glob pattern."""
+    libs = [lib for lib in glob.glob(pattern, recursive=True) if "stub" not in lib]
+    return max(libs, key=os.path.basename, default=None)
+
+
 @functools.lru_cache(maxsize=None)
 def _load_cuda_library_from_system(lib_name: str):
     """
-    Attempts to load shared object file installed via system/cuda-toolkit.
+    Attempts to load shared object file from a location selected by the environment or toolkit.
 
     `lib_name`: Name of library to load without extension or `lib` prefix.
     """
 
+    file_name = f"lib{lib_name}{_get_sys_extension()}"
+
     # Where to look for the shared lib in decreasing order of preference.
     paths = (
         os.environ.get(f"{lib_name.upper()}_HOME"),
@@ -408,16 +412,41 @@ def _load_cuda_library_from_system(lib_name: str):
     for path in paths:
         if path is None:
             continue
-        libs = glob.glob(f"{path}/**/lib{lib_name}{_get_sys_extension()}*", recursive=True)
-        libs = [lib for lib in libs if "stub" not in lib]
-        libs.sort(reverse=True, key=os.path.basename)
-        if libs:
-            return True, ctypes.CDLL(libs[0], mode=ctypes.RTLD_GLOBAL)
+        lib = _newest_cuda_library(f"{path}/**/{file_name}*")
+        if lib:
+            return True, ctypes.CDLL(lib, mode=ctypes.RTLD_GLOBAL)
+
+    origin = os.path.dirname(os.path.realpath(sys.executable))
+    for path in os.environ.get("LD_LIBRARY_PATH", "").split(os.pathsep):
+        if not path:
+            continue
+        # In this variable, $ORIGIN is the executable's directory. dlopen would use the caller's.
+        path = path.replace("${ORIGIN}", origin).replace("$ORIGIN", origin)
+        exact = os.path.join(path, file_name)
+        # dlopen expands the other tokens, such as $LIB, but a file search cannot.
+        if "$" in path:
+            try:
+                return True, ctypes.CDLL(exact, mode=ctypes.RTLD_GLOBAL)
+            except OSError:
+                continue
+        # Follow the unversioned name as the linker does. Python packages only have versioned names.
+        lib = exact if os.path.exists(exact) else _newest_cuda_library(f"{exact}*")
+        if lib:
+            return True, ctypes.CDLL(lib, mode=ctypes.RTLD_GLOBAL)
+
+    return False, None
+
+
+@functools.lru_cache(maxsize=None)
+def _load_cuda_library_from_linker_default(lib_name: str):
+    """
+    Attempts to load shared object file through the dynamic linker's default search.
+
+    `lib_name`: Name of library to load without extension or `lib` prefix.
+    """
 
-    # Search in LD_LIBRARY_PATH.
     try:
-        _lib_handle = ctypes.CDLL(f"lib{lib_name}{_get_sys_extension()}", mode=ctypes.RTLD_GLOBAL)
-        return True, _lib_handle
+        return True, ctypes.CDLL(f"lib{lib_name}{_get_sys_extension()}", mode=ctypes.RTLD_GLOBAL)
     except OSError:
         return False, None
 
@@ -426,11 +455,11 @@ def _load_cuda_library_from_system(lib_name: str):
 def _load_cuda_library(lib_name: str):
     """
     Load given shared library.
-    Prioritize loading from system/toolkit
-    before checking python packages.
+    Prioritize locations selected by the environment or toolkit, then python packages,
+    then the dynamic linker's default search.
     """
 
-    # Attempt to locate library in system.
+    # Attempt to locate library in a selected location.
     found, handle = _load_cuda_library_from_system(lib_name)
     if found:
         return True, handle
@@ -440,6 +469,12 @@ def _load_cuda_library(lib_name: str):
     if found:
         return False, handle
 
+    # A host-wide library comes last: pip-installed frameworks depend on the version in their own
+    # Python environment, and loading another copy globally would replace it in this process.
+    found, handle = _load_cuda_library_from_linker_default(lib_name)
+    if found:
+        return True, handle
+
     raise RuntimeError(f"{lib_name} shared object not found.")
 
 
@@ -503,16 +538,14 @@ def _load_core_library():
     # and an error is thrown if not found.
     _, _CUDNN_LIB_CTYPES = _load_cuda_library("cudnn")
     system_nvrtc, _NVRTC_LIB_CTYPES = _load_cuda_library("nvrtc")
-    system_curand, _CURAND_LIB_CTYPES = _load_cuda_library("curand")
 
     # This additional step is necessary to be able to install TE wheels
     # and import TE (without any guards) in an environment where the cuda
     # toolkit might be absent without being guarded
-    load_libs_for_no_ctk = not system_nvrtc and not system_curand
+    load_libs_for_no_ctk = not system_nvrtc
     if load_libs_for_no_ctk:
-        _CUBLAS_LIB_CTYPES = _load_cuda_library_from_python("cublas", strict=True)
-        _CUDART_LIB_CTYPES = _load_cuda_library_from_python("cudart", strict=True)
-        _CUDNN_ALL_LIB_CTYPES = _load_cuda_library_from_python("cudnn", strict=True)
+        _CUBLAS_LIB_CTYPES = _load_cuda_library_from_python("cublas")
+        _CUDART_LIB_CTYPES = _load_cuda_library_from_python("cudart")
 
     # Prepare CuTeDSL backend for on-demand compilation via TVM-FFI if user enables it.
     if os.environ.get("NVTE_ENABLE_CUTEDSL_BACKEND", "0") != "0":