mooncake: probe pointers via cuda.bindings instead of ctypes libcudart
The ctypes-based probing loaded a second libcudart copy into the process, which corrupted CUDA/runtime state and segfaulted the scheduler processes shortly after registration (reproducible on prefill ranks, intermittent on decode). cuda.bindings.runtime.cudaPointerGetAttributes is a properly typed binding and boots cleanly on both roles.
This commit is contained in:
@@ -101,32 +101,12 @@ FAILED_SESSION_RECOVERIES = Counter(
|
||||
|
||||
import ctypes as _ctypes
|
||||
|
||||
|
||||
class _CudaPointerAttributes(_ctypes.Structure):
|
||||
_fields_ = [
|
||||
("type", _ctypes.c_int),
|
||||
("device", _ctypes.c_int),
|
||||
("devicePointer", _ctypes.c_void_p),
|
||||
("hostPointer", _ctypes.c_void_p),
|
||||
]
|
||||
|
||||
|
||||
_CUDA_MEMORY_TYPE_DEVICE = 2
|
||||
_CUDART = None
|
||||
|
||||
|
||||
def _get_cudart():
|
||||
global _CUDART
|
||||
if _CUDART is None:
|
||||
for name in ("libcudart.so", "libcudart.so.13", "libcudart.so.12"):
|
||||
try:
|
||||
_CUDART = _ctypes.CDLL(name)
|
||||
break
|
||||
except OSError:
|
||||
continue
|
||||
else:
|
||||
_CUDART = False
|
||||
return _CUDART or None
|
||||
try:
|
||||
from cuda.bindings import runtime as _cudart
|
||||
except ImportError: # pragma: no cover - cuda-python is always present in images
|
||||
_cudart = None
|
||||
|
||||
|
||||
def _is_device_pointer(ptr: int) -> bool:
|
||||
@@ -136,34 +116,32 @@ def _is_device_pointer(ptr: int) -> bool:
|
||||
segment addresses). Returns False on any error so the caller falls back
|
||||
to the safe host path.
|
||||
"""
|
||||
cudart = _get_cudart()
|
||||
if cudart is None:
|
||||
if _cudart is None:
|
||||
# Cannot tell; assume device so behavior stays unchanged.
|
||||
return True
|
||||
attr = _CudaPointerAttributes()
|
||||
ret = cudart.cudaPointerGetAttributes(
|
||||
_ctypes.byref(attr), _ctypes.c_void_p(int(ptr))
|
||||
)
|
||||
if ret != 0:
|
||||
err, attr = _cudart.cudaPointerGetAttributes(int(ptr))
|
||||
if int(err) != 0:
|
||||
# Clear the error so subsequent CUDA calls are not poisoned.
|
||||
cudart.cudaGetLastError()
|
||||
_cudart.cudaGetLastError()
|
||||
return False
|
||||
return attr.type == _CUDA_MEMORY_TYPE_DEVICE
|
||||
return int(attr.type) == _CUDA_MEMORY_TYPE_DEVICE
|
||||
|
||||
|
||||
def _read_bytes_from_address(addr: int, length: int) -> Optional[bytes]:
|
||||
if length <= 0:
|
||||
return b""
|
||||
if _is_device_pointer(addr):
|
||||
cudart = _get_cudart()
|
||||
buf = (_ctypes.c_char * length)()
|
||||
buf = bytearray(length)
|
||||
# cudaMemcpyDeviceToHost = 2; synchronous default-stream copy.
|
||||
ret = cudart.cudaMemcpy(
|
||||
buf, _ctypes.c_void_p(int(addr)), _ctypes.c_size_t(length), 2
|
||||
err, = _cudart.cudaMemcpy(
|
||||
_ctypes.addressof((_ctypes.c_char * length).from_buffer(buf)),
|
||||
int(addr),
|
||||
length,
|
||||
2,
|
||||
)
|
||||
if ret != 0:
|
||||
if int(err) != 0:
|
||||
logger.error(
|
||||
f"cudaMemcpy D2H failed (ret={ret}) for addr {hex(addr)} len {length}"
|
||||
f"cudaMemcpy D2H failed (err={err}) for addr {hex(addr)} len {length}"
|
||||
)
|
||||
return None
|
||||
return bytes(buf)
|
||||
@@ -174,15 +152,14 @@ def _write_bytes_to_address(addr: int, data: bytes) -> bool:
|
||||
if not data:
|
||||
return True
|
||||
if _is_device_pointer(addr):
|
||||
cudart = _get_cudart()
|
||||
buf = _ctypes.create_string_buffer(data, len(data))
|
||||
# cudaMemcpyHostToDevice = 1; synchronous default-stream copy.
|
||||
ret = cudart.cudaMemcpy(
|
||||
_ctypes.c_void_p(int(addr)), buf, _ctypes.c_size_t(len(data)), 1
|
||||
err, = _cudart.cudaMemcpy(
|
||||
int(addr), _ctypes.addressof(buf), len(data), 1
|
||||
)
|
||||
if ret != 0:
|
||||
if int(err) != 0:
|
||||
logger.error(
|
||||
f"cudaMemcpy H2D failed (ret={ret}) for addr {hex(addr)} "
|
||||
f"cudaMemcpy H2D failed (err={err}) for addr {hex(addr)} "
|
||||
f"len {len(data)}"
|
||||
)
|
||||
return False
|
||||
|
||||
Reference in New Issue
Block a user