[Data] Add per-stage training-thread blocking attribution to iter_batches

[OneSizeFitsQuorum]

Summary

Decomposes iter_total_blocked_s into per-stage contributions so users can see where the training thread is blocked. Closes #64132, part of RFC #63911.

Design

Each pipeline stage records an independent (start_s, end_s) window via TimeSpan, a lightweight value yielded by Timer.timer(). The training thread records its blocked window around next(batch_iter). Per-stage attribution = overlap of the two windows:

overlap = min(stage.end_s, blocked_end) - max(stage.start_s, blocked_start)

Invariant: sum(iter_blocked_*) approximates iter_total_blocked_s. It is a lower bound in the common case, but split fetch stages (production_wait + data_transfer) may overlap by design for multi-block batches — see _attribute_blocked_time docstring for the full rationale.

The fetch stage is split into two: production_wait spans the upstream wait (blocked on the data pipeline), and data_transfer spans the cross-node transfer via ray.get(). This lets users differentiate data production stall from cross-node transfer stall.

New Prometheus Metrics (9)

Metric	Description
data_iter_blocked_production_wait_seconds	Upstream wait (blocked on data pipeline)
data_iter_blocked_data_transfer_seconds	Cross-node transfer (ray.get)
data_iter_blocked_batching_seconds	Batch creation
data_iter_blocked_format_seconds	Format conversion
data_iter_blocked_collate_seconds	collate_fn
data_iter_blocked_finalize_seconds	finalize_fn
data_iter_batches_total	Batches delivered to training loop
data_iter_rows_total	Rows delivered (exit-side, respects drop_last)
data_iter_total_seconds	Total iterator wall-clock time

Per-rank Prometheus label (rank extracted from dataset tag's trailing split_N) is deferred to a follow-up PR — see TODO in stats.py.

Changes

• interfaces.py — TimeSpan, BlockFetchTiming (per-block: production_wait + data_transfer), BlockFetchResult (block + Optional timing), BatchTimings (per-batch: all 6 stages), merge_fetch / _merge_span
• iter_batches.py — _attribute_blocked_time() overlap computation with docstring documenting known limitations (split-fetch overlap, reorder buffer unattributed gap)
• util.py — resolve_block_refs returns Iterator[BlockFetchResult]; production_wait captured as a bare TimeSpan (no accumulation, to avoid double-counting with prefetch_batches_locally.get_next_ref_bundle); data_transfer accumulated into iter_get_s
• stats.py — 9 Prometheus Gauges, Timer.timer() yields a fresh TimeSpan per call (thread-safe), PipelineStage enum, get_blocked_timer(stage) via match, IterStatsSummary per-stage breakdown
• iterator.py — Final metrics flush on iteration end

Example Output

Per-stage training-thread blocked time breakdown:
    * block fetch (ray.get): 1.74ms
    * batching: 428.36us
    * format: 4.1ms
Total batches consumed: 10
Total rows consumed: 100

Performance

~6 μs/batch overhead. At 10k batches/sec: <0.04% impact on a 10ms training step.

Tests

• 56 unit tests (iter_batches) + 7 unit tests (stats/summary) — all passing
• Integration verified with Ray Train v2 DataParallelTrainer + TorchTrainer

References

Closes #64132 · Part of RFC #63911

cc @JasonLi1909

[gemini-code-assist]

Code Review

This pull request introduces detailed per-stage wall-clock timing and statistics tracking (including fetch, batching, formatting, collating, finalization, and order restoration) for Ray Data batch iteration, along with corresponding Prometheus metrics and unit tests. The review feedback suggests updating the type annotation of dataset_tag in _create_iteration_tags to Optional[str] to align with the underlying function and avoid static type checking errors.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[JasonLi1909]

Thank you for the PR @OneSizeFitsQuorum! The core idea of capturing the start and end of each pipeline stage to compute overlap with the training thread stall makes sense. That said, let's keep the scope of this PR to just the training thread attribution metrics. We can compartmentalize other metrics for a later PR, but at a glance some of them such as data_iter_prefetch_queue_depth are more of an implementation detail, which we should avoid. Also it would be great if you could simplify the PR description so it's shorter and easier to understand. Thanks!

[OneSizeFitsQuorum]

Thanks for the review @JasonLi1909! I've made the following changes:

1. Reduced scope — removed the 5 implementation-detail metrics (data_iter_prefetch_queue_depth, data_iter_restore_order_buffer_peak, data_iter_shuffle_buffer_*). The PR now only contains the 9 core training-thread attribution metrics. The removed metrics are preserved on a separate branch for a future PR.

2. Simplified PR description — cut it down from ~300 lines to ~80 lines, focusing on design, the 9 metrics, and test results.

Also fixed a regex issue flagged by Cursor Bugbot: the rank extraction now uses the last split_N match (instead of the first) to avoid false matches when the user-defined dataset name contains split_<digits>.

[JasonLi1909]

Hey @OneSizeFitsQuorum, thanks for the cleanup. I've left some more comments. Overall, make sure to include docstrings, comments, and that names of data classes/functions communicate their purpose. I'll take another look after you're done, thanks!

[JasonLi1909]

Question: The data model feels off here, an entire block will not populate all of the BatchTimings. The data models should cleanly reflect what is possible and intrinsic to the component, is there a better way to organize this?

[JasonLi1909]

Is there a way to avoid this? This is purely a shim for the typing. Somehow, we should make the timing optional into BatchIterator. If it's too complex and out of scope, leave a detailed TODO comment for the future.

[JasonLi1909]

Two things here:

1. The naming here is confusing as this function is used to merge timings at the block level but the type is BatchTimings. It reads as if batches rather than blocks are being merged. See comment below about the data model of BlockWithTiming.
2. If we split merge into data wait + transfer, we will need to change this to merge the stages properly. This may result in some double counting between the wait and transfer timers in the case of batches that span multiple blocks, so we will need to be careful here. For instance:

Block 1: wait=[0,1], transfer=[1,2]
Block 2: wait=[2,3], transfer=[3,4]

Say a batch spans these two blocks, then when we merge these blocks we get
wait=[0,3], transfer=[1,4], resulting in double counting. We will need to avoid this somehow

[OneSizeFitsQuorum]

@JasonLi1909 I've addressed all your feedback across 4 commits (TimeSpan consolidation, fetch split, data model refactoring, PipelineStage enum). Cursor Bugbot then flagged 3 remaining issues that I'd like to discuss before fixing:

1. Nested timer double-counts iter_get_ref_bundles_s

When prefetch_batches > 0, resolve_block_refs wraps next(block_ref_iter) with _timed(iter_get_ref_bundles_s), but prefetch_batches_locally's get_next_ref_bundle() already times the same next(ref_bundles) call with the same Timer. The upstream wait gets accumulated twice.

Proposed fix: In resolve_block_refs, create the production_wait TimeSpan manually (without _timed()) so it doesn't accumulate to iter_get_ref_bundles_s. The cumulative stats remain handled by get_next_ref_bundle()'s existing timer. The per-batch TimeSpan is only used for overlap attribution, not accumulation. Does this separation make sense?

2. Split fetch spans overlap, violating sum ≤ total

After splitting fetch into production_wait + data_transfer, multi-block span merge creates overlapping windows (stages are interleaved across blocks). For example:
```
Block 1: prod [0,100], xfer [100,200]
Block 2: prod [200,300], xfer [300,400]
Merged: prod [0,300], xfer [100,400] ← overlap [100,300]
```

My view: In a pipelined system, multiple stages can be active simultaneously (data_transfer for block N + production_wait for block N+1). Attributing the same blocked interval to both stages is semantically correct — the training thread was indeed blocked while both stages were running. I'd relax the invariant to allow sum > total and document this behavior. I also considered using interval lists (List[TimeSpan]) for precise non-overlapping merge, but that would significantly increase complexity. Do you agree with the relaxed invariant, or do you prefer a different approach?

3. Reorder buffer skews blocked attribution

With preserve_order=True, stage timings are recorded when format/collate finish, but the batch may sit in restore_original_order's buffer before being yielded. The blocked window includes this reorder wait, but no stage TimeSpan covers it, so it becomes unattributed.

My view: Since you asked to remove restore_order timing, this unattributed gap is expected. If users see high unattributed ratios, they can investigate further. Is this acceptable, or would you prefer a different approach?

[JasonLi1909]

"IterationStage" is a better name for this. Also, document with a description for each stage. "Pipeline stages for blocked attribution" is vague- what pipeline stages? what blocked attribution?

[OneSizeFitsQuorum]

doc added. BTW, I think either name works—they each emphasize a slightly different aspect, and both make sense in this context. Since you seem to prefer it, I've changed it to iteration stage.

[cursor]

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

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