diff --git a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py index f30e6e3838..6e34bc5bcf 100644 --- a/cuda_core/cuda/core/_memory/_virtual_memory_resource.py +++ b/cuda_core/cuda/core/_memory/_virtual_memory_resource.py @@ -317,19 +317,18 @@ 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) - # Register undo for creation - trans.append(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) 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]) ) @@ -382,15 +381,14 @@ 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) - # Register undo for old_handle - trans.append(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 @@ -406,28 +404,26 @@ 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) - - # Register undo for new physical memory - trans.append(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) 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] ) @@ -542,20 +538,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) - # Register undo for physical memory - trans.append(lambda h=handle: raise_if_driver_error(driver.cuMemRelease(h)[0])) + # 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 ---- # 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])) 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) @@ -589,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/cuda/core/_utils/cuda_utils.pyi b/cuda_core/cuda/core/_utils/cuda_utils.pyi index 8706792772..aba411572d 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 exit 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. + Exit callbacks always run: in LIFO order on rollback or FIFO order during commit. Usage: with Transaction() as txn: - txn.append(some_cleanup_function, arg1, arg2) + txn.on_failure(some_cleanup_function, arg1, arg2) + txn.on_exit(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_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: @@ -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_exit(self, fn: Callable[..., Any], /, *args: Any, **kwargs) -> None: + """ + 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 undo actions. After this, exiting the with-block does nothing. + 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 318d4466be..5c507207e2 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 exit 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. + Exit callbacks always run: in LIFO order on rollback or FIFO order during commit. Usage: with Transaction() as txn: - txn.append(some_cleanup_function, arg1, arg2) + txn.on_failure(some_cleanup_function, arg1, arg2) + txn.on_exit(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_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_exit: list[Callable[[], Any]] = [] self._entered = False def __enter__(self): @@ -334,23 +336,38 @@ class Transaction: def __exit__(self, exc_type, exc, tb): # If exit callbacks remain, they'll run in LIFO order. self._entered = False + self._on_exit.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_exit(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> None: + """ + 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_exit()") + callback = partial(fn, *args, **kwargs) + self._stack.callback(callback) + self._on_exit.append(callback) + def commit(self) -> None: """ - Disarm all undo actions. After this, exiting the with-block does nothing. + 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_exit: + fn() + self._on_exit.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 5ef48919a5..1e2b878851 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,11 +1028,43 @@ 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, 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 assert ("set_access", new_ptr, aligned_additional, 1) in calls + assert ("release", NEW_HANDLE) in calls + + +@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") + + mr = VirtualMemoryResource( + device, + config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"), + ) + 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): + allocate_and_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 def test_vmm_allocator_rdma_unsupported_exception():