[Data] Remove _base_dataset from StreamSplitDataIterator

[JasonLi1909]

Summary

StreamSplitDataIterator previously held a direct reference to the base Dataset object, which was used for schema(), get_context(), _get_dataset_tag(), and incrementing _run_index. During init_train_context on RayTrainWorker actors, the iterator is serialized and copied once per training worker, and because it contained the full dataset object, this caused significant head node object store usage proportional to the number of workers. At scale, this led to training ingest instability and head node OOMs. This PR moves all remaining dataset interactions into the SplitCoordinator actor, so the iterator communicates exclusively through the coordinator and no longer serializes the dataset object.

Changes

• Remove _base_dataset from StreamSplitDataIterator and route schema(), get_context(), and _get_dataset_tag() through new coordinator methods (get_dataset_schema(), get_dataset_context(), get_dataset_tag()) via ray.get()
• Remove _run_index increment from _to_ref_bundle_iterator since it is now handled on the SplitCoordinator side via ExecutionPlan.create_executor()

_dataset_state_lock for schema access

get_dataset_schema() on the coordinator is guarded by a _dataset_state_lock. Schema computation can trigger dataset execution, so the lock ensures this only happens once. The first caller computes and caches the result in self._schema, and subsequent callers return the cached value without redundant execution. The lock also prevents race conditions from concurrent access to dataset state.

Schema access during active dataset execution

If a cached schema is available (e.g., computed before iteration started), get_dataset_schema() returns it immediately, even during active execution. Otherwise, if the schema has never been computed and the streaming executor is active, it raises a RuntimeError.

i1, i2 = ds.streaming_split(2, equal=True)

# This works, schema is called before iteration starts:
schema = i1.schema()

# This raises RuntimeError, the executor is already active:
for batch in i1.iter_batches():
    schema = i1.schema()  # RuntimeError: Cannot call schema() during active dataset execution

This guard is necessary because ExecutionPlan.schema() with an unknown output type (e.g., a UDF whose return schema can't be inferred statically) falls back to execute_to_iterator() in plan.py, which creates a second StreamingExecutor on the same plan. That second executor computes its resource budget independently from the cluster via ResourceManager.get_global_limits(), unaware that the first executor has already claimed all available CPUs. The second executor's map tasks are submitted to Ray but can never be scheduled.

Because schema() blocks waiting for the first output bundle (next(gen) in execute_to_iterator), the SplitCoordinator actor becomes permanently blocked. This in turn blocks the consumers of the first executor, creating a circular deadlock where neither executor can make progress.

SplitCoordinator concurrency

get_dataset_context() and get_dataset_tag() are lightweight, non-blocking calls that don't contend with the main execution threads, so the existing max_concurrency=n+1 on the coordinator actor is sufficient.

Tests

• test_streaming_split_schema_before_execution: verifies schema retrieval works before any iteration starts
• test_streaming_split_schema_during_execution: verifies schema() raises RuntimeError when called during active iteration (called from inside the iteration loop to avoid race conditions)
• test_streaming_split_schema_after_execution: verifies schema retrieval works after execution completes
• test_streaming_split_context: verifies get_context() returns a valid DataContext from the coordinator
• test_streaming_split_dataset_tag: verifies _get_dataset_tag() returns correct split-indexed tags (_split_0, _split_1) from the coordinator
• All existing streaming_split tests pass, plus test_context_propagation.py::test_streaming_split

[gemini-code-assist]

Code Review

This pull request effectively decouples StreamSplitDataIterator from direct Dataset access by delegating operations to the SplitCoordinator. This is a significant improvement for reducing memory usage on the head node, particularly in distributed training scenarios. The use of a lock for schema computation is a solid approach to handle concurrency. The accompanying tests are thorough and cover the new functionality well. I have one suggestion to enhance the robustness of the schema() method to align its behavior with other DataIterator implementations.

[justinvyu]

thanks for the fix!!

[justinvyu]

can you try running this without raising the error to see what happens? and add the output to the PR description appendix

[JasonLi1909]

Ran without the safe guard, it can lead to a hang because the second executor will try to schedule tasks but won't have any resources available, added to PR description.

[justinvyu]

thanks for checking

[cursor]

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

[iamjustinhsu]

Just to confirm -- this StreamSplitDataIterator was being serialized 2x?

[JasonLi1909]

No, each worker only gets one iterator so it's only serialized once. The other copy of the dataset object came from the TrainRunContext which was removed in #61953

[iamjustinhsu]

I read the PR description but still a bit confused -- Are you implying there are scenarios where the schema is None, but current executor is not None, and that's why we need this guard? Like, can the executor be running when there is no schema produced yet? I think adding a comment about the if guard would be very helpful for future readers

[JasonLi1909]

self._schema is for caching. If get_dataset_schema was not called previously, then it is very possible for it to be empty when called during execution. This guard is to primary to prevent two executions on the same dataset- which can lead to deadlock. Will update pr description to be more clear