Skip to content

Commit e603e0e

Browse files
committed
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]>
1 parent 0477f61 commit e603e0e

2 files changed

Lines changed: 38 additions & 2 deletions

File tree

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

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -269,7 +269,11 @@ cdef _MemPool MP_from_allocation_handle(cls, alloc_handle):
269269
# Register it.
270270
if uuid is not None:
271271
registered = self.register(uuid)
272-
assert registered is self
272+
if registered is not self:
273+
raise RuntimeError(
274+
"Internal error: register() returned a different memory "
275+
"resource than the one being registered"
276+
)
273277

274278
return self
275279

@@ -292,7 +296,11 @@ cdef _MemPool MP_register(_MemPool self, uuid):
292296
return existing
293297
if not self.is_ipc_enabled:
294298
raise RuntimeError("Memory resource is not IPC-enabled")
295-
assert self.uuid is None or self.uuid == uuid
299+
if self.uuid is not None and self.uuid != uuid:
300+
raise ValueError(
301+
f"Cannot register memory resource with UUID {uuid}: "
302+
f"the resource already has UUID {self.uuid}"
303+
)
296304
registry[uuid] = self
297305
self._ipc_data._alloc_handle._uuid = uuid
298306
return self

‎cuda_core/tests/memory_ipc/test_errors.py‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,34 @@ def test_register_rejects_non_ipc_memory_resource(mempool_device):
7272
DeviceMemoryResource.from_registry(key)
7373

7474

75+
@pytest.mark.human_authored
76+
def test_register_rejects_mismatched_uuid(ipc_memory_resource):
77+
"""register() rejects a UUID that does not match the resource's own.
78+
79+
The check is an explicit ValueError (not a bare ``assert``) so it remains
80+
enforced even when CPython runs with ``-O``, where assertions are removed
81+
(see #2697). Registering under a foreign key would otherwise rewrite the
82+
resource's UUID and silently corrupt the IPC registry.
83+
"""
84+
mr = ipc_memory_resource
85+
assert mr.is_ipc_enabled
86+
87+
# The resource must already carry its own UUID (assigned on creation).
88+
assert mr.uuid is not None
89+
other = uuid.uuid4()
90+
while other == mr.uuid:
91+
other = uuid.uuid4()
92+
93+
with pytest.raises(ValueError, match="already has UUID"):
94+
mr.register(other)
95+
96+
# A rejected registration must not rewrite the resource's own UUID.
97+
assert mr.uuid != other
98+
99+
# Registering under the resource's own UUID still succeeds.
100+
assert mr.register(mr.uuid) is mr
101+
102+
75103
@pytest.mark.skipif(os.name == "nt", reason="IPC allocation handles are not supported on Windows")
76104
@pytest.mark.agent_authored(model="gpt-5.6")
77105
def test_ipc_allocation_handle_state_tracks_close():

0 commit comments

Comments
 (0)