Fix CUDA 13.0 VMM handle type compatibility (#34431)
This commit is contained in:
@@ -281,20 +281,28 @@ def make_device_allocation_prop(
|
|||||||
elif not isinstance(handle_types, int):
|
elif not isinstance(handle_types, int):
|
||||||
raise ValueError("handle_types must be 'auto', an integer, or None")
|
raise ValueError("handle_types must be 'auto', an integer, or None")
|
||||||
|
|
||||||
handle_types = int(handle_types)
|
handle_type_value = int(handle_types)
|
||||||
valid_handle_types = {
|
valid_handle_types = {
|
||||||
int(drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE),
|
int(drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE): (
|
||||||
int(drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR),
|
drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE
|
||||||
int(drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC),
|
),
|
||||||
|
int(drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR): (
|
||||||
|
drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR
|
||||||
|
),
|
||||||
|
int(drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC): (
|
||||||
|
drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC
|
||||||
|
),
|
||||||
}
|
}
|
||||||
if handle_types not in valid_handle_types:
|
if handle_type_value not in valid_handle_types:
|
||||||
raise ValueError(f"invalid CUDA handle-type value: {handle_types}")
|
raise ValueError(f"invalid CUDA handle-type value: {handle_type_value}")
|
||||||
|
|
||||||
prop = drv.CUmemAllocationProp()
|
prop = drv.CUmemAllocationProp()
|
||||||
prop.type = drv.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED
|
prop.type = drv.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED
|
||||||
prop.location.type = drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
|
prop.location.type = drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
|
||||||
prop.location.id = int(device_id)
|
prop.location.id = int(device_id)
|
||||||
prop.requestedHandleTypes = handle_types
|
# cuda-bindings 13.0.x requires the generated enum here; newer releases
|
||||||
|
# also accept a plain int, which previously hid this compatibility issue.
|
||||||
|
prop.requestedHandleTypes = valid_handle_types[handle_type_value]
|
||||||
prop.allocFlags.gpuDirectRDMACapable = int(gpu_direct_rdma)
|
prop.allocFlags.gpuDirectRDMACapable = int(gpu_direct_rdma)
|
||||||
return prop
|
return prop
|
||||||
|
|
||||||
|
|||||||
@@ -159,6 +159,46 @@ def test_default_handle_type_fallback(monkeypatch, rejected, expected) -> None:
|
|||||||
get_device_allocation_handle_type.cache_clear()
|
get_device_allocation_handle_type.cache_clear()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize(
|
||||||
|
("handle_types", "expected"),
|
||||||
|
[
|
||||||
|
(None, drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE),
|
||||||
|
(0, drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_NONE),
|
||||||
|
(_POSIX_FD, _POSIX_FD),
|
||||||
|
(_FABRIC, _FABRIC),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
def test_allocation_prop_assigns_handle_type_enum(
|
||||||
|
monkeypatch, handle_types, expected
|
||||||
|
) -> None:
|
||||||
|
"""CUDA bindings 13.0.x reject integer VMM handle types at assignment."""
|
||||||
|
|
||||||
|
class Fields:
|
||||||
|
pass
|
||||||
|
|
||||||
|
class StrictAllocationProp:
|
||||||
|
def __init__(self):
|
||||||
|
self.location = Fields()
|
||||||
|
self.allocFlags = Fields()
|
||||||
|
self._requested_handle_types = None
|
||||||
|
|
||||||
|
@property
|
||||||
|
def requestedHandleTypes(self):
|
||||||
|
return self._requested_handle_types
|
||||||
|
|
||||||
|
@requestedHandleTypes.setter
|
||||||
|
def requestedHandleTypes(self, value):
|
||||||
|
if not isinstance(value, drv.CUmemAllocationHandleType):
|
||||||
|
raise TypeError("requestedHandleTypes requires a CUDA enum")
|
||||||
|
self._requested_handle_types = value
|
||||||
|
|
||||||
|
monkeypatch.setattr(drv, "CUmemAllocationProp", StrictAllocationProp)
|
||||||
|
|
||||||
|
prop = make_device_allocation_prop(0, handle_types=handle_types)
|
||||||
|
|
||||||
|
assert prop.requestedHandleTypes is expected
|
||||||
|
|
||||||
|
|
||||||
def test_granularity_defaults_to_recommended(monkeypatch) -> None:
|
def test_granularity_defaults_to_recommended(monkeypatch) -> None:
|
||||||
prop = make_device_allocation_prop(0, handle_types=None)
|
prop = make_device_allocation_prop(0, handle_types=None)
|
||||||
seen = []
|
seen = []
|
||||||
|
|||||||
Reference in New Issue
Block a user