[Data] Shut down executor when DataIterator exits early

[xinyuangui2]

When user code breaks out of iter_batches() early (break in a for-loop), the dataset was not marked as finished. The streaming executor's worker thread kept producing blocks that piled up in the object store, and the executor held resources that could starve other datasets (e.g. a validation dataset waiting to run).

The inner _ClosingIterator only shuts down the executor through its __del__, which is non-deterministic — generator references may linger, leaving the executor alive long after iteration stops.

Single deterministic shutdown path: _on_iteration_end

Add a try/finally around yield from batch_iterator in DataIterator._iter_batches, calling a new _on_iteration_end(executor) hook on StopIteration, GeneratorExit (early break), and exceptions.

• The default implementation calls executor.shutdown(force=False) (idempotent), covering the local iter_batches() path.
• StreamSplitDataIterator overrides it to fire client_disengaged(epoch, split_idx) on the remote SplitCoordinator, since its executor lives on the actor and the iterator returns executor=None.

Both run synchronously on the consumer's thread, so cleanup happens the moment the consumer stops pulling.

Why GeneratorExit matters here

_iter_batches is a generator (it yield froms the batch iterator). When the consumer abandons the generator early — either by break-ing out of the for-loop, or by letting the generator go out of scope — Python's generator protocol throws GeneratorExit into it at the current yield point. Without the try/finally, that GeneratorExit unwinds straight out and the cleanup after yield from never runs; the executor keeps producing blocks until _ClosingIterator.__del__ eventually fires, which can be arbitrarily delayed. With the try/finally, GeneratorExit triggers _on_iteration_end synchronously on the consumer's thread, giving us one deterministic shutdown path for all three exit modes.

Why the consumer thread, not gen_blocks's finally

StreamSplitDataIterator's inner gen_blocks generator runs inside make_async_gen's filling worker thread. On early break, that worker exits via the interrupted_event path without explicitly closing the generator — its cleanup is then GC-bound and can be arbitrarily delayed under CI load (10+ s, causing flakes). The consumer-side hook avoids that entirely.

SplitCoordinator cleanup

On each client_disengaged(epoch_id, split_idx) call, the coordinator:

• Clears that split's prefetched-bytes accounting and drops its buffered RefBundle so remaining splits don't suffer stale backpressure or pinned memory.
• Tracks disengaged splits in a set; once every split has disengaged, shuts the executor down. Shutdown is dispatched outside the lock since StreamingExecutor.shutdown joins the scheduling thread for up to 2 s.
• Idempotent per (epoch_id, split_idx) as defense against duplicate/stale notifications (e.g. notifications arriving after the next epoch has started).

StreamSplitDataIterator._active_epoch is a plain attribute (no lock) — split iterators are passed to @ray.remote tasks, and a threading.Lock would make them unpicklable. The single-consumer convention plus the strict happens-before ordering (start_epoch → first yield → consumer exit → _on_iteration_end) makes this safe.

Lingering iter() reference caveat

There is a known case the try/finally cannot cover:

it = iter(ds.iter_batches())
for i, batch in enumerate(it):
    if i == 0:
        break
# the iterator sticks around until ``it`` goes out of scope;
# only ``it.close()`` guarantees the ``finally`` runs.

When the user keeps an explicit reference to the iterator, Python doesn't fire GeneratorExit on break — cleanup is deferred until the reference is dropped. Some libraries (e.g. PyTorch Lightning's batch fetchers) do this internally and the user can't control it. There's no general fix in code, so the docstrings on Dataset.iter_batches and DataIterator.iter_batches document the caveat and recommend calling it.close() for callers who hold the iterator reference themselves. it.close() is Python's built-in generator method — it raises GeneratorExit at the current yield point and runs our finally.

Test plan

Local iterator path (test_iterator.py):

• test_iter_batches_early_exit_shuts_down_executor — break after first batch; assert executor._shutdown is True.
• test_iter_batches_full_iteration_shuts_down_executor — natural exhaustion still shuts down (regression guard).
• test_iter_batches_exception_shuts_down_executor — exception inside the consumer loop still shuts down.
• test_iter_batches_close_on_held_iterator_shuts_down_executor — caller holds an iter() reference; it.close() triggers shutdown (covers the documented escape hatch for libraries like Lightning).

Streaming-split path (test_streaming_integration.py):

• test_streaming_split_early_exit_shuts_down_executor — both splits break early; uses a slow per-row map so natural completion takes ~28 s, proving the executor shuts down via the consumer-thread hook (10 s wait).
• test_streaming_split_partial_early_exit_keeps_executor — only one split breaks early; the other consumes all remaining batches (executor stays alive); shutdown happens once the second split also disengages.

Existing iterator + streaming-integration tests still pass locally.

[gemini-code-assist]

Code Review

This pull request introduces eager shutdown of the streaming executor during batch iteration by wrapping the process in a try-finally block, ensuring resources are released even on early exits or exceptions. Corresponding unit tests were added to validate shutdown behavior across various scenarios. Feedback was provided regarding the placement of the iteration time metric recording, which currently includes the shutdown duration and may lead to inaccurate performance data.

[justinvyu]

Btw, finally doesn't trigger in this case:

it = iter(ds.iter_batches())
for i, batch in enumerate(it):
    if i == 0:
        break

# the iterator sticks around until it goes out of scope
it.close()  # only way to guarantee that finally happens

I can't think of a good way to solve this case so we should just followup with a warning on our docs to not keep a iterator reference around and to just stick with the for loop.

One other problem is that other libraries like lightning could keep an iter() reference around which is outside of user control:

https://github.com/Lightning-AI/pytorch-lightning/blob/0e20e15f2376f4f356470b08875639a945c43334/src/lightning/pytorch/loops/fetchers.py#L79-L83

[xinyuangui2]

I think it is fine to keep this "leak". Since users keep the iterator reference, they might use this later.

[justinvyu]

How does this work if you have N threads on the split coordinator occupied by calls to SplitCoord.get()?

(We technically have N+1 threads to allow the controller to shutdown the data executor in the aborted run case, but are we also relying on this 1 extra thread here?)

[justinvyu]

For the N pending SplitCoord.get() tasks, do we expect those to all error out once the executor shuts down? Can those errors propagate up to the user and possibly crash the user code?

[xinyuangui2]

Discussed with @justinvyu offline, the consumer thread is not likely to propagate the error out.

[justinvyu]

@goutamvenkat-anyscale tried adding a similar "client finished" remote call but it was in the gen_blocks try/finally which is not guaranteed to run immediately if the client breaks out in the main thread because it the gen_blocks thread would only exit once it gets to the next yield call. And after that it might need to wait for gc as well.

The try/finally here should always happen immediately (except for the leaked iterator reference case).

[justinvyu]

Needed to clarify the problem we are solving in this PR. We only solve this following case in the PR, not the pytorch lightning limit_train_batches which falls under the "iterator leak" category which we don't have control over. I removed the mention of limit_train_batches the PR description.

import ray.data
from ray.train import ScalingConfig
import time
from ray.train.v2.api.data_parallel_trainer import DataParallelTrainer

epochs = 2
train_dataloader = ray.data.from_items([{"x": 1} for _ in range(1000)])
val_dataloader = ray.data.from_items([{"x": 2} for _ in range(1000)])


def train_func():
    train_dataloader = ray.train.get_dataset_shard("train").iter_batches(batch_size=100)
    val_dataloader = ray.train.get_dataset_shard("val").iter_batches(batch_size=100)

    for epoch in range(epochs):
        print("Train epoch", epoch)
        for batch in train_dataloader:
            ...
            if epoch == 1:
                time.sleep(10)
        print("Val epoch", epoch)
        for i, batch in enumerate(val_dataloader):
            if i == 1:
                break
        import gc

        gc.collect()


trainer = DataParallelTrainer(
    train_func,
    scaling_config=ScalingConfig(num_workers=2),
    datasets={"train": train_dataloader, "val": val_dataloader},
)
trainer.fit()

# Val epoch 0 executor hangs around for the entire duration of training epoch 1
# Only shuts down when validation epoch 1 starts.
# See that the val epoch 0 execution takes a long time to finish since it needs to wait for the long training epoch 1 to finish.
# (pid=32682) ✔️  Dataset val_4_0 execution finished in 51.27 seconds: : 620 row [00:51, 12.1 row/s]