Fix Buffer.from_handle(mr=...) not calling mr.deallocate()

[Andy-Jost]

Summary

• Fixes[BUG]: Buffer.from_handle does not respect custom MemoryResource deallocation #1619: Buffer.from_handle(ptr, size, mr=mr) now calls mr.deallocate(ptr, size, stream) when the buffer is closed or garbage collected, restoring the pre-RAII behavior.
• Adds deviceptr_create_with_mr() to the resource handle infrastructure, whose shared_ptr deleter invokes a registered callback with the pointer, size, and deallocation stream.
• The callback is implemented in Cython and converts the raw stream handle to a cuda.core.Stream before calling mr.deallocate().
• Removes the unused stream parameter from Buffer._init.

Closes#1619

Changes

• resource_handles.hpp/.cpp: Add MRDeallocCallback type, register_mr_dealloc_callback(), and deviceptr_create_with_mr().
• _resource_handles.pxd/.pyx: Expose the new types and functions to Cython.
• _buffer.pyx: Register the deallocation callback at module init; use deviceptr_create_with_mr when mr is provided in Buffer._init; update docstrings.
• _legacy.py: Remove unused stream argument from Buffer._init calls.
• tests/helpers/buffers.py: Add TrackingMR test helper.
• tests/test_memory.py: Add test_mr_deallocate_called_on_close and test_mr_deallocate_called_on_gc.

Test Plan

• New tests verify mr.deallocate() is called on explicit close() and on GC
• All existing memory and object protocol tests pass (338 passed, 15 skipped)
• CI

Made with Cursor

[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.

[Andy-Jost]

/ok to test c7fdab7f8547c4593f4d705554aa85fc428d5865

[cpcloud]

LGTM, nothing major.

[cpcloud]

Suggested change

	pass# Cannot propagate exceptions from a C++ destructor
	import sys
	print(f"Warning: mr.deallocate() failed during Buffer destruction: {exc}",
	file=sys.stderr)

Just so that something is shown. I think this would mirror the Python interpreter's behavior.

Not a requirement, but would really help debugging I suspect.

[cpcloud]

Seems like these branches should be merged into something like:

Suggested change

	if (gil.acquired() && mr_dealloc_cb) {
	mr_dealloc_cb(mr, b->resource, size, b->h_stream);
	}
	// Decref mr even if callback was skipped (prevent leak when
	// callback is not registered -- should not happen in practice).
	if (gil.acquired()) {
	Py_DECREF(mr);
	}
	if (gil.acquired()) {
	if (mr_dealloc_cb) {
	mr_dealloc_cb(mr, b->resource, size, b->h_stream);
	}
	// Decref mr regardless of whether the callback ran, as long as
	// we have the GIL. Leak is acceptable during interpreter finalization and shouldn't happen in practice.
	Py_DECREF(mr);
	}
	delete b;

[cpcloud]

The two new tests cover close() (no stream) and GC, but neither exercises the path where close(stream=s) forwards the stream through the C++ deleter to mr.deallocate(ptr, size, stream). A third test that verifies the stream argument arrives at mr.deallocate() would close that gap, e.g.:

deftest_mr_deallocate_receives_stream():
device=Device()
device.set_current()
stream=device.create_stream()
received= {}
classStreamCaptureMR(TrackingMR):
defdeallocate(self, ptr, size, stream=None):
received['stream'] =streamsuper().deallocate(ptr, size, stream)
mr=StreamCaptureMR()
buf=mr.allocate(1024)
buf.close(stream)
assertreceived['stream'] isnotNone

[Andy-Jost]

/ok to test 373cb5f5d193a61685b6fa06b44808749ba78c39

[cpcloud]

Thanks!

[Andy-Jost]

/ok to test 4da1c8a

[github-actions]

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