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

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

View File

@ -5,7 +5,6 @@ convenient for use when we just need to call a few functions.
"""
import ctypes
import glob
from dataclasses import dataclass
from typing import Any, Dict, List, Optional
@ -62,29 +61,6 @@ def find_loaded_library(lib_name) -> Optional[str]:
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:
exported_functions = [
# cudaError_t cudaSetDevice ( int device )
@ -131,7 +107,7 @@ class CudaRTLibrary:
if so_file is None:
so_file = find_loaded_library("libcudart")
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, \
(
"libcudart is not loaded in the current process, "