[sgl-kernel] fix runtime error while preloading CUDA runtime (#13089)

This commit is contained in:
Antonin Vidon
2025-12-02 23:03:51 +08:00
committed by GitHub
parent c5947ecd85
commit df1f31241b
+22 -14
View File
@@ -205,20 +205,28 @@ def _find_cuda_home():
def _preload_cuda_library(): def _preload_cuda_library():
"""Preload the CUDA runtime library to help avoid 'libcudart.so.12 not found' issues."""
cuda_home = Path(_find_cuda_home()) cuda_home = Path(_find_cuda_home())
if (cuda_home / "lib").is_dir(): candidate_dirs = [
cuda_path = cuda_home / "lib" cuda_home / "lib",
elif (cuda_home / "lib64").is_dir(): cuda_home / "lib64",
cuda_path = cuda_home / "lib64" Path("/usr/lib/x86_64-linux-gnu"),
else: Path("/usr/lib/aarch64-linux-gnu"),
# Search for 'libcudart.so.12' in subdirectories Path("/usr/lib64"),
for path in cuda_home.rglob("libcudart.so.12"): Path("/usr/lib"),
cuda_path = path.parent ]
break
else:
raise RuntimeError("Could not find CUDA lib directory.")
cuda_include = (cuda_path / "libcudart.so.12").resolve() for base in candidate_dirs:
if cuda_include.exists(): candidate = base / "libcudart.so.12"
ctypes.CDLL(str(cuda_include), mode=ctypes.RTLD_GLOBAL) if candidate.exists():
try:
cuda_runtime_lib = candidate.resolve()
ctypes.CDLL(str(cuda_runtime_lib), mode=ctypes.RTLD_GLOBAL)
logger.debug(f"Preloaded CUDA runtime under {cuda_runtime_lib}")
return
except Exception as e:
logger.debug(f"Failed to load {candidate}: {e}")
continue
logger.debug("[sgl_kernel] Could not preload CUDA runtime library")