Skip to content

Commit 026f513

Browse files
waterWangAndy-Jostclaude
authored
fix(core): reject mismatched UUID in register() even under -O (#2704)
* fix(core): reject mismatched UUID in register() even under -O (#2697) DeviceMemoryResource.register() and PinnedMemoryResource.register() validated the supplied UUID against the resource's own with a bare `assert`, which CPython removes under `-O`. A mismatched UUID was then accepted, registered as a foreign registry key, and rewritten into the resource's own UUID, silently corrupting the IPC registry. Replace the assertion with an explicit ValueError so the check is enforced regardless of the interpreter's assertion mode, and convert the internal register() result check in MP_from_allocation_handle to an explicit RuntimeError as well. Signed-off-by: waterWang <[email protected]> * cuda.core: address the review of the register() UUID check - test_register_rejects_mismatched_uuid skips non-pool resources: VirtualMemoryResource joined the ipc_memory_resource fixture and has no register(). - The internal-invariant error uses the "internal cuda.core error, please report" wording. - register() on DeviceMemoryResource and PinnedMemoryResource documents its RuntimeError and ValueError, and the 1.3.0 release notes record the fix. Co-Authored-By: Claude Opus 5.5 <[email protected]> * test(cuda.core): compare the UUID with its original after a rejected register() A rejected registration that cleared the UUID to None would have passed the old check (uuid != other). CodeRabbit's suggestion. Co-Authored-By: Claude Opus 5.5 <[email protected]> --------- Signed-off-by: waterWang <[email protected]> Co-authored-by: waterWang <[email protected]> Co-authored-by: Andy Jost <[email protected]> Co-authored-by: Claude Fable 5.1 <[email protected]>
1 parent 71e837a commit 026f513

7 files changed

Lines changed: 76 additions & 2 deletions

File tree

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,13 @@ class DeviceMemoryResource(_MemPool):
137137
-------
138138
The registered mapped memory resource. If one was previously registered
139139
with the given key, it is returned.
140+
141+
Raises
142+
------
143+
RuntimeError
144+
If the memory resource is not IPC-enabled.
145+
ValueError
146+
If the memory resource already has a different UUID.
140147
"""
141148
@classmethod
142149
def from_allocation_handle(cls, device_id: Device | int, alloc_handle: int | IPCAllocationHandle) -> DeviceMemoryResource:

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -176,6 +176,13 @@ cdef class DeviceMemoryResource(_MemPool):
176176
-------
177177
The registered mapped memory resource. If one was previously registered
178178
with the given key, it is returned.
179+
180+
Raises
181+
------
182+
RuntimeError
183+
If the memory resource is not IPC-enabled.
184+
ValueError
185+
If the memory resource already has a different UUID.
179186
"""
180187
return <DeviceMemoryResource>(_ipc.MP_register(self, uuid))
181188

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

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -453,7 +453,11 @@ cdef _MemPool MP_from_allocation_handle(cls, alloc_handle):
453453
# Register it.
454454
if uuid is not None:
455455
registered = self.register(uuid)
456-
assert registered is self
456+
if registered is not self:
457+
raise RuntimeError(
458+
"internal cuda.core error, please report: register() returned a "
459+
"different memory resource than the one being registered"
460+
)
457461

458462
return self
459463

@@ -476,7 +480,11 @@ cdef _MemPool MP_register(_MemPool self, uuid):
476480
return existing
477481
if not self.is_ipc_enabled:
478482
raise RuntimeError("Memory resource is not IPC-enabled")
479-
assert self.uuid is None or self.uuid == uuid
483+
if self.uuid is not None and self.uuid != uuid:
484+
raise ValueError(
485+
f"Cannot register memory resource with UUID {uuid}: "
486+
f"the resource already has UUID {self.uuid}"
487+
)
480488
registry[uuid] = self
481489
self._ipc_data._alloc_handle._uuid = uuid
482490
return self

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,13 @@ class PinnedMemoryResource(_MemPool):
126126
-------
127127
The registered mapped memory resource. If one was previously registered
128128
with the given key, it is returned.
129+
130+
Raises
131+
------
132+
RuntimeError
133+
If the memory resource is not IPC-enabled.
134+
ValueError
135+
If the memory resource already has a different UUID.
129136
"""
130137
@classmethod
131138
def from_allocation_handle(cls, alloc_handle: int | IPCAllocationHandle) -> PinnedMemoryResource:

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,13 @@ cdef class PinnedMemoryResource(_MemPool):
179179
-------
180180
The registered mapped memory resource. If one was previously registered
181181
with the given key, it is returned.
182+
183+
Raises
184+
------
185+
RuntimeError
186+
If the memory resource is not IPC-enabled.
187+
ValueError
188+
If the memory resource already has a different UUID.
182189
"""
183190
return <PinnedMemoryResource>(_ipc.MP_register(self, uuid))
184191

‎cuda_core/docs/source/release/1.3.0-notes.rst‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -112,6 +112,13 @@ New features
112112
Fixes and enhancements
113113
----------------------
114114

115+
- ``register()`` on a :class:`~_memory.DeviceMemoryResource` or a
116+
:class:`PinnedMemoryResource` now raises ``ValueError`` when the resource
117+
already has a different UUID. The check used to be an ``assert``, so under
118+
``python -O`` a mismatched UUID silently rewrote the resource's UUID and
119+
corrupted the IPC registry.
120+
(`#2697 <https://github.com/NVIDIA/cuda-python/issues/2697>`__)
121+
115122
- A :class:`~_memory.DeviceMemoryResource` or a :class:`VirtualMemoryResource`
116123
sent through ``multiprocessing`` now names its device, and the
117124
``VirtualMemoryResource`` ``peers``, by UUID. The receiving process maps

‎cuda_core/tests/memory_ipc/test_errors.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,37 @@ def test_register_rejects_non_ipc_memory_resource(mempool_device):
9090
DeviceMemoryResource.from_registry(key)
9191

9292

93+
@pytest.mark.human_authored
94+
def test_register_rejects_mismatched_uuid(ipc_memory_resource):
95+
"""register() rejects a UUID that does not match the resource's own.
96+
97+
The check is an explicit ValueError (not a bare ``assert``) so it remains
98+
enforced even when CPython runs with ``-O``, where assertions are removed
99+
(see #2697). Registering under a foreign key would otherwise rewrite the
100+
resource's UUID and silently corrupt the IPC registry.
101+
"""
102+
mr = ipc_memory_resource
103+
if not is_pool_resource(mr):
104+
pytest.skip("allocation handles and the registry are pool-only")
105+
assert mr.is_ipc_enabled
106+
107+
# The resource must already carry its own UUID (assigned on creation).
108+
original = mr.uuid
109+
assert original is not None
110+
other = uuid.uuid4()
111+
while other == mr.uuid:
112+
other = uuid.uuid4()
113+
114+
with pytest.raises(ValueError, match="already has UUID"):
115+
mr.register(other)
116+
117+
# A rejected registration must leave the resource's own UUID unchanged.
118+
assert mr.uuid == original
119+
120+
# Registering under the resource's own UUID still succeeds.
121+
assert mr.register(mr.uuid) is mr
122+
123+
93124
@pytest.mark.skipif(os.name == "nt", reason="IPC allocation handles are not supported on Windows")
94125
@pytest.mark.agent_authored(model="gpt-5.6")
95126
def test_ipc_allocation_handle_state_tracks_close():

0 commit comments

Comments
 (0)