From b2f881bce11c5483ecffc857da22bf0cdd4881c4 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Wed, 17 Jun 2026 23:48:58 +0200 Subject: [PATCH 1/6] Fix VMM handle leaks in virtual memory paths --- .../core/_memory/_virtual_memory_resource.py | 62 ++- cuda_core/tests/test_memory.py | 488 +++++++++++++++++- 2 files changed, 540 insertions(+), 10 deletions(-) diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py index f30e6e3838d..2e1b26c9d4c 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py @@ -322,8 +322,17 @@ def _grow_allocation_fast_path( ) res, new_handle = driver.cuMemCreate(aligned_additional_size, prop, 0) raise_if_driver_error(res) - # Register undo for creation - trans.append(lambda h=new_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) + new_handle_released = False + + def _release_new_handle() -> None: + nonlocal new_handle_released + if not new_handle_released: + raise_if_driver_error(driver.cuMemRelease(new_handle)[0]) + new_handle_released = True + + # Register undo for creation. Callback is conditional to avoid + # double-release after an explicit successful release. + trans.append(_release_new_handle) # Map the new physical memory to the extended VA range (res,) = driver.cuMemMap(new_ptr, aligned_additional_size, 0, new_handle, 0) @@ -339,6 +348,9 @@ def _grow_allocation_fast_path( (res,) = driver.cuMemSetAccess(new_ptr, aligned_additional_size, descs, len(descs)) raise_if_driver_error(res) + # Release handle ownership now that mapping is stable. + _release_new_handle() + # All succeeded, cancel undo actions trans.commit() @@ -389,8 +401,17 @@ def _grow_allocation_slow_path( # Get the old allocation handle for remapping result, old_handle = driver.cuMemRetainAllocationHandle(buf.handle) raise_if_driver_error(result) - # Register undo for old_handle - trans.append(lambda h=old_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) + old_handle_released = False + + def _release_old_handle() -> None: + nonlocal old_handle_released + if not old_handle_released: + raise_if_driver_error(driver.cuMemRelease(old_handle)[0]) + old_handle_released = True + + # Register undo for old handle. Callback is conditional to avoid + # double-release after explicit success. + trans.append(_release_old_handle) # Unmap the old VA range (aligned previous size) aligned_prev_size = total_aligned_size - aligned_additional_size @@ -419,8 +440,17 @@ def _remap_old() -> None: res, new_handle = driver.cuMemCreate(aligned_additional_size, prop, 0) raise_if_driver_error(res) - # Register undo for new physical memory - trans.append(lambda h=new_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) + new_handle_released = False + + def _release_new_handle() -> None: + nonlocal new_handle_released + if not new_handle_released: + raise_if_driver_error(driver.cuMemRelease(new_handle)[0]) + new_handle_released = True + + # Register undo for new physical memory. Callback is conditional to + # avoid double-release after explicit success. + trans.append(_release_new_handle) # Map the new physical memory to the extended portion (aligned offset) (res,) = driver.cuMemMap(int(new_ptr) + aligned_prev_size, aligned_additional_size, 0, new_handle, 0) @@ -439,6 +469,10 @@ def _remap_old() -> None: (res,) = driver.cuMemSetAccess(new_ptr, total_aligned_size, descs, len(descs)) raise_if_driver_error(res) + # Release handles once all operations that need them have completed. + _release_new_handle() + _release_old_handle() + # All succeeded, cancel undo actions trans.commit() @@ -542,8 +576,17 @@ def allocate(self, size: int, *, stream: Stream | GraphBuilder | None = None) -> # ---- Create physical memory ---- res, handle = driver.cuMemCreate(aligned_size, prop, 0) raise_if_driver_error(res) - # Register undo for physical memory - trans.append(lambda h=handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) + handle_released = False + + def _release_handle() -> None: + nonlocal handle_released + if not handle_released: + raise_if_driver_error(driver.cuMemRelease(handle)[0]) + handle_released = True + + # Register undo for physical memory. Callback is conditional to + # avoid double-release after explicit success. + trans.append(_release_handle) # ---- Reserve VA space ---- # Potentially, use a separate size for the VA reservation from the physical allocation size @@ -563,6 +606,9 @@ def allocate(self, size: int, *, stream: Stream | GraphBuilder | None = None) -> (res,) = driver.cuMemSetAccess(ptr, aligned_size, descs, len(descs)) raise_if_driver_error(res) + # Release handle ownership once map+access setup succeeded. + _release_handle() + trans.commit() # Done — return a Buffer that tracks this VA range diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 5ef48919a50..4c0830d9a30 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1028,11 +1028,495 @@ def __init__(self, size): assert result is buf assert buf._size == new_size - # Successful commit: create, map, set access, and no rollback calls. - assert [c[0] for c in calls] == ["create", "map", "set_access"] + # Successful commit: create, map, set access, explicit handle release. + assert [c[0] for c in calls] == ["create", "map", "set_access", "release"] assert ("create", aligned_additional) in calls assert ("map", new_ptr, aligned_additional, NEW_HANDLE) in calls assert ("set_access", new_ptr, aligned_additional, 1) in calls + assert ("release", NEW_HANDLE) in calls + + +def test_vmm_allocator_allocate_releases_handle_after_map(init_cuda, monkeypatch): + """Exercise VMM allocation with mocked driver calls and explicit handle release. + + On successful map, the handle ownership should be released to avoid leaking + driver-side allocation handles even for the short-lived mock path. + """ + device = Device() + if not device.properties.virtual_memory_management_supported: + pytest.skip("Virtual memory management is not supported on this device") + + vmm_mr = VirtualMemoryResource( + device, + config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"), + ) + + SUCCESS = driver.CUresult.CUDA_SUCCESS + HANDLE = 0xBEEF + calls = [] + + def fake_get_allocation_granularity(_, _granularity_flag): + calls.append(("granularity",)) + return (SUCCESS, 1024) + + def fake_create(size, p, flags): + calls.append(("create", size)) + return (SUCCESS, HANDLE) + + def fake_addr_reserve(size, align, hint, fd): + calls.append(("addr_reserve", size, align, hint, fd)) + return (SUCCESS, 0x20_0000) + + def fake_map(ptr, size, offset, handle, flags): + calls.append(("map", ptr, size, handle)) + return (SUCCESS,) + + def fake_set_access(ptr, size, descs, count): + calls.append(("set_access", ptr, size, count)) + return (SUCCESS,) + + def fake_release(handle): + calls.append(("release", handle)) + return (SUCCESS,) + + def fake_unmap(ptr, size): + calls.append(("unmap", ptr, size)) + return (SUCCESS,) + + def fake_addr_free(ptr, size): + calls.append(("addr_free", ptr, size)) + return (SUCCESS,) + + class FakeBuffer: + def __init__(self, ptr, size, mr): + self.handle = ptr + self.size = size + self.memory_resource = mr + + def fake_from_handle(ptr, size, mr): + calls.append(("from_handle", ptr, size)) + return FakeBuffer(ptr, size, mr) + + monkeypatch.setattr(driver, "cuMemGetAllocationGranularity", fake_get_allocation_granularity) + monkeypatch.setattr(driver, "cuMemCreate", fake_create) + monkeypatch.setattr(driver, "cuMemAddressReserve", fake_addr_reserve) + monkeypatch.setattr(driver, "cuMemMap", fake_map) + monkeypatch.setattr(driver, "cuMemSetAccess", fake_set_access) + monkeypatch.setattr(driver, "cuMemRelease", fake_release) + monkeypatch.setattr(driver, "cuMemUnmap", fake_unmap) + monkeypatch.setattr(driver, "cuMemAddressFree", fake_addr_free) + monkeypatch.setattr(Buffer, "from_handle", fake_from_handle) + + result = vmm_mr.allocate(4096) + + assert isinstance(result, FakeBuffer) + assert result.size == 4096 + assert [c[0] for c in calls] == [ + "granularity", + "create", + "addr_reserve", + "map", + "set_access", + "release", + "from_handle", + ] + assert ("release", HANDLE) in calls + + +def test_vmm_allocator_allocate_rolls_back_if_handle_release_fails(init_cuda, monkeypatch): + device = Device() + if not device.properties.virtual_memory_management_supported: + pytest.skip("Virtual memory management is not supported on this device") + + vmm_mr = VirtualMemoryResource( + device, + config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"), + ) + + SUCCESS = driver.CUresult.CUDA_SUCCESS + ERROR = driver.CUresult.CUDA_ERROR_UNKNOWN + HANDLE = 0xBEEF + calls = [] + release_calls = [0] + + def fake_get_allocation_granularity(_, _granularity_flag): + calls.append(("granularity",)) + return (SUCCESS, 1024) + + def fake_create(size, p, flags): + calls.append(("create", size)) + return (SUCCESS, HANDLE) + + def fake_addr_reserve(size, align, hint, fd): + calls.append(("addr_reserve", size, align, hint, fd)) + return (SUCCESS, 0x20_0000) + + def fake_map(ptr, size, offset, handle, flags): + calls.append(("map", ptr, size, handle)) + return (SUCCESS,) + + def fake_set_access(ptr, size, descs, count): + calls.append(("set_access", ptr, size, count)) + return (SUCCESS,) + + def fake_release(handle): + calls.append(("release", handle)) + release_calls[0] += 1 + if release_calls[0] == 1: + return (ERROR,) + return (SUCCESS,) + + def fake_unmap(ptr, size): + calls.append(("unmap", ptr, size)) + return (SUCCESS,) + + def fake_addr_free(ptr, size): + calls.append(("addr_free", ptr, size)) + return (SUCCESS,) + + class FakeBuffer: + def __init__(self, ptr, size, mr): + self.handle = ptr + self.size = size + self.memory_resource = mr + + def fake_from_handle(ptr, size, mr): + calls.append(("from_handle", ptr, size)) + return FakeBuffer(ptr, size, mr) + + monkeypatch.setattr(driver, "cuMemGetAllocationGranularity", fake_get_allocation_granularity) + monkeypatch.setattr(driver, "cuMemCreate", fake_create) + monkeypatch.setattr(driver, "cuMemAddressReserve", fake_addr_reserve) + monkeypatch.setattr(driver, "cuMemMap", fake_map) + monkeypatch.setattr(driver, "cuMemSetAccess", fake_set_access) + monkeypatch.setattr(driver, "cuMemRelease", fake_release) + monkeypatch.setattr(driver, "cuMemUnmap", fake_unmap) + monkeypatch.setattr(driver, "cuMemAddressFree", fake_addr_free) + monkeypatch.setattr(Buffer, "from_handle", fake_from_handle) + + with pytest.raises(CUDAError): + vmm_mr.allocate(4096) + + assert [c[0] for c in calls] == [ + "granularity", + "create", + "addr_reserve", + "map", + "set_access", + "release", + "unmap", + "addr_free", + ] + assert calls.count(("release", HANDLE)) == 2 + assert "from_handle" not in [c[0] for c in calls] + + +def test_vmm_allocator_grow_allocation_fast_path_rolls_back_if_handle_release_fails(init_cuda, monkeypatch): + """Verify fast-path rollback when the post-map handle release fails.""" + device = Device() + if not device.properties.virtual_memory_management_supported: + pytest.skip("Virtual memory management is not supported on this device") + + vmm_mr = VirtualMemoryResource( + device, + config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"), + ) + + prop = driver.CUmemAllocationProp() + prop.type = driver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + prop.location.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + prop.location.id = device.device_id + + SUCCESS = driver.CUresult.CUDA_SUCCESS + ERROR = driver.CUresult.CUDA_ERROR_UNKNOWN + NEW_HANDLE = 0xBEEF + calls = [] + release_calls = [0] + + def fake_create(size, p, flags): + calls.append(("create", size)) + return (SUCCESS, NEW_HANDLE) + + def fake_map(ptr, size, offset, handle, flags): + calls.append(("map", ptr, size, handle)) + return (SUCCESS,) + + def fake_set_access(ptr, size, descs, count): + calls.append(("set_access", ptr, size, count)) + return (SUCCESS,) + + def fake_release(handle): + calls.append(("release", handle)) + release_calls[0] += 1 + if release_calls[0] == 1: + return (ERROR,) + return (SUCCESS,) + + def fake_unmap(ptr, size): + calls.append(("unmap", ptr, size)) + return (SUCCESS,) + + def fake_addr_free(ptr, size): + calls.append(("addr_free", ptr, size)) + return (SUCCESS,) + + class FakeBuffer: + def __init__(self, size): + self._size = size + + original_size = 2 * 1024 * 1024 + aligned_additional = 2 * 1024 * 1024 + new_size = original_size + aligned_additional + new_ptr = 0x10_0000 + buf = FakeBuffer(original_size) + + monkeypatch.setattr(driver, "cuMemCreate", fake_create) + monkeypatch.setattr(driver, "cuMemMap", fake_map) + monkeypatch.setattr(driver, "cuMemSetAccess", fake_set_access) + monkeypatch.setattr(driver, "cuMemUnmap", fake_unmap) + monkeypatch.setattr(driver, "cuMemRelease", fake_release) + monkeypatch.setattr(driver, "cuMemAddressFree", fake_addr_free) + + with pytest.raises(CUDAError): + vmm_mr._grow_allocation_fast_path(buf, new_size, prop, aligned_additional, new_ptr) + + assert buf._size == original_size + assert [c[0] for c in calls] == [ + "create", + "map", + "set_access", + "release", + "unmap", + "addr_free", + ] + + +def test_vmm_allocator_grow_allocation_slow_path(init_cuda, monkeypatch): + """Exercise the VMM grow slow path with mocked driver calls. + + This verifies that both retained and newly-created handles are released after + successful mapping and commit. + """ + device = Device() + if not device.properties.virtual_memory_management_supported: + pytest.skip("Virtual memory management is not supported on this device") + + vmm_mr = VirtualMemoryResource( + device, + config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"), + ) + + prop = driver.CUmemAllocationProp() + prop.type = driver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + prop.location.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + prop.location.id = device.device_id + + SUCCESS = driver.CUresult.CUDA_SUCCESS + OLD_HANDLE = 0xFEED + NEW_HANDLE = 0xBEEF + calls = [] + + def fake_addr_reserve(size, align, hint, fd): + calls.append(("reserve", size, align, hint, fd)) + return (SUCCESS, 0x30_0000) + + def fake_retain(handle): + calls.append(("retain", handle)) + return (SUCCESS, OLD_HANDLE) + + def fake_create(size, p, flags): + calls.append(("create", size)) + return (SUCCESS, NEW_HANDLE) + + def fake_map(ptr, size, offset, handle, flags): + calls.append(("map", ptr, size, offset, handle)) + return (SUCCESS,) + + def fake_unmap(ptr, size): + calls.append(("unmap", ptr, size)) + return (SUCCESS,) + + def fake_set_access(ptr, size, descs, count): + calls.append(("set_access", ptr, size, count)) + return (SUCCESS,) + + def fake_release(handle): + calls.append(("release", handle)) + return (SUCCESS,) + + def fake_addr_free(ptr, size): + calls.append(("addr_free", ptr, size)) + return (SUCCESS,) + + class FakeInputBuffer: + def __init__(self, size, handle): + self._size = size + self.handle = handle + + def _clear(self): + calls.append(("clear",)) + + class FakeReturnedBuffer: + def __init__(self, ptr, size, mr): + self.handle = ptr + self.size = size + self.memory_resource = mr + + def fake_from_handle(ptr, size, mr): + calls.append(("from_handle", ptr, size)) + return FakeReturnedBuffer(ptr, size, mr) + + monkeypatch.setattr(driver, "cuMemAddressReserve", fake_addr_reserve) + monkeypatch.setattr(driver, "cuMemRetainAllocationHandle", fake_retain) + monkeypatch.setattr(driver, "cuMemUnmap", fake_unmap) + monkeypatch.setattr(driver, "cuMemMap", fake_map) + monkeypatch.setattr(driver, "cuMemCreate", fake_create) + monkeypatch.setattr(driver, "cuMemSetAccess", fake_set_access) + monkeypatch.setattr(driver, "cuMemRelease", fake_release) + monkeypatch.setattr(driver, "cuMemAddressFree", fake_addr_free) + monkeypatch.setattr(Buffer, "from_handle", fake_from_handle) + + original_size = 2 * 1024 * 1024 + aligned_additional = 2 * 1024 * 1024 + result_size = original_size + aligned_additional + original_ptr = 0x10_0000 + buf = FakeInputBuffer(original_size, original_ptr) + + result = vmm_mr._grow_allocation_slow_path(buf, result_size, prop, aligned_additional, result_size, 1024) + + assert isinstance(result, FakeReturnedBuffer) + assert result.size == result_size + assert result.handle == 0x30_0000 + assert [c[0] for c in calls] == [ + "reserve", + "retain", + "unmap", + "map", + "create", + "map", + "set_access", + "release", + "release", + "addr_free", + "clear", + "from_handle", + ] + assert ("release", NEW_HANDLE) in calls + assert ("release", OLD_HANDLE) in calls + + +def test_vmm_allocator_grow_allocation_slow_path_rolls_back_if_handle_release_fails(init_cuda, monkeypatch): + """Verify slow-path rollback when a handle release fails near the end of growth.""" + device = Device() + if not device.properties.virtual_memory_management_supported: + pytest.skip("Virtual memory management is not supported on this device") + + vmm_mr = VirtualMemoryResource( + device, + config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"), + ) + + prop = driver.CUmemAllocationProp() + prop.type = driver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED + prop.location.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE + prop.location.id = device.device_id + + SUCCESS = driver.CUresult.CUDA_SUCCESS + ERROR = driver.CUresult.CUDA_ERROR_UNKNOWN + OLD_HANDLE = 0xFEED + NEW_HANDLE = 0xBEEF + calls = [] + release_calls = [0] + + def fake_addr_reserve(size, align, hint, fd): + calls.append(("reserve", size, align, hint, fd)) + return (SUCCESS, 0x30_0000) + + def fake_retain(handle): + calls.append(("retain", handle)) + return (SUCCESS, OLD_HANDLE) + + def fake_create(size, p, flags): + calls.append(("create", size)) + return (SUCCESS, NEW_HANDLE) + + def fake_map(ptr, size, offset, handle, flags): + calls.append(("map", ptr, size, offset, handle)) + return (SUCCESS,) + + def fake_unmap(ptr, size): + calls.append(("unmap", ptr, size)) + return (SUCCESS,) + + def fake_set_access(ptr, size, descs, count): + calls.append(("set_access", ptr, size, count)) + return (SUCCESS,) + + def fake_release(handle): + calls.append(("release", handle)) + release_calls[0] += 1 + if release_calls[0] == 1: + return (ERROR,) + return (SUCCESS,) + + def fake_addr_free(ptr, size): + calls.append(("addr_free", ptr, size)) + return (SUCCESS,) + + class FakeInputBuffer: + def __init__(self, size, handle): + self._size = size + self.handle = handle + + def _clear(self): + calls.append(("clear",)) + + def fake_from_handle(ptr, size, mr): + calls.append(("from_handle", ptr, size)) + return FakeBuffer(ptr, size, mr) + + class FakeBuffer: + def __init__(self, ptr, size, mr): + self.handle = ptr + self.size = size + self.memory_resource = mr + + original_size = 2 * 1024 * 1024 + aligned_additional = 2 * 1024 * 1024 + result_size = original_size + aligned_additional + original_ptr = 0x10_0000 + buf = FakeInputBuffer(original_size, original_ptr) + + monkeypatch.setattr(driver, "cuMemAddressReserve", fake_addr_reserve) + monkeypatch.setattr(driver, "cuMemRetainAllocationHandle", fake_retain) + monkeypatch.setattr(driver, "cuMemUnmap", fake_unmap) + monkeypatch.setattr(driver, "cuMemMap", fake_map) + monkeypatch.setattr(driver, "cuMemCreate", fake_create) + monkeypatch.setattr(driver, "cuMemSetAccess", fake_set_access) + monkeypatch.setattr(driver, "cuMemRelease", fake_release) + monkeypatch.setattr(driver, "cuMemAddressFree", fake_addr_free) + monkeypatch.setattr(Buffer, "from_handle", fake_from_handle) + + with pytest.raises(CUDAError): + vmm_mr._grow_allocation_slow_path(buf, result_size, prop, aligned_additional, result_size, 1024) + + assert [c[0] for c in calls] == [ + "reserve", + "retain", + "unmap", + "map", + "create", + "map", + "set_access", + "unmap", + "release", + "unmap", + "map", + "release", + "addr_free", + ] + assert calls.count(("release", NEW_HANDLE)) == 2 + assert calls.count(("release", OLD_HANDLE)) == 1 def test_vmm_allocator_rdma_unsupported_exception(): From 98265d606a50c592ef66c7389285e33ed1db4ff7 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Thu, 18 Jun 2026 00:09:46 +0200 Subject: [PATCH 2/6] Adjust VMM rollback tests for release-fail retries --- cuda_core/tests/test_memory.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 4c0830d9a30..d6e100cf49a 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1205,6 +1205,7 @@ def fake_from_handle(ptr, size, mr): "set_access", "release", "unmap", + "release", "addr_free", ] assert calls.count(("release", HANDLE)) == 2 @@ -1287,6 +1288,7 @@ def __init__(self, size): "set_access", "release", "unmap", + "release", "addr_free", ] @@ -1508,6 +1510,7 @@ def __init__(self, ptr, size, mr): "create", "map", "set_access", + "release", "unmap", "release", "unmap", From 76efd95d14a9ccd83f2bde98f81913237f8cb092 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Wed, 15 Jul 2026 06:07:27 +0200 Subject: [PATCH 3/6] Address VMM transaction review feedback --- .../core/_memory/_virtual_memory_resource.py | 80 +-- cuda_core/cuda/core/_utils/cuda_utils.pyi | 29 +- cuda_core/cuda/core/_utils/cuda_utils.pyx | 39 +- cuda_core/tests/test_memory.py | 491 +----------------- 4 files changed, 75 insertions(+), 564 deletions(-) diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py index 2e1b26c9d4c..5cd2b7841e5 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py @@ -317,28 +317,19 @@ def _grow_allocation_fast_path( """ with Transaction() as trans: # Create new physical memory for the additional size - trans.append( + trans.on_failure( lambda np=new_ptr, s=aligned_additional_size: raise_if_driver_error(driver.cuMemAddressFree(np, s)[0]) ) res, new_handle = driver.cuMemCreate(aligned_additional_size, prop, 0) raise_if_driver_error(res) - new_handle_released = False - - def _release_new_handle() -> None: - nonlocal new_handle_released - if not new_handle_released: - raise_if_driver_error(driver.cuMemRelease(new_handle)[0]) - new_handle_released = True - - # Register undo for creation. Callback is conditional to avoid - # double-release after an explicit successful release. - trans.append(_release_new_handle) + trans.on_failure(lambda h=new_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) + trans.on_success(lambda h=new_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) # Map the new physical memory to the extended VA range (res,) = driver.cuMemMap(new_ptr, aligned_additional_size, 0, new_handle, 0) raise_if_driver_error(res) # Register undo for mapping - trans.append( + trans.on_failure( lambda np=new_ptr, s=aligned_additional_size: raise_if_driver_error(driver.cuMemUnmap(np, s)[0]) ) @@ -348,9 +339,6 @@ def _release_new_handle() -> None: (res,) = driver.cuMemSetAccess(new_ptr, aligned_additional_size, descs, len(descs)) raise_if_driver_error(res) - # Release handle ownership now that mapping is stable. - _release_new_handle() - # All succeeded, cancel undo actions trans.commit() @@ -394,24 +382,15 @@ def _grow_allocation_slow_path( res, new_ptr = driver.cuMemAddressReserve(total_aligned_size, addr_align, 0, 0) raise_if_driver_error(res) # Register undo for VA reservation - trans.append( + trans.on_failure( lambda np=new_ptr, s=total_aligned_size: raise_if_driver_error(driver.cuMemAddressFree(np, s)[0]) ) # Get the old allocation handle for remapping result, old_handle = driver.cuMemRetainAllocationHandle(buf.handle) raise_if_driver_error(result) - old_handle_released = False - - def _release_old_handle() -> None: - nonlocal old_handle_released - if not old_handle_released: - raise_if_driver_error(driver.cuMemRelease(old_handle)[0]) - old_handle_released = True - - # Register undo for old handle. Callback is conditional to avoid - # double-release after explicit success. - trans.append(_release_old_handle) + trans.on_failure(lambda h=old_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) + trans.on_success(lambda h=old_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) # Unmap the old VA range (aligned previous size) aligned_prev_size = total_aligned_size - aligned_additional_size @@ -427,37 +406,27 @@ def _remap_old() -> None: # TODO: consider logging this exception pass - trans.append(_remap_old) + trans.on_failure(_remap_old) # Remap the old physical memory to the new VA range (aligned previous size) (res,) = driver.cuMemMap(int(new_ptr), aligned_prev_size, 0, old_handle, 0) raise_if_driver_error(res) # Register undo for mapping - trans.append(lambda np=new_ptr, s=aligned_prev_size: raise_if_driver_error(driver.cuMemUnmap(np, s)[0])) + trans.on_failure(lambda np=new_ptr, s=aligned_prev_size: raise_if_driver_error(driver.cuMemUnmap(np, s)[0])) # Create new physical memory for the additional size res, new_handle = driver.cuMemCreate(aligned_additional_size, prop, 0) raise_if_driver_error(res) - - new_handle_released = False - - def _release_new_handle() -> None: - nonlocal new_handle_released - if not new_handle_released: - raise_if_driver_error(driver.cuMemRelease(new_handle)[0]) - new_handle_released = True - - # Register undo for new physical memory. Callback is conditional to - # avoid double-release after explicit success. - trans.append(_release_new_handle) + trans.on_failure(lambda h=new_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) + trans.on_success(lambda h=new_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) # Map the new physical memory to the extended portion (aligned offset) (res,) = driver.cuMemMap(int(new_ptr) + aligned_prev_size, aligned_additional_size, 0, new_handle, 0) raise_if_driver_error(res) # Register undo for mapping - trans.append( + trans.on_failure( lambda base=int(new_ptr), offs=aligned_prev_size, s=aligned_additional_size: raise_if_driver_error( driver.cuMemUnmap(base + offs, s)[0] ) @@ -469,10 +438,6 @@ def _release_new_handle() -> None: (res,) = driver.cuMemSetAccess(new_ptr, total_aligned_size, descs, len(descs)) raise_if_driver_error(res) - # Release handles once all operations that need them have completed. - _release_new_handle() - _release_old_handle() - # All succeeded, cancel undo actions trans.commit() @@ -576,28 +541,20 @@ def allocate(self, size: int, *, stream: Stream | GraphBuilder | None = None) -> # ---- Create physical memory ---- res, handle = driver.cuMemCreate(aligned_size, prop, 0) raise_if_driver_error(res) - handle_released = False - - def _release_handle() -> None: - nonlocal handle_released - if not handle_released: - raise_if_driver_error(driver.cuMemRelease(handle)[0]) - handle_released = True - - # Register undo for physical memory. Callback is conditional to - # avoid double-release after explicit success. - trans.append(_release_handle) + trans.on_failure(lambda h=handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) + # Once mapped, the physical allocation is kept alive without the creation reference. + trans.on_success(lambda h=handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) # ---- Reserve VA space ---- # Potentially, use a separate size for the VA reservation from the physical allocation size res, ptr = driver.cuMemAddressReserve(aligned_size, addr_align, config.addr_hint, 0) raise_if_driver_error(res) # Register undo for VA reservation - trans.append(lambda p=ptr, s=aligned_size: raise_if_driver_error(driver.cuMemAddressFree(p, s)[0])) + trans.on_failure(lambda p=ptr, s=aligned_size: raise_if_driver_error(driver.cuMemAddressFree(p, s)[0])) # ---- Map physical memory into VA ---- (res,) = driver.cuMemMap(ptr, aligned_size, 0, handle, 0) - trans.append(lambda p=ptr, s=aligned_size: raise_if_driver_error(driver.cuMemUnmap(p, s)[0])) + trans.on_failure(lambda p=ptr, s=aligned_size: raise_if_driver_error(driver.cuMemUnmap(p, s)[0])) raise_if_driver_error(res) # ---- Set access for owner + peers ---- @@ -606,9 +563,6 @@ def _release_handle() -> None: (res,) = driver.cuMemSetAccess(ptr, aligned_size, descs, len(descs)) raise_if_driver_error(res) - # Release handle ownership once map+access setup succeeded. - _release_handle() - trans.commit() # Done — return a Buffer that tracks this VA range diff --git a/cuda_core/cuda/core/_utils/cuda_utils.pyi b/cuda_core/cuda/core/_utils/cuda_utils.pyi index 87067927724..bb368d45cc3 100644 --- a/cuda_core/cuda/core/_utils/cuda_utils.pyi +++ b/cuda_core/cuda/core/_utils/cuda_utils.pyi @@ -19,21 +19,22 @@ class NVRTCError(CUDAError): class Transaction: """ - A context manager for transactional operations with undo capability. + A context manager for transactional operations with failure and success callbacks. - The Transaction class allows you to register undo actions (callbacks) that will be executed - if the transaction is not committed before exiting the context. This is useful for managing - resources or operations that need to be rolled back in case of errors or early exits. + Failure callbacks are executed in LIFO order if the transaction exits without being committed. + Success callbacks are executed in FIFO order when the transaction is committed. Usage: with Transaction() as txn: - txn.append(some_cleanup_function, arg1, arg2) + txn.on_failure(some_cleanup_function, arg1, arg2) + txn.on_success(some_finalize_function, arg1, arg2) # ... perform operations ... - txn.commit() # Disarm undo actions; nothing will be rolled back on exit + txn.commit() Methods: - append(fn, *args, **kwargs): Register an undo action to be called on rollback. - commit(): Disarm all undo actions; nothing will be rolled back on exit. + on_failure(fn, *args, **kwargs): Register a callback to be called on rollback. + on_success(fn, *args, **kwargs): Register a callback to be called on commit. + commit(): Disarm failure callbacks and run success callbacks. """ def __init__(self) -> None: @@ -45,15 +46,21 @@ class Transaction: def __exit__(self, exc_type, exc, tb): ... - def append(self, fn: Callable[..., Any], /, *args: Any, **kwargs) -> None: + def on_failure(self, fn: Callable[..., Any], /, *args: Any, **kwargs) -> None: """ - Register an undo action (runs if the with-block exits without commit()). + Register a failure callback (runs if the with-block exits without commit()). + Values are bound now via partial so late mutations don't bite you. + """ + + def on_success(self, fn: Callable[..., Any], /, *args: Any, **kwargs) -> None: + """ + Register a success callback (runs in FIFO order during commit()). Values are bound now via partial so late mutations don't bite you. """ def commit(self) -> None: """ - Disarm all undo actions. After this, exiting the with-block does nothing. + Disarm all failure callbacks, then run success callbacks in FIFO order. """ _keep_driver_in_stub: 'driver.CUresult' _keep_nvrtc_in_stub: 'nvrtc.nvrtcResult' diff --git a/cuda_core/cuda/core/_utils/cuda_utils.pyx b/cuda_core/cuda/core/_utils/cuda_utils.pyx index 318d4466bee..8cb65811aa5 100644 --- a/cuda_core/cuda/core/_utils/cuda_utils.pyx +++ b/cuda_core/cuda/core/_utils/cuda_utils.pyx @@ -306,24 +306,26 @@ def is_nested_sequence(obj: object) -> bool: class Transaction: """ - A context manager for transactional operations with undo capability. + A context manager for transactional operations with failure and success callbacks. - The Transaction class allows you to register undo actions (callbacks) that will be executed - if the transaction is not committed before exiting the context. This is useful for managing - resources or operations that need to be rolled back in case of errors or early exits. + Failure callbacks are executed in LIFO order if the transaction exits without being committed. + Success callbacks are executed in FIFO order when the transaction is committed. Usage: with Transaction() as txn: - txn.append(some_cleanup_function, arg1, arg2) + txn.on_failure(some_cleanup_function, arg1, arg2) + txn.on_success(some_finalize_function, arg1, arg2) # ... perform operations ... - txn.commit() # Disarm undo actions; nothing will be rolled back on exit + txn.commit() Methods: - append(fn, *args, **kwargs): Register an undo action to be called on rollback. - commit(): Disarm all undo actions; nothing will be rolled back on exit. + on_failure(fn, *args, **kwargs): Register a callback to be called on rollback. + on_success(fn, *args, **kwargs): Register a callback to be called on commit. + commit(): Disarm failure callbacks and run success callbacks. """ def __init__(self) -> None: self._stack = ExitStack() + self._on_success: list[Callable[[], Any]] = [] self._entered = False def __enter__(self): @@ -334,23 +336,36 @@ class Transaction: def __exit__(self, exc_type, exc, tb): # If exit callbacks remain, they'll run in LIFO order. self._entered = False + self._on_success.clear() return self._stack.__exit__(exc_type, exc, tb) - def append(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> None: + def on_failure(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> None: """ - Register an undo action (runs if the with-block exits without commit()). + Register a failure callback (runs if the with-block exits without commit()). Values are bound now via partial so late mutations don't bite you. """ if not self._entered: - raise RuntimeError("Transaction must be entered before append()") + raise RuntimeError("Transaction must be entered before on_failure()") self._stack.callback(partial(fn, *args, **kwargs)) + def on_success(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> None: + """ + Register a success callback (runs in FIFO order during commit()). + Values are bound now via partial so late mutations don't bite you. + """ + if not self._entered: + raise RuntimeError("Transaction must be entered before on_success()") + self._on_success.append(partial(fn, *args, **kwargs)) + def commit(self) -> None: """ - Disarm all undo actions. After this, exiting the with-block does nothing. + Disarm all failure callbacks, then run success callbacks in FIFO order. """ # pop_all() empties this stack so no callbacks are triggered on exit. self._stack.pop_all() + for fn in self._on_success: + fn() + self._on_success.clear() # Track whether we've already warned about fork method diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index d6e100cf49a..3a06ee8415d 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -989,8 +989,8 @@ def fake_set_access(ptr, size, descs, count): calls.append(("set_access", ptr, size, count)) return (SUCCESS,) - # Rollback-only entry points: registered as undo actions but, on a - # successful commit, must never be invoked. + # Cleanup entry points. Release runs on commit; unmap and address free + # remain rollback-only. def fake_unmap(ptr, size): calls.append(("unmap", ptr, size)) return (SUCCESS,) @@ -1028,7 +1028,7 @@ def __init__(self, size): assert result is buf assert buf._size == new_size - # Successful commit: create, map, set access, explicit handle release. + # Successful commit: create, map, set access, and release the creation handle. assert [c[0] for c in calls] == ["create", "map", "set_access", "release"] assert ("create", aligned_additional) in calls assert ("map", new_ptr, aligned_additional, NEW_HANDLE) in calls @@ -1036,490 +1036,25 @@ def __init__(self, size): assert ("release", NEW_HANDLE) in calls -def test_vmm_allocator_allocate_releases_handle_after_map(init_cuda, monkeypatch): - """Exercise VMM allocation with mocked driver calls and explicit handle release. - - On successful map, the handle ownership should be released to avoid leaking - driver-side allocation handles even for the short-lived mock path. - """ - device = Device() - if not device.properties.virtual_memory_management_supported: - pytest.skip("Virtual memory management is not supported on this device") - - vmm_mr = VirtualMemoryResource( - device, - config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"), - ) - - SUCCESS = driver.CUresult.CUDA_SUCCESS - HANDLE = 0xBEEF - calls = [] - - def fake_get_allocation_granularity(_, _granularity_flag): - calls.append(("granularity",)) - return (SUCCESS, 1024) - - def fake_create(size, p, flags): - calls.append(("create", size)) - return (SUCCESS, HANDLE) - - def fake_addr_reserve(size, align, hint, fd): - calls.append(("addr_reserve", size, align, hint, fd)) - return (SUCCESS, 0x20_0000) - - def fake_map(ptr, size, offset, handle, flags): - calls.append(("map", ptr, size, handle)) - return (SUCCESS,) - - def fake_set_access(ptr, size, descs, count): - calls.append(("set_access", ptr, size, count)) - return (SUCCESS,) - - def fake_release(handle): - calls.append(("release", handle)) - return (SUCCESS,) - - def fake_unmap(ptr, size): - calls.append(("unmap", ptr, size)) - return (SUCCESS,) - - def fake_addr_free(ptr, size): - calls.append(("addr_free", ptr, size)) - return (SUCCESS,) - - class FakeBuffer: - def __init__(self, ptr, size, mr): - self.handle = ptr - self.size = size - self.memory_resource = mr - - def fake_from_handle(ptr, size, mr): - calls.append(("from_handle", ptr, size)) - return FakeBuffer(ptr, size, mr) - - monkeypatch.setattr(driver, "cuMemGetAllocationGranularity", fake_get_allocation_granularity) - monkeypatch.setattr(driver, "cuMemCreate", fake_create) - monkeypatch.setattr(driver, "cuMemAddressReserve", fake_addr_reserve) - monkeypatch.setattr(driver, "cuMemMap", fake_map) - monkeypatch.setattr(driver, "cuMemSetAccess", fake_set_access) - monkeypatch.setattr(driver, "cuMemRelease", fake_release) - monkeypatch.setattr(driver, "cuMemUnmap", fake_unmap) - monkeypatch.setattr(driver, "cuMemAddressFree", fake_addr_free) - monkeypatch.setattr(Buffer, "from_handle", fake_from_handle) - - result = vmm_mr.allocate(4096) - - assert isinstance(result, FakeBuffer) - assert result.size == 4096 - assert [c[0] for c in calls] == [ - "granularity", - "create", - "addr_reserve", - "map", - "set_access", - "release", - "from_handle", - ] - assert ("release", HANDLE) in calls - - -def test_vmm_allocator_allocate_rolls_back_if_handle_release_fails(init_cuda, monkeypatch): - device = Device() - if not device.properties.virtual_memory_management_supported: - pytest.skip("Virtual memory management is not supported on this device") - - vmm_mr = VirtualMemoryResource( - device, - config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"), - ) - - SUCCESS = driver.CUresult.CUDA_SUCCESS - ERROR = driver.CUresult.CUDA_ERROR_UNKNOWN - HANDLE = 0xBEEF - calls = [] - release_calls = [0] - - def fake_get_allocation_granularity(_, _granularity_flag): - calls.append(("granularity",)) - return (SUCCESS, 1024) - - def fake_create(size, p, flags): - calls.append(("create", size)) - return (SUCCESS, HANDLE) - - def fake_addr_reserve(size, align, hint, fd): - calls.append(("addr_reserve", size, align, hint, fd)) - return (SUCCESS, 0x20_0000) - - def fake_map(ptr, size, offset, handle, flags): - calls.append(("map", ptr, size, handle)) - return (SUCCESS,) - - def fake_set_access(ptr, size, descs, count): - calls.append(("set_access", ptr, size, count)) - return (SUCCESS,) - - def fake_release(handle): - calls.append(("release", handle)) - release_calls[0] += 1 - if release_calls[0] == 1: - return (ERROR,) - return (SUCCESS,) - - def fake_unmap(ptr, size): - calls.append(("unmap", ptr, size)) - return (SUCCESS,) - - def fake_addr_free(ptr, size): - calls.append(("addr_free", ptr, size)) - return (SUCCESS,) - - class FakeBuffer: - def __init__(self, ptr, size, mr): - self.handle = ptr - self.size = size - self.memory_resource = mr - - def fake_from_handle(ptr, size, mr): - calls.append(("from_handle", ptr, size)) - return FakeBuffer(ptr, size, mr) - - monkeypatch.setattr(driver, "cuMemGetAllocationGranularity", fake_get_allocation_granularity) - monkeypatch.setattr(driver, "cuMemCreate", fake_create) - monkeypatch.setattr(driver, "cuMemAddressReserve", fake_addr_reserve) - monkeypatch.setattr(driver, "cuMemMap", fake_map) - monkeypatch.setattr(driver, "cuMemSetAccess", fake_set_access) - monkeypatch.setattr(driver, "cuMemRelease", fake_release) - monkeypatch.setattr(driver, "cuMemUnmap", fake_unmap) - monkeypatch.setattr(driver, "cuMemAddressFree", fake_addr_free) - monkeypatch.setattr(Buffer, "from_handle", fake_from_handle) - - with pytest.raises(CUDAError): - vmm_mr.allocate(4096) - - assert [c[0] for c in calls] == [ - "granularity", - "create", - "addr_reserve", - "map", - "set_access", - "release", - "unmap", - "release", - "addr_free", - ] - assert calls.count(("release", HANDLE)) == 2 - assert "from_handle" not in [c[0] for c in calls] - - -def test_vmm_allocator_grow_allocation_fast_path_rolls_back_if_handle_release_fails(init_cuda, monkeypatch): - """Verify fast-path rollback when the post-map handle release fails.""" +def test_vmm_allocate_close_does_not_leak(init_cuda): device = Device() if not device.properties.virtual_memory_management_supported: pytest.skip("Virtual memory management is not supported on this device") - vmm_mr = VirtualMemoryResource( + mr = VirtualMemoryResource( device, config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"), ) + buf = mr.allocate(8 * 1024 * 1024) + aligned_size = buf.size + buf.close() - prop = driver.CUmemAllocationProp() - prop.type = driver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED - prop.location.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE - prop.location.id = device.device_id - - SUCCESS = driver.CUresult.CUDA_SUCCESS - ERROR = driver.CUresult.CUDA_ERROR_UNKNOWN - NEW_HANDLE = 0xBEEF - calls = [] - release_calls = [0] - - def fake_create(size, p, flags): - calls.append(("create", size)) - return (SUCCESS, NEW_HANDLE) - - def fake_map(ptr, size, offset, handle, flags): - calls.append(("map", ptr, size, handle)) - return (SUCCESS,) - - def fake_set_access(ptr, size, descs, count): - calls.append(("set_access", ptr, size, count)) - return (SUCCESS,) - - def fake_release(handle): - calls.append(("release", handle)) - release_calls[0] += 1 - if release_calls[0] == 1: - return (ERROR,) - return (SUCCESS,) - - def fake_unmap(ptr, size): - calls.append(("unmap", ptr, size)) - return (SUCCESS,) - - def fake_addr_free(ptr, size): - calls.append(("addr_free", ptr, size)) - return (SUCCESS,) - - class FakeBuffer: - def __init__(self, size): - self._size = size - - original_size = 2 * 1024 * 1024 - aligned_additional = 2 * 1024 * 1024 - new_size = original_size + aligned_additional - new_ptr = 0x10_0000 - buf = FakeBuffer(original_size) - - monkeypatch.setattr(driver, "cuMemCreate", fake_create) - monkeypatch.setattr(driver, "cuMemMap", fake_map) - monkeypatch.setattr(driver, "cuMemSetAccess", fake_set_access) - monkeypatch.setattr(driver, "cuMemUnmap", fake_unmap) - monkeypatch.setattr(driver, "cuMemRelease", fake_release) - monkeypatch.setattr(driver, "cuMemAddressFree", fake_addr_free) - - with pytest.raises(CUDAError): - vmm_mr._grow_allocation_fast_path(buf, new_size, prop, aligned_additional, new_ptr) - - assert buf._size == original_size - assert [c[0] for c in calls] == [ - "create", - "map", - "set_access", - "release", - "unmap", - "release", - "addr_free", - ] - - -def test_vmm_allocator_grow_allocation_slow_path(init_cuda, monkeypatch): - """Exercise the VMM grow slow path with mocked driver calls. - - This verifies that both retained and newly-created handles are released after - successful mapping and commit. - """ - device = Device() - if not device.properties.virtual_memory_management_supported: - pytest.skip("Virtual memory management is not supported on this device") - - vmm_mr = VirtualMemoryResource( - device, - config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"), - ) - - prop = driver.CUmemAllocationProp() - prop.type = driver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED - prop.location.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE - prop.location.id = device.device_id - - SUCCESS = driver.CUresult.CUDA_SUCCESS - OLD_HANDLE = 0xFEED - NEW_HANDLE = 0xBEEF - calls = [] - - def fake_addr_reserve(size, align, hint, fd): - calls.append(("reserve", size, align, hint, fd)) - return (SUCCESS, 0x30_0000) - - def fake_retain(handle): - calls.append(("retain", handle)) - return (SUCCESS, OLD_HANDLE) - - def fake_create(size, p, flags): - calls.append(("create", size)) - return (SUCCESS, NEW_HANDLE) - - def fake_map(ptr, size, offset, handle, flags): - calls.append(("map", ptr, size, offset, handle)) - return (SUCCESS,) - - def fake_unmap(ptr, size): - calls.append(("unmap", ptr, size)) - return (SUCCESS,) - - def fake_set_access(ptr, size, descs, count): - calls.append(("set_access", ptr, size, count)) - return (SUCCESS,) - - def fake_release(handle): - calls.append(("release", handle)) - return (SUCCESS,) - - def fake_addr_free(ptr, size): - calls.append(("addr_free", ptr, size)) - return (SUCCESS,) - - class FakeInputBuffer: - def __init__(self, size, handle): - self._size = size - self.handle = handle - - def _clear(self): - calls.append(("clear",)) - - class FakeReturnedBuffer: - def __init__(self, ptr, size, mr): - self.handle = ptr - self.size = size - self.memory_resource = mr - - def fake_from_handle(ptr, size, mr): - calls.append(("from_handle", ptr, size)) - return FakeReturnedBuffer(ptr, size, mr) - - monkeypatch.setattr(driver, "cuMemAddressReserve", fake_addr_reserve) - monkeypatch.setattr(driver, "cuMemRetainAllocationHandle", fake_retain) - monkeypatch.setattr(driver, "cuMemUnmap", fake_unmap) - monkeypatch.setattr(driver, "cuMemMap", fake_map) - monkeypatch.setattr(driver, "cuMemCreate", fake_create) - monkeypatch.setattr(driver, "cuMemSetAccess", fake_set_access) - monkeypatch.setattr(driver, "cuMemRelease", fake_release) - monkeypatch.setattr(driver, "cuMemAddressFree", fake_addr_free) - monkeypatch.setattr(Buffer, "from_handle", fake_from_handle) - - original_size = 2 * 1024 * 1024 - aligned_additional = 2 * 1024 * 1024 - result_size = original_size + aligned_additional - original_ptr = 0x10_0000 - buf = FakeInputBuffer(original_size, original_ptr) - - result = vmm_mr._grow_allocation_slow_path(buf, result_size, prop, aligned_additional, result_size, 1024) - - assert isinstance(result, FakeReturnedBuffer) - assert result.size == result_size - assert result.handle == 0x30_0000 - assert [c[0] for c in calls] == [ - "reserve", - "retain", - "unmap", - "map", - "create", - "map", - "set_access", - "release", - "release", - "addr_free", - "clear", - "from_handle", - ] - assert ("release", NEW_HANDLE) in calls - assert ("release", OLD_HANDLE) in calls - - -def test_vmm_allocator_grow_allocation_slow_path_rolls_back_if_handle_release_fails(init_cuda, monkeypatch): - """Verify slow-path rollback when a handle release fails near the end of growth.""" - device = Device() - if not device.properties.virtual_memory_management_supported: - pytest.skip("Virtual memory management is not supported on this device") - - vmm_mr = VirtualMemoryResource( - device, - config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"), - ) - - prop = driver.CUmemAllocationProp() - prop.type = driver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED - prop.location.type = driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE - prop.location.id = device.device_id - - SUCCESS = driver.CUresult.CUDA_SUCCESS - ERROR = driver.CUresult.CUDA_ERROR_UNKNOWN - OLD_HANDLE = 0xFEED - NEW_HANDLE = 0xBEEF - calls = [] - release_calls = [0] - - def fake_addr_reserve(size, align, hint, fd): - calls.append(("reserve", size, align, hint, fd)) - return (SUCCESS, 0x30_0000) - - def fake_retain(handle): - calls.append(("retain", handle)) - return (SUCCESS, OLD_HANDLE) - - def fake_create(size, p, flags): - calls.append(("create", size)) - return (SUCCESS, NEW_HANDLE) - - def fake_map(ptr, size, offset, handle, flags): - calls.append(("map", ptr, size, offset, handle)) - return (SUCCESS,) - - def fake_unmap(ptr, size): - calls.append(("unmap", ptr, size)) - return (SUCCESS,) - - def fake_set_access(ptr, size, descs, count): - calls.append(("set_access", ptr, size, count)) - return (SUCCESS,) - - def fake_release(handle): - calls.append(("release", handle)) - release_calls[0] += 1 - if release_calls[0] == 1: - return (ERROR,) - return (SUCCESS,) - - def fake_addr_free(ptr, size): - calls.append(("addr_free", ptr, size)) - return (SUCCESS,) - - class FakeInputBuffer: - def __init__(self, size, handle): - self._size = size - self.handle = handle - - def _clear(self): - calls.append(("clear",)) - - def fake_from_handle(ptr, size, mr): - calls.append(("from_handle", ptr, size)) - return FakeBuffer(ptr, size, mr) - - class FakeBuffer: - def __init__(self, ptr, size, mr): - self.handle = ptr - self.size = size - self.memory_resource = mr + baseline = handle_return(driver.cuMemGetInfo())[0] + for _ in range(8): + mr.allocate(8 * 1024 * 1024).close() + free = handle_return(driver.cuMemGetInfo())[0] - original_size = 2 * 1024 * 1024 - aligned_additional = 2 * 1024 * 1024 - result_size = original_size + aligned_additional - original_ptr = 0x10_0000 - buf = FakeInputBuffer(original_size, original_ptr) - - monkeypatch.setattr(driver, "cuMemAddressReserve", fake_addr_reserve) - monkeypatch.setattr(driver, "cuMemRetainAllocationHandle", fake_retain) - monkeypatch.setattr(driver, "cuMemUnmap", fake_unmap) - monkeypatch.setattr(driver, "cuMemMap", fake_map) - monkeypatch.setattr(driver, "cuMemCreate", fake_create) - monkeypatch.setattr(driver, "cuMemSetAccess", fake_set_access) - monkeypatch.setattr(driver, "cuMemRelease", fake_release) - monkeypatch.setattr(driver, "cuMemAddressFree", fake_addr_free) - monkeypatch.setattr(Buffer, "from_handle", fake_from_handle) - - with pytest.raises(CUDAError): - vmm_mr._grow_allocation_slow_path(buf, result_size, prop, aligned_additional, result_size, 1024) - - assert [c[0] for c in calls] == [ - "reserve", - "retain", - "unmap", - "map", - "create", - "map", - "set_access", - "release", - "unmap", - "release", - "unmap", - "map", - "release", - "addr_free", - ] - assert calls.count(("release", NEW_HANDLE)) == 2 - assert calls.count(("release", OLD_HANDLE)) == 1 + assert baseline - free < aligned_size def test_vmm_allocator_rdma_unsupported_exception(): From 2e8eaf76c22e5618fea3191e4a02a55c23200063 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Wed, 15 Jul 2026 06:09:56 +0200 Subject: [PATCH 4/6] Use exactly-once VMM handle cleanup hooks --- .../core/_memory/_virtual_memory_resource.py | 12 +++----- cuda_core/cuda/core/_utils/cuda_utils.pyi | 16 +++++----- cuda_core/cuda/core/_utils/cuda_utils.pyx | 30 ++++++++++--------- 3 files changed, 28 insertions(+), 30 deletions(-) diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py index 5cd2b7841e5..b0f90f213dc 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py @@ -322,8 +322,7 @@ def _grow_allocation_fast_path( ) res, new_handle = driver.cuMemCreate(aligned_additional_size, prop, 0) raise_if_driver_error(res) - trans.on_failure(lambda h=new_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) - trans.on_success(lambda h=new_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) + trans.on_exit(lambda h=new_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) # Map the new physical memory to the extended VA range (res,) = driver.cuMemMap(new_ptr, aligned_additional_size, 0, new_handle, 0) @@ -389,8 +388,7 @@ def _grow_allocation_slow_path( # Get the old allocation handle for remapping result, old_handle = driver.cuMemRetainAllocationHandle(buf.handle) raise_if_driver_error(result) - trans.on_failure(lambda h=old_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) - trans.on_success(lambda h=old_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) + trans.on_exit(lambda h=old_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) # Unmap the old VA range (aligned previous size) aligned_prev_size = total_aligned_size - aligned_additional_size @@ -418,8 +416,7 @@ def _remap_old() -> None: # Create new physical memory for the additional size res, new_handle = driver.cuMemCreate(aligned_additional_size, prop, 0) raise_if_driver_error(res) - trans.on_failure(lambda h=new_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) - trans.on_success(lambda h=new_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) + trans.on_exit(lambda h=new_handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) # Map the new physical memory to the extended portion (aligned offset) (res,) = driver.cuMemMap(int(new_ptr) + aligned_prev_size, aligned_additional_size, 0, new_handle, 0) @@ -541,9 +538,8 @@ def allocate(self, size: int, *, stream: Stream | GraphBuilder | None = None) -> # ---- Create physical memory ---- res, handle = driver.cuMemCreate(aligned_size, prop, 0) raise_if_driver_error(res) - trans.on_failure(lambda h=handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) # Once mapped, the physical allocation is kept alive without the creation reference. - trans.on_success(lambda h=handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) + trans.on_exit(lambda h=handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) # ---- Reserve VA space ---- # Potentially, use a separate size for the VA reservation from the physical allocation size diff --git a/cuda_core/cuda/core/_utils/cuda_utils.pyi b/cuda_core/cuda/core/_utils/cuda_utils.pyi index bb368d45cc3..aba411572df 100644 --- a/cuda_core/cuda/core/_utils/cuda_utils.pyi +++ b/cuda_core/cuda/core/_utils/cuda_utils.pyi @@ -19,22 +19,22 @@ class NVRTCError(CUDAError): class Transaction: """ - A context manager for transactional operations with failure and success callbacks. + A context manager for transactional operations with failure and exit callbacks. Failure callbacks are executed in LIFO order if the transaction exits without being committed. - Success callbacks are executed in FIFO order when the transaction is committed. + Exit callbacks always run: in LIFO order on rollback or FIFO order during commit. Usage: with Transaction() as txn: txn.on_failure(some_cleanup_function, arg1, arg2) - txn.on_success(some_finalize_function, arg1, arg2) + txn.on_exit(some_finalize_function, arg1, arg2) # ... perform operations ... txn.commit() Methods: on_failure(fn, *args, **kwargs): Register a callback to be called on rollback. - on_success(fn, *args, **kwargs): Register a callback to be called on commit. - commit(): Disarm failure callbacks and run success callbacks. + on_exit(fn, *args, **kwargs): Register a callback to be called on rollback or commit. + commit(): Disarm failure callbacks and run exit callbacks. """ def __init__(self) -> None: @@ -52,15 +52,15 @@ class Transaction: Values are bound now via partial so late mutations don't bite you. """ - def on_success(self, fn: Callable[..., Any], /, *args: Any, **kwargs) -> None: + def on_exit(self, fn: Callable[..., Any], /, *args: Any, **kwargs) -> None: """ - Register a success callback (runs in FIFO order during commit()). + Register an exit callback (runs exactly once, on rollback or during commit()). Values are bound now via partial so late mutations don't bite you. """ def commit(self) -> None: """ - Disarm all failure callbacks, then run success callbacks in FIFO order. + Disarm all failure callbacks, then run exit callbacks in FIFO order. """ _keep_driver_in_stub: 'driver.CUresult' _keep_nvrtc_in_stub: 'nvrtc.nvrtcResult' diff --git a/cuda_core/cuda/core/_utils/cuda_utils.pyx b/cuda_core/cuda/core/_utils/cuda_utils.pyx index 8cb65811aa5..5c507207e26 100644 --- a/cuda_core/cuda/core/_utils/cuda_utils.pyx +++ b/cuda_core/cuda/core/_utils/cuda_utils.pyx @@ -306,26 +306,26 @@ def is_nested_sequence(obj: object) -> bool: class Transaction: """ - A context manager for transactional operations with failure and success callbacks. + A context manager for transactional operations with failure and exit callbacks. Failure callbacks are executed in LIFO order if the transaction exits without being committed. - Success callbacks are executed in FIFO order when the transaction is committed. + Exit callbacks always run: in LIFO order on rollback or FIFO order during commit. Usage: with Transaction() as txn: txn.on_failure(some_cleanup_function, arg1, arg2) - txn.on_success(some_finalize_function, arg1, arg2) + txn.on_exit(some_finalize_function, arg1, arg2) # ... perform operations ... txn.commit() Methods: on_failure(fn, *args, **kwargs): Register a callback to be called on rollback. - on_success(fn, *args, **kwargs): Register a callback to be called on commit. - commit(): Disarm failure callbacks and run success callbacks. + on_exit(fn, *args, **kwargs): Register a callback to be called on rollback or commit. + commit(): Disarm failure callbacks and run exit callbacks. """ def __init__(self) -> None: self._stack = ExitStack() - self._on_success: list[Callable[[], Any]] = [] + self._on_exit: list[Callable[[], Any]] = [] self._entered = False def __enter__(self): @@ -336,7 +336,7 @@ class Transaction: def __exit__(self, exc_type, exc, tb): # If exit callbacks remain, they'll run in LIFO order. self._entered = False - self._on_success.clear() + self._on_exit.clear() return self._stack.__exit__(exc_type, exc, tb) def on_failure(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> None: @@ -348,24 +348,26 @@ class Transaction: raise RuntimeError("Transaction must be entered before on_failure()") self._stack.callback(partial(fn, *args, **kwargs)) - def on_success(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> None: + def on_exit(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> None: """ - Register a success callback (runs in FIFO order during commit()). + Register an exit callback (runs exactly once, on rollback or during commit()). Values are bound now via partial so late mutations don't bite you. """ if not self._entered: - raise RuntimeError("Transaction must be entered before on_success()") - self._on_success.append(partial(fn, *args, **kwargs)) + raise RuntimeError("Transaction must be entered before on_exit()") + callback = partial(fn, *args, **kwargs) + self._stack.callback(callback) + self._on_exit.append(callback) def commit(self) -> None: """ - Disarm all failure callbacks, then run success callbacks in FIFO order. + Disarm all failure callbacks, then run exit callbacks in FIFO order. """ # pop_all() empties this stack so no callbacks are triggered on exit. self._stack.pop_all() - for fn in self._on_success: + for fn in self._on_exit: fn() - self._on_success.clear() + self._on_exit.clear() # Track whether we've already warned about fork method From d91108daa684565b01e12d2834816a8a1f53e0cc Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Wed, 15 Jul 2026 06:14:57 +0200 Subject: [PATCH 5/6] Simplify VMM deallocation ownership --- cuda_core/cuda/core/_memory/_virtual_memory_resource.py | 5 +---- cuda_core/tests/test_memory.py | 3 ++- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py index b0f90f213dc..c1812412e79 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py @@ -585,14 +585,11 @@ def deallocate(self, ptr: DevicePointerType, size: int, *, stream: Stream | Grap from cuda.core._stream import Stream_accept Stream_accept(stream) - result, handle = driver.cuMemRetainAllocationHandle(ptr) - raise_if_driver_error(result) + # The mapping owns the allocation; unmapping frees its backing memory when no external references remain. (result,) = driver.cuMemUnmap(ptr, size) raise_if_driver_error(result) (result,) = driver.cuMemAddressFree(ptr, size) raise_if_driver_error(result) - (result,) = driver.cuMemRelease(handle) - raise_if_driver_error(result) @property def is_device_accessible(self) -> bool: diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index 3a06ee8415d..f28e2eb0afa 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1045,7 +1045,7 @@ def test_vmm_allocate_close_does_not_leak(init_cuda): device, config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"), ) - buf = mr.allocate(8 * 1024 * 1024) + buf = mr.allocate(8 * 1024 * 1024) # Warm up and learn the aligned allocation size. aligned_size = buf.size buf.close() @@ -1054,6 +1054,7 @@ def test_vmm_allocate_close_does_not_leak(init_cuda): mr.allocate(8 * 1024 * 1024).close() free = handle_return(driver.cuMemGetInfo())[0] + # Current main leaks aligned_size per iteration; the fixed path stays near baseline. assert baseline - free < aligned_size From 6de0c3643343fd00dd23c95f2ff01d883b6dbd00 Mon Sep 17 00:00:00 2001 From: Minh Vu Date: Wed, 15 Jul 2026 06:21:04 +0200 Subject: [PATCH 6/6] Cover VMM grow leaks and preserve map errors --- .../core/_memory/_virtual_memory_resource.py | 4 ++-- cuda_core/tests/test_memory.py | 19 ++++++++++++++----- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py index c1812412e79..6e34bc5bcfc 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py @@ -538,7 +538,7 @@ def allocate(self, size: int, *, stream: Stream | GraphBuilder | None = None) -> # ---- Create physical memory ---- res, handle = driver.cuMemCreate(aligned_size, prop, 0) raise_if_driver_error(res) - # Once mapped, the physical allocation is kept alive without the creation reference. + # Drop the creation reference on either outcome; a successful mapping keeps the allocation alive. trans.on_exit(lambda h=handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) # ---- Reserve VA space ---- @@ -550,8 +550,8 @@ def allocate(self, size: int, *, stream: Stream | GraphBuilder | None = None) -> # ---- Map physical memory into VA ---- (res,) = driver.cuMemMap(ptr, aligned_size, 0, handle, 0) - trans.on_failure(lambda p=ptr, s=aligned_size: raise_if_driver_error(driver.cuMemUnmap(p, s)[0])) raise_if_driver_error(res) + trans.on_failure(lambda p=ptr, s=aligned_size: raise_if_driver_error(driver.cuMemUnmap(p, s)[0])) # ---- Set access for owner + peers ---- descs = self._build_access_descriptors(prop) diff --git a/cuda_core/tests/test_memory.py b/cuda_core/tests/test_memory.py index f28e2eb0afa..1e2b878851c 100644 --- a/cuda_core/tests/test_memory.py +++ b/cuda_core/tests/test_memory.py @@ -1036,7 +1036,8 @@ def __init__(self, size): assert ("release", NEW_HANDLE) in calls -def test_vmm_allocate_close_does_not_leak(init_cuda): +@pytest.mark.parametrize("grow", [False, True], ids=["allocate", "grow"]) +def test_vmm_allocate_close_does_not_leak(init_cuda, grow): device = Device() if not device.properties.virtual_memory_management_supported: pytest.skip("Virtual memory management is not supported on this device") @@ -1045,13 +1046,21 @@ def test_vmm_allocate_close_does_not_leak(init_cuda): device, config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"), ) - buf = mr.allocate(8 * 1024 * 1024) # Warm up and learn the aligned allocation size. - aligned_size = buf.size - buf.close() + requested_size = 8 * 1024 * 1024 + + def allocate_and_close(): + buf = mr.allocate(requested_size) + if grow: + buf = mr.modify_allocation(buf, 2 * buf.size) + aligned_size = buf.size + buf.close() + return aligned_size + + aligned_size = allocate_and_close() # Warm up and learn the aligned allocation size. baseline = handle_return(driver.cuMemGetInfo())[0] for _ in range(8): - mr.allocate(8 * 1024 * 1024).close() + allocate_and_close() free = handle_return(driver.cuMemGetInfo())[0] # Current main leaks aligned_size per iteration; the fixed path stays near baseline.