From 685682032c0a88a4c8938bc8e7efb1aebb546b85 Mon Sep 17 00:00:00 2001 From: Bruno Volpato Date: Mon, 2 Nov 2026 19:32:42 -0500 Subject: [PATCH 1/4] [Common] Search Python packages before the linker default paths A pip-installed framework depends on the CUDA libraries in its own Python environment. Transformer Engine searched the dynamic linker's default paths first, so a host-wide cuDNN replaced the one that torch needs. With torch 2.15 and an older system cuDNN, this raises "cuDNN version incompatibility". Keep explicit locations first: the library-specific variables, CUDA_HOME, CUDA_PATH, /usr/local/cuda, and LD_LIBRARY_PATH. Then search the Python environment, and use the linker default last. LD_LIBRARY_PATH is now read directly, so it also finds versioned libraries that have no unversioned symlink. Signed-off-by: Bruno Volpato --- docs/troubleshooting.rst | 4 ++ tests/pytorch/test_sanity.py | 53 +++++++++++++++++++++++++++ transformer_engine/common/__init__.py | 40 ++++++++++++++------ 3 files changed, 86 insertions(+), 11 deletions(-) diff --git a/docs/troubleshooting.rst b/docs/troubleshooting.rst index cd5bdd18381..abc136f5b78 100644 --- a/docs/troubleshooting.rst +++ b/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 ``_HOME``, + ``_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 diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 2627de023ac..2ed7fcdb362 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -8,8 +8,10 @@ import pytest import os import copy +import ctypes 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 +1653,54 @@ def check_weights(): with autocast(enabled=with_quantization, recipe=quantization_recipe): y = module(x, **kwargs) check_weights() + + +@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.""" + lib_name = "tesearchordertest" + locations = { + "python": tmp_path / "python" / "nvidia" / lib_name / "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: + location.mkdir(parents=True) + (location / f"lib{lib_name}.so.1").touch() + + loaded = [] + + def fake_cdll(name, mode=0): + if name == f"lib{lib_name}.so" and "linker" not in available: + 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"{lib_name.upper()}_HOME"): + monkeypatch.delenv(variable, raising=False) + monkeypatch.setenv("LD_LIBRARY_PATH", str(locations["ld_library_path"])) + monkeypatch.setenv(f"{lib_name.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(lib_name) + + assert found_in_system == (expected != "python") + if expected == "linker": + assert loaded == [f"lib{lib_name}.so"] + else: + assert loaded == [str(locations[expected] / f"lib{lib_name}.so.1")] diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index 6237c6732b9..292b324ac71 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -391,11 +391,13 @@ def _load_cuda_library_from_python(lib_name: str, strict: bool = False): @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. """ + ext = _get_sys_extension() + # Where to look for the shared lib in decreasing order of preference. paths = ( os.environ.get(f"{lib_name.upper()}_HOME"), @@ -404,20 +406,30 @@ def _load_cuda_library_from_system(lib_name: str): os.environ.get("CUDA_PATH"), "/usr/local/cuda", ) + patterns = [f"{path}/**/lib{lib_name}{ext}*" for path in paths if path is not None] + ld_library_path = os.environ.get("LD_LIBRARY_PATH", "").split(os.pathsep) + patterns += [f"{path}/lib{lib_name}{ext}*" for path in ld_library_path if path] - for path in paths: - if path is None: - continue - libs = glob.glob(f"{path}/**/lib{lib_name}{_get_sys_extension()}*", recursive=True) + for pattern in patterns: + libs = glob.glob(pattern, 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) - # Search in LD_LIBRARY_PATH. + 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. + """ + 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 +438,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 +452,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.") From 5e01f300167ce0c57226542b5be33c2255e8064e Mon Sep 17 00:00:00 2001 From: Bruno Volpato Date: Mon, 2 Nov 2026 19:57:57 -0500 Subject: [PATCH 2/4] [Common] Do not require the cuRAND library at import Importing Transformer Engine raises "curand shared object not found" on a host without the CUDA toolkit when PyTorch 2.15 is installed. PyTorch 2.15 removed curand from its cuda-toolkit extras, and JAX does not install it. Transformer Engine only uses cuRAND's device-side headers. The core library does not link libcurand and calls none of its host functions, so the import does not need to load it. Decide whether the toolkit is present from NVRTC alone. Signed-off-by: Bruno Volpato --- tests/pytorch/test_sanity.py | 26 ++++++++++++++++++++++++++ transformer_engine/common/__init__.py | 3 +-- 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 2ed7fcdb362..c4523afc9ef 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -9,6 +9,9 @@ import os import copy import ctypes +import subprocess +import sys +import textwrap import transformer_engine import transformer_engine.common as te_common @@ -1704,3 +1707,26 @@ def fake_cdll(name, mode=0): assert loaded == [f"lib{lib_name}.so"] else: assert loaded == [str(locations[expected] / f"lib{lib_name}.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) diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index 292b324ac71..0867bbcc5af 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -521,12 +521,11 @@ 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) From 6238e1689374ece82107084c4eb07984a9451140 Mon Sep 17 00:00:00 2001 From: Bruno Volpato Date: Tue, 3 Nov 2026 14:50:05 -0500 Subject: [PATCH 3/4] [Common] Follow linker behavior in the CUDA library search Keep three cases that the previous search changes altered: - Only the dynamic linker can expand tokens such as $ORIGIN in LD_LIBRARY_PATH, so leave those entries to the linker. Otherwise, the Python packages take priority over the selected location. - Follow the unversioned name in an LD_LIBRARY_PATH entry when it exists, because it selects a version. Use the newest versioned name only when it does not exist. - When NVRTC comes from a Python package, the other libraries can still come from the system. Load cuBLAS and cudart from Python packages only if they are there, and leave the rest to the dynamic linker. Remove the second cuDNN load, which is redundant for a Python cuDNN and otherwise maps another copy. Signed-off-by: Bruno Volpato --- tests/pytorch/test_sanity.py | 94 +++++++++++++++++++-------- transformer_engine/common/__init__.py | 49 +++++++++----- 2 files changed, 99 insertions(+), 44 deletions(-) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index c4523afc9ef..b40814799c4 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -1658,6 +1658,41 @@ def check_weights(): 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): + if name == f"lib{_SEARCH_LIB}.so" and not linker_finds_library: + 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", [ @@ -1669,44 +1704,51 @@ def check_weights(): ) def test_cuda_library_search_order(monkeypatch, tmp_path, available, expected): """Selected locations outrank the Python environment, which outranks a host-wide library.""" - lib_name = "tesearchordertest" locations = { - "python": tmp_path / "python" / "nvidia" / lib_name / "lib", + "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: - location.mkdir(parents=True) - (location / f"lib{lib_name}.so.1").touch() + _create_libraries(location, ".1") - loaded = [] + found_in_system, loaded = _run_cuda_library_search( + monkeypatch, + tmp_path, + str(locations["ld_library_path"]), + linker_finds_library="linker" in available, + ) - def fake_cdll(name, mode=0): - if name == f"lib{lib_name}.so" and "linker" not in available: - raise OSError(f"{name}: cannot open shared object file") - loaded.append(name) - return object() + 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")] - 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"{lib_name.upper()}_HOME"): - monkeypatch.delenv(variable, raising=False) - monkeypatch.setenv("LD_LIBRARY_PATH", str(locations["ld_library_path"])) - monkeypatch.setenv(f"{lib_name.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() +@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) - found_in_system, _ = te_common._load_cuda_library(lib_name) + _, loaded = _run_cuda_library_search(monkeypatch, tmp_path, str(directory)) - assert found_in_system == (expected != "python") - if expected == "linker": - assert loaded == [f"lib{lib_name}.so"] - else: - assert loaded == [str(locations[expected] / f"lib{lib_name}.so.1")] + assert loaded == [str(directory / f"lib{_SEARCH_LIB}.so{expected_suffix}")] + + +def test_cuda_library_ld_library_path_token_uses_linker(monkeypatch, tmp_path): + """Only the linker can expand $ORIGIN, so that entry must outrank the Python environment.""" + _create_libraries(tmp_path / "python" / "nvidia" / _SEARCH_LIB / "lib", ".1") + + found_in_system, loaded = _run_cuda_library_search(monkeypatch, tmp_path, "$ORIGIN/lib") + + assert found_in_system + assert loaded == [f"lib{_SEARCH_LIB}.so"] def test_import_does_not_require_curand_library(): diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index 0867bbcc5af..b092bdf5d00 100644 --- a/transformer_engine/common/__init__.py +++ b/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,12 +381,15 @@ 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): """ @@ -396,7 +398,7 @@ def _load_cuda_library_from_system(lib_name: str): `lib_name`: Name of library to load without extension or `lib` prefix. """ - ext = _get_sys_extension() + file_name = f"lib{lib_name}{_get_sys_extension()}" # Where to look for the shared lib in decreasing order of preference. paths = ( @@ -406,16 +408,28 @@ def _load_cuda_library_from_system(lib_name: str): os.environ.get("CUDA_PATH"), "/usr/local/cuda", ) - patterns = [f"{path}/**/lib{lib_name}{ext}*" for path in paths if path is not None] - ld_library_path = os.environ.get("LD_LIBRARY_PATH", "").split(os.pathsep) - patterns += [f"{path}/lib{lib_name}{ext}*" for path in ld_library_path if path] - for pattern in patterns: - libs = glob.glob(pattern, 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) + for path in paths: + if path is None: + continue + lib = _newest_cuda_library(f"{path}/**/{file_name}*") + if lib: + return True, ctypes.CDLL(lib, mode=ctypes.RTLD_GLOBAL) + + for path in os.environ.get("LD_LIBRARY_PATH", "").split(os.pathsep): + if not path: + continue + # Only the dynamic linker can expand tokens such as $ORIGIN. + if "$" in path: + found, handle = _load_cuda_library_from_linker_default(lib_name) + if found: + return True, handle + continue + # Follow the unversioned name as the linker does. Python packages only have versioned names. + exact = os.path.join(path, file_name) + 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 @@ -527,9 +541,8 @@ def _load_core_library(): # toolkit might be absent without being guarded 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": From acfbca92a24c1652c8afa17119754eac43c44c3a Mon Sep 17 00:00:00 2001 From: Bruno Volpato Date: Tue, 3 Nov 2026 15:04:02 -0500 Subject: [PATCH 4/4] [Common] Limit LD_LIBRARY_PATH token handling to the entry An LD_LIBRARY_PATH entry with a token such as $ORIGIN sent the search to the dynamic linker's bare-name lookup. That lookup also searches the default locations, so an unrelated token entry could load a host-wide cuDNN ahead of the Python package and bring back the version mismatch. Expand $ORIGIN to the executable's directory, as the linker does for this variable, and pass entries with other tokens to dlopen as a path. If an entry does not have the library, the search continues to the next location. Signed-off-by: Bruno Volpato --- tests/pytorch/test_sanity.py | 27 ++++++++++++++++++++++----- transformer_engine/common/__init__.py | 15 +++++++++------ 2 files changed, 31 insertions(+), 11 deletions(-) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index b40814799c4..d73a0891e91 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -1666,7 +1666,8 @@ def _run_cuda_library_search(monkeypatch, tmp_path, ld_library_path, linker_find loaded = [] def fake_cdll(name, mode=0): - if name == f"lib{_SEARCH_LIB}.so" and not linker_finds_library: + 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() @@ -1741,14 +1742,30 @@ def test_cuda_library_ld_library_path_name_choice(monkeypatch, tmp_path, suffixe assert loaded == [str(directory / f"lib{_SEARCH_LIB}.so{expected_suffix}")] -def test_cuda_library_ld_library_path_token_uses_linker(monkeypatch, tmp_path): - """Only the linker can expand $ORIGIN, so that entry must outrank the Python environment.""" +@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, "$ORIGIN/lib") + found_in_system, loaded = _run_cuda_library_search(monkeypatch, tmp_path, entry) assert found_in_system - assert loaded == [f"lib{_SEARCH_LIB}.so"] + 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(): diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index b092bdf5d00..cbd6a373cfc 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -416,17 +416,20 @@ def _load_cuda_library_from_system(lib_name: str): 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 - # Only the dynamic linker can expand tokens such as $ORIGIN. + # 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: - found, handle = _load_cuda_library_from_linker_default(lib_name) - if found: - return True, handle - continue + 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. - exact = os.path.join(path, file_name) lib = exact if os.path.exists(exact) else _newest_cuda_library(f"{exact}*") if lib: return True, ctypes.CDLL(lib, mode=ctypes.RTLD_GLOBAL)