Internal change.

PiperOrigin-RevId: 724053185

grain/_src/python/dataset/stats.py

@@ -19,6 +19,7 @@
 from collections.abc import Sequence
 import contextlib
 import dataclasses
+import functools
 import pprint
 import sys
 import threading
@@ -51,6 +52,20 @@
     bucketer=monitoring.Bucketer.PowersOf(2.0),
 )
 
+_next_duration_ns_histogram = monitoring.EventMetric(
+    "/grain/python/dataset/next_duration_ns",
+    metadata=monitoring.Metadata(
+        description=(
+            "Histogram of durations of every `__next__` call on the output"
+            " iterator. Each data point is the duration value of `__next__`"
+            " call."
+        ),
+        units=monitoring.Units.NANOSECONDS,
+    ),
+    root=grain_monitoring.get_monitoring_root(),
+    bucketer=monitoring.Bucketer.PowersOf(2.0),
+)
+
 T = TypeVar("T")
 # Time between two consecutive monitoring reports.
 _REPORTING_PERIOD_SEC = 10
@@ -267,6 +282,37 @@ def _pretty_format_summary(
   return table.get_pretty_wrapped_summary()  # pylint: disable=protected-access
 
 
+def record_next_duration_if_output(next_fn):
+  """Records the duration of the `__next__` call on the output iterator node.
+
+  Expected to be used as follows:
+  ```
+  class MyMapDatasetIterator(DatasetIterator):
+    ...
+    @stats.record_next_duration_if_output
+    def __next__(self):
+      ...
+  ```
+
+  Args:
+    next_fn: The `__next__` function to wrap.
+
+  Returns:
+    The wrapped `next_fn`.
+  """
+
+  @functools.wraps(next_fn)
+  def wrapper(iterator):
+    start_time = time.perf_counter_ns()
+    result = next_fn(iterator)
+    if iterator._stats._is_output:  # pylint:disable=protected-access
+      next_duration_ns = time.perf_counter_ns() - start_time
+      _next_duration_ns_histogram.Record(next_duration_ns)
+    return result
+
+  return wrapper
+
+
 class _Table:
   """Table class for pretty printing tabular data."""
 

grain/_src/python/dataset/transformations/BUILD

@@ -227,7 +227,10 @@ py_library(
     name = "interleave",
     srcs = ["interleave.py"],
     srcs_version = "PY3",
-    deps = ["//grain/_src/python/dataset"],
+    deps = [
+        "//grain/_src/python/dataset",
+        "//grain/_src/python/dataset:stats",
+    ],
 )
 
 py_test(
@@ -247,7 +250,10 @@ py_library(
     name = "limit",
     srcs = ["limit.py"],
     srcs_version = "PY3",
-    deps = ["//grain/_src/python/dataset"],
+    deps = [
+        "//grain/_src/python/dataset",
+        "//grain/_src/python/dataset:stats",
+    ],
 )
 
 py_test(

grain/_src/python/dataset/transformations/batch.py

@@ -22,6 +22,7 @@
 
 from grain._src.core import tree_lib
 from grain._src.python.dataset import dataset
+from grain._src.python.dataset import stats
 import numpy as np
 
 T = TypeVar("T")
@@ -61,6 +62,7 @@ def __init__(
     self._drop_remainder = drop_remainder
     self._batch_fn = batch_fn
 
+  @stats.record_next_duration_if_output
   def __next__(self) -> T:
     values = []
     for _ in range(self._batch_size):

grain/_src/python/dataset/transformations/filter.py

@@ -147,6 +147,7 @@ def _threshold_checker(self):
         raise_threshold=self._ctx.dataset_options.filter_raise_threshold_ratio,
     )
 
+  @dataset_stats.record_next_duration_if_output
   def __next__(self):
     value = None
     passed_filter = False

grain/_src/python/dataset/transformations/flatmap.py

@@ -96,6 +96,7 @@ def __init__(
   def _has_consumed_all_buffer_elements(self):
     return self._next_index_in_buffer >= len(self._buffer)
 
+  @dataset_stats.record_next_duration_if_output
   def __next__(self):
     timer = dataset_stats.Timer()
     while self._has_consumed_all_buffer_elements():

grain/_src/python/dataset/transformations/interleave.py

@@ -17,6 +17,7 @@
 from typing import TypeVar
 
 from grain._src.python.dataset import dataset
+from grain._src.python.dataset import stats
 from grain._src.python.dataset.transformations import prefetch
 
 T = TypeVar("T")
@@ -53,6 +54,7 @@ def __init__(
         None
     ] * self._cycle_length
 
+  @stats.record_next_duration_if_output
   def __next__(self) -> T:
     while True:
       if iterator_to_use := self._iterators_in_use[self._next_index_in_cycle]:

grain/_src/python/dataset/transformations/limit.py

@@ -16,6 +16,7 @@
 from typing import Any, TypeVar
 
 from grain._src.python.dataset import dataset
+from grain._src.python.dataset import stats
 
 Element = Any
 T = TypeVar("T")  # pylint: disable=invalid-name
@@ -33,6 +34,7 @@ def __init__(
     self._count = count
     self._count_elements_read = 0
 
+  @stats.record_next_duration_if_output
   def __next__(self):
     if self._count_elements_read >= self._count:
       raise StopIteration

grain/_src/python/dataset/transformations/map.py

@@ -19,6 +19,7 @@
 
 from grain._src.core import transforms
 from grain._src.python.dataset import dataset
+from grain._src.python.dataset import stats
 import numpy as np
 
 
@@ -215,6 +216,7 @@ def __init__(
     # TODO: Move users away from this and remove.
     self._index_for_rng = 0
 
+  @stats.record_next_duration_if_output
   def __next__(self):
     element = next(self._parent)
     with self._stats.record_self_time():
@@ -254,6 +256,7 @@ def __init__(
     self._rng = np.random.Generator(np.random.Philox(seed))
     self._transform_name = transform_name
 
+  @stats.record_next_duration_if_output
   def __next__(self):
     element = next(self._parent)
     with self._stats.record_self_time():

grain/_src/python/dataset/transformations/mix.py

@@ -24,6 +24,7 @@
 from grain._src.core import exceptions
 from grain._src.python.dataset import base
 from grain._src.python.dataset import dataset
+from grain._src.python.dataset import stats
 from typing_extensions import override
 
 Element = Any
@@ -139,6 +140,7 @@ def __init__(
     self._index = 0
     self._stop = False
 
+  @stats.record_next_duration_if_output
   def __next__(self):
     if self._stop:
       # Although there may be elements available in some parent datasets, do not

grain/_src/python/dataset/transformations/packing.py

@@ -193,6 +193,7 @@ def _finalize_current_batch(self, element_for_shapes):
           meta_features=self._meta_features,
       )
 
+  @dataset_stats.record_next_duration_if_output
   def __next__(self):
     timer = dataset_stats.Timer()
     if self._packed_batch is not None:

grain/_src/python/dataset/transformations/packing_concat_then_split.py

@@ -400,6 +400,7 @@ def _maybe_add_to_buffer(
       tokens_in_buffer[k] += v
     return remainder
 
+  @stats.record_next_duration_if_output
   def __next__(self):
     if self._packed_elements:
       self._state.elements_from_buffer_after_checkpoint += 1

grain/_src/python/dataset/transformations/prefetch.py

@@ -135,6 +135,7 @@ def _threshold_checker(self):
         raise_threshold=self._ctx.dataset_options.filter_raise_threshold_ratio,
     )
 
+  @dataset_stats.record_next_duration_if_output
   def __next__(self) -> T:
     # The time recorded here is the time spent in prefetch node to return an
     # element, including the time spent in parent node.
@@ -496,6 +497,7 @@ def _stats(self):
   def __iter__(self) -> dataset.DatasetIterator[T]:
     return self
 
+  @dataset_stats.record_next_duration_if_output
   def __next__(self) -> T:
     self._ensure_iterator_initialized()
     result, state = next(self._iterator)
@@ -714,6 +716,7 @@ def _producer(
     except Exception as e:  # pylint: disable=broad-except
       output_buffer.put((None, None, e))
 
+  @dataset_stats.record_next_duration_if_output
   def __next__(self):
     self.start_prefetch()
     assert self._buffer is not None

grain/_src/python/dataset/transformations/shuffle.py

@@ -19,6 +19,7 @@
 from typing import TypeVar
 
 from grain._src.python.dataset import dataset
+from grain._src.python.dataset import stats
 from grain._src.python.experimental.index_shuffle.python import index_shuffle_module as index_shuffle
 
 
@@ -188,6 +189,7 @@ def _fill_and_shuffle_window(self):
         seed=self._global_seed + self._window_index, window=self._window
     )
 
+  @stats.record_next_duration_if_output
   def __next__(self):
     # Window is empty, fill up the next window.
     if not self._window: