Fix VMM handle leaks in virtual memory paths

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:

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'

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

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():