Repository navigation
cuda.core: allow Buffer deallocation stream to be updated before teardown #2600
Description
Summary
Add a public method for changing the stream that will order a Buffer's eventual deallocation without closing the buffer:
Buffer.set_deallocation_stream(stream: Stream | GraphBuilder) -> None
This is a follow-up to #2526.
Problem
A Buffer records a deallocation recipe when it acquires an allocation. Buffer.close(stream=...) can replace that recipe, but it immediately closes the Buffer.
Applications may transfer responsibility for a live allocation from work on stream A to work on stream B while retaining the Buffer or relying on automatic/deferred destruction. After establishing the necessary stream dependency, there is currently no public way to make stream B order the eventual deallocation without also closing the buffer.
This matters when the final release occurs through garbage collection or through another retained DevicePtrHandle, such as a graph attachment.
Proposed API
buffer.set_deallocation_stream(stream)The method should:
- Accept the same
Stream | GraphBuildervalues asBuffer.close(stream=...). - Replace the stored deallocation recipe without freeing or otherwise invalidating the buffer.
- Use the context-binding and PTDS handling introduced by cuda.core: capture complete Buffer deallocation recipe at creation #2526.
- Reject
None. - Raise if the buffer is already closed.
- Return
None.
Buffer.close(stream=stream) should remain equivalent to updating the deallocation stream and then closing the buffer.
Ordering semantics
This method does not enqueue synchronization, transfer ownership to a CUDA stream, or establish dependencies between streams. The caller must ensure that the new deallocation stream is ordered after allocation and all uses of the memory, using events or other CUDA synchronization mechanisms as appropriate.
Concurrent mutation and destruction of the same buffer remain the caller's responsibility.
Reflection
Do not add a public getter in this initial change.
The stored DeallocationStream is a teardown recipe rather than an ordinary cuda.core.Stream: default-stream recipes include a bound context, and PTDS recipes also retain host-thread metadata. Returning only a Stream would not faithfully represent or round-trip the complete recipe.
A separate immutable diagnostic representation can be considered later if a concrete introspection use case arises.
Acceptance criteria
- A live buffer can have its deallocation stream changed without being closed.
- Subsequent
close()or automatic destruction uses the replacement recipe. - The buffer remains usable after
set_deallocation_stream(). - Context-bound default-stream and PTDS behavior matches cuda.core: capture complete Buffer deallocation recipe at creation #2526.
- Existing
close(stream=...)behavior remains unchanged.
Related
- [BUG]: Buffer cleanup must capture a complete deallocation recipe at creation #2497
- cuda.core: capture complete Buffer deallocation recipe at creation #2526
- cuda.core: reject stream=None — require explicit stream everywhere #2001
- Inconsistent lifetime management across
from_handle/from_*APIs #1989
Activity