Address VMM transaction review feedback

Author: fallintoplace

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

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'

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