[BugFix] fix cuda-python requirement (#5261)

* fix cuda-python requirement

* update

* fix
This commit is contained in:
Yuanle Liu
2025-11-27 13:58:18 +08:00
committed by GitHub
parent 84c7fa49a5
commit ef5aa5c03b
3 changed files with 20 additions and 2 deletions
+11 -1
View File
@@ -22,7 +22,11 @@ from typing import List, Tuple
import numpy as np
import paddle
from cuda import cudart
try:
from cuda import cudart
except ImportError:
cudart = None
from fastdeploy.config import EPLBConfig
@@ -90,6 +94,12 @@ def create_mmap(model_name: List, ep_rank: int, ep_size: int, shm_uuid: str, epl
shm_ptr = ctypes.cast(shm_ptr, ctypes.POINTER(ctypes.c_int8))
addr = ctypes.addressof(shm_ptr.contents)
if cudart is None:
raise ImportError(
"cuda-python not installed. Install the version matching your CUDA toolkit:\n"
" CUDA 12.x → pip install cuda-python==12.*\n"
)
# Register memory with CUDA
(ret,) = cudart.cudaHostRegister(addr, shm_size, 0)
if ret != cudart.cudaError_t.cudaSuccess: