diff --git a/CHANGELOG.md b/CHANGELOG.md index f4b519c01..d97f7e21e 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/grain/_src/core/config.py b/grain/_src/core/config.py index 8a5e7c5e0..3421fdfa6 100644 --- a/grain/_src/core/config.py +++ b/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( diff --git a/grain/_src/python/dataset/dataset.py b/grain/_src/python/dataset/dataset.py index 1bf6ca2e1..9c6ff8bf8 100644 --- a/grain/_src/python/dataset/dataset.py +++ b/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 ] diff --git a/grain/_src/python/dataset/stats.py b/grain/_src/python/dataset/stats.py index de6c389de..6ce17c5fd 100644 --- a/grain/_src/python/dataset/stats.py +++ b/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" ) diff --git a/grain/_src/python/dataset/stats_test.py b/grain/_src/python/dataset/stats_test.py index a4f57ae99..8e8a2cd6a 100644 --- a/grain/_src/python/dataset/stats_test.py +++ b/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):