Prefer Python CUDA libs over implicit system paths

[sbhavani]

Description

Prefer CUDA runtime libraries from the active Python environment before falling back to implicit system locations such as /usr/local/cuda or dynamic linker lookup.

Explicit user configuration via CUDNN_HOME, CUDNN_PATH, CUDA_HOME, or CUDA_PATH still takes highest priority.

Fixes #2147.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Fixed CUDA runtime path precedence
• Ran a GPU smoke test to verify loader precedence:
  • explicit CUDNN_HOME over Python package
  • Python package over dynamic-linker fallback
  • dynamic-linker fallback when no Python package is present

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• [] I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR refactors library loading in transformer_engine/common/__init__.py to change the precedence of CUDA library discovery. The monolithic _load_cuda_library_from_system function is split into clearly-named helper functions, and Python-installed nvidia packages are now consulted before the implicit /usr/local/cuda path and dynamic linker.

• Old priority: Explicit env vars → /usr/local/cuda + dynamic linker → Python dist-packages
• New priority: Explicit env vars → Python dist-packages → /usr/local/cuda → dynamic linker
• The load_libs_for_no_ctk boolean logic in the module-level loader (lines 389–400) correctly interprets the new return-value semantics; False is still returned only for Python-package sources, so cublas/cudart are still explicitly loaded when needed.

Confidence Score: 5/5

Safe to merge — the change is a well-scoped refactor of library discovery order with no functional regressions in the existing load_libs_for_no_ctk logic.

The priority inversion (Python packages before /usr/local/cuda and the dynamic linker) is the intended behavioral change, and the return-value contract (True = system source, False = Python package) is preserved correctly across all four load paths. The module-level consumer of that boolean (load_libs_for_no_ctk) will therefore behave identically to before for every combination of installed libraries.

No files require special attention.

Important Files Changed

Filename	Overview
transformer_engine/common/init.py	Refactors CUDA library loading into four focused helper functions and inverts the Python-package vs. system-default priority; return-value semantics and load_libs_for_no_ctk logic remain consistent with the previous implementation.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["_load_cuda_library(lib_name)"] --> B["_try_load_cuda_library_from_envvars\n(LIBNAME_HOME/PATH, CUDA_HOME/PATH)"]
    B -- "found" --> R1["return True, handle"]
    B -- "None" --> C["_load_cuda_library_from_python\n(nvidia dist-packages)"]
    C -- "found" --> R2["return False, handle"]
    C -- "not found" --> D["_try_load_cuda_library_from_default_path\n(/usr/local/cuda)"]
    D -- "found" --> R3["return True, handle"]
    D -- "None" --> E["_try_load_library_from_dynamic_linker\n(LD_LIBRARY_PATH / system linker)"]
    E -- "found" --> R4["return True, handle"]
    E -- "None" --> F["raise RuntimeError"]

    style R2 fill:#d4edda,stroke:#28a745
    style F fill:#f8d7da,stroke:#dc3545

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A["_load_cuda_library(lib_name)"] --> B["_try_load_cuda_library_from_envvars\n(LIBNAME_HOME/PATH, CUDA_HOME/PATH)"]
    B -- "found" --> R1["return True, handle"]
    B -- "None" --> C["_load_cuda_library_from_python\n(nvidia dist-packages)"]
    C -- "found" --> R2["return False, handle"]
    C -- "not found" --> D["_try_load_cuda_library_from_default_path\n(/usr/local/cuda)"]
    D -- "found" --> R3["return True, handle"]
    D -- "None" --> E["_try_load_library_from_dynamic_linker\n(LD_LIBRARY_PATH / system linker)"]
    E -- "found" --> R4["return True, handle"]
    E -- "None" --> F["raise RuntimeError"]

    style R2 fill:#d4edda,stroke:#28a745
    style F fill:#f8d7da,stroke:#dc3545

Loading

_{Reviews (2): Last reviewed commit: "refactor(common): split CUDA library sea..." | Re-trigger Greptile}

[timmoon10]

This design is kind of convoluted and overtuned to the use-case in _load_cuda_library. Now that we've separated the envvar search from the /usr/local/cuda+ dynamic linker path, there's no reason for them to live in the same function. I'd suggest splitting the function up into fine-grained helper functions that _load_cuda_library can call in any order:

def _try_load_library_from_paths(file_name: str, paths: Iterable[str]) -> Optional[ctypes.CDLL]:
    for path in paths:
        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 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]:
    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 = [p for p in paths if p is not None]
    return _try_load_library_from_paths(
        f"lib{lib_name}{_get_sys_extension()}",
        paths,
    )

@functools.lru_cache(maxsize=None)
def _try_load_cuda_library_from_default_path(lib_name: str) -> Optional[ctypes.CDLL]:
    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:
        return ctypes.CDLL(file_name, mode=ctypes.RTLD_GLOBAL)
    except OSError:
        return None

[sbhavani]

thanks for the suggestion! Just refactored the code into helpers