Add py_permanently_disable_stats config option to Grain.

CHANGELOG.md

@@ -6,6 +6,9 @@ changes. Best viewed [here](https://google-grain.readthedocs.io/en/latest/change
 ## Unreleased
 
 * New features:
+  * Added the `py_permanently_disable_stats` config option
+    (`--grain_py_permanently_disable_stats`) to permanently disable dataset
+    stats collection and avoid stats tree initialization overhead.
   * Added `HFIterDataset` to wrap Hugging Face streaming datasets, supporting sharding, slicing (`set_slice`), and state checkpointing.
   * `MapDataset.slice` now accepts a `Sequence[int]` (including a
     `MapDataset[int]`) in addition to a `slice` object, enabling arbitrary index

grain/_src/core/config.py

@@ -130,6 +130,15 @@
     "How Grain filters internal stack frames from the traceback upon errors.",
 )
 
+_PERMANENTLY_DISABLE_STATS = flags.DEFINE_bool(
+    "grain_py_permanently_disable_stats",
+    False,
+    (
+        "If True, permanently disables Grain dataset statistics collection and"
+        " on-demand profiling to avoid initialization overhead."
+    ),
+)
+
 _GRAIN_FLAGS = (
     _INTERLEAVED_SHUFFLE,
     _INTERLEAVED_SHUFFLE_BLOCK_SIZE,
@@ -143,6 +152,7 @@
     _DATASET_VISUALIZATION_OUTPUT_DIR,
     _RELAY_SIGTERM_TO_MAIN,
     _TRACEBACK_FILTERING,
+    _PERMANENTLY_DISABLE_STATS,
 )
 
 _grain_experiment_metric = monitoring.Metric(

grain/_src/python/dataset/dataset.py

@@ -934,6 +934,9 @@ def _initialize_stats(
     Returns:
       The initialized stats object.
     """
+    if dataset_stats.stats_are_permanently_disabled():
+      self._stats = dataset_stats._NULL_STATS
+      return self._stats
     # There may be parent `MapDataset` nodes introduced by users that did not
     # call super init and thus don't have `_parents`.
     parents_stats = []
@@ -1680,6 +1683,9 @@ def _initialize_stats(
     Returns:
       The initialized stats object.
     """
+    if dataset_stats.stats_are_permanently_disabled():
+      self._stats = dataset_stats._NULL_STATS
+      return self._stats
     parent_stats = [
         p._initialize_stats(execution_tracking_mode) for p in self._parents
     ]

grain/_src/python/dataset/stats.py

@@ -761,11 +761,14 @@ def _running_in_colab() -> bool:
   return "google.colab" in sys.modules
 
 
+_NULL_CONTEXT = contextlib.nullcontext()
+
+
 class _DefaultStats(Stats):
   """Default implementation for statistics collection that does nothing."""
 
   def record_self_time(self, *, num_elements: int = 1, offset_ns: int = 0):  # pyrefly: ignore[bad-override]
-    return contextlib.nullcontext()
+    return _NULL_CONTEXT
 
   def record_output_spec(self, element: T) -> T:
     return element
@@ -780,6 +783,24 @@ def record_bytes_produced(self, element: Any) -> Any:
     return element
 
 
+class _NullStats(_DefaultStats):
+  """Singleton no-op stats object used when stats are permanently disabled.
+
+  Unlike `_DefaultStats`, this does not support on-demand profiling.
+  """
+
+  def __init__(self):  # pylint: disable=super-init-not-called
+    # Intentionally skip `super().__init__()` to avoid the overhead of
+    # registering iterators and starting the on-demand profiling monitor.
+    self._config = StatsConfig(name="")
+    self._self_output_spec = None
+    self._parents = ()
+    self._is_output = False
+
+
+_NULL_STATS = _NullStats()
+
+
 class _VisualizationStats(Stats):
   """Produces Dataset Visualization Graph."""
 
@@ -1145,6 +1166,13 @@ def _build_execution_summary(
     return execution_summary, node_id
 
 
+def stats_are_permanently_disabled() -> bool:
+  """Returns whether dataset statistics collection is permanently disabled."""
+  return bool(
+      grain_config.config.get_or_default("py_permanently_disable_stats")
+  )
+
+
 def make_stats(
     config: StatsConfig,
     parents: Sequence[Stats],
@@ -1153,6 +1181,8 @@ def make_stats(
     ),
 ) -> Stats:
   """Produces statistics instance according to the current execution mode."""
+  if stats_are_permanently_disabled():
+    return _NULL_STATS
   vis_output_dir = grain_config.config.get_or_default(
       "py_dataset_visualization_output_dir"
   )

grain/_src/python/dataset/stats_test.py

@@ -254,6 +254,15 @@ def test_report(self):
     s = s._parents[0]
     s.report()
 
+  @flagsaver.flagsaver(grain_py_permanently_disable_stats=True)
+  def test_permanently_disable_stats(self):
+    ds = dataset.MapDataset.range(5).map(_identity)
+    it = ds.to_iter_dataset().batch(2).__iter__()
+    self.assertEqual([list(x) for x in it], [[0, 1], [2, 3], [4]])
+    self.assertIs(ds._stats, stats._NULL_STATS)
+    self.assertIs(it._stats, stats._NULL_STATS)
+    self.assertEmpty(stats._iter_weakref_registry)
+
 
 class DebugModeStatsTest(absltest.TestCase):