Cap memory pool size in some tests

Author: juenglin

cuda_core/tests/test_memory.py

@@ -1452,11 +1452,13 @@ def test_pinned_mr_numa_id_default_no_ipc(init_cuda):
     device = Device()
     skip_if_pinned_memory_unsupported(device)
 
-    mr = create_pinned_memory_resource_or_xfail(PinnedMemoryResourceOptions(), xfail_device=device)
+    mr = create_pinned_memory_resource_or_xfail(PinnedMemoryResourceOptions(max_size=POOL_SIZE), xfail_device=device)
     assert mr.numa_id == -1
     mr.close()
 
-    mr = create_pinned_memory_resource_or_xfail(PinnedMemoryResourceOptions(ipc_enabled=False), xfail_device=device)
+    mr = create_pinned_memory_resource_or_xfail(
+        PinnedMemoryResourceOptions(ipc_enabled=False, max_size=POOL_SIZE), xfail_device=device
+    )
     assert mr.numa_id == -1
     mr.close()
 
@@ -1491,7 +1493,9 @@ def test_pinned_mr_numa_id_explicit(init_cuda):
     if host_numa_id < 0:
         pytest.skip("System does not support NUMA")
 
-    mr = create_pinned_memory_resource_or_xfail(PinnedMemoryResourceOptions(numa_id=host_numa_id), xfail_device=device)
+    mr = create_pinned_memory_resource_or_xfail(
+        PinnedMemoryResourceOptions(numa_id=host_numa_id, max_size=POOL_SIZE), xfail_device=device
+    )
     assert mr.numa_id == host_numa_id
     mr.close()
 

cuda_core/tests/test_memory_peer_access.py

@@ -11,6 +11,14 @@
 from cuda.core._utils.cuda_utils import CUDAError
 
 NBYTES = 1024
+# Every owned pool below holds at most NBYTES, but a pool created without an
+# explicit max_size reserves a system-dependent window that scales with device
+# memory -- hundreds of GiB on large-memory GPUs. The per-process virtual
+# address budget is bounded (~1 TB on Windows MCDM), and reservations are not
+# returned until a pool is torn down and its stream-ordered frees retire, so
+# oversized windows accumulate across a session and eventually starve later
+# pool creations with CUDA_ERROR_OUT_OF_MEMORY (issue #2381). Cap them.
+POOL_SIZE = 2097152  # 2MB size
 
 pytestmark = pytest.mark.thread_unsafe(reason="peer access tests mutate process-global CUDA memory-pool access state")
 
@@ -22,7 +30,7 @@ def test_peer_access_basic(mempool_device_x2):
     one_on_dev0 = make_scratch_buffer(dev0, 1, NBYTES)
     stream_on_dev0 = dev0.create_stream()
     # Use owned pool to ensure clean initial state (no stale peer access).
-    dmr_on_dev1 = DeviceMemoryResource(dev1, DeviceMemoryResourceOptions())
+    dmr_on_dev1 = DeviceMemoryResource(dev1, DeviceMemoryResourceOptions(max_size=POOL_SIZE))
     buf_on_dev1 = dmr_on_dev1.allocate(NBYTES, stream=dev1.default_stream)
 
     # No access at first.
@@ -73,7 +81,7 @@ def test_peer_access_transitions(mempool_device_x3):
     pgens = [PatternGen(devs[i], NBYTES, streams[i]) for i in range(3)]
     # Use owned pools (with options) to ensure clean initial state.
     # Default pools are shared and may have stale peer access from prior tests.
-    dmrs = [DeviceMemoryResource(dev, DeviceMemoryResourceOptions()) for dev in devs]
+    dmrs = [DeviceMemoryResource(dev, DeviceMemoryResourceOptions(max_size=POOL_SIZE)) for dev in devs]
     bufs = [dmr.allocate(NBYTES, stream=dev.default_stream) for dmr, dev in zip(dmrs, devs)]
 
     def verify_state(state, pattern_seed):
@@ -163,7 +171,7 @@ def isolated_dmr_x2(mempool_device_x2):
     proxy tests are not polluted by other tests sharing a default pool.
     """
     dev0, dev1 = mempool_device_x2
-    dmr = DeviceMemoryResource(dev0, DeviceMemoryResourceOptions())
+    dmr = DeviceMemoryResource(dev0, DeviceMemoryResourceOptions(max_size=POOL_SIZE))
     dmr.peer_accessible_by = []
     return dmr, dev0, dev1
 
@@ -273,7 +281,7 @@ def test_peer_accessible_by_no_cache_across_proxies(mempool_device_x2):
 def test_peer_accessible_by_iteration_order_is_sorted(mempool_device_x2):
     """``__iter__`` yields peers in ascending device-ordinal order."""
     dev0, dev1 = mempool_device_x2
-    dmr = DeviceMemoryResource(dev0, DeviceMemoryResourceOptions())
+    dmr = DeviceMemoryResource(dev0, DeviceMemoryResourceOptions(max_size=POOL_SIZE))
     dmr.peer_accessible_by = [dev1]
     devices = list(dmr.peer_accessible_by)
     ids = [d.device_id for d in devices]