Reset plugin EP stream chunks before release

[Akshay Sonawane (apsonawane)]

This pull request adds cleanup logic to ensure that memory arena back-pointers and stream-tagged chunks are reset before stream objects are deleted. This helps prevent resource leaks and dangling references when CUDA streams are destroyed or released.

Memory management improvements:

• In CudaSyncStream::~CudaSyncStream(), before the stream object is destroyed, the code now calls factory_.ResetDeviceArenaChunksUsingStream to clear arena back-pointers associated with the CUDA stream, and releases any returned status.
• In StreamImpl::ReleaseImpl, before deleting the stream implementation, the code checks for an associated arena allocator and calls ResetChunksUsingStream to reset any stream-tagged chunks.

[Copilot]

Pull request overview

This PR adds additional cleanup steps when plugin EP stream objects are released/destroyed, aiming to clear arena stream-tagged chunk bookkeeping to avoid dangling back-pointers and leaks when CUDA streams go away.

Changes:

• Reset stream-tagged arena chunks in the example plugin EP stream ReleaseImpl prior to deleting the stream object.
• Reset CUDA plugin device-arena stream-tagged chunk bookkeeping in CudaSyncStream::~CudaSyncStream() prior to stream teardown.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
onnxruntime/test/autoep/library/example_plugin_ep/ep_stream_support.cc	Adds arena chunk reset in stream ReleaseImpl before deleting the stream impl.
onnxruntime/core/providers/cuda/plugin/cuda_stream_plugin.cc	Adds device-arena chunk reset in CudaSyncStream destructor during stream teardown.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Tianlei Wu (tianleiwu)]

The underlying stale stream back-pointer is a real bug, but the blocking teardown-order issue is already tracked in the existing CUDA stream thread. Chunks must not become globally reusable until prior GPU work is known complete; if synchronization fails, the affected chunks need to be quarantined or the arena invalidated rather than having their stream tags cleared.

Please also add focused regression coverage for release without a successful OnSessionRunEnd. It should prove both that no dangling stream pointer is consulted and that another stream cannot reuse the chunk before the original stream's work completes.

[Tianlei Wu (tianleiwu)]

Requesting changes for one additional teardown race in the example plugin. The CUDA failure-lifetime concern is continued in the existing review thread so it is not duplicated here.

[Ti-Tai Wang (titaiwangms)]

Major review finding

onnxruntime/core/providers/cuda/plugin/cuda_arena.cc:493-524

The new exception rollback can permanently orphan a free chunk if the stream-map insertion in SplitChunk throws.

SplitFreeChunkFromBin first removes the original chunk h from its bin, then calls SplitChunk. SplitChunk now attempts:

stream_to_chunks_.at(new_chunk->stream).insert(h_new_chunk);

If that allocation throws, its catch block only deallocates the new remainder handle. The original h is not restored to its bin, and the caller's catch only rolls back the destination stream mapping. Consequently h remains not in use but absent from every free bin, permanently reducing usable arena capacity and potentially causing later allocation failures.

Please make the operation transactional: either perform all potentially throwing bookkeeping before removing h, or catch the split failure in SplitFreeChunkFromBin and reinsert the unchanged original chunk into its bin.

[Tianlei Wu (tianleiwu)]

Reviewed the current head after the teardown-lifetime follow-ups. The undrained/capture path now detaches the outer stream pointer, abandons the arena before any backing memory can be reused, and unregisters the native handle; the normal path synchronizes on the owning device before resetting mappings. The example-plugin allocator race and the free-bin rollback issue are also addressed. I found no remaining blocking issues, and the required CUDA plugin checks are green.