Cover VMM grow leaks and preserve map errors

Author: fallintoplace

cuda_core/cuda/core/_memory/_virtual_memory_resource.py

@@ -538,7 +538,7 @@ 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)
-            # Once mapped, the physical allocation is kept alive without the creation reference.
+            # 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 ----
@@ -550,8 +550,8 @@ def allocate(self, size: int, *, stream: Stream | GraphBuilder | None = None) ->
 
             # ---- Map physical memory into VA ----
             (res,) = driver.cuMemMap(ptr, aligned_size, 0, handle, 0)
-            trans.on_failure(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)

cuda_core/tests/test_memory.py

@@ -1036,7 +1036,8 @@ def __init__(self, size):
     assert ("release", NEW_HANDLE) in calls
 
 
-def test_vmm_allocate_close_does_not_leak(init_cuda):
+@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")
@@ -1045,13 +1046,21 @@ def test_vmm_allocate_close_does_not_leak(init_cuda):
         device,
         config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd"),
     )
-    buf = mr.allocate(8 * 1024 * 1024)  # Warm up and learn the aligned allocation size.
-    aligned_size = buf.size
-    buf.close()
+    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):
-        mr.allocate(8 * 1024 * 1024).close()
+        allocate_and_close()
     free = handle_return(driver.cuMemGetInfo())[0]
 
     # Current main leaks aligned_size per iteration; the fixed path stays near baseline.