Repository: vllm-project/vllm
File: vllm/env_override.py (Python)

Code before:
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os


# The CUDA compatibility library path must be set before importing
# torch, because torch loads CUDA shared libraries at import time and the
# dynamic linker only consults LD_LIBRARY_PATH when a library is first
# loaded. Setting it afterwards has no effect.
# ruff: noqa: E402
def _maybe_set_cuda_compatibility_path():
    enable_cuda_compatibility = os.environ.get(
        "VLLM_ENABLE_CUDA_COMPATIBILITY", "0"
    ).strip().lower() in ("1", "true")
    if not enable_cuda_compatibility:
        return
    cuda_compat_path = os.environ.get("VLLM_CUDA_COMPATIBILITY_PATH", None)
    if not cuda_compat_path or not os.path.isdir(cuda_compat_path):
        return
    existing_ld_library = os.environ.get("LD_LIBRARY_PATH", None)
    if not existing_ld_library:
        ld_library_paths = []
    else:
        ld_library_paths = existing_ld_library.split(os.pathsep)
    norm_cuda_compat_path = os.path.normpath(cuda_compat_path)
    if (
        ld_library_paths
        and ld_library_paths[0]
        and os.path.normpath(ld_library_paths[0]) == norm_cuda_compat_path
    ):
        return
    new_ld_library_paths = [norm_cuda_compat_path] + [
        ld_library_path
        for ld_library_path in ld_library_paths
        if not ld_library_path
        or os.path.normpath(ld_library_path) != norm_cuda_compat_path
    ]
    os.environ["LD_LIBRARY_PATH"] = os.pathsep.join(new_ld_library_paths)


_maybe_set_cuda_compatibility_path()

import torch


Review comment (bug):
When `VLLM_ENABLE_CUDA_COMPATIBILITY` is set but `VLLM_CUDA_COMPATIBILITY_PATH` is missing/invalid, the code silently returns. This can make CUDA import failures much harder to diagnose. Consider logging a warning (or raising a clear error) when compatibility is explicitly enabled but the path is not usable.

Code after:
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# ruff: noqa: E402
import importlib.util
import os


def _get_torch_cuda_version():
    """Peripheral function to _maybe_set_cuda_compatibility_path().
    PyTorch version must not be determined by importing directly
    because it will trigger the CUDA initialization, losing the
    chance to set the LD_LIBRARY_PATH beforehand.
    """
    try:
        spec = importlib.util.find_spec("torch")
        if not spec:
            return None
        if spec.origin:
            torch_root = os.path.dirname(spec.origin)
        elif spec.submodule_search_locations:
            torch_root = spec.submodule_search_locations[0]
        else:
            return None
        version_path = os.path.join(torch_root, "version.py")
        if not os.path.exists(version_path):
            return None
        # Load the version module without importing torch
        ver_spec = importlib.util.spec_from_file_location("torch.version", version_path)
        if not ver_spec or not ver_spec.loader:
            return None
        module = importlib.util.module_from_spec(ver_spec)
        # Avoid registering in sys.modules to not confuse future imports
        ver_spec.loader.exec_module(module)
        return getattr(module, "cuda", None)
    except Exception:
        return None


def _maybe_set_cuda_compatibility_path():
    """Set LD_LIBRARY_PATH for CUDA forward compatibility if enabled.

    Must run before 'import torch' since torch loads CUDA shared libraries
    at import time and the dynamic linker only consults LD_LIBRARY_PATH when
    a library is first loaded.