[Data] Add iter_prefetched_bytes stats (#58900)

> Thank you for contributing to Ray! 🚀
> Please review the [Ray Contribution
Guide](https://docs.ray.io/en/master/ray-contribute/getting-involved.html)
before opening a pull request.

> ⚠️ Remove these instructions before submitting your PR.

> 💡 Tip: Mark as draft if you want early feedback, or ready for review
when it's complete.

## Description
> Briefly describe what this PR accomplishes and why it's needed.

### [Data] Add iter_prefetched_blocks stats

Report prefetched bytes per iterator as stats.


## Related issues
> Link related issues: "Fixes #1234", "Closes #1234", or "Related to
#1234".

## Additional information
> Optional: Add implementation details, API changes, usage examples,
screenshots, etc.

---------

Signed-off-by: Srinath Krishnamachari <[email protected]>
Signed-off-by: Srinath Krishnamachari <[email protected]>

Author: srinathk10

python/ray/dashboard/modules/metrics/dashboards/data_dashboard_panels.py

@@ -1050,6 +1050,21 @@
     stack=False,
 )
 
+ITERATION_PREFETCHED_BYTES_PANEL = Panel(
+    id=90,
+    title="Iteration Prefetched Bytes",
+    description="Current bytes of prefetched blocks in the iterator",
+    unit="bytes",
+    targets=[
+        Target(
+            expr="sum(ray_data_iter_prefetched_bytes{{{global_filters}}}) by (dataset)",
+            legend="Prefetched Bytes: {{dataset}}",
+        )
+    ],
+    fill=0,
+    stack=False,
+)
+
 ITERATION_TIME_TO_FIRST_BATCH_PANEL = Panel(
     id=120,
     title="Iteration Time to First Batch",
@@ -1420,6 +1435,7 @@
             ITERATION_BLOCKS_LOCAL_PANEL,
             ITERATION_BLOCKS_REMOTE_PANEL,
             ITERATION_BLOCKS_UNKNOWN_LOCATION_PANEL,
+            ITERATION_PREFETCHED_BYTES_PANEL,
             ITERATION_TIME_TO_FIRST_BATCH_PANEL,
             ITERATION_GET_REF_BUNDLES_PANEL,
         ],

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

@@ -373,6 +373,8 @@ def get_next_ref_bundle() -> RefBundle:
     current_window_size = 0
 
     if num_batches_to_prefetch <= 0:
+        if stats:
+            stats.iter_prefetched_bytes = 0
         for ref_bundle in ref_bundles:
             for block_ref in ref_bundle.block_refs:
                 yield block_ref
@@ -398,6 +400,10 @@ def get_next_ref_bundle() -> RefBundle:
             break
 
     prefetcher.prefetch_blocks([block_ref for block_ref, _ in list(sliding_window)])
+    if stats:
+        stats.iter_prefetched_bytes = sum(
+            metadata.size_bytes or 0 for _, metadata in sliding_window
+        )
 
     while sliding_window:
         block_ref, metadata = sliding_window.popleft()
@@ -413,6 +419,10 @@ def get_next_ref_bundle() -> RefBundle:
                 )
             except StopIteration:
                 pass
+        if stats:
+            stats.iter_prefetched_bytes = sum(
+                metadata.size_bytes or 0 for _, metadata in sliding_window
+            )
         yield block_ref
         trace_deallocation(block_ref, loc="iter_batches", free=eager_free)
     prefetcher.stop()

python/ray/data/_internal/stats.py

@@ -381,6 +381,11 @@ def __init__(self, max_stats=1000):
             description="Number of blocks that have unknown locations",
             tag_keys=iter_tag_keys,
         )
+        self.iter_prefetched_bytes = Gauge(
+            "data_iter_prefetched_bytes",
+            description="Current bytes of prefetched blocks in the iterator",
+            tag_keys=iter_tag_keys,
+        )
 
         # === Dataset and Operator Metadata Metrics ===
         dataset_tags = ("dataset", "job_id", "start_time")
@@ -568,6 +573,7 @@ def update_iteration_metrics(
         self.iter_blocks_local.set(stats.iter_blocks_local, tags)
         self.iter_blocks_remote.set(stats.iter_blocks_remote, tags)
         self.iter_unknown_location.set(stats.iter_unknown_location, tags)
+        self.iter_prefetched_bytes.set(stats.iter_prefetched_bytes, tags)
 
         self.iter_block_fetching_s.set(stats.iter_get_s.get(), tags)
         self.iter_batch_shaping_s.set(stats.iter_next_batch_s.get(), tags)
@@ -976,6 +982,7 @@ def __init__(
         self.iter_blocks_local: int = 0
         self.iter_blocks_remote: int = 0
         self.iter_unknown_location: int = 0
+        self.iter_prefetched_bytes: int = 0
 
         # Memory usage stats
         self.global_bytes_spilled: int = 0
@@ -1018,6 +1025,7 @@ def to_summary(self) -> "DatasetStatsSummary":
             self.iter_blocks_local,
             self.iter_blocks_remote,
             self.iter_unknown_location,
+            self.iter_prefetched_bytes,
         )
 
         stats_summary_parents = []
@@ -1730,6 +1738,8 @@ class IterStatsSummary:
     iter_blocks_remote: int
     # Num of blocks with unknown locations
     iter_unknown_location: int
+    # Current bytes of prefetched blocks in the iterator
+    iter_prefetched_bytes: int
 
     def __str__(self) -> str:
         return self.to_string()
@@ -1830,6 +1840,8 @@ def to_string(self) -> str:
                 out += "    * Num blocks unknown location: {}\n".format(
                     self.iter_unknown_location
                 )
+            if self.iter_prefetched_bytes:
+                out += "    * Prefetched bytes: {}\n".format(self.iter_prefetched_bytes)
             if self.streaming_split_coord_time.get() != 0:
                 out += "Streaming split coordinator overhead time: "
                 out += f"{fmt(self.streaming_split_coord_time.get())}\n"
@@ -1846,6 +1858,7 @@ def __repr__(self, level=0) -> str:
             f"{indent}   iter_blocks_local={self.iter_blocks_local or None},\n"
             f"{indent}   iter_blocks_remote={self.iter_blocks_remote or None},\n"
             f"{indent}   iter_unknown_location={self.iter_unknown_location or None},\n"
+            f"{indent}   iter_prefetched_bytes={self.iter_prefetched_bytes or None},\n"
             f"{indent}   next_time={fmt(self.next_time.get()) or None},\n"
             f"{indent}   format_time={fmt(self.format_time.get()) or None},\n"
             f"{indent}   user_time={fmt(self.user_time.get()) or None},\n"

python/ray/data/tests/test_stats.py

@@ -791,6 +791,7 @@ def test_dataset__repr__(ray_start_regular_shared, restore_data_context):
         "      iter_blocks_local=None,\n"
         "      iter_blocks_remote=None,\n"
         "      iter_unknown_location=None,\n"
+        "      iter_prefetched_bytes=None,\n"
         "      next_time=T,\n"
         "      format_time=T,\n"
         "      user_time=T,\n"
@@ -813,6 +814,7 @@ def test_dataset__repr__(ray_start_regular_shared, restore_data_context):
         "            iter_blocks_local=None,\n"
         "            iter_blocks_remote=None,\n"
         "            iter_unknown_location=None,\n"
+        "            iter_prefetched_bytes=None,\n"
         "            next_time=T,\n"
         "            format_time=T,\n"
         "            user_time=T,\n"
@@ -933,6 +935,7 @@ def check_stats():
         "      iter_blocks_local=None,\n"
         "      iter_blocks_remote=None,\n"
         "      iter_unknown_location=N,\n"
+        "      iter_prefetched_bytes=None,\n"
         "      next_time=T,\n"
         "      format_time=T,\n"
         "      user_time=T,\n"
@@ -1029,6 +1032,7 @@ def check_stats():
         "            iter_blocks_local=None,\n"
         "            iter_blocks_remote=None,\n"
         "            iter_unknown_location=None,\n"
+        "            iter_prefetched_bytes=None,\n"
         "            next_time=T,\n"
         "            format_time=T,\n"
         "            user_time=T,\n"
@@ -1051,6 +1055,7 @@ def check_stats():
         "                  iter_blocks_local=None,\n"
         "                  iter_blocks_remote=None,\n"
         "                  iter_unknown_location=None,\n"
+        "                  iter_prefetched_bytes=None,\n"
         "                  next_time=T,\n"
         "                  format_time=T,\n"
         "                  user_time=T,\n"