test(cuda.core): pop the context when the init_cuda leftover warning aborts setup

cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp

@@ -39,9 +39,12 @@ struct VaReservationBox {
     size_t size;                          // exact reserved size; cuMemAddressFree needs both
 };
 
-// Process-wide count of live VA reservations (cuMemAddressReserve), the leak
-// quantity issue #2381 is about; read by the test suite to gate the per-test
-// gc.collect(). Atomic for free-threaded builds.
+// Process-wide count of live VA reservations created here with
+// cuMemAddressReserve. Memory-pool VA windows (issue #2381) are not included:
+// owned pools are tracked by the owned-pool count, and the device default pool
+// (cuDeviceGetMemPool) is a borrowed reference that neither count tracks.
+// Read by the test suite to gate the per-test gc.collect(). Atomic for
+// free-threaded builds.
 std::atomic<long long> g_live_va_reservations{0};
 
 struct VaMappingBox {

cuda_core/tests/conftest.py

@@ -231,36 +231,42 @@ def session_setup():
 def init_cuda(request):
     # TODO: rename this to e.g. init_context
     device = Device(0)
+    # Everything that can raise after set_current() runs inside the try, so the
+    # finally pops this context; the leftover warning below is escalated to an
+    # error. A failed set_current() leaves nothing for the finally to undo.
+    # live_at_setup stays None until the setup snapshot is recorded.
     device.set_current()
-
-    # Set option to avoid spin-waiting on synchronization.
-    if int(os.environ.get("CUDA_CORE_TEST_BLOCKING_SYNC", 0)) != 0:
-        handle_return(
-            driver.cuDevicePrimaryCtxSetFlags(device.device_id, driver.CUctx_flags.CU_CTX_SCHED_BLOCKING_SYNC)
-        )
-
-    # Leftover check: growth since the previous init_cuda test means something
-    # in between left an owned resource alive. Collect first to sweep a cyclic
-    # leftover, and before the counts at setup are recorded so they are reflected there.
-    if _mr_expected_live is not None and any(now > exp for now, exp in zip(_mr_live_at_item_start, _mr_expected_live)):
-        gc.collect()
-        now_counts = _live_resource_counts()
-        held = tuple(max(n - e, 0) for n, e in zip(now_counts, _mr_expected_live))
-        excess = tuple(s - e for s, e in zip(_mr_live_at_item_start, _mr_expected_live))
-        seed = getattr(request.config.option, "randomly_seed", None)
-        seed_msg = f" (pytest-randomly seed: {seed})" if seed is not None else ""
-        warnings.warn(
-            f"Owned resources were left alive by tests that ran since {_mr_last_init_cuda_nodeid} "
-            f"({_format_counts(excess)}; still alive after a collect: {_format_counts(held)}). "
-            f"A test or fixture that constructs a memory resource must request init_cuda.{seed_msg}",
-            UnreleasedMemoryResourceWarning,
-            stacklevel=2,
-        )
-
-    # Counts at setup. The device default pool is a borrowed reference and is
-    # not counted, so no cached-default bookkeeping is needed.
-    live_at_setup = _live_resource_counts()
+    live_at_setup = None
     try:
+        # Set option to avoid spin-waiting on synchronization.
+        if int(os.environ.get("CUDA_CORE_TEST_BLOCKING_SYNC", 0)) != 0:
+            handle_return(
+                driver.cuDevicePrimaryCtxSetFlags(device.device_id, driver.CUctx_flags.CU_CTX_SCHED_BLOCKING_SYNC)
+            )
+
+        # Leftover check: growth since the previous init_cuda test means something
+        # in between left an owned resource alive. Collect first to sweep a cyclic
+        # leftover, and before the counts at setup are recorded so they are reflected there.
+        if _mr_expected_live is not None and any(
+            now > exp for now, exp in zip(_mr_live_at_item_start, _mr_expected_live)
+        ):
+            gc.collect()
+            now_counts = _live_resource_counts()
+            held = tuple(max(n - e, 0) for n, e in zip(now_counts, _mr_expected_live))
+            excess = tuple(s - e for s, e in zip(_mr_live_at_item_start, _mr_expected_live))
+            seed = getattr(request.config.option, "randomly_seed", None)
+            seed_msg = f" (pytest-randomly seed: {seed})" if seed is not None else ""
+            warnings.warn(
+                f"Owned resources were left alive by tests that ran since {_mr_last_init_cuda_nodeid} "
+                f"({_format_counts(excess)}; still alive after a collect: {_format_counts(held)}). "
+                f"A test or fixture that constructs a memory resource must request init_cuda.{seed_msg}",
+                UnreleasedMemoryResourceWarning,
+                stacklevel=2,
+            )
+
+        # Counts at setup. The device default pool is a borrowed reference and is
+        # not counted, so no cached-default bookkeeping is needed.
+        live_at_setup = _live_resource_counts()
         yield device
     finally:
         # pytest holds fixture values in item.funcargs (a pytest-internal dict)
@@ -271,21 +277,24 @@ def init_cuda(request):
         # Only an owned resource in a reference cycle outlives the test frame, so
         # collect only if the counts differ from their values at setup. A lower
         # count means an older resource was released, which could mask a leak
-        # in this test: collect in that case too.
-        live_after = _live_resource_counts()
-        if live_after != live_at_setup:
+        # in this test: collect in that case too. Skip the comparison when setup
+        # aborted before the snapshot (the leftover warning raises).
+        live_after = _live_resource_counts() if live_at_setup is not None else None
+        if live_after is not None and live_after != live_at_setup:
             gc.collect()
         # Force any pool/allocation whose only remaining reference was a local
         # in this test's frame to actually get destroyed now, then drain the
         # context so the stream-ordered frees that destruction enqueues retire
         # before the next test runs (issue #2381). gc.collect() above must run
         # first: cuCtxSynchronize alone cannot drain frees that were never
         # enqueued because their owning object had not been collected yet.
-        # With pytest-run-parallel this runs after worker join.
+        # With pytest-run-parallel this runs after worker join. This also runs
+        # when setup raises, so an escalated leftover warning cannot leave the
+        # context current.
         driver.cuCtxSynchronize()
         _ = _device_unset_current()
         # After the cleanup, so an escalated warning cannot skip it.
-        if any(now < snap for now, snap in zip(live_after, live_at_setup)):
+        if live_after is not None and any(now < snap for now, snap in zip(live_after, live_at_setup)):
             released = tuple(s - n for s, n in zip(live_at_setup, live_after))
             warnings.warn(
                 f"{request.node.nodeid} released owned resources that existed before it started "

cuda_core/tests/memory/test_init_cuda_collect_gate.py

@@ -24,13 +24,17 @@
 
 import gc
 import inspect
+import warnings
 
 import pytest
+from helpers.contexts import current_context_handle
 from helpers.memory import (
     BaselineMemoryResourceReleasedWarning,
     UnreleasedMemoryResourceWarning,
 )
 
+from cuda.core._utils.cuda_utils import driver, handle_return
+
 
 class _ScriptedCounts:
     """Callable stand-in for ``conftest._live_resource_counts``.
@@ -184,6 +188,52 @@ def test_setup_warns_when_resources_left_alive_since_last_init_cuda(request, mon
     _finish_teardown(gen)  # live_after == live_at_setup -> no baseline warning
 
 
+def _pop_context_stack():
+    """Pop every current context. Returns the handles, top last, for restoration."""
+    saved = []
+    while current_context_handle() != 0:
+        saved.append(handle_return(driver.cuCtxPopCurrent()))
+    return saved
+
+
+def _push_context_stack(saved):
+    for ctx in reversed(saved):
+        handle_return(driver.cuCtxPushCurrent(ctx))
+
+
+@_THREAD_UNSAFE_MARKER
+@pytest.mark.agent_authored(model="grok-4.7")
+def test_escalated_setup_warning_still_pops_the_context(request, monkeypatch):
+    """An error-filtered leftover warning still runs teardown and clears the context.
+
+    pytest.ini escalates ``UnreleasedMemoryResourceWarning`` to an error. The
+    warning is raised after ``set_current()``, so the fixture finally has to pop
+    that context before the error propagates. ``pytest.warns`` hides this: it
+    forces the warning filter to ``always`` and the generator yields. The test
+    sets the error filter itself so it does not depend on ``pytest.ini`` or
+    ``-W`` / ``-p no:warnings`` on the command line.
+    """
+    func = _init_cuda_func(request)
+    _patch_conftest_global(monkeypatch, func, "_mr_expected_live", (0, 0))
+    _patch_conftest_global(monkeypatch, func, "_mr_live_at_item_start", (1, 0))
+    _patch_conftest_global(monkeypatch, func, "_mr_last_init_cuda_nodeid", "tests/prev::test_prev")
+    _patch_conftest_global(monkeypatch, func, "_live_resource_counts", _ScriptedCounts((1, 0)))
+
+    # set_current replaces the top of a non-empty stack, so start from empty.
+    # Otherwise a successful pop does not restore the handle this test observed.
+    saved = _pop_context_stack()
+    gen = func(_FakeRequest())
+    try:
+        with warnings.catch_warnings():
+            warnings.simplefilter("error", UnreleasedMemoryResourceWarning)
+            with pytest.raises(UnreleasedMemoryResourceWarning):
+                next(gen)
+        assert current_context_handle() == 0
+    finally:
+        gen.close()
+        _push_context_stack(saved)
+
+
 @_THREAD_UNSAFE_MARKER
 @pytest.mark.agent_authored(model="claude-opus-4.8")
 def test_teardown_warns_on_baseline_release(request, monkeypatch):