[data][1/n] Add the push-based streaming_split coordinator by xinyuangui2 · Pull Request #66465 · ray-project/ray

xinyuangui2 · 2026-09-24T18:03:49Z

Refactor the splitcoordinator to push based. Passing the ref as a task arg forces Ray to dereference and materialize a secondary copy on the worker node. So by the time the worker touches the data, it's already materialized in the worker’s local object store and is not affected by a remote death.

Why are these changes needed?

First PR of the push-based streaming_split for Ray Train ingestion. It adds the coordinator actor only; the pushers and the consumer-side iterator follow in #66466.

Epoch lifecycle: runs one streaming executor per epoch behind a barrier that all n splits must reach. Waiters sleep on a Condition rather than polling.
Row window: _SplitFlow tracks each split's target_rows, rows_pushed, and rows_consumed. A sender may send while target_rows - (rows_pushed - rows_consumed) > 0. Flow state is recreated every epoch, so a stale report can't affect the new one.
Pacing: bytes sent but not yet consumed feed the executor's external-consumer backpressure.
Scope: internal only. There's no public API change and the pull path is untouched. Dataset.streaming_split will switch over in a later PR.

Split out of #65842 (supersedes it), so this doesn't duplicate other work. Written with AI assistance (Claude Code) and reviewed by a human. Tests: pytest python/ray/data/tests/test_push_split_coordinator.py (10 passed).

TODO:

After mid-epoch resumption PRs are merged, this coordinator needs minor changes.

🤖 Generated with Claude Code

https://claude.ai/code/session_01QfaKctvP9gCsC2EBTpdqnv

Add Dataset.streaming_split_push_based: a coordinator actor runs the streaming executor and pushes each split's blocks, by value, to the actor hosting that split's PushBasedDataIterator, which iterates them from a local queue. Each consumer declares a prefetch_batches * batch_size row window and reports consumption; the coordinator pushes whole blocks while target_rows - (rows_pushed - rows_consumed) > 0, and feeds pushed-minus-consumed bytes to the executor's external-consumer backpressure. Deliveries are sequence-numbered and reordered on arrival, so ordering doesn't depend on the hosting actor's concurrency. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QfaKctvP9gCsC2EBTpdqnv Signed-off-by: xgui <xgui@anyscale.com>

gemini-code-assist

Code Review

This pull request introduces a push-based streaming split mechanism for Ray Data, including a coordinator actor that pushes blocks to consumers, a local receiver that reorders out-of-order deliveries, and a custom data iterator. Feedback on these changes highlights several critical improvements: resolving a permanent memory leak of receiver objects in the global registry, replacing an inefficient busy-wait loop in the epoch barrier with a condition variable, ignoring duplicate out-of-order deliveries to prevent memory leaks, validating the split index during registration to avoid pusher thread crashes, and adding missing API decorators to the new dataset method.

Drop Dataset.streaming_split_push_based. The push-based split is built by an internal streaming_split_push_based(dataset, n, ...) helper instead, so a later PR can switch Dataset.streaming_split over directly. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QfaKctvP9gCsC2EBTpdqnv Signed-off-by: xgui <xgui@anyscale.com>

Move the coordinator into its own module, push_split_coordinator.py: the epoch barrier and per-epoch executor, the per-split row window (_SplitFlow), and the external-consumer-bytes pacing feed. The pusher threads, the consumer-side receiver and PushBasedDataIterator move to the follow-up PR. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QfaKctvP9gCsC2EBTpdqnv Signed-off-by: xgui <xgui@anyscale.com>

Wake barrier waiters with a Condition instead of polling every 0.1s, and reject out-of-range split indices. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QfaKctvP9gCsC2EBTpdqnv Signed-off-by: xgui <xgui@anyscale.com>

cursor

Stale Bugbot comment from a previous run.

…efly If the last arrival's teardown of the previous epoch raised, the other splits waited forever and a retry drove the arrival count negative. Release the barrier in a finally block and log the teardown failure. Also make the coordinator tests type-check under pyrefly. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QfaKctvP9gCsC2EBTpdqnv Signed-off-by: xgui <xgui@anyscale.com>

cursor

Cursor Bugbot has reviewed your changes and found 3 potential issues.

^{Reviewed by Cursor Bugbot for commit 6ad886e. Configure here.}

- Ignore consumer reports that arrive after a split finished, so a late report can't reopen its send window. - Tag metrics with the next run's dataset id, as SplitCoordinator does: the tag is fetched before the epoch's executor is created. - Report coordinator overhead in stats() (streaming_split_coordinator_s), log per-split finish counts, and restore the full schema() error. - Document why consumer reports don't refresh the pacing feed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QfaKctvP9gCsC2EBTpdqnv Signed-off-by: xgui <xgui@anyscale.com>

…cstring Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QfaKctvP9gCsC2EBTpdqnv Signed-off-by: xgui <xgui@anyscale.com>

Rename _create_split_dataset to create_streaming_split_dataset and move it to the shared Ray Data utils, so the push-based split and, later, Dataset.streaming_split can share it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QfaKctvP9gCsC2EBTpdqnv Signed-off-by: xgui <xgui@anyscale.com>

ayushk7102

Can we add some more information in the PR description about why these changes are needed? we can motivate some things about the push model e.g.:

We can handle reconstruction for lost blocks before passing refs to the workers
Any implications of this design for mid epoch resumption, etc.

ayushk7102 · 2026-10-09T19:23:54Z

+            except Exception:
+                logger.warning(
+                    f"Failed to tear down epoch {starting_epoch}; starting the "
+                    "next epoch anyway.",
+                    exc_info=True,
+                )


It looks like if there was a failure during the teardown of the epoch in the original split coordinator, we would re-raise for all threads. But with the new push split coordinator, we swallow the error and move on. Is this expected?

https://github.com/ray-project/ray/blob/master/python/ray/data/_internal/iterator/stream_split_iterator.py#L315-L337

Good catch — this now matches SplitCoordinator: a teardown failure is raised to every split for that epoch. Unlike before, the barrier is still released first, so no split hangs (the bug the earlier change fixed) and a later start_epoch can start the next epoch. Covered by test_teardown_failure_fails_epoch_for_all_splits.

…nt flow invariant Match SplitCoordinator: if tearing down the previous epoch fails, the next epoch fails for every split instead of starting anyway. The barrier is still released, so no split hangs and a retry can start a later epoch. Also document why bytes_consumed never exceeds bytes_pushed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QfaKctvP9gCsC2EBTpdqnv Signed-off-by: xgui <xgui@anyscale.com>

…helper Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QfaKctvP9gCsC2EBTpdqnv Signed-off-by: xgui <xgui@anyscale.com>

A sender that was mid-fetch when its split finished would otherwise add bytes no consumer reads back, inflating the pacing feed for the rest of the epoch. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QfaKctvP9gCsC2EBTpdqnv Signed-off-by: xgui <xgui@anyscale.com>

Blocks are sent by value, so the consumer holds no ObjectRef and Ray Data's BlockRefCounter stopped counting a block as soon as its delivery task finished, while the consumer still pinned it. The producer's budget and output throttles then under-counted by about a prefetch window per consumer, letting it buffer more upstream (peak object store 1.11x pull's on training_ingest_regression_test). _SplitFlow now keeps each sent block's ref until the consumer reports its bytes consumed, matching how a pull consumer's prefetched refs are counted. No extra object-store memory: the consumer pins the object anyway. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QfaKctvP9gCsC2EBTpdqnv Signed-off-by: xgui <xgui@anyscale.com>

The tests start actors and wait on epoch barriers with up-to-60 s timeouts; size "large" (900 s) leaves room for slower CI machines. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QfaKctvP9gCsC2EBTpdqnv Signed-off-by: xgui <xgui@anyscale.com>

goutamvenkat-anyscale · 2026-10-13T20:57:04Z

+
+    def rows_to_send(self) -> int:
+        with self.cond:
+            return self.target_rows - (self.rows_pushed - self.rows_consumed)


I'm confused. If rows_pushed > rows_consumed why are we sending more?

goutamvenkat-anyscale · 2026-10-13T20:58:05Z

+            "target_rows": {i: f.target_rows for i, f in flows.items()},
+            "rows_pushed": {i: f.rows_pushed for i, f in flows.items()},
+            "rows_consumed": {i: f.rows_consumed for i, f in flows.items()},
+            "bytes_pushed": {i: f.bytes_pushed for i, f in flows.items()},
+            "bytes_consumed": {i: f.bytes_consumed for i, f in flows.items()},
+            "unconsumed_refs": {i: f.num_unconsumed_refs() for i, f in flows.items()},


Construct this dict from a single loop

goutamvenkat-anyscale · 2026-10-13T21:09:02Z

+    def _finish_split(self, epoch_id: int, split_idx: int) -> None:
+        executor_to_shutdown = None
+        with self._lock:
+            if epoch_id != self._cur_epoch:


In what case can the split not correspond to the current epoch? Also shouldn't this be an error since all splits of an epoch need to finish before moving to the next epoch?

goutamvenkat-anyscale · 2026-10-13T21:16:16Z

+        with self._barrier_cond:
+            while self._cur_epoch == starting_epoch and (
+                self._num_unarrived_splits_at_barrier != 0
+                or self._teardown_complete_for != starting_epoch
+            ):
+                if time.time() - start_time > BLOCKED_CLIENT_WARN_TIMEOUT:
+                    if log_once(f"push_split_blocked_{split_idx}_{starting_epoch}"):
+                        logger.warning(
+                            f"Push-based split (epoch={starting_epoch}, "
+                            f"split={split_idx}) blocked waiting on other "
+                            f"clients for more than "
+                            f"{BLOCKED_CLIENT_WARN_TIMEOUT}s. All clients must "
+                            "iterate their splits at the same time."
+                        )
+                # Woken by the last arrival; the timeout only paces the
+                # blocked-client warning.
+                self._barrier_cond.wait(timeout=1.0)


Why not use https://docs.python.org/3/library/threading.html#threading.Condition.wait_for

xinyuangui2 mentioned this pull request Sep 24, 2026

[data][train][2/n] Push-based streaming_split: pushers, iterator, Ray Train integration, release tests #66466

Open

xinyuangui2 requested a review from justinvyu September 24, 2026 20:18

xinyuangui2 added the go add ONLY when ready to merge, run all tests label Sep 24, 2026

xinyuangui2 marked this pull request as ready for review September 24, 2026 20:20

xinyuangui2 requested a review from a team as a code owner September 24, 2026 20:21

gemini-code-assist Bot reviewed Sep 24, 2026

View reviewed changes

xinyuangui2 and others added 3 commits September 24, 2026 20:25

xinyuangui2 changed the title ~~[data][1/n] Add push-based streaming_split~~ [data][1/n] Add the push-based streaming_split coordinator Sep 24, 2026

cursor Bot reviewed Sep 24, 2026

View reviewed changes

Comment thread python/ray/data/_internal/iterator/push_split_coordinator.py Outdated

ray-gardener Bot added the data Ray Data-related issues label Sep 25, 2026

JasonLi1909 mentioned this pull request Sep 30, 2026

[data] Fix off-by-one run_index in Ray Data iteration metrics dataset_id tag #66586

Merged

Merge branch 'master' into push-split-coordinator

6ad886e

cursor Bot reviewed Oct 8, 2026

View reviewed changes

Comment thread python/ray/data/_internal/iterator/push_split_coordinator.py

Comment thread python/ray/data/_internal/iterator/push_split_coordinator.py Outdated

Comment thread python/ray/data/_internal/iterator/push_split_coordinator.py Outdated

xinyuangui2 requested review from JasonLi1909 and ayushk7102 October 8, 2026 20:26

xinyuangui2 and others added 3 commits October 8, 2026 20:27

[data] Put the push-based split overview in the coordinator module do…

163ef0e

…cstring Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QfaKctvP9gCsC2EBTpdqnv Signed-off-by: xgui <xgui@anyscale.com>

ayushk7102 reviewed Oct 9, 2026

View reviewed changes

ayushk7102 approved these changes Oct 9, 2026

View reviewed changes

xinyuangui2 and others added 6 commits October 9, 2026 19:36

Merge branch 'master' into push-split-coordinator

8eed180

[data] Build Dataset.streaming_split's split dataset with the shared …

0157c00

…helper Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QfaKctvP9gCsC2EBTpdqnv Signed-off-by: xgui <xgui@anyscale.com>

goutamvenkat-anyscale reviewed Oct 13, 2026

View reviewed changes

Conversation

xinyuangui2 commented Sep 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why are these changes needed?

TODO:

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cursor Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

cursor Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ayushk7102 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ayushk7102 Oct 9, 2026

Choose a reason for hiding this comment

Uh oh!

xinyuangui2 Oct 9, 2026

Choose a reason for hiding this comment

Uh oh!

goutamvenkat-anyscale Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

goutamvenkat-anyscale Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

goutamvenkat-anyscale Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

goutamvenkat-anyscale Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

xinyuangui2 commented Sep 24, 2026 •

edited

Loading

cursor Bot left a comment •

edited

Loading