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

[fallintoplace]

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

[leofang]

Thanks for the quick turnaround — I've verified the fix locally. cc @benhg for vis

One pre-existing gap surfaced while testing (reproduced on main, not introduced here): when a slow-path grow fails after the old range was remapped, _remap_old restores the mapping but not the access grants, so the rolled-back buffer faults on next access. I'll file a separate issue; no action needed in this PR.

-- Leo's bot

[leofang]

Non-blocking nit before merge: on_failure and on_exit now differ by a single line, so please extract the shared logic:

def _register(self, callback: Callable[[], Any], on_commit: bool) -> None:
    if not self._entered:
        raise RuntimeError("Transaction must be entered before registering callbacks")
    # The ExitStack copy runs on rollback (LIFO, interleaved with the failure
    # callbacks); the _on_exit copy runs at commit(). commit() disarms the stack
    # before running _on_exit, so exactly one of the two ever fires.
    self._stack.callback(callback)
    if on_commit:
        self._on_exit.append(callback)

def on_failure(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> None:
    """Register a failure callback (runs if the with-block exits without commit())."""
    self._register(partial(fn, *args, **kwargs), on_commit=False)

def on_exit(self, fn: Callable[..., Any], /, *args: Any, **kwargs: Any) -> None:
    """Register an exit callback (runs exactly once, on rollback or during commit())."""
    self._register(partial(fn, *args, **kwargs), on_commit=True)

Besides the dedup, the comment at the dual-registration site is the important part — a reviewer already tripped on why on_exit writes to both registries.

[fallintoplace]

Thanks, I can fix this in a few hours. @leofang

[leofang]

Note to myself: append() is gone but this is not public API.