ray-project--ray/python/ray/data/tests/block_batching/test_iter_batches.py

import queue
import threading
import time
from typing import Iterator, List, Optional
from unittest.mock import patch

import pandas as pd
import pyarrow as pa
import pytest

import ray
from ray.data._internal.block_batching.interfaces import (
    Batch,
    BatchMetadata,
    BatchStageTimings,
    BlockPrefetcher,
)
from ray.data._internal.block_batching.iter_batches import (
    BatchIterator,
    prefetch_batches_locally,
    restore_original_order,
)
from ray.data._internal.block_batching.util import (
    WaitBlockPrefetcher,
)
from ray.data._internal.execution.interfaces.ref_bundle import BlockEntry, RefBundle
from ray.data._internal.stats import DatasetStats, TimeSpan
from ray.data.block import Block, BlockAccessor, BlockMetadata
from ray.types import ObjectRef

# Sleep duration injected into each scenario's bottleneck stage. Picked to be
# large enough to dominate scheduling/measurement noise but small enough to
# keep the test fast (5 batches × 0.3s ≈ 1.5s per scenario).
SLEEP_S = 0.3


def ref_bundle_generator(num_rows: int, num_blocks: int) -> Iterator[RefBundle]:
    for i in range(num_blocks):
        block = pa.table({"foo": [i] * num_rows})
        metadata = BlockMetadata(
            num_rows=num_rows,
            size_bytes=0,
            input_files=[],
            exec_stats=None,
        )
        schema = block.schema
        yield RefBundle(
            blocks=(BlockEntry(ray.put(block), metadata),),
            owns_blocks=True,
            schema=schema,
        )


@pytest.mark.parametrize("num_batches_to_prefetch", [1, 2])
@pytest.mark.parametrize("batch_size", [None, 1, 4])
def test_prefetch_batches_locally(
    ray_start_regular_shared, num_batches_to_prefetch, batch_size
):
    class DummyPrefetcher(BlockPrefetcher):
        def __init__(self):
            self.windows = []

        def prefetch_blocks(self, block_refs: List[ObjectRef[Block]]):
            if batch_size is None:
                assert len(block_refs) == num_batches_to_prefetch
            else:
                assert (
                    sum(len(ray.get(block_ref)) for block_ref in block_refs)
                    >= batch_size * num_batches_to_prefetch
                )
            self.windows.append(block_refs)

    num_blocks = 10
    num_rows = 2
    prefetcher = DummyPrefetcher()
    ref_bundles = list(ref_bundle_generator(num_blocks=num_blocks, num_rows=num_rows))
    prefetch_block_iter = prefetch_batches_locally(
        iter(ref_bundles),
        prefetcher=prefetcher,
        num_batches_to_prefetch=num_batches_to_prefetch,
        batch_size=batch_size,
    )

    block_count = 0
    prefetched_blocks = []
    previous_num_windows = 1

    for block in prefetch_block_iter:
        prefetched_blocks.append(block)
        block_count += 1
        remaining_rows = (num_blocks - block_count) * num_rows
        if batch_size is None and block_count < num_blocks - num_batches_to_prefetch:
            # Test that we are actually prefetching in advance if this is not the last
            # block.
            assert len(prefetcher.windows) == previous_num_windows + 1
            previous_num_windows = len(prefetcher.windows)
        elif (
            batch_size is not None
            and remaining_rows > batch_size * num_batches_to_prefetch
        ):
            # Test that we are actually prefetching in advance if this is not the last
            # batch.
            assert len(prefetcher.windows) == previous_num_windows + 1
            previous_num_windows = len(prefetcher.windows)

    # Test that original blocks are unchanged.
    expected_blocks = []
    for ref_bundle in ref_bundles:
        expected_blocks.extend(ref_bundle.block_refs)
    assert prefetched_blocks == expected_blocks


def test_restore_from_original_order():
    base_iterator = [
        Batch(BatchMetadata(batch_idx=1), None),
        Batch(BatchMetadata(batch_idx=0), None),
        Batch(BatchMetadata(batch_idx=3), None),
        Batch(BatchMetadata(batch_idx=2), None),
    ]

    ordered = list(restore_original_order(iter(base_iterator)))
    idx = [batch.metadata.batch_idx for batch in ordered]
    assert idx == [0, 1, 2, 3]


def test_attribute_blocked_time_overlap_attribution():
    stats = DatasetStats(metadata={}, parent=None)
    batch_iterator = BatchIterator(iter([]), stats=stats)
    timings = BatchStageTimings()
    timings.production_wait.append(TimeSpan(start_s=10.0, end_s=20.0))
    timings.batching = TimeSpan(start_s=20.0, end_s=30.0)
    timings.format = TimeSpan(start_s=30.0, end_s=40.0)
    timings.finalize = TimeSpan(start_s=50.0, end_s=60.0)
    batch = Batch(BatchMetadata(batch_idx=0, num_rows=8, stage_timings=timings), None)

    batch_iterator._attribute_blocked_time(
        batch, blocked_start_s=15.0, blocked_end_s=35.0
    )

    assert stats.iter_blocked_production_wait_s.get() == pytest.approx(5.0)
    assert stats.iter_blocked_batching_s.get() == pytest.approx(10.0)
    assert stats.iter_blocked_format_s.get() == pytest.approx(5.0)
    assert stats.iter_blocked_collate_s.get() == 0
    assert stats.iter_blocked_finalize_s.get() == 0
    assert stats.iter_batches_total == 1
    assert stats.iter_rows_total == 8


def _make_span(start: Optional[float], end: Optional[float]) -> Optional[TimeSpan]:
    """Create a TimeSpan, or None if the stage didn't run."""
    if start is None or end is None:
        return None
    return TimeSpan(start_s=start, end_s=end)


def _make_batch_with_timings(
    production_wait_start: Optional[float] = None,
    production_wait_end: Optional[float] = None,
    data_transfer_start: Optional[float] = None,
    data_transfer_end: Optional[float] = None,
    batching_start: Optional[float] = None,
    batching_end: Optional[float] = None,
    format_start: Optional[float] = None,
    format_end: Optional[float] = None,
    collate_start: Optional[float] = None,
    collate_end: Optional[float] = None,
    finalize_start: Optional[float] = None,
    finalize_end: Optional[float] = None,
    num_rows: int = 0,
):
    """Helper to construct a Batch with specific stage timing windows."""
    timings = BatchStageTimings()
    pw = _make_span(production_wait_start, production_wait_end)
    if pw is not None:
        timings.production_wait.append(pw)
    dt = _make_span(data_transfer_start, data_transfer_end)
    if dt is not None:
        timings.data_transfer.append(dt)
    timings.batching = _make_span(batching_start, batching_end)
    timings.format = _make_span(format_start, format_end)
    timings.collate = _make_span(collate_start, collate_end)
    timings.finalize = _make_span(finalize_start, finalize_end)
    return Batch(
        BatchMetadata(batch_idx=0, num_rows=num_rows, stage_timings=timings), None
    )


def _make_test_iterator(stats):
    """Create a BatchIterator wired to the given stats without a real pipeline."""
    it = BatchIterator.__new__(BatchIterator)
    it._stats = stats
    return it


class TestAttributeBlockedTimeEdgeCases:
    """Edge case tests for overlap-based blocked attribution."""

    def test_zero_overlap_stage_finished_before_blocked(self):
        """Fetch [0, 1.5] finished before training blocked at t=2 → 0 attribution."""
        stats = DatasetStats(metadata={}, parent=None)
        it = _make_test_iterator(stats)
        batch = _make_batch_with_timings(
            production_wait_start=0.0, production_wait_end=1.5
        )
        it._attribute_blocked_time(batch, blocked_start_s=2.0, blocked_end_s=3.0)
        assert stats.iter_blocked_production_wait_s.get() == 0.0

    def test_zero_overlap_blocked_before_stage(self):
        """Training blocked [0, 1], stage ran [2, 3] → 0 attribution."""
        stats = DatasetStats(metadata={}, parent=None)
        it = _make_test_iterator(stats)
        batch = _make_batch_with_timings(format_start=2.0, format_end=3.0)
        it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=1.0)
        assert stats.iter_blocked_format_s.get() == 0.0

    def test_partial_overlap(self):
        """Fetch [0, 2], blocked [1, 3] → overlap = min(2,3)-max(0,1) = 1.0."""
        stats = DatasetStats(metadata={}, parent=None)
        it = _make_test_iterator(stats)
        batch = _make_batch_with_timings(
            production_wait_start=0.0, production_wait_end=2.0
        )
        it._attribute_blocked_time(batch, blocked_start_s=1.0, blocked_end_s=3.0)
        assert stats.iter_blocked_production_wait_s.get() == pytest.approx(1.0)

    def test_full_overlap_stage_inside_blocked(self):
        """Stage [1, 2] entirely inside blocked [0, 3] → full 1.0 credit."""
        stats = DatasetStats(metadata={}, parent=None)
        it = _make_test_iterator(stats)
        batch = _make_batch_with_timings(batching_start=1.0, batching_end=2.0)
        it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=3.0)
        assert stats.iter_blocked_batching_s.get() == pytest.approx(1.0)

    def test_no_collate_fn_zero_attribution(self):
        """collate stage has start_s=0 → skipped, 0 attribution."""
        stats = DatasetStats(metadata={}, parent=None)
        it = _make_test_iterator(stats)
        batch = _make_batch_with_timings(format_start=1.0, format_end=2.0)
        it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=3.0)
        assert stats.iter_blocked_format_s.get() == pytest.approx(1.0)
        assert stats.iter_blocked_collate_s.get() == 0.0

    def test_no_finalize_fn_zero_attribution(self):
        """finalize stage has start_s=0 → skipped, 0 attribution."""
        stats = DatasetStats(metadata={}, parent=None)
        it = _make_test_iterator(stats)
        batch = _make_batch_with_timings(collate_start=1.0, collate_end=2.0)
        it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=3.0)
        assert stats.iter_blocked_collate_s.get() == pytest.approx(1.0)
        assert stats.iter_blocked_finalize_s.get() == 0.0

    def test_prefetch_hides_fetch_from_training(self):
        """Effective prefetch: fetch done before training blocks → 0 fetch attribution."""
        stats = DatasetStats(metadata={}, parent=None)
        it = _make_test_iterator(stats)
        batch = _make_batch_with_timings(
            production_wait_start=0.0,
            production_wait_end=1.5,
            collate_start=2.3,
            collate_end=2.6,
        )
        # Training only starts blocking at t=2 (prefetch worked)
        it._attribute_blocked_time(batch, blocked_start_s=2.0, blocked_end_s=2.6)
        assert stats.iter_blocked_production_wait_s.get() == 0.0
        assert stats.iter_blocked_collate_s.get() == pytest.approx(0.3)

    def test_accumulation_across_batches(self):
        """Two batches each contribute to fetch — values accumulate."""
        stats = DatasetStats(metadata={}, parent=None)
        it = _make_test_iterator(stats)
        # Batch 1: fetch [0,1], blocked [0,2] → overlap 1.0
        b1 = _make_batch_with_timings(
            production_wait_start=0.0, production_wait_end=1.0, num_rows=10
        )
        it._attribute_blocked_time(b1, blocked_start_s=0.0, blocked_end_s=2.0)
        # Batch 2: fetch [5,6], blocked [5,7] → overlap 1.0
        b2 = _make_batch_with_timings(
            production_wait_start=5.0, production_wait_end=6.0, num_rows=20
        )
        it._attribute_blocked_time(b2, blocked_start_s=5.0, blocked_end_s=7.0)

        assert stats.iter_blocked_production_wait_s.get() == pytest.approx(2.0)
        assert stats.iter_batches_total == 2
        assert stats.iter_rows_total == 30

    def test_overlap_invariant_sum_leq_total(self):
        """sum(iter_blocked_*) <= iter_total_blocked_s holds for non-overlapping stages."""
        stats = DatasetStats(metadata={}, parent=None)
        it = _make_test_iterator(stats)
        stats.iter_total_blocked_s.add(5.0)
        batch = _make_batch_with_timings(
            production_wait_start=0.0,
            production_wait_end=1.0,
            batching_start=1.0,
            batching_end=2.0,
            format_start=2.0,
            format_end=3.0,
            num_rows=5,
        )
        it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=5.0)

        total = stats.iter_total_blocked_s.get()
        sum_stages = (
            stats.iter_blocked_production_wait_s.get()
            + stats.iter_blocked_batching_s.get()
            + stats.iter_blocked_format_s.get()
            + stats.iter_blocked_collate_s.get()
            + stats.iter_blocked_finalize_s.get()
        )
        assert sum_stages <= total + 1e-9

    def test_blocked_inside_stage(self):
        """Stage [0, 10] fully contains blocked [3, 5] → overlap = 2.0."""
        stats = DatasetStats(metadata={}, parent=None)
        it = _make_test_iterator(stats)
        batch = _make_batch_with_timings(
            production_wait_start=0.0, production_wait_end=10.0
        )
        it._attribute_blocked_time(batch, blocked_start_s=3.0, blocked_end_s=5.0)
        assert stats.iter_blocked_production_wait_s.get() == pytest.approx(2.0)

    def test_all_stages_simultaneous_overlap(self):
        """Multiple stages overlap with blocked window simultaneously."""
        stats = DatasetStats(metadata={}, parent=None)
        it = _make_test_iterator(stats)
        batch = _make_batch_with_timings(
            production_wait_start=0.0,
            production_wait_end=1.0,
            batching_start=1.0,
            batching_end=2.0,
            format_start=2.0,
            format_end=3.0,
            collate_start=3.0,
            collate_end=4.0,
            finalize_start=4.0,
            finalize_end=5.0,
            num_rows=100,
        )
        # Blocked window covers all stages
        it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=5.0)
        assert stats.iter_blocked_production_wait_s.get() == pytest.approx(1.0)
        assert stats.iter_blocked_batching_s.get() == pytest.approx(1.0)
        assert stats.iter_blocked_format_s.get() == pytest.approx(1.0)
        assert stats.iter_blocked_collate_s.get() == pytest.approx(1.0)
        assert stats.iter_blocked_finalize_s.get() == pytest.approx(1.0)
        assert stats.iter_batches_total == 1
        assert stats.iter_rows_total == 100

    def test_overlapping_spans_not_double_counted(self):
        """Two overlapping production_wait spans: union, not sum."""
        stats = DatasetStats(metadata={}, parent=None)
        it = _make_test_iterator(stats)
        # Block 1: prod [0, 100], Block 2: prod [50, 150] — overlap [50, 100]
        # Blocked [0, 200] covers both
        batch = _make_batch_with_timings(
            production_wait_start=0.0,
            production_wait_end=100.0,
            num_rows=10,
        )
        # Add a second production_wait span (multi-block batch)
        batch.metadata.stage_timings.production_wait.append(
            TimeSpan(start_s=50.0, end_s=150.0)
        )
        it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=200.0)
        # Union of [0,100] and [50,150] = [0,150] = 150, NOT 100+100=200
        assert stats.iter_blocked_production_wait_s.get() == pytest.approx(150.0)


def test_attribute_blocked_time_all_stages_full_overlap():
    """All stages with realistic timing, full overlap with blocked window."""
    stats = DatasetStats(metadata={}, parent=None)
    it = _make_test_iterator(stats)
    stats.iter_total_blocked_s.add(5.0)

    batch = _make_batch_with_timings(
        production_wait_start=0.0,
        production_wait_end=0.5,
        batching_start=0.5,
        batching_end=1.0,
        format_start=1.0,
        format_end=2.0,
        collate_start=2.0,
        collate_end=2.5,
        finalize_start=2.5,
        finalize_end=3.0,
        num_rows=256,
    )

    # Blocked window covers all stages
    it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=5.0)

    # Each stage gets its full duration
    assert stats.iter_blocked_production_wait_s.get() == pytest.approx(0.5)
    assert stats.iter_blocked_batching_s.get() == pytest.approx(0.5)
    assert stats.iter_blocked_format_s.get() == pytest.approx(1.0)
    assert stats.iter_blocked_collate_s.get() == pytest.approx(0.5)
    assert stats.iter_blocked_finalize_s.get() == pytest.approx(0.5)
    assert stats.iter_batches_total == 1
    assert stats.iter_rows_total == 256

    # Invariant: sum = 3.0 <= total_blocked = 5.0
    sum_stages = (
        stats.iter_blocked_production_wait_s.get()
        + stats.iter_blocked_batching_s.get()
        + stats.iter_blocked_format_s.get()
        + stats.iter_blocked_collate_s.get()
        + stats.iter_blocked_finalize_s.get()
    )
    assert sum_stages == pytest.approx(3.0)
    assert sum_stages <= stats.iter_total_blocked_s.get() + 1e-9


def test_finalize_fn_uses_single_thread(ray_start_regular_shared):
    """Tests that finalize_fn is not run with multiple threads."""
    ref_bundles_iter = ref_bundle_generator(num_blocks=20, num_rows=2)

    q = queue.Queue()
    semaphore = threading.Semaphore(value=1)

    def finalize_enforce_single_thread(batch):
        already_acquired = not semaphore.acquire(blocking=False)
        if already_acquired:
            e = AssertionError("finalize_fn is being run concurrently.")
            q.put(e, block=True)
        semaphore.release()
        return batch

    # Test that finalize_fn is called in a single thread,
    # even if prefetch_batches is set.
    output_batches = BatchIterator(
        ref_bundles_iter,
        collate_fn=lambda batch: batch,
        finalize_fn=finalize_enforce_single_thread,
        prefetch_batches=4,
    )

    # Force execution of the iterator.
    # This step should not raise an exception.
    list(output_batches)

    try:
        e = q.get(block=False, timeout=0.1)
        raise e
    except queue.Empty:
        pass


# Test for 3 cases
# 1. Batch size is less than block size
# 2. Batch size is more than block size
# 3. Block size is not divisble by batch size
@pytest.mark.parametrize("batch_size", [1, 4, 3])
@pytest.mark.parametrize("drop_last", [True, False])
@pytest.mark.parametrize("prefetch_batches", [0, 1])
def test_iter_batches_e2e(
    ray_start_regular_shared, batch_size, drop_last, prefetch_batches
):
    def collate_fn(batch: pd.DataFrame):
        return batch + 1

    ref_bundles_iter = ref_bundle_generator(num_blocks=4, num_rows=2)

    output_batches = BatchIterator(
        ref_bundles_iter,
        batch_size=batch_size,
        prefetch_batches=prefetch_batches,
        batch_format="pandas",
        collate_fn=collate_fn,
        drop_last=drop_last,
        preserve_order=True,
    )

    output_batches = list(output_batches)

    assert len(output_batches) > 0
    for df in output_batches:
        # Check batch formatting.
        assert isinstance(df, pd.DataFrame)
        # Check batch size.
        if batch_size == 3 and not drop_last:
            assert len(df) in {2, 3}
        else:
            assert len(df) == batch_size

    concat_df = pd.concat(output_batches)
    # Test that collate_fn is applied.
    assert concat_df["foo"].iloc[0] == 1
    # Make sure order is preserved.
    for i in range(len(concat_df) - 1):
        assert concat_df["foo"].iloc[i + 1] >= concat_df["foo"].iloc[i]


def test_iter_batches_counts_rows_at_pipeline_exit(ray_start_regular_shared):
    stats = DatasetStats(metadata={}, parent=None)
    ref_bundles_iter = ref_bundle_generator(num_blocks=4, num_rows=2)

    output_batches = list(
        BatchIterator(
            ref_bundles_iter,
            stats=stats,
            batch_size=3,
            prefetch_batches=0,
            batch_format="pandas",
            drop_last=True,
        )
    )

    assert len(output_batches) == 2
    assert [len(batch) for batch in output_batches] == [3, 3]
    assert stats.iter_batches_total == 2
    assert stats.iter_rows_total == 6


def test_iter_batches_e2e_async(ray_start_regular_shared):
    """We add time.sleep in 3 places:
    1. In the base generator to simulate streaming executor blocking on next results.
    2. In the collate_fn to simulate expensive slicing/formatting/collation
    3. In the user thread to simulate training.
    """

    def collate_fn(batch):
        time.sleep(2)
        return batch

    ref_bundles = ref_bundle_generator(num_blocks=20, num_rows=2)
    start_time = time.time()
    output_batches = BatchIterator(
        ref_bundles,
        batch_size=None,
        collate_fn=collate_fn,
        prefetch_batches=4,
    )
    batches = []
    for batch in output_batches:
        time.sleep(1.5)
        batches.append(batch)
    end_time = time.time()

    # 20 batches, 1.5 second sleep. Should be less than 45 seconds, even with some
    # overhead.
    # If there was no overlap, then we would expect this to take at least 20*2.5 = 50
    assert end_time - start_time < 45, end_time - start_time

    assert len(batches) == 20
    assert all(len(batch) == 2 for batch in batches)


@pytest.mark.parametrize("preserve_order", [True, False])
def test_iter_batches_preserve_order_flag(
    ray_start_regular_shared, preserve_order, restore_data_context
):
    """When `execution_options.preserve_order` is True, batches must come
    out in input order even with a multi-worker format threadpool. When
    False, ordering is not guaranteed (but the full set of batches must
    still be produced)."""
    # Variable per-batch collate cost makes worker-completion order
    # arbitrary so the reorder path actually does work when enabled.
    def collate_fn(batch):
        idx = int(batch["foo"][0])
        time.sleep(0.05 * (idx % 4))
        return batch

    num_blocks = 16
    ref_bundles = ref_bundle_generator(num_blocks=num_blocks, num_rows=1)
    output_batches = list(
        BatchIterator(
            ref_bundles,
            batch_size=1,
            collate_fn=collate_fn,
            batch_format="pandas",
            prefetch_batches=4,
            preserve_order=preserve_order,
        )
    )

    indices = [int(df["foo"].iloc[0]) for df in output_batches]
    assert sorted(indices) == list(range(num_blocks))
    if preserve_order:
        assert indices == list(range(num_blocks)), indices


def test_finalize_fn_runs_after_restore_original_order(ray_start_regular_shared):
    """When preserve_order=True, finalize_fn must run after the reorder
    buffer so that the buffer holds CPU batches rather than finalize_fn
    outputs (e.g., GPU tensors). Asserts finalize_fn sees batches in
    monotonically increasing order even when the format/collate threadpool
    completes them out of order."""

    def collate_fn(batch):
        # Variable per-batch cost so worker-completion order is arbitrary.
        idx = int(batch["foo"].iloc[0])
        time.sleep(0.05 * (idx % 4))
        return batch

    seen_by_finalize = []
    seen_lock = threading.Lock()

    def finalize_fn(batch):
        idx = int(batch["foo"].iloc[0])
        with seen_lock:
            seen_by_finalize.append(idx)
        return batch

    num_blocks = 16
    ref_bundles = ref_bundle_generator(num_blocks=num_blocks, num_rows=1)
    list(
        BatchIterator(
            ref_bundles,
            batch_size=1,
            collate_fn=collate_fn,
            finalize_fn=finalize_fn,
            batch_format="pandas",
            prefetch_batches=4,
            preserve_order=True,
        )
    )

    assert seen_by_finalize == list(range(num_blocks)), seen_by_finalize


def _ref_bundles_with_size(
    num_blocks: int, num_rows: int, size_bytes_per_block: int
) -> Iterator[RefBundle]:
    """Create ref bundles with explicit size_bytes for testing."""
    for i in range(num_blocks):
        block = pa.table({"foo": [i] * num_rows})
        metadata = BlockMetadata(
            num_rows=num_rows,
            size_bytes=size_bytes_per_block,
            input_files=[],
            exec_stats=None,
        )
        schema = block.schema
        yield RefBundle(
            blocks=(BlockEntry(ray.put(block), metadata),),
            owns_blocks=True,
            schema=schema,
        )


@pytest.mark.parametrize(
    "num_batches_to_prefetch,expected_bytes_sequence",
    [
        # No prefetching: all 5 blocks report 0 prefetched bytes
        (0, [0, 0, 0, 0, 0]),
        # prefetch 2 blocks: with 5 blocks of 100 bytes each
        # After yield block 0: window has 1,2 -> 200 (added block 2)
        # After yield block 1: window has 2,3 -> 200 (added block 3)
        # After yield block 2: window has 3,4 -> 200 (added block 4)
        # After yield block 3: window has 4 -> 100 (no more to add)
        # After yield block 4: window empty -> 0
        (2, [200, 200, 200, 100, 0]),
    ],
)
def test_prefetch_bytes_tracking(
    ray_start_regular_shared, num_batches_to_prefetch, expected_bytes_sequence
):
    """Test iter_prefetched_bytes is set correctly during prefetching.

    Tests prefetch_batches_locally directly to verify exact values,
    bypassing async BatchIterator which has non-deterministic timing.
    """
    stats = DatasetStats(metadata={}, parent=None)

    # Create 5 ref bundles, each with size_bytes=100
    num_blocks = 5
    ref_bundles = list(
        _ref_bundles_with_size(num_blocks, num_rows=2, size_bytes_per_block=100)
    )

    prefetcher = WaitBlockPrefetcher()
    block_iter = prefetch_batches_locally(
        iter(ref_bundles),
        prefetcher=prefetcher,
        num_batches_to_prefetch=num_batches_to_prefetch,
        batch_size=None,
        stats=stats,
    )

    # Track iter_prefetched_bytes after each block is yielded
    recorded_bytes = []
    for _ in block_iter:
        recorded_bytes.append(stats.iter_prefetched_bytes)

    assert recorded_bytes == expected_bytes_sequence, f"Got {recorded_bytes}"


@pytest.mark.parametrize("prefetch_batches", [0, 2])
def test_prefetch_bytes_callback(ray_start_regular_shared, prefetch_batches):
    """Test prefetch_bytes_callback is invoked correctly by BatchIterator."""
    reported_bytes = []

    def prefetch_callback(num_bytes: int):
        reported_bytes.append(num_bytes)

    stats = DatasetStats(metadata={}, parent=None)

    # Create 5 ref bundles
    num_blocks = 5
    ref_bundles = list(
        _ref_bundles_with_size(num_blocks, num_rows=2, size_bytes_per_block=100)
    )

    output_batches = BatchIterator(
        iter(ref_bundles),
        stats=stats,
        batch_size=None,
        prefetch_batches=prefetch_batches,
        prefetch_bytes_callback=prefetch_callback,
    )

    # Consume all batches
    batches = list(output_batches)

    assert len(batches) == 5

    # Callback is called 5 times (per batch) + 1 time at epoch end
    assert len(reported_bytes) == 6, f"Expected 6, got {len(reported_bytes)}"

    # All values should be non-negative
    assert all(b >= 0 for b in reported_bytes), f"Negative: {reported_bytes}"

    # Last value should be 0 (after_epoch_end)
    assert reported_bytes[-1] == 0, f"Last should be 0: {reported_bytes}"


@pytest.mark.parametrize(
    "scenario,bound_stage",
    [
        ("production", "iter_blocked_production_wait_s"),
        ("data_transfer", "iter_blocked_data_transfer_s"),
        ("batching", "iter_blocked_batching_s"),
        ("collate", "iter_blocked_collate_s"),
        ("format", "iter_blocked_format_s"),
        ("finalize", "iter_blocked_finalize_s"),
    ],
)
def test_e2e_blocked_attribution_by_scenario(
    ray_start_regular_shared, scenario, bound_stage
):
    """E2e: when a specific stage is the bottleneck, its blocked metric
    should be the largest among all stages, and at least SLEEP_S."""
    from ray.data._internal.stats import _StatsManager

    iter_kwargs = {"batch_size": 10, "prefetch_batches": 0}
    patches = []

    if scenario == "production":
        # Slow upstream map → production_wait dominates.
        def slow_map(batch):
            time.sleep(SLEEP_S)
            return batch

        ds = ray.data.range(50, override_num_blocks=5).map(slow_map)

    elif scenario == "data_transfer":
        # Patch ray.get ONLY in util.resolve_block_refs (not globally) so the
        # streaming executor's own ray.get calls aren't slowed (which would
        # inflate production_wait). We replace util_mod.ray with a proxy that
        # has a slow `get` but delegates everything else to the real ray.
        from ray.data._internal.block_batching import util as util_mod

        orig_get = ray.get

        class _SlowGetRayProxy:
            """Proxy that sleeps on `.get` but delegates everything else."""

            def __getattr__(self, name):
                return getattr(ray, name)

            @staticmethod
            def get(ref):
                time.sleep(SLEEP_S)
                return orig_get(ref)

        patches.append(patch.object(util_mod, "ray", _SlowGetRayProxy()))
        ds = ray.data.range(50, override_num_blocks=5)

    elif scenario == "batching":
        # Patch Batcher.next_batch to inject slow batching.
        from ray.data._internal.batcher import Batcher

        orig_next_batch = Batcher.next_batch

        def slow_next_batch(self):
            time.sleep(SLEEP_S)
            return orig_next_batch(self)

        patches.append(patch.object(Batcher, "next_batch", slow_next_batch))
        ds = ray.data.range(50, override_num_blocks=5)

    elif scenario == "collate":
        # Pass _collate_fn via _iter_batches (private signature accepts it;
        # public iter_batches does not).
        def slow_collate(batch):
            time.sleep(SLEEP_S)
            return batch

        iter_kwargs["_collate_fn"] = slow_collate
        ds = ray.data.range(50, override_num_blocks=5)

    elif scenario == "format":
        # Patch BlockAccessor.to_batch_format — it's called INSIDE
        # _format_batch's _maybe_time context, so the sleep is captured by
        # the format timing span.
        orig_to_batch_format = BlockAccessor.to_batch_format

        def slow_to_batch_format(self, batch_format):
            time.sleep(SLEEP_S)
            return orig_to_batch_format(self, batch_format)

        patches.append(
            patch.object(BlockAccessor, "to_batch_format", slow_to_batch_format)
        )
        ds = ray.data.range(50, override_num_blocks=5)

    elif scenario == "finalize":
        # Pass _finalize_fn via _iter_batches (private signature accepts it).
        def slow_finalize(data):
            time.sleep(SLEEP_S)
            return data

        iter_kwargs["_finalize_fn"] = slow_finalize
        ds = ray.data.range(50, override_num_blocks=5)

    it = ds.iterator()
    captured = []
    orig = _StatsManager.update_iteration_metrics

    def spy(stats, dataset_tag):
        captured.append(stats)
        return orig(stats, dataset_tag)

    patches.append(patch.object(_StatsManager, "update_iteration_metrics", spy))

    import contextlib

    with contextlib.ExitStack() as stack:
        for p in patches:
            stack.enter_context(p)
        # Use _iter_batches (private) so we can pass _collate_fn / _finalize_fn
        # which the public iter_batches signature does not expose.
        for _ in it._iter_batches(**iter_kwargs):
            pass

    stats = captured[-1]
    all_stages = [
        stats.iter_blocked_production_wait_s.get(),
        stats.iter_blocked_data_transfer_s.get(),
        stats.iter_blocked_batching_s.get(),
        stats.iter_blocked_format_s.get(),
        stats.iter_blocked_collate_s.get(),
        stats.iter_blocked_finalize_s.get(),
    ]
    bound_value = getattr(stats, bound_stage).get()
    # The bottleneck stage should be at least the sleep time we injected,
    # proving the timing capture is actually recording the stall.
    assert bound_value >= SLEEP_S, (
        f"{scenario}-bound: {bound_stage}={bound_value} < SLEEP_S={SLEEP_S}; "
        "timing capture missed the injected stall"
    )
    # The bottleneck stage should be strictly greater than all others.
    for v in all_stages:
        if v == bound_value:
            continue
        assert (
            bound_value > v
        ), f"{scenario}-bound: {bound_stage}={bound_value} not > {v}"


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))