diff --git a/cuda_core/cuda/core/_memory/_device_memory_resource.pyi b/cuda_core/cuda/core/_memory/_device_memory_resource.pyi index 6e2a5d7bb87..ff16f8bef92 100644 --- a/cuda_core/cuda/core/_memory/_device_memory_resource.pyi +++ b/cuda_core/cuda/core/_memory/_device_memory_resource.pyi @@ -137,6 +137,13 @@ class DeviceMemoryResource(_MemPool): ------- The registered mapped memory resource. If one was previously registered with the given key, it is returned. + + Raises + ------ + RuntimeError + If the memory resource is not IPC-enabled. + ValueError + If the memory resource already has a different UUID. """ @classmethod def from_allocation_handle(cls, device_id: Device | int, alloc_handle: int | IPCAllocationHandle) -> DeviceMemoryResource: diff --git a/cuda_core/cuda/core/_memory/_device_memory_resource.pyx b/cuda_core/cuda/core/_memory/_device_memory_resource.pyx index f078c3ac7d3..b7df931cdac 100644 --- a/cuda_core/cuda/core/_memory/_device_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_device_memory_resource.pyx @@ -176,6 +176,13 @@ cdef class DeviceMemoryResource(_MemPool): ------- The registered mapped memory resource. If one was previously registered with the given key, it is returned. + + Raises + ------ + RuntimeError + If the memory resource is not IPC-enabled. + ValueError + If the memory resource already has a different UUID. """ return (_ipc.MP_register(self, uuid)) diff --git a/cuda_core/cuda/core/_memory/_ipc.pyx b/cuda_core/cuda/core/_memory/_ipc.pyx index 5091d39f0fc..37e08729dae 100644 --- a/cuda_core/cuda/core/_memory/_ipc.pyx +++ b/cuda_core/cuda/core/_memory/_ipc.pyx @@ -453,7 +453,11 @@ cdef _MemPool MP_from_allocation_handle(cls, alloc_handle): # Register it. if uuid is not None: registered = self.register(uuid) - assert registered is self + if registered is not self: + raise RuntimeError( + "internal cuda.core error, please report: register() returned a " + "different memory resource than the one being registered" + ) return self @@ -476,7 +480,11 @@ cdef _MemPool MP_register(_MemPool self, uuid): return existing if not self.is_ipc_enabled: raise RuntimeError("Memory resource is not IPC-enabled") - assert self.uuid is None or self.uuid == uuid + if self.uuid is not None and self.uuid != uuid: + raise ValueError( + f"Cannot register memory resource with UUID {uuid}: " + f"the resource already has UUID {self.uuid}" + ) registry[uuid] = self self._ipc_data._alloc_handle._uuid = uuid return self diff --git a/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyi b/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyi index 323ae7c18a9..244553778b9 100644 --- a/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyi +++ b/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyi @@ -126,6 +126,13 @@ class PinnedMemoryResource(_MemPool): ------- The registered mapped memory resource. If one was previously registered with the given key, it is returned. + + Raises + ------ + RuntimeError + If the memory resource is not IPC-enabled. + ValueError + If the memory resource already has a different UUID. """ @classmethod def from_allocation_handle(cls, alloc_handle: int | IPCAllocationHandle) -> PinnedMemoryResource: diff --git a/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyx b/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyx index 5762c940301..62ff20bf8a8 100644 --- a/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyx +++ b/cuda_core/cuda/core/_memory/_pinned_memory_resource.pyx @@ -179,6 +179,13 @@ cdef class PinnedMemoryResource(_MemPool): ------- The registered mapped memory resource. If one was previously registered with the given key, it is returned. + + Raises + ------ + RuntimeError + If the memory resource is not IPC-enabled. + ValueError + If the memory resource already has a different UUID. """ return (_ipc.MP_register(self, uuid)) diff --git a/cuda_core/docs/source/release/1.3.0-notes.rst b/cuda_core/docs/source/release/1.3.0-notes.rst index 5d6521d30a8..c4de5454b21 100644 --- a/cuda_core/docs/source/release/1.3.0-notes.rst +++ b/cuda_core/docs/source/release/1.3.0-notes.rst @@ -104,6 +104,13 @@ New features Fixes and enhancements ---------------------- +- ``register()`` on a :class:`~_memory.DeviceMemoryResource` or a + :class:`PinnedMemoryResource` now raises ``ValueError`` when the resource + already has a different UUID. The check used to be an ``assert``, so under + ``python -O`` a mismatched UUID silently rewrote the resource's UUID and + corrupted the IPC registry. + (`#2697 `__) + - A :class:`~_memory.DeviceMemoryResource` or a :class:`VirtualMemoryResource` sent through ``multiprocessing`` now names its device, and the ``VirtualMemoryResource`` ``peers``, by UUID. The receiving process maps diff --git a/cuda_core/tests/memory_ipc/test_errors.py b/cuda_core/tests/memory_ipc/test_errors.py index dffb410134f..aec81c82c6e 100644 --- a/cuda_core/tests/memory_ipc/test_errors.py +++ b/cuda_core/tests/memory_ipc/test_errors.py @@ -90,6 +90,37 @@ def test_register_rejects_non_ipc_memory_resource(mempool_device): DeviceMemoryResource.from_registry(key) +@pytest.mark.human_authored +def test_register_rejects_mismatched_uuid(ipc_memory_resource): + """register() rejects a UUID that does not match the resource's own. + + The check is an explicit ValueError (not a bare ``assert``) so it remains + enforced even when CPython runs with ``-O``, where assertions are removed + (see #2697). Registering under a foreign key would otherwise rewrite the + resource's UUID and silently corrupt the IPC registry. + """ + mr = ipc_memory_resource + if not is_pool_resource(mr): + pytest.skip("allocation handles and the registry are pool-only") + assert mr.is_ipc_enabled + + # The resource must already carry its own UUID (assigned on creation). + original = mr.uuid + assert original is not None + other = uuid.uuid4() + while other == mr.uuid: + other = uuid.uuid4() + + with pytest.raises(ValueError, match="already has UUID"): + mr.register(other) + + # A rejected registration must leave the resource's own UUID unchanged. + assert mr.uuid == original + + # Registering under the resource's own UUID still succeeds. + assert mr.register(mr.uuid) is mr + + @pytest.mark.skipif(os.name == "nt", reason="IPC allocation handles are not supported on Windows") @pytest.mark.agent_authored(model="gpt-5.6") def test_ipc_allocation_handle_state_tracks_close():