Skip to content

Commit

Permalink
[core][misc] update libcudart finding (vllm-project#7620)
Browse files Browse the repository at this point in the history
Co-authored-by: cjackal <[email protected]>
  • Loading branch information
youkaichao and cjackal authored Aug 17, 2024
1 parent 5bf45db commit d95cc0a
Showing 1 changed file with 7 additions and 2 deletions.
9 changes: 7 additions & 2 deletions vllm/distributed/device_communicators/cuda_wrapper.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,8 +49,13 @@ def find_loaded_library(lib_name) -> Optional[str]:
if not found:
# the library is not loaded in the current process
return None
# if lib_name is libcudart, we need to match a line with:
# address /path/to/libcudart-hash.so.11.0
start = line.index("/")
path = line[start:].strip()
filename = path.split("/")[-1]
assert filename.rpartition(".so")[0].startswith(lib_name), \
f"Unexpected filename: {filename} for library {lib_name}"
return path


Expand Down Expand Up @@ -98,9 +103,9 @@ class CudaRTLibrary:

def __init__(self, so_file: Optional[str] = None):
if so_file is None:
so_file = find_loaded_library("libcudart.so")
so_file = find_loaded_library("libcudart")
assert so_file is not None, \
"libcudart.so is not loaded in the current process"
"libcudart is not loaded in the current process"
if so_file not in CudaRTLibrary.path_to_library_cache:
lib = ctypes.CDLL(so_file)
CudaRTLibrary.path_to_library_cache[so_file] = lib
Expand Down

0 comments on commit d95cc0a

Please sign in to comment.