[Data] [stats] Add RefBundle retrieval time metric to iterator dataset stats (#58422)

## Why These Changes Are Needed

This PR adds a new metric to track the time spent retrieving `RefBundle`
objects during dataset iteration. This metric provides better visibility
into the performance breakdown of batch iteration, specifically
capturing the time spent in `get_next_ref_bundle()` calls within the
`prefetch_batches_locally` function.

## Related Issue Number

N/A

## Example

```
  dataloader/train = {'producer_throughput': 8361.841782656593, 'iter_stats': {'prefetch_block-avg': inf, 'prefetch_block-min': inf, 'prefetch_block-max': 0, 'prefetch_block-total': 0, 'get_ref_bundles-avg': 0.05172277254545271, 'get_ref_bundles-min': 1.1991999997462699e-05, 'get_ref_bundles-max': 11.057470971999976, 'get_ref_bundles-total': 15.361663445999454, 'fetch_block-avg': 0.31572694455743233, 'fetch_block-min': 0.0006362799999806157, 'fetch_block-max': 2.1665870369999993, 'fetch_block-total': 93.45517558899996, 'block_to_batch-avg': 0.001048687573988573, 'block_to_batch-min': 2.10620000302697e-05, 'block_to_batch-max': 0.049948245999985375, 'block_to_batch-total': 2.048086831999683, 'format_batch-avg': 0.0001013781433686053, 'format_batch-min': 1.415700000961806e-05, 'format_batch-max': 0.009682661999988795, 'format_batch-total': 0.19799151399888615, 'collate-avg': 0.01303446213312943, 'collate-min': 0.00025646699998560507, 'collate-max': 0.9855495820000328, 'collate-total': 25.456304546001775, 'finalize-avg': 0.012211385266257683, 'finalize-min': 0.004209667999987232, 'finalize-max': 0.3785081949999949, 'finalize-total': 23.848835425001255, 'time_spent_blocked-avg': 0.04783407008137157, 'time_spent_blocked-min': 1.2316999971062614e-05, 'time_spent_blocked-max': 12.46102861700001, 'time_spent_blocked-total': 93.46777293900004, 'time_spent_training-avg': 0.015053571562211652, 'time_spent_training-min': 1.3704999958008557e-05, 'time_spent_training-max': 1.079616685000019, 'time_spent_training-total': 29.399625260999358}}
```

## Checks

- [ ] I've signed off every commit(by using the -s flag, i.e., `git
commit -s`) in this PR.
- [ ] I've run `scripts/format.sh` to lint the changes in this PR.
- [ ] I've included any doc changes needed for
https://docs.ray.io/en/master/.
- [ ] I've added any new APIs to the API Reference. For example, if I
added a
method in Tune, I've added it in `doc/source/tune/api/` under the
corresponding `.rst` file.
- [ ] I've made sure the tests are passing. Note that there might be a
few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [x] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(

---------

Signed-off-by: xgui <[email protected]>
Signed-off-by: Xinyuan <[email protected]>

Author: xinyuangui2

python/ray/data/_internal/block_batching/iter_batches.py

@@ -146,6 +146,7 @@ def _prefetch_blocks(
             num_batches_to_prefetch=self._prefetch_batches,
             batch_size=self._batch_size,
             eager_free=self._eager_free,
+            stats=self._stats,
         )
 
     def _resolve_block_refs(
@@ -322,6 +323,7 @@ def prefetch_batches_locally(
     num_batches_to_prefetch: int,
     batch_size: Optional[int],
     eager_free: bool = False,
+    stats: Optional[DatasetStats] = None,
 ) -> Iterator[ObjectRef[Block]]:
     """Given an iterator of batched RefBundles, returns an iterator over the
     corresponding block references while prefetching `num_batches_to_prefetch`
@@ -334,8 +336,13 @@ def prefetch_batches_locally(
             current batch during the scan.
         batch_size: User specified batch size, or None to let the system pick.
         eager_free: Whether to eagerly free the object reference from the object store.
+        stats: Dataset stats object used to store ref bundle retrieval time.
     """
 
+    def get_next_ref_bundle() -> RefBundle:
+        with stats.iter_get_ref_bundles_s.timer() if stats else nullcontext():
+            return next(ref_bundles)
+
     sliding_window = collections.deque()
     current_window_size = 0
 
@@ -358,7 +365,7 @@ def prefetch_batches_locally(
         batch_size is None and len(sliding_window) < num_batches_to_prefetch
     ):
         try:
-            next_ref_bundle = next(ref_bundles)
+            next_ref_bundle = get_next_ref_bundle()
             sliding_window.extend(next_ref_bundle.blocks)
             current_window_size += next_ref_bundle.num_rows()
         except StopIteration:
@@ -371,7 +378,7 @@ def prefetch_batches_locally(
         current_window_size -= metadata.num_rows
         if batch_size is None or current_window_size < num_rows_to_prefetch:
             try:
-                next_ref_bundle = next(ref_bundles)
+                next_ref_bundle = get_next_ref_bundle()
                 for block_ref_and_md in next_ref_bundle.blocks:
                     sliding_window.append(block_ref_and_md)
                     current_window_size += block_ref_and_md[1].num_rows

python/ray/data/_internal/stats.py

@@ -327,6 +327,11 @@ def __init__(self, max_stats=1000):
             description="Seconds spent in iterator initialization code",
             tag_keys=iter_tag_keys,
         )
+        self.iter_get_ref_bundles_s = Gauge(
+            "data_iter_get_ref_bundles_seconds",
+            description="Seconds spent getting RefBundles from the dataset iterator",
+            tag_keys=iter_tag_keys,
+        )
         self.iter_get_s = Gauge(
             "data_iter_get_seconds",
             description="Seconds spent in ray.get() while resolving block references",
@@ -565,6 +570,7 @@ def update_iteration_metrics(
         tags = self._create_tags(dataset_tag)
 
         self.iter_initialize_s.set(stats.iter_initialize_s.get(), tags)
+        self.iter_get_ref_bundles_s.set(stats.iter_get_ref_bundles_s.get(), tags)
         self.iter_get_s.set(stats.iter_get_s.get(), tags)
         self.iter_next_batch_s.set(stats.iter_next_batch_s.get(), tags)
         self.iter_format_batch_s.set(stats.iter_format_batch_s.get(), tags)
@@ -1098,6 +1104,7 @@ def __init__(
 
         # Iteration stats, filled out if the user iterates over the dataset.
         self.iter_wait_s: Timer = Timer()
+        self.iter_get_ref_bundles_s: Timer = Timer()
         self.iter_get_s: Timer = Timer()
         self.iter_next_batch_s: Timer = Timer()
         self.iter_format_batch_s: Timer = Timer()
@@ -1146,6 +1153,7 @@ def to_summary(self) -> "DatasetStatsSummary":
 
         iter_stats = IterStatsSummary(
             self.iter_wait_s,
+            self.iter_get_ref_bundles_s,
             self.iter_get_s,
             self.iter_next_batch_s,
             self.iter_format_batch_s,
@@ -1843,6 +1851,8 @@ def __repr__(self, level=0) -> str:
 class IterStatsSummary:
     # Time spent in actor based prefetching, in seconds.
     wait_time: Timer
+    # Time spent getting RefBundles from the dataset iterator, in seconds
+    get_ref_bundles_time: Timer
     # Time spent in `ray.get()`, in seconds
     get_time: Timer
     # Time spent in batch building, in seconds
@@ -1880,6 +1890,7 @@ def to_string(self) -> str:
             self.block_time.get()
             or self.time_to_first_batch.get()
             or self.total_time.get()
+            or self.get_ref_bundles_time.get()
             or self.get_time.get()
             or self.next_time.get()
             or self.format_time.get()
@@ -1911,6 +1922,13 @@ def to_string(self) -> str:
             out += (
                 "* Batch iteration time breakdown (summed across prefetch threads):\n"
             )
+            if self.get_ref_bundles_time.get():
+                out += "    * In get RefBundles: {} min, {} max, {} avg, {} total\n".format(
+                    fmt(self.get_ref_bundles_time.min()),
+                    fmt(self.get_ref_bundles_time.max()),
+                    fmt(self.get_ref_bundles_time.avg()),
+                    fmt(self.get_ref_bundles_time.get()),
+                )
             if self.get_time.get():
                 out += "    * In ray.get(): {} min, {} max, {} avg, {} total\n".format(
                     fmt(self.get_time.min()),
@@ -1973,6 +1991,7 @@ def __repr__(self, level=0) -> str:
         return (
             f"IterStatsSummary(\n"
             f"{indent}   wait_time={fmt(self.wait_time.get()) or None},\n"
+            f"{indent}   get_ref_bundles_time={fmt(self.get_ref_bundles_time.get()) or None},\n"
             f"{indent}   get_time={fmt(self.get_time.get()) or None},\n"
             f"{indent}   iter_blocks_local={self.iter_blocks_local or None},\n"
             f"{indent}   iter_blocks_remote={self.iter_blocks_remote or None},\n"

python/ray/data/tests/test_stats.py

@@ -501,6 +501,7 @@ def test_streaming_split_stats(ray_start_regular_shared, restore_data_context):
     * Total time spent waiting for the first batch after starting iteration: T
     * Total execution time for user thread: T
 * Batch iteration time breakdown (summed across prefetch threads):
+    * In get RefBundles: T min, T max, T avg, T total
     * In ray.get(): T min, T max, T avg, T total
     * In batch creation: T min, T max, T avg, T total
     * In batch formatting: T min, T max, T avg, T total
@@ -696,6 +697,7 @@ def test_dataset_stats_basic(
         f"    * Total time spent waiting for the first batch after starting iteration: T\n"
         f"    * Total execution time for user thread: T\n"
         f"* Batch iteration time breakdown (summed across prefetch threads):\n"
+        f"    * In get RefBundles: T min, T max, T avg, T total\n"
         f"    * In ray.get(): T min, T max, T avg, T total\n"
         f"    * In batch creation: T min, T max, T avg, T total\n"
         f"    * In batch formatting: T min, T max, T avg, T total\n"
@@ -740,6 +742,7 @@ def test_block_location_nums(ray_start_regular_shared, restore_data_context):
         f"    * Total time spent waiting for the first batch after starting iteration: T\n"
         f"    * Total execution time for user thread: T\n"
         f"* Batch iteration time breakdown (summed across prefetch threads):\n"
+        f"    * In get RefBundles: T min, T max, T avg, T total\n"
         f"    * In ray.get(): T min, T max, T avg, T total\n"
         f"    * In batch creation: T min, T max, T avg, T total\n"
         f"    * In batch formatting: T min, T max, T avg, T total\n"
@@ -845,6 +848,7 @@ def test_dataset__repr__(ray_start_regular_shared, restore_data_context):
         "   ],\n"
         "   iter_stats=IterStatsSummary(\n"
         "      wait_time=T,\n"
+        "      get_ref_bundles_time=T,\n"
         "      get_time=T,\n"
         "      iter_blocks_local=None,\n"
         "      iter_blocks_remote=None,\n"
@@ -866,6 +870,7 @@ def test_dataset__repr__(ray_start_regular_shared, restore_data_context):
         "         operators_stats=[],\n"
         "         iter_stats=IterStatsSummary(\n"
         "            wait_time=T,\n"
+        "            get_ref_bundles_time=T,\n"
         "            get_time=T,\n"
         "            iter_blocks_local=None,\n"
         "            iter_blocks_remote=None,\n"
@@ -985,6 +990,7 @@ def check_stats():
         "   ],\n"
         "   iter_stats=IterStatsSummary(\n"
         "      wait_time=T,\n"
+        "      get_ref_bundles_time=T,\n"
         "      get_time=T,\n"
         "      iter_blocks_local=None,\n"
         "      iter_blocks_remote=None,\n"
@@ -1080,6 +1086,7 @@ def check_stats():
         "         ],\n"
         "         iter_stats=IterStatsSummary(\n"
         "            wait_time=T,\n"
+        "            get_ref_bundles_time=T,\n"
         "            get_time=T,\n"
         "            iter_blocks_local=None,\n"
         "            iter_blocks_remote=None,\n"
@@ -1101,6 +1108,7 @@ def check_stats():
         "               operators_stats=[],\n"
         "               iter_stats=IterStatsSummary(\n"
         "                  wait_time=T,\n"
+        "                  get_ref_bundles_time=T,\n"
         "                  get_time=T,\n"
         "                  iter_blocks_local=None,\n"
         "                  iter_blocks_remote=None,\n"
@@ -1537,6 +1545,7 @@ def test_streaming_stats_full(ray_start_regular_shared, restore_data_context):
     * Total time spent waiting for the first batch after starting iteration: T
     * Total execution time for user thread: T
 * Batch iteration time breakdown (summed across prefetch threads):
+    * In get RefBundles: T min, T max, T avg, T total
     * In ray.get(): T min, T max, T avg, T total
     * In batch creation: T min, T max, T avg, T total
     * In batch formatting: T min, T max, T avg, T total

release/train_tests/benchmark/ray_dataloader_factory.py

@@ -123,6 +123,10 @@ def get_metrics(self) -> Dict[str, Any]:
                     "prefetch_block-min": iter_stats.wait_time.min(),
                     "prefetch_block-max": iter_stats.wait_time.max(),
                     "prefetch_block-total": iter_stats.wait_time.get(),
+                    "get_ref_bundles-avg": iter_stats.get_ref_bundles_time.avg(),
+                    "get_ref_bundles-min": iter_stats.get_ref_bundles_time.min(),
+                    "get_ref_bundles-max": iter_stats.get_ref_bundles_time.max(),
+                    "get_ref_bundles-total": iter_stats.get_ref_bundles_time.get(),
                     "fetch_block-avg": iter_stats.get_time.avg(),
                     "fetch_block-min": iter_stats.get_time.min(),
                     "fetch_block-max": iter_stats.get_time.max(),