From dd365a55d2646b8c604c76696133ae63fc18cdb5 Mon Sep 17 00:00:00 2001 From: Santosh Bhavani Date: Mon, 6 Jul 2026 08:02:38 -0700 Subject: [PATCH 1/2] Prefer Python CUDA libs over implicit system paths Signed-off-by: Santosh Bhavani --- transformer_engine/common/__init__.py | 48 +++++++++++++++++++-------- 1 file changed, 34 insertions(+), 14 deletions(-) diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index 231680321e..cfabc4bfa3 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -300,21 +300,33 @@ 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): +def _load_cuda_library_from_system( + lib_name: str, + include_env_paths: bool = True, + include_default_paths: bool = True, + include_linker_paths: bool = True, +): """ Attempts to load shared object file installed via system/cuda-toolkit. `lib_name`: Name of library to load without extension or `lib` prefix. + `include_env_paths`: Search paths set with environment variables. + `include_default_paths`: Search default system CUDA installation paths. + `include_linker_paths`: Search paths known to the dynamic linker. """ - # Where to look for the shared lib in decreasing order of preference. - paths = ( - os.environ.get(f"{lib_name.upper()}_HOME"), - os.environ.get(f"{lib_name.upper()}_PATH"), - os.environ.get("CUDA_HOME"), - os.environ.get("CUDA_PATH"), - "/usr/local/cuda", - ) + paths = [] + if include_env_paths: + paths.extend( + ( + os.environ.get(f"{lib_name.upper()}_HOME"), + os.environ.get(f"{lib_name.upper()}_PATH"), + os.environ.get("CUDA_HOME"), + os.environ.get("CUDA_PATH"), + ) + ) + if include_default_paths: + paths.append("/usr/local/cuda") for path in paths: if path is None: @@ -325,7 +337,10 @@ def _load_cuda_library_from_system(lib_name: str): if libs: return True, ctypes.CDLL(libs[0], mode=ctypes.RTLD_GLOBAL) - # Search in LD_LIBRARY_PATH. + # Search paths known to the dynamic linker. + if not include_linker_paths: + return False, None + try: _lib_handle = ctypes.CDLL(f"lib{lib_name}{_get_sys_extension()}", mode=ctypes.RTLD_GLOBAL) return True, _lib_handle @@ -337,12 +352,12 @@ 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. """ - # Attempt to locate library in system. - found, handle = _load_cuda_library_from_system(lib_name) + # Preferred order: explicit user paths, active Python environment, implicit system fallback. + found, handle = _load_cuda_library_from_system( + lib_name, include_default_paths=False, include_linker_paths=False + ) if found: return True, handle @@ -351,6 +366,11 @@ def _load_cuda_library(lib_name: str): if found: return False, handle + # Fall back to the default system installation and dynamic linker paths. + found, handle = _load_cuda_library_from_system(lib_name, include_env_paths=False) + if found: + return True, handle + raise RuntimeError(f"{lib_name} shared object not found.") From 46dbba4f8331a1543da4c4ec12c4bc0b15618e54 Mon Sep 17 00:00:00 2001 From: Santosh Bhavani Date: Fri, 10 Jul 2026 12:53:24 -0700 Subject: [PATCH 2/2] refactor(common): split CUDA library search helpers Signed-off-by: Santosh Bhavani --- transformer_engine/common/__init__.py | 94 +++++++++++++-------------- 1 file changed, 47 insertions(+), 47 deletions(-) diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index cfabc4bfa3..4c1143b4ca 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -15,7 +15,7 @@ import subprocess import sys import sysconfig -from typing import Optional, Tuple +from typing import Iterable, Optional, Tuple @functools.lru_cache(maxsize=None) @@ -299,66 +299,62 @@ def _load_cuda_library_from_python(lib_name: str, strict: bool = False): return path_found, ctypes_handles -@functools.lru_cache(maxsize=None) -def _load_cuda_library_from_system( - lib_name: str, - include_env_paths: bool = True, - include_default_paths: bool = True, - include_linker_paths: bool = True, -): - """ - Attempts to load shared object file installed via system/cuda-toolkit. - - `lib_name`: Name of library to load without extension or `lib` prefix. - `include_env_paths`: Search paths set with environment variables. - `include_default_paths`: Search default system CUDA installation paths. - `include_linker_paths`: Search paths known to the dynamic linker. - """ - - paths = [] - if include_env_paths: - paths.extend( - ( - os.environ.get(f"{lib_name.upper()}_HOME"), - os.environ.get(f"{lib_name.upper()}_PATH"), - os.environ.get("CUDA_HOME"), - os.environ.get("CUDA_PATH"), - ) - ) - if include_default_paths: - paths.append("/usr/local/cuda") - +def _try_load_library_from_paths(file_name: str, paths: Iterable[str]) -> Optional[ctypes.CDLL]: + """Try loading a library from a sequence of filesystem paths.""" for path in paths: - if path is None: - continue - libs = glob.glob(f"{path}/**/lib{lib_name}{_get_sys_extension()}*", recursive=True) + libs = glob.glob(f"{path}/**/{file_name}*", 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) + return ctypes.CDLL(libs[0], mode=ctypes.RTLD_GLOBAL) + return None + + +@functools.lru_cache(maxsize=None) +def _try_load_cuda_library_from_envvars(lib_name: str) -> Optional[ctypes.CDLL]: + """Try loading a CUDA library from paths configured by environment variables.""" + paths = [ + os.environ.get(f"{lib_name.upper()}_HOME"), + os.environ.get(f"{lib_name.upper()}_PATH"), + os.environ.get("CUDA_HOME"), + os.environ.get("CUDA_PATH"), + ] + paths = [path for path in paths if path is not None] + return _try_load_library_from_paths( + f"lib{lib_name}{_get_sys_extension()}", + paths, + ) - # Search paths known to the dynamic linker. - if not include_linker_paths: - return False, None +@functools.lru_cache(maxsize=None) +def _try_load_cuda_library_from_default_path(lib_name: str) -> Optional[ctypes.CDLL]: + """Try loading a CUDA library from the default toolkit installation path.""" + return _try_load_library_from_paths( + f"lib{lib_name}{_get_sys_extension()}", + ["/usr/local/cuda"], + ) + + +def _try_load_library_from_dynamic_linker(file_name: str) -> Optional[ctypes.CDLL]: + """Try loading a library from paths known to the dynamic linker.""" try: - _lib_handle = ctypes.CDLL(f"lib{lib_name}{_get_sys_extension()}", mode=ctypes.RTLD_GLOBAL) - return True, _lib_handle + return ctypes.CDLL(file_name, mode=ctypes.RTLD_GLOBAL) except OSError: - return False, None + return None @functools.lru_cache(maxsize=None) def _load_cuda_library(lib_name: str): """ Load given shared library. + Prioritize explicit system paths, then python packages, then system defaults. """ - # Preferred order: explicit user paths, active Python environment, implicit system fallback. - found, handle = _load_cuda_library_from_system( - lib_name, include_default_paths=False, include_linker_paths=False - ) - if found: + file_name = f"lib{lib_name}{_get_sys_extension()}" + + # Attempt to locate library in explicitly configured system paths. + handle = _try_load_cuda_library_from_envvars(lib_name) + if handle is not None: return True, handle # Attempt to locate library in Python dist-packages. @@ -367,8 +363,12 @@ def _load_cuda_library(lib_name: str): return False, handle # Fall back to the default system installation and dynamic linker paths. - found, handle = _load_cuda_library_from_system(lib_name, include_env_paths=False) - if found: + handle = _try_load_cuda_library_from_default_path(lib_name) + if handle is not None: + return True, handle + + handle = _try_load_library_from_dynamic_linker(file_name) + if handle is not None: return True, handle raise RuntimeError(f"{lib_name} shared object not found.")