Skip to content

Commit 001192a

Browse files
committed
fix(cuda.core): address fabric IPC review feedback
1 parent 26698eb commit 001192a

8 files changed

Lines changed: 70 additions & 59 deletions

File tree

‎cuda_core/cuda/core/_memory/_ipc.pxd‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,9 +56,9 @@ cdef class IPCAllocationHandle:
5656
# Descriptor of a VirtualMemoryResource buffer: one exported handle per
5757
# physical allocation that backs the buffer, in address order.
5858
cdef class _VirtualMemoryDescriptor(IPCBufferDescriptor):
59-
cdef readonly int _handle_type
6059
cdef readonly tuple _chunk_sizes
6160

61+
cdef int driver_handle_type(self) except -1
6262
cdef void* handle_ptr(self, Py_ssize_t i, cydriver.CUmemFabricHandle* scratch) except? NULL
6363

6464

‎cuda_core/cuda/core/_memory/_ipc.pyi‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,6 @@ class IPCAllocationHandle:
6060

6161
class _VirtualMemoryDescriptor(IPCBufferDescriptor):
6262
"""Common allocation layout and driver-handle interface for VMM descriptors."""
63-
_handle_type: int
6463
_chunk_sizes: tuple[Any, ...]
6564

6665
@property
@@ -106,7 +105,7 @@ class VirtualMemoryFdBufferDescriptor(_VirtualMemoryDescriptor):
106105
_handles: tuple[Any, ...]
107106

108107
@staticmethod
109-
def _from_exports(handle_type: int, chunk_sizes: tuple[int, ...], handles: tuple[Any, ...], size: int) -> VirtualMemoryFdBufferDescriptor: ...
108+
def _from_exports(chunk_sizes: tuple[int, ...], handles: tuple[Any, ...], size: int) -> VirtualMemoryFdBufferDescriptor: ...
110109
@classmethod
111110
def from_fds(cls, fds: Iterable[int | IPCAllocationHandle], sizes: Iterable[int], *, size: int | None=None) -> VirtualMemoryFdBufferDescriptor:
112111
"""Build a descriptor from POSIX file descriptors received outside ``multiprocessing``.

‎cuda_core/cuda/core/_memory/_ipc.pyx‎

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -187,16 +187,20 @@ cdef class _VirtualMemoryDescriptor(IPCBufferDescriptor):
187187
"""The export handle type; the importing resource must use the same type."""
188188
from cuda.core._memory._virtual_memory_resource import VirtualMemoryResourceOptions
189189

190+
handle_type = self.driver_handle_type()
190191
for spec, value in VirtualMemoryResourceOptions._handle_types.items():
191-
if spec is not None and int(value) == self._handle_type:
192+
if spec is not None and int(value) == handle_type:
192193
return spec
193-
raise ValueError(f"the descriptor carries an unknown handle type ({self._handle_type})")
194+
raise ValueError(f"the descriptor carries an unknown handle type ({handle_type})")
194195

195196
@property
196197
def sizes(self) -> tuple[int, ...]:
197198
"""Size of each physical allocation, in address order; the sum is at least ``size``."""
198199
return self._chunk_sizes
199200

201+
cdef int driver_handle_type(self) except -1:
202+
raise NotImplementedError("a concrete VMM descriptor is required")
203+
200204
cdef void* handle_ptr(self, Py_ssize_t i, cydriver.CUmemFabricHandle* scratch) except? NULL:
201205
raise NotImplementedError("a concrete VMM descriptor is required")
202206

@@ -237,13 +241,12 @@ cdef class VirtualMemoryFdBufferDescriptor(_VirtualMemoryDescriptor):
237241

238242
@staticmethod
239243
def _from_exports(
240-
handle_type: int, chunk_sizes: tuple[int, ...], handles: tuple, size: int
244+
chunk_sizes: tuple[int, ...], handles: tuple, size: int
241245
) -> VirtualMemoryFdBufferDescriptor:
242246
cdef VirtualMemoryFdBufferDescriptor self = VirtualMemoryFdBufferDescriptor.__new__(
243247
VirtualMemoryFdBufferDescriptor)
244248
self._payload = b""
245249
self._size = size
246-
self._handle_type = handle_type
247250
self._chunk_sizes = tuple(chunk_sizes)
248251
self._handles = tuple(handles)
249252
return self
@@ -288,7 +291,6 @@ cdef class VirtualMemoryFdBufferDescriptor(_VirtualMemoryDescriptor):
288291
OSError
289292
If an ``int`` is not an open file descriptor.
290293
"""
291-
cdef int driver_type = <int>cydriver.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR
292294
cdef tuple sizes_t = tuple(operator.index(n) for n in sizes)
293295
cdef list given = list(fds)
294296
if len(given) != len(sizes_t):
@@ -315,7 +317,10 @@ cdef class VirtualMemoryFdBufferDescriptor(_VirtualMemoryDescriptor):
315317
raise
316318
# The handle owns the duplicate now; if append fails, the handle closes it once.
317319
handles.append(handle)
318-
return VirtualMemoryFdBufferDescriptor._from_exports(driver_type, sizes_t, tuple(handles), size)
320+
return VirtualMemoryFdBufferDescriptor._from_exports(sizes_t, tuple(handles), size)
321+
322+
cdef int driver_handle_type(self) except -1:
323+
return <int>cydriver.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR
319324

320325
cdef void* handle_ptr(self, Py_ssize_t i, cydriver.CUmemFabricHandle* scratch) except? NULL:
321326
# fd zero is valid, so NULL is a valid result unless an exception is set.
@@ -344,7 +349,7 @@ cdef class VirtualMemoryFdBufferDescriptor(_VirtualMemoryDescriptor):
344349
def __reduce__(self) -> tuple[object, ...]:
345350
# multiprocessing duplicates the IPCAllocationHandle fds into the receiver; plain pickle raises.
346351
return VirtualMemoryFdBufferDescriptor._from_exports, (
347-
self._handle_type, self._chunk_sizes, self._handles, self._size)
352+
self._chunk_sizes, self._handles, self._size)
348353

349354

350355
cdef class VirtualMemoryFabricBufferDescriptor(_VirtualMemoryDescriptor):
@@ -419,7 +424,6 @@ cdef class VirtualMemoryFabricBufferDescriptor(_VirtualMemoryDescriptor):
419424
raise ValueError(f"size {size} is outside 0..{total}, the sum of the allocation sizes")
420425
cdef VirtualMemoryFabricBufferDescriptor self = VirtualMemoryFabricBufferDescriptor.__new__(cls)
421426
self._payload = b""
422-
self._handle_type = <int>cydriver.CU_MEM_HANDLE_TYPE_FABRIC
423427
self._chunk_sizes = sizes_t
424428
self._handles = handles_t
425429
self._size = size
@@ -436,6 +440,9 @@ cdef class VirtualMemoryFabricBufferDescriptor(_VirtualMemoryDescriptor):
436440
"""The exporting device's UUID, if supplied by the exporter."""
437441
return self._owner_uuid
438442

443+
cdef int driver_handle_type(self) except -1:
444+
return <int>cydriver.CU_MEM_HANDLE_TYPE_FABRIC
445+
439446
cdef void* handle_ptr(self, Py_ssize_t i, cydriver.CUmemFabricHandle* scratch) except? NULL:
440447
cdef bytes handle = self._handles[i]
441448
memcpy(scratch.data, <const char*>handle, 64)

‎cuda_core/cuda/core/_memory/_virtual_memory_resource.pyx‎

Lines changed: 22 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -263,37 +263,36 @@ cdef int _raise_last_error() except -1:
263263
)
264264

265265

266-
cdef inline int _attach_note(exc, str detail) except -1:
267-
# PEP 678 notes (Python 3.11+) keep the detail separable from the message;
268-
# older interpreters get it appended to the message instead.
269-
add_note = getattr(exc, "add_note", None)
270-
if add_note is not None:
271-
add_note(detail)
272-
else:
273-
exc.args = (f"{exc.args[0]} ({detail})", *exc.args[1:])
274-
return 0
275-
276-
277266
cdef int _raise_last_error_noting(str detail) except -1:
278267
"""Raise the status a handle factory recorded, with ``detail`` attached as a note."""
279268
try:
280269
_raise_last_error()
281270
except Exception as exc:
282-
_attach_note(exc, detail)
271+
exc.add_note(detail)
283272
raise
284273
return 0
285274

286275

287-
cdef int _raise_allocation_error(const cydriver.CUmemAllocationProp* prop) except -1:
288-
"""Raise the status of a failed physical allocation; a fabric resource gets the IMEX note."""
289-
if prop.requestedHandleTypes == cydriver.CU_MEM_HANDLE_TYPE_FABRIC:
290-
_raise_last_error_noting(
276+
cdef int _raise_allocation_error(
277+
cydriver.CUresult status, cydriver.CUmemAllocationHandleType requested_handle_types
278+
) except -1:
279+
"""Raise a physical-allocation failure, noting possible fabric configuration errors."""
280+
if (requested_handle_types == cydriver.CU_MEM_HANDLE_TYPE_FABRIC
281+
and status in (cydriver.CUDA_ERROR_NOT_PERMITTED,
282+
cydriver.CUDA_ERROR_NOT_SUPPORTED,
283+
cydriver.CUDA_ERROR_INVALID_VALUE)):
284+
_check_noting(
285+
status,
291286
"while creating a physical allocation with a fabric handle; fabric handles need an NVLink fabric "
292287
"and an accessible IMEX channel, and a device can report support while the allocation is refused "
293288
"(the attribute does not check for an NVLink fabric, and older drivers do not check for an IMEX "
294289
"channel)"
295290
)
296-
return _raise_last_error()
291+
HANDLE_RETURN(status)
292+
raise RuntimeError(
293+
"internal cuda.core error, please report: a physical allocation factory "
294+
"returned an empty handle without recording a CUDA error"
295+
)
297296

298297

299298
cdef int _check_noting(cydriver.CUresult status, str detail) except -1:
@@ -303,7 +302,7 @@ cdef int _check_noting(cydriver.CUresult status, str detail) except -1:
303302
try:
304303
HANDLE_RETURN(status)
305304
except Exception as exc:
306-
_attach_note(exc, detail)
305+
exc.add_note(detail)
307306
raise
308307
return 0
309308

@@ -740,7 +739,7 @@ cdef class VirtualMemoryResource(MemoryResource):
740739
with nogil:
741740
h_alloc = create_mem_allocation_handle(aligned, prop, descs.data(), descs.size())
742741
if not h_alloc:
743-
_raise_allocation_error(&prop)
742+
_raise_allocation_error(get_last_error(), prop.requestedHandleTypes)
744743
with nogil:
745744
h_res = create_va_reservation_handle(aligned, addr_align, hint)
746745
if not h_res:
@@ -893,7 +892,7 @@ cdef class VirtualMemoryResource(MemoryResource):
893892
with nogil:
894893
h_alloc = create_mem_allocation_handle(add, prop, descs.data(), descs.size())
895894
if not h_alloc:
896-
_raise_allocation_error(&prop)
895+
_raise_allocation_error(get_last_error(), prop.requestedHandleTypes)
897896
with nogil:
898897
h_map = create_va_mapping_handle(base + total, h_alloc, h_res)
899898
if not h_map:
@@ -926,7 +925,7 @@ cdef class VirtualMemoryResource(MemoryResource):
926925
with nogil:
927926
h_alloc = create_mem_allocation_handle(add, prop, descs.data(), descs.size())
928927
if not h_alloc:
929-
_raise_allocation_error(&prop)
928+
_raise_allocation_error(get_last_error(), prop.requestedHandleTypes)
930929
with nogil:
931930
m2 = create_va_mapping_handle(base_new + offset, h_alloc, h_res_new)
932931
if not m2:
@@ -1044,7 +1043,7 @@ cdef class VirtualMemoryResource(MemoryResource):
10441043
if not self.is_ipc_enabled:
10451044
raise RuntimeError("Memory resource is not IPC-enabled")
10461045
handle_type = int(VirtualMemoryResourceOptions._handle_type_to_driver(cfg.handle_type))
1047-
if desc._handle_type != handle_type:
1046+
if desc.driver_handle_type() != handle_type:
10481047
raise ValueError(
10491048
"the descriptor's handle type does not match this resource's "
10501049
f"handle_type={str(cfg.handle_type)!r}; import it with a resource configured like the exporter's"
@@ -1282,4 +1281,4 @@ cdef _VirtualMemoryDescriptor _export_ipc_descriptor(Buffer buf):
12821281
if handle_type == <int>cydriver.CU_MEM_HANDLE_TYPE_FABRIC:
12831282
return VirtualMemoryFabricBufferDescriptor._from_exports(
12841283
tuple(sizes), tuple(handles), buf._size, mr.device.uuid)
1285-
return VirtualMemoryFdBufferDescriptor._from_exports(handle_type, tuple(sizes), tuple(handles), buf._size)
1284+
return VirtualMemoryFdBufferDescriptor._from_exports(tuple(sizes), tuple(handles), buf._size)

‎cuda_core/tests/conftest.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -454,9 +454,9 @@ def _require_vmm_fabric_device(device):
454454
except CUDAError as exc:
455455
# Only a refused physical allocation means "no fabric here" (NOT_PERMITTED on the CI runners,
456456
# INVALID_VALUE on an H200 without an IMEX channel). VirtualMemoryResource attaches its note about
457-
# the fabric and IMEX requirements at that stage alone (on Python 3.10 it lands in the message);
457+
# the fabric and IMEX requirements at that stage alone;
458458
# a failure of the reservation, the mapping, or the stream setup is a real error and propagates.
459-
if "fabric handle" not in " ".join((str(exc), *getattr(exc, "__notes__", ()))):
459+
if "fabric handle" not in " ".join(getattr(exc, "__notes__", ())):
460460
raise
461461
_fabric_allocation_error[device.device_id] = str(exc).split(":", 1)[0]
462462
finally:

‎cuda_core/tests/memory_ipc/test_errors.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -216,7 +216,7 @@ def PARENT_ACTION(self, queue):
216216
else:
217217
# The VMM descriptor carries the exported handles and their sizes;
218218
# forge only the buffer size.
219-
oversized = type(desc)._from_exports(desc._handle_type, desc._chunk_sizes, desc._handles, desc.size * 100)
219+
oversized = type(desc)._from_exports(desc._chunk_sizes, desc._handles, desc.size * 100)
220220
stream.sync()
221221
queue.put(oversized)
222222

‎cuda_core/tests/memory_ipc/test_vmm_fabric.py‎

Lines changed: 8 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -40,15 +40,11 @@ def _read(buf):
4040
return bytes(host)
4141

4242

43-
def _notes(exc):
44-
# Python 3.10 appends the note to the message instead of using PEP 678.
45-
return " ".join((str(exc), *getattr(exc, "__notes__", ())))
46-
47-
4843
def _check_import_error(exc):
4944
# The driver's status for a bad handle varies with the IMEX setup; the note is cuda.core's and does not.
45+
notes = " ".join(exc.__notes__)
5046
for detail in ("stale", "exporting buffer was closed", "forged", "different IMEX channel"):
51-
assert detail in _notes(exc)
47+
assert detail in notes
5248

5349

5450
def _exchange(owner_uuid, payload):
@@ -137,7 +133,7 @@ def test_export_shape(vmm_fabric_device):
137133
prefix = VirtualMemoryFabricBufferDescriptor.from_handles(desc.handles, desc.sizes, size=1)
138134
with Buffer.from_ipc_descriptor(mr, prefix, stream=device.default_stream) as alias:
139135
assert alias.size == 1 < sum(prefix.sizes)
140-
assert not any(hasattr(desc, name) for name in ("fds", "close", "is_closed"))
136+
assert not hasattr(desc, "fds")
141137
with mr.allocate(0) as empty:
142138
desc = empty.ipc_descriptor
143139
assert (desc.handles, desc.sizes, desc.size) == ((), (), 0)
@@ -228,8 +224,10 @@ def test_second_gpu_and_wrong_owner(vmm_fabric_device_x2):
228224

229225

230226
@pytest.mark.agent_authored(model="gpt-6.1-sol")
231-
def test_handle_type_mismatch(vmm_fabric_device):
232-
device = vmm_fabric_device
227+
def test_handle_type_mismatch(vmm_ipc_device):
228+
device = vmm_ipc_device
229+
if not device.properties.handle_type_fabric_supported:
230+
pytest.skip("Device does not support fabric handles")
233231
fabric, fd = _resource(device), _resource(device, handle_type="posix_fd")
234232
# Empty descriptors would need no import call; even these must reject a mismatch.
235233
for exporting, importing in ((fabric, fd), (fd, fabric)):
@@ -279,13 +277,9 @@ def test_reexport_imported_derivatives(vmm_fabric_device):
279277

280278

281279
@pytest.mark.agent_authored(model="gpt-6.1-sol")
282-
def test_configuration_and_buffer_pickle(vmm_fabric_device):
280+
def test_buffer_pickle(vmm_fabric_device):
283281
device = vmm_fabric_device
284282
mr = _resource(device)
285-
assert mr.is_ipc_enabled
286-
for location in ("host", "host_numa", "host_numa_current"):
287-
host = _resource(device, location_type=location, handle_type=None if location == "host" else "fabric")
288-
assert not host.is_ipc_enabled
289283
with mr.allocate(1) as buf:
290284
_fill(buf, 0x55)
291285
with pickle.loads(pickle.dumps(buf)) as alias: # noqa: S301 our own live buffer

‎cuda_core/tests/memory_ipc/test_vmm_ipc.py‎

Lines changed: 20 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,12 @@
3030
from cuda.bindings import driver
3131
from cuda.core import Buffer, Device, MemoryResource, VirtualMemoryResource, VirtualMemoryResourceOptions
3232
from cuda.core._dlpack import DLDeviceType
33-
from cuda.core._memory._ipc import IPCAllocationHandle, IPCBufferDescriptor, VirtualMemoryFdBufferDescriptor
33+
from cuda.core._memory._ipc import (
34+
IPCAllocationHandle,
35+
IPCBufferDescriptor,
36+
VirtualMemoryFabricBufferDescriptor,
37+
VirtualMemoryFdBufferDescriptor,
38+
)
3439
from cuda.core._memory._virtual_memory_resource import VirtualMemoryBuffer
3540
from cuda.core._utils.cuda_utils import CUDAError, handle_return
3641
from cuda.core.utils import StridedMemoryView
@@ -109,13 +114,10 @@ def _has_note(exc, text):
109114

110115
def _forge(desc, **fields):
111116
"""A copy of ``desc`` with some fields replaced, as a hostile or buggy peer could send."""
112-
fields.setdefault("handle_type", desc._handle_type)
113117
fields.setdefault("chunk_sizes", desc._chunk_sizes)
114118
fields.setdefault("handles", desc._handles)
115119
fields.setdefault("size", desc.size)
116-
return VirtualMemoryFdBufferDescriptor._from_exports(
117-
fields["handle_type"], fields["chunk_sizes"], fields["handles"], fields["size"]
118-
)
120+
return VirtualMemoryFdBufferDescriptor._from_exports(fields["chunk_sizes"], fields["handles"], fields["size"])
119121

120122

121123
class _SharedStub(MemoryResource):
@@ -177,16 +179,26 @@ def _on_queue_feeder_error(e, _obj):
177179
("options", "expected"),
178180
[
179181
pytest.param({"handle_type": "posix_fd"}, True, id="posix_fd"),
182+
pytest.param({"handle_type": "fabric"}, True, id="fabric"),
180183
pytest.param({"handle_type": None}, False, id="none"),
181184
pytest.param({"handle_type": "win32_kmt"}, False, id="win32_kmt"),
182185
pytest.param({"location_type": "host", "handle_type": None}, False, id="host"),
183186
pytest.param(
184187
{"location_type": "host_numa", "numa_id": 0, "handle_type": "posix_fd"}, False, id="host_numa-posix_fd"
185188
),
189+
pytest.param(
190+
{"location_type": "host_numa", "numa_id": 0, "handle_type": "fabric"}, False, id="host_numa-fabric"
191+
),
186192
],
187193
)
188194
def test_is_ipc_enabled(vmm_ipc_device, options, expected):
189-
"""POSIX file descriptors on device memory have a transport; host-located memory is not shared yet."""
195+
"""POSIX fd and fabric device memory have transports; host-located memory is not shared yet."""
196+
if (
197+
options.get("location_type", "device") == "device"
198+
and options.get("handle_type") == "fabric"
199+
and not vmm_ipc_device.properties.handle_type_fabric_supported
200+
):
201+
pytest.skip("Device does not support fabric handles")
190202
mr = VirtualMemoryResource(vmm_ipc_device, config=VirtualMemoryResourceOptions(**options))
191203
assert mr.is_ipc_enabled is expected
192204

@@ -352,9 +364,9 @@ def test_rejected_before_import(self, vmm_ipc_device):
352364
with mr.allocate(1) as buf:
353365
gran = buf.size
354366
desc = buf.ipc_descriptor
355-
fabric = int(driver.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC)
367+
fabric = VirtualMemoryFabricBufferDescriptor.from_handles([bytes(64)], [gran])
356368
cases = [
357-
(_forge(desc, handle_type=fabric), ValueError, "handle type does not match"),
369+
(fabric, ValueError, "handle type does not match"),
358370
(_forge(desc, chunk_sizes=(gran, gran)), ValueError, "malformed"),
359371
(_forge(desc, chunk_sizes=(0,), size=0), ValueError, "positive multiple"),
360372
(_forge(desc, chunk_sizes=(gran + 1,)), ValueError, "positive multiple"),

0 commit comments

Comments
 (0)