Fix VMM handle leaks in virtual memory paths

[fallintoplace]

Fixes #2344

Summary

VirtualMemoryResource kept the creation reference returned by cuMemCreate after a successful mapping, so repeated allocation and growth leaked physical memory. This drops that reference on both transaction outcomes.

Changes

• Rename the rollback hook to on_failure and add on_exit for exactly-once obligations. on_exit callbacks run LIFO during rollback or FIFO during commit.
• Use on_exit for every created or retained VMM handle in allocate and both grow paths. VA-free, unmap, and remap callbacks remain failure-only.
• Register allocation rollback only after cuMemMap succeeds, preserving the original map error when mapping fails.
• Simplify deallocate to unmap the allocation and free its VA reservation; the mapping now owns the physical allocation.
• Replace the added driver-mock tests with one focused cuMemGetInfo regression covering allocation and growth.

Validation

• Relevant pre-commit checks pass locally, including formatting, generated stubs, type checking, and Cython lint.
• Transaction callbacks were checked directly for FIFO commit and interleaved LIFO rollback ordering.
• pre-commit.ci passes on the pushed branch.

[copy-pr-bot]

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.

[leofang]

@fallintoplace what issue are we fixing here? Would you kindly file an issue first with a reproducer, per our contribution guideline? Without any context it is hard to judge what's going on. For example, the transaction mechanism is meant to take care of the handle cleanup. This seems like opening a backdoor to me?

[leofang]

Thanks for filing #2344 with a reproducer — the leak is real. I confirmed on an RTX A6000 against main: 16 × allocate(64 MiB) + close() leaks exactly 64 MiB per iteration, the grow path leaks 128 MiB per grow/close cycle, and with this PR's source change applied both loops leak 0 MiB. The analysis is also right: the transaction only guards the failure path; on success nothing drops the cuMemCreate reference, and deallocate()'s retain+release nets to zero.

So the fix direction is correct, but two things need to be reworked before this can merge:

1. Extend the transaction abstraction instead of working around it. The repeated nonlocal-flag release helpers are hard to reason about and duplicate state the transaction should own. Please add a success-side hook to Transaction and use it at all three sites — see the inline comment in allocate() for the concrete shape. Transaction (cuda_core/cuda/core/_utils/cuda_utils.pyx) is only used by this file, so the change is fully contained in this PR.

2. Replace the mocked tests with a single real-GPU regression test. We don't accept driver-monkeypatch tests in cuda.core — details in the inline comments, including one added test that asserts an impossible rollback order and fails on real hardware. One compact cuMemGetInfo-based leak test is the right size for this fix; sketch inline. Please run the suite locally on a VMM-capable GPU and confirm before pushing updates.

Optional while you're here: once "the mapping owns the physical allocation" holds, deallocate() can shrink to just cuMemUnmap + cuMemAddressFree — its retain/release pair only exists to compensate for the reference this PR stops leaking.

-- Leo's bot

[leofang]

This conditional-release pattern (repeated three times in this file) fights the transaction abstraction. What the fix actually needs is a success-side counterpart to the undo list, scope-guard style. Since Transaction is only used by this file, please extend it as part of this PR:

• rename append(fn) → on_failure(fn): runs LIFO only if the block exits without commit() (behavior unchanged)
• add on_success(fn): runs FIFO inside commit(), right after the undos are disarmed

def commit(self):
    self._stack.pop_all()
    for fn in self._on_success:   # registered via on_success(); FIFO
        fn()
    self._on_success.clear()

(plus the trivial __init__/on_success plumbing and docstring updates, mirroring on_failure). Then every handle site becomes two symmetric, declarative lines — no flags, no nested defs — and a double release is impossible by construction, because commit and rollback are mutually exclusive:

res, handle = driver.cuMemCreate(aligned_size, prop, 0)
raise_if_driver_error(res)
trans.on_failure(lambda h=handle: raise_if_driver_error(driver.cuMemRelease(h)[0]))
# Once cuMemMap succeeds, the mapping keeps the pages alive (cuMemRelease documents
# that the handle "can be freed when there are still outstanding mappings"), so on
# success we drop the now-redundant creation reference: the mapping owns the memory.
trans.on_success(lambda h=handle: raise_if_driver_error(driver.cuMemRelease(h)[0]))

This also fixes a semantic wrinkle in the current version: if the explicit pre-commit release fails, this code tears down a fully-constructed allocation and re-runs the release that just failed as part of the rollback. A commit-time release failure (not a realistic event for a valid handle) should propagate without destroying a working allocation.

[leofang]

Addendum, after I thought about it further: every handle site would pass the identical callable to both hooks (a reference must be dropped exactly once, on whichever path), while the VA-free/unmap/remap undos remain failure-only. So instead of a public on_success, please add on_exit(fn): register the callback both as a LIFO undo and in an internal commit-time list; commit() disarms the undos first and then runs that list, so the callback fires exactly once on either path. Handle sites then become a single call:

trans.on_exit(lambda h=handle: raise_if_driver_error(driver.cuMemRelease(h)[0]))

Note this distinction is exactly where the original leak came from: the release was registered as a failure-only undo, so commit() cancelled it on success and the creation reference was never dropped. With on_exit, commit() runs it instead of cancelling it. Net API: on_failure (failure-only undos) + on_exit (exactly-once obligations) + commit().

-- Leo's bot

[leofang]

Same as in allocate(): with on_failure/on_success this whole block is two lines, and the explicit _release_new_handle() call at L352 disappears.

[leofang]

Both handles in the slow path follow the same two-line pattern (this one, the created chunk at L443–453, and the explicit calls at L472–474 all go away). Note that on_success gives the retained old_handle exactly the timing it needs, for free: it must stay referenced until every step has succeeded, because the _remap_old undo needs a live handle to remap the old range on rollback (the LIFO undo order already guarantees _remap_old runs before the release) — and "inside commit()" is precisely "after everything succeeded".

[leofang]

Please drop all of the mocked tests in this PR — and don't extend the pattern of the existing ones from #2130 (we intend to revisit those for the same reasons):

1. They don't verify the fix. Every test here re-asserts the call sequence the implementation makes, so they pass with or without the leak; only a real-memory check fails on the broken code.
2. One of them asserts an impossible rollback order and fails on a real GPU (see the comment at L1200).
3. We plan to lower this module to Cython, at which point monkeypatching driver.* at the Python level stops working and all of these tests break.

The suite is already large and VMM is a niche path, so one focused real-GPU test is the right size:

def test_vmm_allocate_close_does_not_leak(init_cuda):
    device = Device()
    if not device.properties.virtual_memory_management_supported:
        pytest.skip("Virtual memory management is not supported on this device")
    mr = VirtualMemoryResource(
        device, config=VirtualMemoryResourceOptions(handle_type="win32_kmt" if IS_WINDOWS else "posix_fd")
    )
    buf = mr.allocate(8 * 1024 * 1024)  # warm-up; also learn the aligned size
    aligned_size = buf.size
    buf.close()
    baseline = handle_return(driver.cuMemGetInfo())[0]
    for _ in range(8):
        mr.allocate(8 * 1024 * 1024).close()
    free = handle_return(driver.cuMemGetInfo())[0]
    # on current main this leaks aligned_size per iteration; with the fix it's ~0
    assert baseline - free < aligned_size

(IS_WINDOWS and handle_return are already imported in this file.) If you want to cover the grow paths too, add one variant that calls mr.modify_allocation(buf, 2 * size) before close() in the loop — same assertion. I've verified both versions fail on current main and pass with your source change, on an RTX A6000.

[leofang]

This assertion can never pass: Transaction is an ExitStack, so undo actions run in LIFO order. In allocate() the undos are registered as release → addr_free → unmap, so a rollback runs unmap → addr_free → release. Running this exact scenario against your branch on an RTX A6000 gives:

['granularity', 'create', 'addr_reserve', 'map', 'set_access', 'release', 'unmap', 'addr_free', 'release']

i.e. this test fails on any VMM-capable machine, which tells me the suite wasn't executed on real hardware. Please always run the tests locally on a GPU before pushing — community PRs don't get CI until a maintainer vets each push, so failures like this otherwise surface very late.

[fallintoplace]

Thank you, let me have a look and push the commit.