test(cuda.core): gate the init_cuda teardown gc.collect() on live owned-resource counts by juenglin · Pull Request #3053 · NVIDIA/cuda-python

juenglin · 2026-10-08T20:33:09Z

Description

The init_cuda teardown added by #2484 runs gc.collect() after every test to return memory-pool VA reservations before the next test (issue #2381). With gc.freeze() already in place, a measured full-suite run still spends ~18.0 s of ~86.1 s (20.9%) in that unconditional collect, and the cost tracks session age rather than what a test allocates. The collect is only needed when a test leaves an owned pool or VA reservation alive in a reference cycle; a resource whose last reference is a local dies by refcount at function return and needs no collect.

This PR replaces the always-on collect with a runtime gate driven by process-wide counts of live owned driver resources, plus the test-hygiene changes it requires. Logical pieces:

test(cuda.core): request init_cuda for memory resource tests — route tests that construct or allocate from a memory resource without init_cuda onto it, so they pop the context and clear the Device's cached default resource. Documents the rule in tests/AGENTS.md.
test(cuda.core): release baseline-retained memory resources at test boundaries — fix three tests (ChildErrorHarness, the escalated-cudawarning test, the fill_env fixture) that kept a closed resource alive past teardown.
test(cuda.core): gate the init_cuda teardown gc.collect() on live count — init_cuda snapshots a live count at setup and collects at teardown only when it changed, after dropping the item's funcargs. Adds UnreleasedMemoryResourceWarning and BaselineMemoryResourceReleasedWarning (the latter escalated to an error in pytest.ini), counter unit tests, and tests/AGENTS.md updates.
test(cuda.core): count live owned pools and VA reservations in C++ — per review, move the counting out of a Python MemoryResource object counter into the C++ RAII layer, split into two process-wide atomic counts: owned memory pools (cuMemPoolCreate/imported) and VA reservations (cuMemAddressReserve), exposed through _live_owned_mempool_count/_live_va_reservation_count. init_cuda snapshots a (owned_pools, va_reservations) tuple; borrowed pool references (the device default) and MemoryResource subclasses that hold no pool or reservation are no longer counted, so the per-Device cached-default bookkeeping is removed.

Approach

A runtime count is preferred over a per-test owns_pool marker plus a lint rule (as in #3045): it needs no annotation to stay in sync, covers resources constructed indirectly through fixtures/helpers, and cannot silently skip a needed collect from a forgotten marker. Counting owned pools and VA reservations at their C++ driver create/destroy sites — rather than using the Python MemoryResource object as a proxy — measures the resources #2381 is about, ignores borrowed pool references, and still counts an owned pool or reservation that outlives its Python resource (for example through a graph in a reference cycle holding allocation handles), which the object counter missed. The cost is test-only instrumentation: two atomic counters and two underscore-prefixed accessors, no per-object field and no layout change.

Measured results (RTX 6000 Ada, Linux, pinned `pytest-randomly` seed 42)

Baseline: 18.04 s of an 86.13 s suite in the unconditional collect, 3849 passed / 0 failed.
Gate projection (collect still on, instrumented): skips 1955/1974 teardowns (99.0%), saves 18.64 s (~21% of suite wall), with parity.

These numbers were taken with the earlier per-object counter (piece 3); the C++ counts (piece 4) track strictly fewer resources, so the gate skips a superset of those teardowns and a full-suite re-measurement is a follow-up. Correctness under the C++ counter was re-verified on test_memory.py, test_error_handling.py, and memory_ipc/ with both warnings escalated to errors, with no misfires.

Soundness and scope

The snapshot gate is sound as long as no older owned resource is released during a test (which could cancel a new leak); piece 2 removes the known retention sources, the gate collects on any count change (not just growth), and it warns when a count falls below the snapshot. Still not covered: cyclic non-resource garbage that holds device memory without its own owned pool or VA reservation (for example a cyclic Buffer on the borrowed device-default pool), and resources from module/class-scoped fixtures. Verification here was single-GPU Linux; the Windows MCDM/WDDM CI rows remain the real OOM check.

Checklist

New or existing tests cover these changes.
The documentation is up to date with these changes.

Tests that allocate through Device.memory_resource / Device.allocate or construct any memory resource without init_cuda did Device() + set_current() by hand, leaving a stale cached default MemoryResource on the thread-local Device: no teardown gc.collect(), no cuCtxSynchronize, and the context never popped. The stale slot is the aliasing route for the planned live-resource-count gc.collect() gate, and the tests violate the tests/AGENTS.md rule that constructing a memory resource counts as touching CUDA. Route both groups onto init_cuda (using the yielded device) so the fixture teardown clears the cached slot via del _tls.devices: - device.memory_resource allocators: 6 test_from_buffer* functions in test_utils.py and test_buffer_external_device in test_memory.py. - resource constructors: 28 functions in test_memory.py (buffers, dlpack, vmm, pinned/managed ipc, legacy pinned, abstract stubs) and 4 launch-scalar tests in test_launcher.py. test_buffer_external_device and test_buffer_external_pinned_alloc still select Device(n - 1) as before; init_cuda only supplies the context and teardown. The no-context tests (test_close_with_default_stream_requires_context, test_from_handle_mr_default_stream_requires_context) keep their no_current_context() body. Assertions are unchanged. Add a sentence to tests/AGENTS.md that allocating through device.memory_resource or Device.allocate, or constructing any memory resource, counts as touching CUDA and requires init_cuda.

…oundaries The snapshot gc.collect() gate is sound only if no older MemoryResource is released during a test, otherwise a mid-test baseline release could mask a leak. Three tests retained a closed (but not destroyed) memory resource past their teardown, which would ratchet the baseline and later cancel a new leak: - tests/memory_ipc/test_errors.py: ChildErrorHarness keeps the fixture's buffer/mr/device references on the pytest-retained item._instance for the session; drop them from __dict__ so the closed MR objects do not persist. - tests/test_error_handling.py::test_escalated_cudawarning_from_cleanup_is_not_a_crash: the unraisable hook args hold exc_traceback whose frame holds ``mr`` (never closed); clear the list, drop the locals, and close the pool. - tests/test_memory.py::fill_env: the fixture returns the MR to the test via a tuple; convert it to a generator that dels ``mr`` in its teardown so the closed MR does not persist in the pytest fixture cache. These are the baseline-retention fixes prerequisite for the planned live-resource-count gc.collect() gate (see plan mr-instance-count-gc-gate, "Soundness conditions of the snapshot gate").

…oryResource count Replace the always-on gc.collect() in the init_cuda teardown with a runtime gate. MemoryResource keeps a process-wide atomic count of live instances (incremented in __cinit__, decremented in __dealloc__) exposed through a private accessor. init_cuda snapshots the count at setup and collects at teardown only when the count has grown. Also warn when a test leaves a MemoryResource alive past teardown (UnreleasedMemoryResourceWarning) or releases one that was alive at setup (BaselineMemoryResourceReleasedWarning), add tests for the counter, and update tests/AGENTS.md.

copy-pr-bot · 2026-10-08T20:33:14Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

juenglin · 2026-10-08T21:35:00Z

/ok to test c4f065b

juenglin · 2026-10-08T21:41:08Z

+        # in this test: collect in that case too.
+        _release_device_default_memory_resources()
+        live_after_release = _live_memory_resource_count()
+        if live_after_release != live_at_setup:


Conditional gc.collect. This where the cost savings come from.

github-actions · 2026-10-08T21:54:16Z

Doc Preview CI
Preview removed because the pull request was closed or merged.

juenglin · 2026-10-08T21:56:36Z

/ok to test fc62237

juenglin · 2026-10-08T22:10:10Z

/ok to test b19d0ae

test_latchkernel constructed 3 DummyUnifiedMemoryResource instances (via make_scratch_buffer) plus a LegacyPinnedMemoryResource (via LatchKernel) without init_cuda. On Windows and aarch64 the ctypes-pointer/Buffer/MR chains form cycles that survive refcounting, so the live-MemoryResource count stays elevated past the test and the next init_cuda test (test_oom_diagnostics_probe_basics_is_live_and_cheap) is blamed by the leftover check (UnreleasedMemoryResourceWarning, escalated to an error). Migrating to init_cuda lets the teardown gate collect the cyclic leftovers, matching the pattern from 68f832c for the other MR-constructing tests.

juenglin · 2026-10-08T22:39:54Z

/ok to test b102d26

coderabbitai · 2026-10-08T23:53:00Z

📝 Summary

Summary by CodeRabbit

Tests
- Improved detection of memory resources left unreleased between CUDA tests. Test runs now warn when resources remain alive or when tests release resources that were already present.
- Reduced unnecessary garbage collection during CUDA test cleanup by collecting only when tracked resource counts change.
- Updated CUDA tests to share initialized devices and release memory resources more consistently during cleanup.

Walkthrough

The PR adds process-wide counters for live owned memory pools and VA reservations. The init_cuda fixture uses these counts to check resource lifetimes and control garbage collection. CUDA tests adopt the fixture’s device and explicitly release selected resources.

Changes

Memory Resource Lifecycle Tracking

Layer / File(s)	Summary
Live resource counters `cuda_core/cuda/core/_cpp/rt/api.hpp`, `cuda_core/cuda/core/_cpp/rt/memory.cpp`, `cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp`, `cuda_core/cuda/core/_rt.pxd`, `cuda_core/cuda/core/_rt.pyx`, `cuda_core/cuda/core/_memory/_buffer.pyx`, `cuda_core/cuda/core/_memory/_buffer.pyi`, `cuda_core/tests/memory/test_resource_live_count.py`	The runtime counts live owned memory pools and VA reservations. Python wrappers expose the counts to test instrumentation. Tests cover owned and borrowed pools, cyclic references, VA allocation, and concurrent pool creation and release.
Fixture lifecycle checks `cuda_core/tests/conftest.py`, `cuda_core/tests/helpers/memory.py`, `cuda_core/pytest.ini`, `cuda_core/tests/AGENTS.md`, `cuda_core/tests/memory/test_init_cuda_collect_gate.py`	The pytest protocol and `init_cuda` fixture track resource counts. The fixture conditionally collects garbage and warns when counts indicate unreleased resources or release of pre-existing resources. Warning filters, test guidance, and fixture tests cover these checks.
CUDA test fixture adoption `cuda_core/tests/test_helpers.py`, `cuda_core/tests/test_launcher.py`, `cuda_core/tests/test_memory.py`, `cuda_core/tests/test_utils.py`, `cuda_core/tests/graph/test_graph_definition_lifetime.py`	CUDA-dependent tests use the `init_cuda` fixture’s device instead of creating and setting a device locally. VMM tests use a shared platform-specific handle type.
Explicit resource cleanup `cuda_core/tests/memory_ipc/test_errors.py`, `cuda_core/tests/test_error_handling.py`, `cuda_core/tests/test_green_context.py`, `cuda_core/tests/test_memory.py`	Tests clear retained references or explicitly close buffers and resources during cleanup. The `fill_env` fixture also deletes its local memory-resource reference at teardown.

Priority: ⬇️ Low

Change: Other

Merge Risk: 🟡 Moderate · up to bc94e

Rare allocation and driver-cleanup failures can leave GPU resources retained while the lifecycle gate reports clean state. These localized fixes should be applied before merge.

Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

Actionable comments posted: 1

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository: NVIDIA/cuda-python/.coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 758efc2b-aff0-4124-8160-a9d263def74b

📥 Commits

Reviewing files that changed from the base of the PR and between 106da2f and b102d26.

📒 Files selected for processing (15)

cuda_core/cuda/core/_memory/_buffer.pyi
cuda_core/cuda/core/_memory/_buffer.pyx
cuda_core/pytest.ini
cuda_core/tests/AGENTS.md
cuda_core/tests/conftest.py
cuda_core/tests/helpers/memory.py
cuda_core/tests/memory/test_init_cuda_collect_gate.py
cuda_core/tests/memory/test_resource_live_count.py
cuda_core/tests/memory_ipc/test_errors.py
cuda_core/tests/test_error_handling.py
cuda_core/tests/test_green_context.py
cuda_core/tests/test_helpers.py
cuda_core/tests/test_launcher.py
cuda_core/tests/test_memory.py
cuda_core/tests/test_utils.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

Move the init_cuda teardown gc.collect() gate's instrumentation out of a Python MemoryResource object counter and into the C++ RAII layer, split into two process-wide atomic counts: owned memory pools (cuMemPoolCreate or imported, not borrowed references) and VA reservations (cuMemAddressReserve). These are the resources issue NVIDIA#2381 is about; borrowed pool references (the device default pool) and MemoryResource subclasses that hold no pool or reservation are no longer counted. - Add std::atomic counts incremented at the owning RAII create sites (wrap_mempool_owned in memory.cpp, create_va_reservation_handle in virtual_memory.cpp) and decremented in their deleters, with live_owned_mempool_count / live_va_reservation_count accessors declared in api.hpp and bridged through _rt.pyx / _rt.pxd. - _buffer.pyx: drop _mr_live_count and MemoryResource.__cinit__/__dealloc__; expose _live_owned_mempool_count / _live_va_reservation_count. - conftest.py: snapshot a (owned_pools, va_reservations) tuple; collect when either component differs; warn per-component. Drop the device-default bookkeeping, since the default pool is a borrowed reference and is no longer counted. - Rewrite the counter tests for the two counters (owned pool, borrowed pool not counted, cycle, VA reservation, concurrency); the gate tests script _live_resource_counts for deterministic skip/collect/warning coverage.

juenglin · 2026-10-09T21:19:50Z

/ok to test 0651f45

- test_owned_pool_in_cycle_stays_counted_until_collect: create the cyclic resource and assert it is still counted inside the gc.disable() block, so an automatic collection cannot break the cycle before the "+1" assertion. - test_concurrent_owned_pool_construct_and_drop_returns_to_start: record worker-thread exceptions and fail on them; a raising worker only kills its own thread, leaving the count balanced and the test passing silently.

- helpers/memory.py: define VMM_HANDLE_TYPE once from IS_WINDOWS and use it in test_memory.py, graph/test_graph_definition_lifetime.py and memory/test_resource_live_count.py, replacing four local copies (one of which used platform.system()). - test_init_cuda_collect_gate.py: patch conftest globals through _patch_conftest_global, which asserts the name exists. monkeypatch.setitem silently adds a missing key, so a rename in conftest would have left the gate unpatched and a test passing without exercising it. - tests/AGENTS.md: state that the leftover and baseline-release checks track only owned memory pools and VA reservations, name both warnings escalated in pytest.ini, and note that default-pool wrappers, pool-less MemoryResource subclasses and their buffers are not counted.

…wording - pytest.ini: escalate BaselineMemoryResourceReleasedWarning to an error alongside UnreleasedMemoryResourceWarning, matching tests/AGENTS.md. - conftest.py: print labelled counts (owned_pools=..., va_reservations=...) in both warnings; replace the stale GraphMemoryResource example in the leftover-check comment. - helpers/memory.py: describe both warnings in terms of owned memory pools and VA reservations. - test_resource_live_count.py: release the VA test buffer with a context manager so a failed assertion does not leak the reservation. - test_init_cuda_collect_gate.py: assert on the warning class only.

juenglin · 2026-10-12T16:17:56Z

/ok to test bc94ed4

coderabbitai

Caution

Some comments are outside the diff and can’t be posted inline due to GitHub limitations.

⚠️ Outside diff range comments (3)

🟠 Major · important: Preserve the owned-pool count when destruction fails. · memory.cpp:82-87

cuda_core/cuda/core/_cpp/rt/memory.cpp:82-87
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

important: Preserve the owned-pool count when destruction fails.

pw_cuMemPoolDestroy returns the driver status, but this deleter discards it. If destruction fails, the pool can remain live. The unconditional decrement then makes init_cuda see unchanged counts, so it can skip gc.collect() and lose the signal used by the unreleased-resource check. Decrement only after successful destruction, or after CUDA_ERROR_DEINITIALIZED when the driver can no longer retain the pool.
Suggested fix
-            pw_cuMemPoolDestroy(b->resource);
-            g_live_owned_mempools.fetch_sub(1, std::memory_order_relaxed);
+            const auto status = pw_cuMemPoolDestroy(b->resource);
+            if (status == CUDA_SUCCESS || status == CUDA_ERROR_DEINITIALIZED) {
+                g_live_owned_mempools.fetch_sub(1, std::memory_order_relaxed);
+            }

🟡 Minor · important: Add rollback before allocating MemoryPoolBox. · memory.cpp:98-105

cuda_core/cuda/core/_cpp/rt/memory.cpp:98-105
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

important: Add rollback before allocating MemoryPoolBox.

cuMemPoolCreate and cuMemPoolImportFromShareableHandle return a live pool before wrap_mempool_owned allocates MemoryPoolBox. If that host allocation throws, the custom deleter is not installed and g_live_owned_mempools is not incremented. The pool remains undisposed and invisible to the teardown gate. Keep a rollback guard active until raw is allocated. Release it before constructing std::shared_ptr; the existing deleter handles control-block allocation failures.
Suggested fix
 static MemoryPoolHandle wrap_mempool_owned(CUmemoryPool pool, int owner_device) {
+    struct PoolRollback {
+        CUmemoryPool pool;
+        int owner_device;
+
+        ~PoolRollback() noexcept {
+            if (pool != nullptr) {
+                clear_mempool_peer_access(pool, owner_device);
+                pw_cuMemPoolDestroy(pool);
+            }
+        }
+    } rollback{pool, owner_device};
+
     auto* raw = new MemoryPoolBox{pool};
+    rollback.pool = nullptr;
     // Count before constructing the shared_ptr: if its control-block allocation
     // throws, it runs the deleter on raw, whose decrement then balances this.

🟡 Minor · important: Roll back the VA reservation when VaReservationBox… · virtual_memory.cpp:99-105

cuda_core/cuda/core/_cpp/rt/virtual_memory.cpp:99-105
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

important: Roll back the VA reservation when VaReservationBox allocation fails.

cuMemAddressReserve succeeds before new VaReservationBox{{ptr}, size} allocates host memory. If that allocation throws, h_res remains empty, so the caller's handle cleanup cannot run. The reservation remains allocated and g_live_va_reservations remains unchanged.
Suggested fix
-    auto* raw = new VaReservationBox{{ptr}, size};
+    VaReservationBox* raw;
+    try {
+        raw = new VaReservationBox{{ptr}, size};
+    } catch (...) {
+        pw_cuMemAddressFree(ptr, size);
+        throw;
+    }

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository: NVIDIA/cuda-python/.coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c43e98f0-baaa-47ff-8788-5a7ce887f21e

📥 Commits

Reviewing files that changed from the base of the PR and between 4b4e769 and bc94ed4.

📒 Files selected for processing (5)

cuda_core/pytest.ini
cuda_core/tests/conftest.py
cuda_core/tests/helpers/memory.py
cuda_core/tests/memory/test_init_cuda_collect_gate.py
cuda_core/tests/memory/test_resource_live_count.py

🚧 Files skipped from review as they are similar to previous changes (2)

cuda_core/tests/helpers/memory.py
cuda_core/tests/conftest.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 9 remain after this review.

Removed preview folders for the following PRs: - PR #3045 - PR #3053 - PR #3054 - PR #3055 - PR #3078 - PR #3079 - PR #3080

juenglin added 3 commits October 8, 2026 11:41

juenglin added this to the cuda.core next milestone Oct 8, 2026

juenglin added test Improvements or additions to tests cuda.core Everything related to the cuda.core module labels Oct 8, 2026

juenglin self-assigned this Oct 8, 2026

address review feedback

c4f065b

juenglin commented Oct 8, 2026

View reviewed changes

address TCC failure

f50ec59

juenglin force-pushed the lighter-init-cuda branch 2 times, most recently from fd9a0c7 to b19d0ae Compare October 8, 2026 22:08

juenglin force-pushed the lighter-init-cuda branch from b19d0ae to b102d26 Compare October 8, 2026 22:39

juenglin mentioned this pull request Oct 8, 2026

test(cuda.core): close leftover owned pools and run the teardown gc.collect() only for pool-owning tests #3045

Closed

2 tasks done

juenglin marked this pull request as ready for review October 8, 2026 23:44

coderabbitai Bot reviewed Oct 8, 2026

View reviewed changes

Comment thread cuda_core/tests/conftest.py Outdated

juenglin requested a review from Andy-Jost October 8, 2026 23:59

juenglin changed the title ~~test(cuda.core): gate the init_cuda teardown gc.collect() on live MemoryResource count~~ test(cuda.core): gate the init_cuda teardown gc.collect() on live owned-resource counts Oct 9, 2026

Andy-Jost approved these changes Oct 12, 2026

View reviewed changes

juenglin added 2 commits October 12, 2026 09:06

coderabbitai Bot reviewed Oct 12, 2026

View reviewed changes

juenglin merged commit 0271e3a into NVIDIA:main Oct 12, 2026
226 of 228 checks passed

This was referenced Oct 12, 2026

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

Open

ci: run cuda.core tests in heap-bounded batches, merge JUnit XML #3083

Closed

github-actions Bot pushed a commit that referenced this pull request Oct 13, 2026

Clean up PR preview folders for 7 closed/merged PRs

e8a08af

Removed preview folders for the following PRs: - PR #3045 - PR #3053 - PR #3054 - PR #3055 - PR #3078 - PR #3079 - PR #3080

Conversation

juenglin commented Oct 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Approach

Measured results (RTX 6000 Ada, Linux, pinned pytest-randomly seed 42)

Soundness and scope

Checklist

Uh oh!

copy-pr-bot Bot commented Oct 8, 2026

Uh oh!

juenglin commented Oct 8, 2026

Uh oh!

juenglin Oct 8, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Oct 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

juenglin commented Oct 8, 2026

Uh oh!

juenglin commented Oct 8, 2026

Uh oh!

juenglin commented Oct 8, 2026

Uh oh!

coderabbitai Bot commented Oct 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

juenglin commented Oct 9, 2026

Uh oh!

juenglin commented Oct 12, 2026

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

juenglin commented Oct 8, 2026 •

edited

Loading

Measured results (RTX 6000 Ada, Linux, pinned `pytest-randomly` seed 42)

github-actions Bot commented Oct 8, 2026 •

edited

Loading

coderabbitai Bot commented Oct 8, 2026 •

edited

Loading