fix(core): restore access descriptors when a VMM grow rolls back

cuda_core/cuda/core/_memory/_virtual_memory_resource.py

@@ -377,6 +377,8 @@ def _grow_allocation_slow_path(
         Returns:
             Buffer: The buffer object updated with the new pointer and size.
         """
+        descs = self._build_access_descriptors(prop)
+
         with Transaction() as trans:
             # Reserve a completely new, larger VA range
             res, new_ptr = driver.cuMemAddressReserve(total_aligned_size, addr_align, 0, 0)
@@ -402,6 +404,9 @@ def _remap_old() -> None:
                 try:
                     (res,) = driver.cuMemMap(int(buf.handle), aligned_prev_size, 0, old_handle, 0)
                     raise_if_driver_error(res)
+                    if descs:
+                        (res,) = driver.cuMemSetAccess(int(buf.handle), aligned_prev_size, descs, len(descs))
+                        raise_if_driver_error(res)
                 except Exception:  # noqa: S110
                     # TODO: consider logging this exception
                     pass
@@ -434,7 +439,6 @@ def _remap_old() -> None:
             )
 
             # Set access permissions for the entire new range
-            descs = self._build_access_descriptors(prop)
             if descs:
                 (res,) = driver.cuMemSetAccess(new_ptr, total_aligned_size, descs, len(descs))
                 raise_if_driver_error(res)

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

@@ -18,6 +18,12 @@ Fixes and enhancements
   (`#2357 <https://github.com/NVIDIA/cuda-python/pull/2357>`__,
   `#2371 <https://github.com/NVIDIA/cuda-python/pull/2371>`__)
 
+- Rolling back a failed :meth:`VirtualMemoryResource.modify_allocation` now
+  restores access to the original range. Previously the rollback remapped the
+  old physical memory to its original address but left the access descriptors
+  unset, so the recovered buffer faulted on its next access.
+  (`#2440 <https://github.com/NVIDIA/cuda-python/pull/2440>`__)
+
 Deprecation Notices
 -------------------
 

cuda_core/tests/test_memory.py

@@ -1035,6 +1035,109 @@ def __init__(self, size):
     assert ("set_access", new_ptr, aligned_additional, 1) in calls
 
 
+@pytest.mark.agent_authored(model="claude-opus-5")
+def test_vmm_allocator_grow_slow_path_rollback_restores_access(init_cuda, monkeypatch):
+    """Fail the slow path after the old range is remapped and check the rollback.
+
+    Rolling back a failed grow has to leave the original range usable, so
+    restoring the mapping is not enough on its own: the access descriptors that
+    the unmap dropped have to be reapplied to the same range.
+    """
+    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
+    FAILURE = driver.CUresult.CUDA_ERROR_INVALID_VALUE
+    OLD_PTR = 0x20_0000
+    NEW_PTR = 0x40_0000
+    OLD_HANDLE = 0xF00D
+    NEW_HANDLE = 0xBEEF
+    calls = []
+
+    def fake_addr_reserve(size, align, addr, flags):
+        return (SUCCESS, NEW_PTR)
+
+    def fake_retain(ptr):
+        return (SUCCESS, OLD_HANDLE)
+
+    def fake_create(size, p, flags):
+        return (SUCCESS, NEW_HANDLE)
+
+    def fake_map(ptr, size, offset, handle, flags):
+        calls.append(("map", ptr, size, 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))
+        # Fail the forward call over the new range; the rollback call must still
+        # be allowed to succeed.
+        return (FAILURE,) if ptr == NEW_PTR else (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,)
+
+    monkeypatch.setattr(driver, "cuMemAddressReserve", fake_addr_reserve)
+    monkeypatch.setattr(driver, "cuMemRetainAllocationHandle", fake_retain)
+    monkeypatch.setattr(driver, "cuMemCreate", fake_create)
+    monkeypatch.setattr(driver, "cuMemMap", fake_map)
+    monkeypatch.setattr(driver, "cuMemUnmap", fake_unmap)
+    monkeypatch.setattr(driver, "cuMemSetAccess", fake_set_access)
+    monkeypatch.setattr(driver, "cuMemRelease", fake_release)
+    monkeypatch.setattr(driver, "cuMemAddressFree", fake_addr_free)
+
+    class FakeBuffer:
+        def __init__(self, ptr, size):
+            self.handle = ptr
+            self._size = size
+
+    aligned_prev_size = 2 * 1024 * 1024
+    aligned_additional_size = 2 * 1024 * 1024
+    total_aligned_size = aligned_prev_size + aligned_additional_size
+    addr_align = 2 * 1024 * 1024
+    buf = FakeBuffer(OLD_PTR, aligned_prev_size)
+
+    with pytest.raises(CUDAError):
+        vmm_mr._grow_allocation_slow_path(
+            buf,
+            total_aligned_size,
+            prop,
+            aligned_additional_size,
+            total_aligned_size,
+            addr_align,
+        )
+
+    # The rollback remapped the old physical memory back to the original VA.
+    assert ("map", OLD_PTR, aligned_prev_size, OLD_HANDLE) in calls
+    # ... and reapplied access to that same range, which is the regression here.
+    assert ("set_access", OLD_PTR, aligned_prev_size, 1) in calls
+    # Access is restored only after the mapping exists again.
+    assert calls.index(("map", OLD_PTR, aligned_prev_size, OLD_HANDLE)) < calls.index(
+        ("set_access", OLD_PTR, aligned_prev_size, 1)
+    )
+    # The failed reservation is still released.
+    assert ("addr_free", NEW_PTR, total_aligned_size) in calls
+
+
 def test_vmm_allocator_rdma_unsupported_exception():
     """Test that VirtualMemoryResource throws an exception when RDMA is requested but device doesn't support it.