cuda.core: close the old buffer when a VMM grow moves the mapping

The slow path of VirtualMemoryResource.modify_allocation unmapped the
old VA range, remapped the physical memory into the new range, freed
the old reservation by hand and then reset the old buffer's handle with
Buffer._clear(). That reset runs the handle's deleter, which calls
mr.deallocate() on the range that was just freed. The failing
cuMemRetainAllocationHandle is reported as a CUDAWarning since #2759
and was silently swallowed before.

Map the old physical memory into the new range as a second mapping
instead, which the VMM APIs allow (virtual aliasing), and then close the
old buffer. Closing runs deallocate() on a range that is still mapped,
so the old range is unmapped, its reservation freed and its handle
reference released through the normal path. The remap-on-rollback undo
step, which swallowed its own errors, is no longer needed because the
old mapping is never removed before the transaction commits.

Issue #2877

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

Author: Andy-Jost

cuda_core/cuda/core/_memory/_virtual_memory_resource.py

@@ -371,10 +371,11 @@ def _grow_allocation_slow_path(
         Slow path for growing a virtual memory allocation when the new region cannot be
         reserved contiguously after the existing buffer.
 
-        This function reserves a new, larger virtual address (VA) range, remaps the old
-        physical memory to the beginning of the new VA range, creates and maps new physical
-        memory for the additional size, sets access permissions, and updates the buffer's
-        pointer and size.
+        This function reserves a new, larger virtual address (VA) range, maps the old
+        physical memory to the beginning of the new VA range as a second mapping, creates
+        and maps new physical memory for the additional size, sets access permissions, and
+        then closes the old buffer, which releases the old VA range through
+        :meth:`deallocate`.
 
         Args:
             buf (Buffer): The buffer to grow.
@@ -385,8 +386,9 @@ def _grow_allocation_slow_path(
             addr_align (int): The required address alignment for the new VA range.
 
         Returns:
-            Buffer: The buffer object updated with the new pointer and size.
+            Buffer: A new buffer for the new VA range. ``buf`` is closed.
         """
+        aligned_prev_size = total_aligned_size - aligned_additional_size
         with Transaction() as trans:
             # Reserve a completely new, larger VA range
             res, new_ptr = driver.cuMemAddressReserve(total_aligned_size, addr_align, 0, 0)
@@ -396,28 +398,17 @@ def _grow_allocation_slow_path(
                 lambda np=new_ptr, s=total_aligned_size: raise_if_driver_error(driver.cuMemAddressFree(np, s)[0])
             )
 
-            # Get the old allocation handle for remapping
+            # Retain the old allocation handle to map it a second time. The mappings
+            # keep the memory alive, so the retained reference is dropped again on
+            # either outcome.
             result, old_handle = driver.cuMemRetainAllocationHandle(buf.handle)
             raise_if_driver_error(result)
             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
-            (result,) = driver.cuMemUnmap(int(buf.handle), aligned_prev_size)
-            raise_if_driver_error(result)
-
-            def _remap_old() -> None:
-                # Try to remap the old physical memory back to the original VA range
-                try:
-                    (res,) = driver.cuMemMap(int(buf.handle), aligned_prev_size, 0, old_handle, 0)
-                    raise_if_driver_error(res)
-                except Exception:  # noqa: S110
-                    # TODO: consider logging this exception
-                    pass
-
-            trans.on_failure(_remap_old)
-
-            # Remap the old physical memory to the new VA range (aligned previous size)
+            # Map the old physical memory to the new VA range (aligned previous size).
+            # The old VA range stays mapped too (virtual aliasing), so the old buffer
+            # is untouched if anything below fails and is released as a whole by
+            # buf.close() once the new mapping is complete.
             (res,) = driver.cuMemMap(int(new_ptr), aligned_prev_size, 0, old_handle, 0)
             raise_if_driver_error(res)
 
@@ -449,12 +440,10 @@ def _remap_old() -> None:
             # All succeeded, cancel undo actions
             trans.commit()
 
-        # Free the old VA range (aligned previous size)
-        (res2,) = driver.cuMemAddressFree(int(buf.handle), aligned_prev_size)
-        raise_if_driver_error(res2)
-
-        # Invalidate the old buffer so its destructor won't try to free again
-        buf._clear()
+        # Release the old VA range through the resource: closing the buffer runs
+        # deallocate(), which unmaps the old range and frees its reservation. The
+        # physical memory stays alive through the new mapping.
+        buf.close()
 
         # Return a new Buffer for the new mapping
         return Buffer.from_handle(ptr=new_ptr, size=new_size, mr=self)

cuda_core/docs/source/release/1.3.0-notes.rst

@@ -34,6 +34,14 @@ New features
 Fixes and enhancements
 ----------------------
 
+- Growing a :class:`~_memory.VirtualMemoryResource` allocation with
+  :meth:`~_memory.VirtualMemoryResource.modify_allocation` when the address
+  range cannot be extended in place no longer reports a spurious
+  :class:`CUDAWarning`. The old buffer is now closed through the resource's
+  :meth:`~_memory.VirtualMemoryResource.deallocate` instead of having its
+  range freed by hand and then released a second time.
+  (`#2877 <https://github.com/NVIDIA/cuda-python/issues/2877>`__)
+
 - ``Graph.__getitem__`` now declares an overload for each node type that has
   an executable view, so type checkers and editors see the precise view type:
   indexing with a :class:`~graph.KernelNode` yields an

cuda_core/tests/test_memory.py

@@ -1537,6 +1537,45 @@ def allocate_and_close():
     assert baseline - free < aligned_size
 
 
+@pytest.mark.agent_authored(model="claude-fable-5-1")
+@pytest.mark.thread_unsafe(reason="warning capture is process-global")
+def test_vmm_allocator_grow_allocation_slow_path_closes_old_buffer(init_cuda):
+    """The slow grow path closes the old buffer without a spurious CUDAWarning (#2877).
+
+    It used to free the old VA range by hand and then reset the old buffer's
+    handle, whose deleter called deallocate() on the freed range a second time.
+    """
+    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"),
+    )
+    buf = vmm_mr.allocate(2 * 1024 * 1024)
+    old_ptr, old_size = int(buf.handle), buf.size
+    handle_return(driver.cuMemsetD8(old_ptr, 7, old_size))
+
+    # Occupy the address range right after buf so the adjacent reservation cannot
+    # be honored and modify_allocation has to take the slow path.
+    decoy = handle_return(driver.cuMemAddressReserve(old_size, 0, old_ptr + old_size, 0))
+    try:
+        with assert_no_cuda_warning():
+            grown = vmm_mr.modify_allocation(buf, 2 * old_size)
+    finally:
+        handle_return(driver.cuMemAddressFree(decoy, old_size))
+
+    assert buf.is_closed
+    assert int(grown.handle) != old_ptr
+    assert grown.size == 2 * old_size
+    # The old contents are reachable through the new mapping.
+    host = (ctypes.c_ubyte * old_size)()
+    handle_return(driver.cuMemcpyDtoH(ctypes.addressof(host), int(grown.handle), old_size))
+    assert bytes(host) == bytes([7]) * old_size
+    grown.close()
+
+
 def test_vmm_allocator_rdma_unsupported_exception():
     """Test that VirtualMemoryResource throws an exception when RDMA is requested but device doesn't support it.