Internal

PiperOrigin-RevId: 837296133

Author: Grain Team

grain/_src/python/dataset/dataset_test.py

@@ -1392,6 +1392,40 @@ def test_execution_summary_with_no_logging(self):
     log_value = "Grain Dataset Execution Summary"
     self.assertNotIn(log_value, "".join(logs.output))
 
+  @flagsaver.flagsaver(grain_py_debug_mode=True)
+  @mock.patch.object(dataset_stats, "_REPORTING_PERIOD_SEC", 0.05)
+  def test_execution_summary_with_mp_prefetch(self):
+    def worker_init_fn_wrapper(worker_index, worker_count):
+      del worker_index, worker_count
+      dataset_stats._REPORTING_PERIOD_SEC = 0.05
+
+    ds = dataset.MapDataset.range(10000).map(MapTransformAddingOne())
+    ds = ds.to_iter_dataset()
+    ds = ds.mp_prefetch(
+        options.MultiprocessingOptions(num_workers=1),
+        worker_init_fn=worker_init_fn_wrapper,
+    )
+    it = ds.__iter__()
+    _ = list(it)
+    all_nodes_present = False
+    while not all_nodes_present:
+      time.sleep(1)
+      all_nodes_present = True
+      summary = dataset.get_execution_summary(it)
+      node_names = {node.name for node in summary.nodes.values()}
+      all_nodes_present = all_nodes_present and any(
+          "RangeMapDataset" in name for name in node_names
+      )
+      all_nodes_present = all_nodes_present and any(
+          "MapMapDataset" in name for name in node_names
+      )
+      all_nodes_present = all_nodes_present and any(
+          "PrefetchDatasetIterator" in name for name in node_names
+      )
+      all_nodes_present = all_nodes_present and any(
+          "MultiprocessPrefetchDatasetIterator" in name for name in node_names
+      )
+
 
 class GetElementSpecTest(parameterized.TestCase):
 

grain/_src/python/dataset/stats.py

@@ -818,10 +818,6 @@ def __init__(self, config: StatsConfig, parents: Sequence[Stats]):
     self._last_update_time = 0
     self._last_report_time = 0
     self._summary_dispatcher = None
-    if self._config.stats_out_queue:
-      self._summary_dispatcher = self._send_stats_to_main_process_loop
-    elif self._config.log_summary:
-      self._summary_dispatcher = self._logging_execution_summary_loop
 
   def __reduce__(self):
     return _ExecutionStats, (self._config, self._parents)
@@ -939,14 +935,18 @@ def record_self_time(self, num_elements: int = 1, offset_ns: int = 0):
                   target=self._reporting_loop, daemon=True
               )
               self._reporting_thread.start()
-        if (
-            self._summary_dispatcher_thread is None
-            and self._summary_dispatcher is not None
+        if self._summary_dispatcher_thread is None and (
+            self._config.stats_out_queue or self._config.log_summary
         ):
           with self._summary_dispatcher_thread_init_lock:
             # Check above together with update would not be atomic -- another
             # thread may have started the logging thread.
             if self._summary_dispatcher_thread is None:
+              self._summary_dispatcher = (
+                  self._send_stats_to_main_process_loop
+                  if self._config.stats_out_queue
+                  else self._logging_execution_summary_loop
+              )
               self._summary_dispatcher_thread = threading.Thread(
                   target=self._summary_dispatcher, daemon=True
               )