[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 StageTiming, a lightweight context manager. 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_*) ≤ iter_total_blocked_s

The fetch stage spans from upstream wait (blocked on data pipeline) through ray.get() completion, capturing both production delays and cross-node transfer.

New Prometheus Metrics (8)

Metric	Description
data_iter_blocked_fetch_seconds	Upstream wait + 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

Also: rank extracted from dataset tag (last split_N) for per-rank Prometheus querying.

Changes

• interfaces.py — StageTiming (context manager), BatchTimings, BlockWithTiming
• iter_batches.py — _report_batch_timings() overlap computation with docstring
• util.py — resolve_block_refs always returns BlockWithTiming (no record_timings flag, no Union/isinstance); nested context managers replace redundant perf_counter() calls; fetch window includes upstream wait
• stats.py — 8 Prometheus Gauges, Timer gains start_s/end_s, IterStatsSummary per-stage breakdown, rank extraction
• 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

• 51 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]

Instead of having tow timers for the same thing, can we consolidate the capture of start_s and end_s into the Timer class (in stats.py)? Double check for any backwards compatibility issues with other consumers of Timer

[OneSizeFitsQuorum]

Done. StageTiming now supports context manager protocol (__enter__/__exit__), and Timer gained start_s/end_s fields. Pipeline functions use nested context managers instead of redundant perf_counter() calls. No backwards compatibility issues — Timer is only used internally within ray/data.

[JasonLi1909]

Let's also capture the blocked time here. It will be useful to indicate how we much we are blocked on the upstream data pipeline and cross node transfer.

[OneSizeFitsQuorum]

Done. The fetch window in resolve_block_refs now spans from next(block_ref_iter) (upstream wait) through ray.get() completion, capturing both production delays and cross-node transfer.

[JasonLi1909]

Let's remove the optionality here and always record the timings to avoid downstream isinstance(block, BlockWithTimings) checks, Union types, and other type branching logic. And if we can consolidate the (start, end) capture into Timer, then we can still toggle this timing via the presence of stats

[OneSizeFitsQuorum]

Done. resolve_block_refs always returns BlockWithTiming (no record_timings parameter). batch_blocks() wraps raw blocks in BlockWithTiming with zero timing before passing to blocks_to_batches(), so _BatchingIterator receives a uniform type. All isinstance checks, Union types, and type branching logic removed.

[JasonLi1909]

Let's hold off on timing restore order and surfacing data_iter_blocked_restore_order_seconds. This is also more of an implementation detail. Instead we should focus on surfacing actionable metrics for users.

[OneSizeFitsQuorum]

Done. Removed restore_order stage, data_iter_blocked_restore_order_seconds metric, and all related code. restore_original_order() reverted to the original simple for-loop. PR now exposes 8 metrics (5 blocked stages + batches/rows/total).

[JasonLi1909]

nit: docstring

[OneSizeFitsQuorum]

added

[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!

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 2 potential issues.

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

[OneSizeFitsQuorum]

Re: merge_fetch idle gaps (Cursor Bugbot)

After deeper analysis, we've reverted to the span-based approach (commit aa41de5). Here's why:

Semantic correctness: From the training thread's perspective, when it calls next(batch_iter) for a multi-block batch, it's blocked for the entire span from when the first block starts fetching to when the last block finishes. Even if there are gaps between consecutive block fetches, the training thread is still waiting.

The "idle gaps" aren't really idle: The gaps between block fetches are actually pipeline overhead (batching logic, scheduling, etc.). They're part of the blocking experience from the training thread's viewpoint.

Sum underestimates: If we sum durations, we'd report less blocking time than the training thread actually experienced, which defeats the purpose of observability.

Example:

Block 1: [0ms, 100ms]
Pipeline overhead: [100ms, 150ms]
Block 2: [150ms, 250ms]

Training thread blocked: 0ms → 250ms = 250ms (span) ✓
Sum of fetch work: 100ms + 100ms = 200ms (underestimates) ✗

The span approach correctly answers "how long did the training thread wait for this batch?", which is what users care about for performance debugging.

Cursor Bugbot's concern about idle gaps is theoretically valid, but in practice those gaps are microseconds of pipeline overhead and represent real blocking time.