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:
        return
    existing_ld_library = os.environ.get("LD_LIBRARY_PATH", None)
    if not existing_ld_library:
        ld_library_paths = []
    else:
        ld_library_paths = [
            ld_library_path
            for ld_library_path in existing_ld_library.split(":")
            if ld_library_path
        ]
    norm_cuda_compat_path = os.path.normpath(cuda_compat_path)
    if (
        ld_library_paths
        and os.path.normpath(ld_library_paths[0]) == norm_cuda_compat_path
    ):
        return
    new_ld_library_paths = [cuda_compat_path] + [
        ld_library_path
        for ld_library_path in ld_library_paths
        if os.path.normpath(ld_library_path) != norm_cuda_compat_path
    ]
    os.environ["LD_LIBRARY_PATH"] = ":".join(new_ld_library_paths)


_maybe_set_cuda_compatibility_path()


Review comment (suggestion):
If VLLM_CUDA_COMPATIBILITY_PATH is set to an empty string, it will be treated as falsy and the function will return early (line 18). However, if a user explicitly sets it to an empty string (e.g., to disable a default value), this silent behavior might be confusing. Consider adding explicit handling for empty strings, either by stripping whitespace and checking, or by logging that an empty path was provided.

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.