Add TraceMe in Grain Pipeline to include main process nodes in host-side analysis.

grain/_src/python/dataset/stats.py

@@ -300,52 +300,48 @@ def pretty_format_summary(
   return table.get_pretty_wrapped_summary()
 
 
-def record_next_duration_if_output(next_fn):
-  """Records the duration of the `__next__` call on the output iterator node.
+# Input pipeline stage categories.
+IPL_CAT_PREPROCESSING = "preprocessing"
+IPL_CAT_READ = "read"
+IPL_CAT_ENQUEUE = "enqueue"
+IPL_CAT_UNKNOWN = "unknown"
 
-  Expected to be used as follows:
-  ```
-  class MyMapDatasetIterator(DatasetIterator):
-    ...
-    @stats.record_next_duration_if_output
-    def __next__(self):
-      ...
-  ```
+
+def record_next_duration_if_output(stage_category: str = IPL_CAT_ENQUEUE):
+  """Records the duration of the `__next__` call on the output iterator node.
 
   Args:
-    next_fn: The `__next__` function to wrap.
+    stage_category: The category of the input pipeline stage.
 
   Returns:
-    The wrapped `next_fn`.
+    A decorator that records the duration of the `__next__` call.
   """
 
-  @functools.wraps(next_fn)
-  def wrapper(iterator):
-    if _TRACE_ANNOTATION and _TRACE_ANNOTATION.is_enabled():
-      with _TRACE_ANNOTATION(
-          f"{iterator.__class__.__name__}.{next_fn.__name__}",
-          _ipl_stage_name=str(iterator),
-          _ipl_stage_id=id(iterator),
-      ):
+  def inner_wrapper(next_fn):
+
+    @functools.wraps(next_fn)
+    def wrapper(iterator):
+      if _TRACE_ANNOTATION and _TRACE_ANNOTATION.is_enabled():
+        with _TRACE_ANNOTATION(
+            f"{iterator.__class__.__name__}.{next_fn.__name__}",
+            _ipl_stage_name=str(iterator),
+            _ipl_stage_id=id(iterator),
+            _ipl_stage_cat=stage_category,
+        ):
+          start_time = time.perf_counter_ns()
+          result = next_fn(iterator)
+      else:
         start_time = time.perf_counter_ns()
         result = next_fn(iterator)
-    else:
-      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
+      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
 
-# Input pipeline stage categories.
-IPL_CAT_PREPROCESSING = "preprocessing"
-IPL_CAT_READ = "read"
-IPL_CAT_ENQUEUE = "enqueue"
-IPL_CAT_UNKNOWN = "unknown"
+  return inner_wrapper
 
 
 def trace_input_pipeline(stage_category: str = IPL_CAT_UNKNOWN, **trace_kwargs):

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

@@ -122,7 +122,7 @@ def __init__(
     self._drop_remainder = drop_remainder
     self._batch_fn = batch_fn
 
-  @stats.record_next_duration_if_output
+  @stats.record_next_duration_if_output(stage_category=stats.IPL_CAT_ENQUEUE)
   def __next__(self) -> T:
     values = []
     for _ in range(self._batch_size):

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

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

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

@@ -94,7 +94,9 @@ 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
+  @dataset_stats.record_next_duration_if_output(
+      stage_category=dataset_stats.IPL_CAT_ENQUEUE
+  )
   def __next__(self):
     timer = dataset_stats.Timer()
     while self._has_consumed_all_buffer_elements():

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

@@ -82,7 +82,7 @@ def __init__(
         None
     ] * self._cycle_length
 
-  @stats.record_next_duration_if_output
+  @stats.record_next_duration_if_output(stage_category=stats.IPL_CAT_ENQUEUE)
   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

@@ -34,7 +34,7 @@ def __init__(
     self._count = count
     self._count_elements_read = 0
 
-  @stats.record_next_duration_if_output
+  @stats.record_next_duration_if_output(stage_category=stats.IPL_CAT_ENQUEUE)
   def __next__(self):
     if self._count_elements_read >= self._count:
       raise StopIteration

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

@@ -210,7 +210,7 @@ def __init__(
     self._map_fn = map_fn
     self._transform_name = transform_name
 
-  @stats.record_next_duration_if_output
+  @stats.record_next_duration_if_output(stage_category=stats.IPL_CAT_ENQUEUE)
   def __next__(self):
     element = next(self._parent)
     with self._stats.record_self_time():
@@ -245,7 +245,7 @@ def __init__(
     self._rng = np.random.Generator(np.random.Philox(seed))
     self._transform_name = transform_name
 
-  @stats.record_next_duration_if_output
+  @stats.record_next_duration_if_output(stage_category=stats.IPL_CAT_ENQUEUE)
   def __next__(self):
     element = next(self._parent)
     with self._stats.record_self_time():
@@ -291,7 +291,7 @@ def __init__(
     self._transform_name = transform_name
     self._counter = 0
 
-  @stats.record_next_duration_if_output
+  @stats.record_next_duration_if_output(stage_category=stats.IPL_CAT_ENQUEUE)
   def __next__(self):
     element = next(self._parent)
     with self._stats.record_self_time():

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

@@ -141,7 +141,7 @@ def __init__(
     self._index = 0
     self._stop = False
 
-  @stats.record_next_duration_if_output
+  @stats.record_next_duration_if_output(stage_category=stats.IPL_CAT_ENQUEUE)
   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

@@ -225,7 +225,9 @@ def _finalize_current_batch(self, element_for_shapes):
           padding_struct=self._padding_struct,
       )
 
-  @dataset_stats.record_next_duration_if_output
+  @dataset_stats.record_next_duration_if_output(
+      stage_category=dataset_stats.IPL_CAT_ENQUEUE
+  )
   def __next__(self):
     timer = dataset_stats.Timer()
     if self._packed_batch is not None:

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

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

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

@@ -167,7 +167,9 @@ def _threshold_checker(self):
         raise_threshold=self._ctx.dataset_options.filter_raise_threshold_ratio,
     )
 
-  @dataset_stats.record_next_duration_if_output
+  @dataset_stats.record_next_duration_if_output(
+      stage_category=dataset_stats.IPL_CAT_ENQUEUE
+  )
   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.
@@ -670,7 +672,9 @@ def _stats(self):
   def __iter__(self) -> dataset.DatasetIterator[T]:
     return self
 
-  @dataset_stats.record_next_duration_if_output
+  @dataset_stats.record_next_duration_if_output(
+      stage_category=dataset_stats.IPL_CAT_ENQUEUE
+  )
   def __next__(self) -> T:
     self._ensure_iterator_initialized()
     # The time recorded here is the time spent in prefetch node to return an
@@ -881,7 +885,9 @@ def start_prefetch(self):
     )
     self._prefetch_thread.start()
 
-  @dataset_stats.record_next_duration_if_output
+  @dataset_stats.record_next_duration_if_output(
+      stage_category=dataset_stats.IPL_CAT_ENQUEUE
+  )
   def __next__(self):
     self.start_prefetch()
     element, state, err = self._buffer.get()

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

@@ -96,7 +96,7 @@ def _flatten(self, pytree_element):
       self._treedef = pytree_element
     return tree_lib.flatten(pytree_element)
 
-  @stats.record_next_duration_if_output
+  @stats.record_next_duration_if_output(stage_category=stats.IPL_CAT_ENQUEUE)
   def __next__(self):
     timer = stats.Timer()
 

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

@@ -191,7 +191,7 @@ def _fill_and_shuffle_window(self):
         seed=self._global_seed + self._window_index, window=self._window
     )
 
-  @stats.record_next_duration_if_output
+  @stats.record_next_duration_if_output(stage_category=stats.IPL_CAT_ENQUEUE)
   def __next__(self):
     # Window is empty, fill up the next window.
     if not self._window:

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

@@ -80,7 +80,9 @@ def __init__(
     super().__init__([p.__iter__() for p in parents])
     self._strict = strict
 
-  @dataset_stats.record_next_duration_if_output
+  @dataset_stats.record_next_duration_if_output(
+      stage_category=dataset_stats.IPL_CAT_ENQUEUE
+  )
   def __next__(self) -> tuple[T, ...]:
     with self._stats.record_self_time():
       # Can't use for a `for` loop because we need to raise StopIteration from