20231088/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Simplify logic of locating CUDART so file path (#13203)

Browse Source 
Signed-off-by: Lu Fang <lufang@fb.com>
This commit is contained in:
Lu Fang 2025-02-12 21:52:41 -08:00 committed by GitHub
parent bc55d13070
commit 009439caeb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
vllm/distributed/device_communicators/cuda_wrapper.py
View File

@ -5,7 +5,6 @@ convenient for use when we just need to call a few functions.
""" """
import ctypes import ctypes
import glob
from dataclasses import dataclass from dataclasses import dataclass
from typing import Any, Dict, List, Optional from typing import Any, Dict, List, Optional
@ -62,29 +61,6 @@ def find_loaded_library(lib_name) -> Optional[str]:
return path return path
def get_cudart_lib_path_from_env() -> Optional[str]:
"""
In some system, find_loaded_library() may not work. So we allow users to
specify the path through environment variable VLLM_CUDART_SO_PATH.
"""
cudart_so_env = envs.VLLM_CUDART_SO_PATH
if cudart_so_env is not None:
cudart_paths = [
cudart_so_env,
]
for path in cudart_paths:
file_paths = glob.glob(path)
if len(file_paths) > 0:
logger.info(
"Found cudart library at %s through env var"
"VLLM_CUDART_SO_PATH=%s",
file_paths[0],
cudart_so_env,
)
return file_paths[0]
return None
class CudaRTLibrary: class CudaRTLibrary:
exported_functions = [ exported_functions = [
# ​cudaError_t cudaSetDevice ( int device ) # ​cudaError_t cudaSetDevice ( int device )
@ -131,7 +107,7 @@ class CudaRTLibrary:
if so_file is None: if so_file is None:
so_file = find_loaded_library("libcudart") so_file = find_loaded_library("libcudart")
if so_file is None: if so_file is None:
so_file = get_cudart_lib_path_from_env() so_file = envs.VLLM_CUDART_SO_PATH # fallback to env var
assert so_file is not None, \ assert so_file is not None, \
( (
"libcudart is not loaded in the current process, " "libcudart is not loaded in the current process, "