Skip to content

cuda.core: allow updating Buffer deallocation streams - #2602

Open
Andy-Jost wants to merge 9 commits into
NVIDIA:mainfrom
Andy-Jost:ajost/buffer-set-deallocation-stream
Open

cuda.core: allow updating Buffer deallocation streams#2602
Andy-Jost wants to merge 9 commits into
NVIDIA:mainfrom
Andy-Jost:ajost/buffer-set-deallocation-stream

Conversation

@Andy-Jost

@Andy-Jost Andy-Jost commented Aug 11, 2026

Copy link
Copy Markdown
Contributor

Important

Please hold review until #2526 merges. This PR is based on that work and will be rebased onto main afterward.

Summary

Closes #2600.

Add Buffer.set_deallocation_stream so callers can replace the stream that orders eventual deallocation without closing the buffer. This supports transferring an allocation's lifetime to another stream while preserving the existing Buffer object.

Changes

  • Add the public Buffer.set_deallocation_stream API and generated type stub.
  • Route initialization, explicit close overrides, and the new setter through the same validation and native DeallocationStream capture path.
  • Document the caller's responsibility to order prior allocation and access work before deallocation on the replacement stream.
  • Add focused tests covering initialization, setter, and close configuration followed by explicit close or garbage collection.

Related Work

Record a DeallocationStream at device-pointer creation so default-stream tokens pin the allocation context (and PTDS the allocating thread) instead of relying on ambient state at free time.
Make the deallocation stream's context current around free/unmap/MR cleanup so destruction no longer depends on ambient CUDA context, and wire cuCtxSetCurrent into the resource-handles driver table.
Add keyword-only stream= on Buffer/ManagedBuffer.from_handle when mr owns the pointer, bind it at construction, and cover teardown with no or foreign current context.
Stop treating CUDA_ERROR_INVALID_CONTEXT as a successful pool free, and let
explicit mr.deallocate() raise; destruction still contains errors in the
callback. Document PTDS deallocation ordering on the stream parameters and
note the context-safe Buffer teardown fix in the 1.2.0 release notes.
Require default deallocation streams to bind a current context at creation so teardown never relies on an ambiguous ambient token. Expand coverage and documentation for context-independent cleanup and failure reporting.
Ensure spawned children can bind the imported buffer's default deallocation stream before their process target starts.
@Andy-Jost Andy-Jost added this to the cuda.core 1.2.0 milestone Aug 11, 2026
@Andy-Jost Andy-Jost added enhancement Any code-related improvements P1 Medium priority - Should do cuda.core Everything related to the cuda.core module labels Aug 11, 2026
@Andy-Jost Andy-Jost self-assigned this Aug 11, 2026
@Andy-Jost
Andy-Jost requested a review from juenglin August 11, 2026 18:58
@github-actions

Copy link
Copy Markdown

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cuda.core Everything related to the cuda.core module enhancement Any code-related improvements P1 Medium priority - Should do

Projects

None yet

Development

Successfully merging this pull request may close these issues.

cuda.core: allow Buffer deallocation stream to be updated before teardown

1 participant