Unverified Commit 009439ca authored by Lu Fang's avatar Lu Fang Committed by GitHub
Browse files

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


Signed-off-by: default avatarLu Fang <lufang@fb.com>
parent bc55d130
...@@ -5,7 +5,6 @@ convenient for use when we just need to call a few functions. ...@@ -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]: ...@@ -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: ...@@ -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, "
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment