[Data] Release pinned blocks after dataset execution (executor shutdown, terminal ops, OutputSplitter)

[marwan116]

Summary

Dataset.streaming_split() can keep ObjectRefs alive across epochs because the split coordinator holds an iterator that closes over the executor topology, and because OutputSplitter can leave the metrics “shadow” input queue out of sync with the real buffer after equal=True truncation. This PR fixes both issues.

What was going wrong

1. Topology + queues across epochs
   The coordinator keeps _output_iterator until the next epoch replaces it. Anything reachable from the old executor’s topology (including OpState queues) can pin the object store until that chain is dropped and queues are drained.

2. Truncated remainder + metrics
   After equalizing splits, a bare _buffer.clear() emptied the real buffer without on_input_dequeued, so OpRuntimeMetrics could still hold RefBundles for the truncated remainder.

3. Terminal OutputSplitter and mark_execution_finished
mark_execution_finished() (which clears internal queues via the mixin) was not run for operators with no downstream deps, so terminal splitters did not follow the same lifecycle path as other ops.

What we changed

• OutputSplitter.all_inputs_done: replace _buffer.clear() with clear_internal_input_queue() so truncation drains through metrics and releases refs consistently.
• update_operator_states: treat fully completed terminal ops like split sinks as finished for execution and run mark_execution_finished() when output_dependencies is empty.
• StreamingExecutor.shutdown: after op.shutdown(), call _clear_topology_queues_post_shutdown(force) to clear internal queues and topology OpState queues; for the DAG sink with num_output_splits > 1, skip clearing the external output queues on cooperative (force=False) shutdown so other splits keep their bundles; force=True still clears everything (e.g. epoch rollover in SplitCoordinator._try_start_new_epoch).
• Tests: test_output_split_shutdown_preserves_sibling_split_queues covers fast vs slow split consumers.

Links (for reviewers)

• Cooperative shutdown from first exhausted split: streaming_executor.py (_ClosingIterator.get_next)
• Forced teardown on epoch advance: stream_split_iterator.py (_try_start_new_epoch)
• Barrier / when the next epoch starts (not tied to first consumer finishing): start_epoch / _barrier, _barrier → _try_start_new_epoch

[gemini-code-assist]

Code Review

This pull request updates the shutdown method in StreamingExecutor to clear internal and operator queues, which helps release ObjectRef references and prevent memory leaks in the object store. A review comment suggests also clearing state.active_tasks to ensure that references held by completed tasks are also promptly released.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 1b3ec6b68f74358ca2e3e81e37e1b716a07fcf37. Configure here.}

[iamjustinhsu]

thanks! lgtm