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)]

Confirmed this is reachable, not just theoretical. When synchronization is disabled, InferenceSession calls DeviceStreamCollection::CleanUp(false), which skips the plugin OnSessionRunEnd; additionally, CudaSyncStream::OnSessionRunEndImpl returns before the arena reset if cudaStreamSynchronize fails. In both cases the destructor can be the first reset. Since ResetChunksUsingStream clears the tag and FindChunkPtr treats an untagged chunk as globally reusable, the reset must follow successful synchronization. On synchronization failure, the affected chunks need quarantine/invalidation or another stable-lifetime strategy; clearing the tag is unsafe.

[Tianlei Wu (tianleiwu)]

The current head avoids clearing arena tags when draining fails, but the proposed quarantine still does not provide the stable lifetime requested here. ArenaImpl stores the outer OrtSyncStream* in each chunk; it does not store CudaSyncStream*. Keeping the implementation alive in ReleaseImpl does not stop plugin_ep::Stream::~Stream from completing and destroying the outer stream after the callback returns, so a later FindChunkPtr can still pass a dangling chunk->stream to GetSyncIdForLastWaitOnSyncStream. Also, active capture is treated as a successful drain without resetting the mapping, and arena-reset errors are released/ignored before the drain reports success. Please make failure/capture preserve or invalidate the object actually retained by the arena, or represent quarantine without retaining a dereferenceable released stream pointer.

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