Account for missing mempool support and harden test

Author: juenglin

cuda_core/tests/test_mempool_oom_retry.py

@@ -3,10 +3,10 @@
 # SPDX-License-Identifier: Apache-2.0
 """Regression test for the deferred-release retry in ``get_device_mempool``.
 
-Issue #2381: the driver reserves virtual address space per memory pool -- on
-Windows MCDM roughly 2x device memory against a 40-bit (1 TiB) cap -- so pools
-that are awaiting teardown can starve the *default* pool's reservation. cuda.core
-then reports ``CUDA_ERROR_OUT_OF_MEMORY`` on a device with ample free memory.
+Issue #2381: each memory pool reserves virtual address space scaling with device
+memory, and the per-process budget is bounded (notably ~1 TB on Windows MCDM), so
+pools awaiting teardown can starve the *default* pool's reservation. cuda.core then
+reports ``CUDA_ERROR_OUT_OF_MEMORY`` on a device with ample free memory.
 
 A pool's address space comes back only once the pool is destroyed, and
 ``cuMemPoolDestroy`` waits on the stream-ordered frees of the pool's outstanding
@@ -22,6 +22,8 @@
 from __future__ import annotations
 
 import multiprocessing as _mp
+import queue
+import traceback
 
 import pytest
 from helpers.child_processes import child_timeout_sec, kill_subprocesses
@@ -39,14 +41,17 @@
 NOT_DEFERRED = "not-deferred"
 RECOVERED = "recovered"
 NOT_RECOVERED = "not-recovered"
+UNSUPPORTED = "unsupported"
+CRASHED = "crashed"
+TIMED_OUT = "timed-out"
 
 
-def _worker_deferred_release(result_queue):
+def _run_deferred_release():
     """Drive a default-pool lookup into deferred-release failure, then retry it.
 
-    Reports one of the module-level outcome strings so the parent can tell a
-    genuine regression apart from a machine whose address space is too large to
-    exhaust.
+    Returns an ``(outcome, detail)`` pair using the module-level outcome strings,
+    so the parent can tell a genuine regression apart from a machine that cannot
+    run this at all.
     """
     import gc
     import time
@@ -57,15 +62,18 @@ def _worker_deferred_release(result_queue):
     from cuda.core import Device, DeviceMemoryResource, DeviceMemoryResourceOptions
     from cuda.core._utils.cuda_utils import CUDAError
 
-    def default_pool_available():
-        err, _pool = driver.cuDeviceGetMemPool(dev)
-        return err == driver.CUresult.CUDA_SUCCESS
-
     device = Device(0)
     device.set_current()
+    if not device.properties.memory_pools_supported:
+        return UNSUPPORTED, "Device does not support mempool operations"
+
     err, dev = driver.cuDeviceGet(0)
     assert err == driver.CUresult.CUDA_SUCCESS, err
 
+    def default_pool_available():
+        err, _pool = driver.cuDeviceGetMemPool(dev)
+        return err == driver.CUresult.CUDA_SUCCESS
+
     # Build everything before the address space is gone; afterwards even kernel
     # compilation could fail for unrelated reasons.
     sleeper = NanosleepKernel(device, sleep_duration_ms=BLOCK_MS)
@@ -80,19 +88,22 @@ def reserve_until_oom(pool_bytes):
         for _ in range(MAX_POOLS_PER_PASS):
             try:
                 mr = DeviceMemoryResource(device, options=options)
+                # An empty pool is torn down immediately, which is the one case
+                # that cannot defer anything, so hold an allocation to force
+                # teardown to wait. Exhaustion frequently surfaces here rather
+                # than at construction, so both must share this guard -- an
+                # unguarded allocate kills the whole child process.
+                buffer = mr.allocate(1024, stream=dealloc_stream)
             except (CUDAError, RuntimeError):
                 return
-            # An empty pool is torn down immediately, which is the one case that
-            # cannot defer anything. Hold an allocation so teardown must wait.
-            buffers.append(mr.allocate(1024, stream=dealloc_stream))
             pools.append(mr)
+            buffers.append(buffer)
 
     reserve_until_oom(POOL_BYTES_COARSE)
     reserve_until_oom(POOL_BYTES_FINE)
 
     if default_pool_available():
-        result_queue.put((NO_EXHAUSTION, len(pools)))
-        return
+        return NO_EXHAUSTION, len(pools)
 
     # Stall the deallocation stream behind a slow kernel on another stream, so
     # the frees below cannot retire until the context is drained.
@@ -105,8 +116,7 @@ def reserve_until_oom(pool_bytes):
 
     if default_pool_available():
         # Release outran us; there was no deferred window to recover from.
-        result_queue.put((NOT_DEFERRED, len(pools)))
-        return
+        return NOT_DEFERRED, None
 
     # The raw lookup just failed, so anything the retried lookup achieves below
     # is attributable to the retry itself -- no cross-build comparison needed.
@@ -116,9 +126,20 @@ def reserve_until_oom(pool_bytes):
     try:
         DeviceMemoryResource(device)
     except (CUDAError, RuntimeError) as exc:
-        result_queue.put((NOT_RECOVERED, repr(exc)))
-        return
-    result_queue.put((RECOVERED, time.perf_counter() - started))
+        return NOT_RECOVERED, repr(exc)
+    return RECOVERED, time.perf_counter() - started
+
+
+def _worker_deferred_release(result_queue):
+    """Always report an outcome, so a crash surfaces as a diagnosis.
+
+    Without this the parent would block until its timeout and raise a bare
+    ``queue.Empty``, hiding whatever actually went wrong in the child.
+    """
+    try:
+        result_queue.put(_run_deferred_release())
+    except BaseException:
+        result_queue.put((CRASHED, traceback.format_exc()))
 
 
 @pytest.mark.agent_authored(model="claude-opus-5")
@@ -130,16 +151,20 @@ def test_default_mempool_lookup_recovers_from_deferred_release():
     proc.start()
     try:
         outcome, detail = result_queue.get(timeout=child_timeout_sec())
+    except queue.Empty:
+        outcome, detail = TIMED_OUT, f"child produced no result within {child_timeout_sec()}s"
     finally:
         proc.join(timeout=child_timeout_sec())
         survivors = kill_subprocesses(proc)
     assert not survivors, "child process did not exit"
 
+    if outcome == UNSUPPORTED:
+        pytest.skip(detail)
     if outcome == NO_EXHAUSTION:
         pytest.skip(f"could not exhaust the address space; reserved {detail} pools")
     if outcome == NOT_DEFERRED:
         pytest.skip("pool release completed too quickly to leave a deferred window")
-    assert outcome == RECOVERED, f"default pool lookup did not recover: {detail}"
+    assert outcome == RECOVERED, f"{outcome}: {detail}"
     # Recovery had to wait out the blocking kernel. Returning much faster would
     # mean the lookup succeeded for some reason other than draining the context,
     # leaving this test passing without exercising the retry.