Skip to content

Commit 90eb8b0

Browse files
authored
test(cuda.core): pop the context when the leftover warning aborts init_cuda setup (#3082)
UnreleasedMemoryResourceWarning is escalated to an error and was raised before the fixture try, so teardown never popped the context. Raise it inside the try, and correct the VA-count comment that attributed the #2381 pool leaks to cuMemAddressReserve.
1 parent a864219 commit 90eb8b0

3 files changed

Lines changed: 98 additions & 36 deletions

File tree

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

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -39,9 +39,12 @@ struct VaReservationBox {
3939
size_t size; // exact reserved size; cuMemAddressFree needs both
4040
};
4141

42-
// Process-wide count of live VA reservations (cuMemAddressReserve), the leak
43-
// quantity issue #2381 is about; read by the test suite to gate the per-test
44-
// gc.collect(). Atomic for free-threaded builds.
42+
// Process-wide count of live VA reservations created here with
43+
// cuMemAddressReserve. Memory-pool VA windows (issue #2381) are not included:
44+
// owned pools are tracked by the owned-pool count, and the device default pool
45+
// (cuDeviceGetMemPool) is a borrowed reference that neither count tracks.
46+
// Read by the test suite to gate the per-test gc.collect(). Atomic for
47+
// free-threaded builds.
4548
std::atomic<long long> g_live_va_reservations{0};
4649

4750
struct VaMappingBox {

‎cuda_core/tests/conftest.py‎

Lines changed: 42 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -231,36 +231,42 @@ def session_setup():
231231
def init_cuda(request):
232232
# TODO: rename this to e.g. init_context
233233
device = Device(0)
234+
# Everything that can raise after set_current() runs inside the try, so the
235+
# finally pops this context; the leftover warning below is escalated to an
236+
# error. A failed set_current() leaves nothing for the finally to undo.
237+
# live_at_setup stays None until the setup snapshot is recorded.
234238
device.set_current()
235-
236-
# Set option to avoid spin-waiting on synchronization.
237-
if int(os.environ.get("CUDA_CORE_TEST_BLOCKING_SYNC", 0)) != 0:
238-
handle_return(
239-
driver.cuDevicePrimaryCtxSetFlags(device.device_id, driver.CUctx_flags.CU_CTX_SCHED_BLOCKING_SYNC)
240-
)
241-
242-
# Leftover check: growth since the previous init_cuda test means something
243-
# in between left an owned resource alive. Collect first to sweep a cyclic
244-
# leftover, and before the counts at setup are recorded so they are reflected there.
245-
if _mr_expected_live is not None and any(now > exp for now, exp in zip(_mr_live_at_item_start, _mr_expected_live)):
246-
gc.collect()
247-
now_counts = _live_resource_counts()
248-
held = tuple(max(n - e, 0) for n, e in zip(now_counts, _mr_expected_live))
249-
excess = tuple(s - e for s, e in zip(_mr_live_at_item_start, _mr_expected_live))
250-
seed = getattr(request.config.option, "randomly_seed", None)
251-
seed_msg = f" (pytest-randomly seed: {seed})" if seed is not None else ""
252-
warnings.warn(
253-
f"Owned resources were left alive by tests that ran since {_mr_last_init_cuda_nodeid} "
254-
f"({_format_counts(excess)}; still alive after a collect: {_format_counts(held)}). "
255-
f"A test or fixture that constructs a memory resource must request init_cuda.{seed_msg}",
256-
UnreleasedMemoryResourceWarning,
257-
stacklevel=2,
258-
)
259-
260-
# Counts at setup. The device default pool is a borrowed reference and is
261-
# not counted, so no cached-default bookkeeping is needed.
262-
live_at_setup = _live_resource_counts()
239+
live_at_setup = None
263240
try:
241+
# Set option to avoid spin-waiting on synchronization.
242+
if int(os.environ.get("CUDA_CORE_TEST_BLOCKING_SYNC", 0)) != 0:
243+
handle_return(
244+
driver.cuDevicePrimaryCtxSetFlags(device.device_id, driver.CUctx_flags.CU_CTX_SCHED_BLOCKING_SYNC)
245+
)
246+
247+
# Leftover check: growth since the previous init_cuda test means something
248+
# in between left an owned resource alive. Collect first to sweep a cyclic
249+
# leftover, and before the counts at setup are recorded so they are reflected there.
250+
if _mr_expected_live is not None and any(
251+
now > exp for now, exp in zip(_mr_live_at_item_start, _mr_expected_live)
252+
):
253+
gc.collect()
254+
now_counts = _live_resource_counts()
255+
held = tuple(max(n - e, 0) for n, e in zip(now_counts, _mr_expected_live))
256+
excess = tuple(s - e for s, e in zip(_mr_live_at_item_start, _mr_expected_live))
257+
seed = getattr(request.config.option, "randomly_seed", None)
258+
seed_msg = f" (pytest-randomly seed: {seed})" if seed is not None else ""
259+
warnings.warn(
260+
f"Owned resources were left alive by tests that ran since {_mr_last_init_cuda_nodeid} "
261+
f"({_format_counts(excess)}; still alive after a collect: {_format_counts(held)}). "
262+
f"A test or fixture that constructs a memory resource must request init_cuda.{seed_msg}",
263+
UnreleasedMemoryResourceWarning,
264+
stacklevel=2,
265+
)
266+
267+
# Counts at setup. The device default pool is a borrowed reference and is
268+
# not counted, so no cached-default bookkeeping is needed.
269+
live_at_setup = _live_resource_counts()
264270
yield device
265271
finally:
266272
# pytest holds fixture values in item.funcargs (a pytest-internal dict)
@@ -271,21 +277,24 @@ def init_cuda(request):
271277
# Only an owned resource in a reference cycle outlives the test frame, so
272278
# collect only if the counts differ from their values at setup. A lower
273279
# count means an older resource was released, which could mask a leak
274-
# in this test: collect in that case too.
275-
live_after = _live_resource_counts()
276-
if live_after != live_at_setup:
280+
# in this test: collect in that case too. Skip the comparison when setup
281+
# aborted before the snapshot (the leftover warning raises).
282+
live_after = _live_resource_counts() if live_at_setup is not None else None
283+
if live_after is not None and live_after != live_at_setup:
277284
gc.collect()
278285
# Force any pool/allocation whose only remaining reference was a local
279286
# in this test's frame to actually get destroyed now, then drain the
280287
# context so the stream-ordered frees that destruction enqueues retire
281288
# before the next test runs (issue #2381). gc.collect() above must run
282289
# first: cuCtxSynchronize alone cannot drain frees that were never
283290
# enqueued because their owning object had not been collected yet.
284-
# With pytest-run-parallel this runs after worker join.
291+
# With pytest-run-parallel this runs after worker join. This also runs
292+
# when setup raises, so an escalated leftover warning cannot leave the
293+
# context current.
285294
driver.cuCtxSynchronize()
286295
_ = _device_unset_current()
287296
# After the cleanup, so an escalated warning cannot skip it.
288-
if any(now < snap for now, snap in zip(live_after, live_at_setup)):
297+
if live_after is not None and any(now < snap for now, snap in zip(live_after, live_at_setup)):
289298
released = tuple(s - n for s, n in zip(live_at_setup, live_after))
290299
warnings.warn(
291300
f"{request.node.nodeid} released owned resources that existed before it started "

‎cuda_core/tests/memory/test_init_cuda_collect_gate.py‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,13 +24,17 @@
2424

2525
import gc
2626
import inspect
27+
import warnings
2728

2829
import pytest
30+
from helpers.contexts import current_context_handle
2931
from helpers.memory import (
3032
BaselineMemoryResourceReleasedWarning,
3133
UnreleasedMemoryResourceWarning,
3234
)
3335

36+
from cuda.core._utils.cuda_utils import driver, handle_return
37+
3438

3539
class _ScriptedCounts:
3640
"""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
184188
_finish_teardown(gen) # live_after == live_at_setup -> no baseline warning
185189

186190

191+
def _pop_context_stack():
192+
"""Pop every current context. Returns the handles, top last, for restoration."""
193+
saved = []
194+
while current_context_handle() != 0:
195+
saved.append(handle_return(driver.cuCtxPopCurrent()))
196+
return saved
197+
198+
199+
def _push_context_stack(saved):
200+
for ctx in reversed(saved):
201+
handle_return(driver.cuCtxPushCurrent(ctx))
202+
203+
204+
@_THREAD_UNSAFE_MARKER
205+
@pytest.mark.agent_authored(model="grok-4.7")
206+
def test_escalated_setup_warning_still_pops_the_context(request, monkeypatch):
207+
"""An error-filtered leftover warning still runs teardown and clears the context.
208+
209+
pytest.ini escalates ``UnreleasedMemoryResourceWarning`` to an error. The
210+
warning is raised after ``set_current()``, so the fixture finally has to pop
211+
that context before the error propagates. ``pytest.warns`` hides this: it
212+
forces the warning filter to ``always`` and the generator yields. The test
213+
sets the error filter itself so it does not depend on ``pytest.ini`` or
214+
``-W`` / ``-p no:warnings`` on the command line.
215+
"""
216+
func = _init_cuda_func(request)
217+
_patch_conftest_global(monkeypatch, func, "_mr_expected_live", (0, 0))
218+
_patch_conftest_global(monkeypatch, func, "_mr_live_at_item_start", (1, 0))
219+
_patch_conftest_global(monkeypatch, func, "_mr_last_init_cuda_nodeid", "tests/prev::test_prev")
220+
_patch_conftest_global(monkeypatch, func, "_live_resource_counts", _ScriptedCounts((1, 0)))
221+
222+
# set_current replaces the top of a non-empty stack, so start from empty.
223+
# Otherwise a successful pop does not restore the handle this test observed.
224+
saved = _pop_context_stack()
225+
gen = func(_FakeRequest())
226+
try:
227+
with warnings.catch_warnings():
228+
warnings.simplefilter("error", UnreleasedMemoryResourceWarning)
229+
with pytest.raises(UnreleasedMemoryResourceWarning):
230+
next(gen)
231+
assert current_context_handle() == 0
232+
finally:
233+
gen.close()
234+
_push_context_stack(saved)
235+
236+
187237
@_THREAD_UNSAFE_MARKER
188238
@pytest.mark.agent_authored(model="claude-opus-4.8")
189239
def test_teardown_warns_on_baseline_release(request, monkeypatch):

0 commit comments

Comments
 (0)