From eae281640bf655cf9ce7672db5d5987bbd0885f8 Mon Sep 17 00:00:00 2001 From: Phani Aenugula Date: Thu, 23 Jul 2026 15:35:19 -0700 Subject: [PATCH] Internal PiperOrigin-RevId: 952991982 --- grain/_src/python/dataset/BUILD | 2 +- grain/_src/python/dataset/dataset.py | 63 ++++++++++++++++++- grain/_src/python/dataset/elastic_iterator.py | 20 +++--- .../transformations/interleave_test.py | 46 +++++++++----- .../dataset/transformations/prefetch.py | 29 --------- 5 files changed, 100 insertions(+), 60 deletions(-) diff --git a/grain/_src/python/dataset/BUILD b/grain/_src/python/dataset/BUILD index 404295d44..e5046876a 100644 --- a/grain/_src/python/dataset/BUILD +++ b/grain/_src/python/dataset/BUILD @@ -166,7 +166,6 @@ py_library( ":dataset", "//grain/_src/core:sharding", "//grain/_src/python:options", - "//grain/_src/python/dataset/transformations:zip", ], ) @@ -186,6 +185,7 @@ py_test( "//grain/_src/core:sharding", "//grain/_src/python:options", "//grain/_src/python/checkpoint:handler", + "//grain/_src/python/dataset/transformations:zip", "//grain/_src/python/testing:experimental", "@abseil-py//absl/testing:absltest", "@abseil-py//absl/testing:parameterized", diff --git a/grain/_src/python/dataset/dataset.py b/grain/_src/python/dataset/dataset.py index 2b67ad561..200ccfc54 100644 --- a/grain/_src/python/dataset/dataset.py +++ b/grain/_src/python/dataset/dataset.py @@ -66,7 +66,6 @@ from grain.proto import execution_summary_pb2 import numpy as np - _api_usage_counter = monitoring.Counter( "/grain/python/lazy_dataset/api", metadata=monitoring.Metadata( @@ -1895,6 +1894,35 @@ def set_slice(self, sl: slice, sequential_slice: bool = False) -> None: ... +@runtime_checkable +class SupportsSlicedStateManagement(Protocol): + """Iterators that support setting a sliced state. + + This protocol is used to support elastic resizing of iterators. + """ + + def get_shard_states(self) -> Sequence[Any]: + """Returns the states of all shards managed by this iterator. + + Used for elastic resizing to capture the current progress of each + shard. + """ + ... + + def set_shard_states(self, shard_states: Sequence[Any]) -> None: + """Sets the states of all shards managed by this iterator. + + Used for elastic resizing to restore the progress of each shard. + + Args: + shard_states: A sequence of dictionaries, one for each shard. Each dict + must contain 'exhausted' key with value bool indicating if the shard is + exhausted and 'state' key with value Any representing the state of the + shard. + """ + ... + + def set_slice( ds: MapDataset | IterDataset, sl: slice, @@ -1921,3 +1949,36 @@ def set_slice( raise ValueError(f"Cannot slice `IterDataset` source. {type(ds)}") for parent in ds.parents: set_slice(parent, sl, sequential_slice) + + +def find_shard_states(it: DatasetIterator) -> Sequence[Any]: + """Returns the shard states for the given dataset iterator.""" + if isinstance(it, SupportsSlicedStateManagement): + return it.get_shard_states() + if not hasattr(it, "_parents"): + return [] + parent_shard_states = [] + # Access the iterators parents to get the parents of the iterator. + for parent in it._parents: # pylint: disable=protected-access + parent_shard_states.append(find_shard_states(parent)) + return list(zip(*parent_shard_states)) + + +def set_shard_states(it: DatasetIterator, shard_states: Sequence[Any]) -> None: + """Sets the shard states for the given dataset iterator.""" + if isinstance(it, SupportsSlicedStateManagement): + it.set_shard_states(shard_states) + return + if not hasattr(it, "_parents"): + raise ValueError(f"Iterator {it} does not support elastic resizing.") + if not shard_states: + return + unzipped_states = list(zip(*shard_states)) + # Access the iterators parents to set the parents of the iterator. + if len(it._parents) != len(unzipped_states): # pylint: disable=protected-access + raise ValueError( + f"Iterator has {len(it._parents)} parents but shard states have" # pylint: disable=protected-access + f" {len(unzipped_states)} elements." + ) + for parent, states in zip(it._parents, unzipped_states): # pylint: disable=protected-access + set_shard_states(parent, states) diff --git a/grain/_src/python/dataset/elastic_iterator.py b/grain/_src/python/dataset/elastic_iterator.py index deacf4097..ff8690517 100644 --- a/grain/_src/python/dataset/elastic_iterator.py +++ b/grain/_src/python/dataset/elastic_iterator.py @@ -26,9 +26,6 @@ from grain._src.python.dataset.transformations import interleave from grain._src.python.dataset.transformations import mix from grain._src.python.dataset.transformations import prefetch -from grain._src.python.dataset.transformations import ( - zip as zip_dataset, -) T = TypeVar("T") @@ -50,14 +47,13 @@ def _verify_transformations_supported(ds: dataset.IterDataset) -> None: if isinstance( next_ds, ( - zip_dataset.ZipIterDataset, prefetch.PrefetchIterDataset, mix.MixedIterDataset, ), ): raise ValueError( - "ElasticIterator for IterDataset does not support zip, mix or" - " prefetch transformation yet." + "ElasticIterator for IterDataset does not support mix or prefetch" + " transformation yet." ) to_check.extend(next_ds.parents) # pyrefly: ignore[bad-argument-type] @@ -91,7 +87,7 @@ def _find_sliceable_iterator( """Finds the first sliceable iterator in the iterator graph. This function recursively searches through the parents of the given iterator - to find an iterator that supports `prefetch.SupportsSlicedStateManagement`. + to find an iterator that supports `dataset.SupportsSlicedStateManagement`. Args: it: The starting DatasetIterator. @@ -100,7 +96,7 @@ def _find_sliceable_iterator( The first sliceable DatasetIterator found, or None if no such iterator is found in the graph. """ - if isinstance(it, prefetch.SupportsSlicedStateManagement): + if isinstance(it, dataset.SupportsSlicedStateManagement): return it if not hasattr(it, "_parents"): return None @@ -207,11 +203,11 @@ def get_shard_states(self) -> dict[int, Any]: A dictionary mapping global shard indices to their states. """ if not isinstance( - self._sliceable_iterator, prefetch.SupportsSlicedStateManagement + self._sliceable_iterator, dataset.SupportsSlicedStateManagement ): raise ValueError( "ElasticIterDatasetIterator does not support" - " prefetch.SupportsSlicedStateManagement." + " dataset.SupportsSlicedStateManagement." ) iter_shard_states = self._sliceable_iterator.get_shard_states() state_by_shard_index = {} @@ -249,11 +245,11 @@ def set_shard_states(self, state: dict[int, Any]) -> None: self._closed = False if not isinstance( - self._sliceable_iterator, prefetch.SupportsSlicedStateManagement + self._sliceable_iterator, dataset.SupportsSlicedStateManagement ): raise ValueError( "ElasticIterDatasetIterator only supports" - " prefetch.SupportsSlicedStateManagement." + " dataset.SupportsSlicedStateManagement." ) ds_iterator_states = {int(k): v for k, v in state.items()} # We need to sort the states by the global shard index to ensure that the diff --git a/grain/_src/python/dataset/transformations/interleave_test.py b/grain/_src/python/dataset/transformations/interleave_test.py index 361f1e6f3..e35998453 100644 --- a/grain/_src/python/dataset/transformations/interleave_test.py +++ b/grain/_src/python/dataset/transformations/interleave_test.py @@ -23,10 +23,9 @@ from grain._src.python.dataset import base from grain._src.python.dataset import dataset from grain._src.python.dataset.transformations import interleave -from grain._src.python.dataset.transformations import prefetch from grain._src.python.dataset.transformations import repeat from grain._src.python.dataset.transformations import zip as zip_dataset -from grain._src.python.testing.experimental import assert_equal_output_after_checkpoint +from grain._src.python.testing import experimental import numpy as np _INTERLEAVE_TEST_CASES = ( @@ -212,7 +211,7 @@ def test_checkpointing_comprehensive(self): ] ds = self._create_dataset(ds, cycle_length=5) ds = self._maybe_wrap_ds(ds) - assert_equal_output_after_checkpoint(ds) + experimental.assert_equal_output_after_checkpoint(ds) def test_set_state_does_not_recreate_iterators_if_not_needed(self): cycle_length = 5 @@ -392,17 +391,21 @@ def test_slice_state_management_checkpoints_correctly( next(it) # Get the shard state. - assert isinstance(it, prefetch.SupportsSlicedStateManagement) - shard_state = it.get_shard_states() + assert isinstance(it, dataset.SupportsSlicedStateManagement) + shard_state = cast( + dataset.SupportsSlicedStateManagement, it + ).get_shard_states() self.assertEqual(shard_state, expected_shard_state) # Create a new iterator and restore state. it2 = ds.__iter__() - assert isinstance(it2, prefetch.SupportsSlicedStateManagement) - it2.set_shard_states(shard_state) + assert isinstance(it2, dataset.SupportsSlicedStateManagement) + cast(dataset.SupportsSlicedStateManagement, it2).set_shard_states( + shard_state + ) # Verify it continues from the correct position. - self.assertSequenceEqual(list(it2), expected_remaining) # pyrefly: ignore[bad-argument-type] + self.assertSequenceEqual(list(it2), expected_remaining) @parameterized.named_parameters( dict( @@ -461,20 +464,27 @@ def test_correct_interleave_state_after_setting_shards( for _ in range(2): next(it) - assert isinstance(it, prefetch.SupportsSlicedStateManagement) - shard_state = it.get_shard_states() + assert isinstance(it, dataset.SupportsSlicedStateManagement) + shard_state = cast( + dataset.SupportsSlicedStateManagement, it + ).get_shard_states() self.assertEqual(shard_state, expected_shard_state) # Create a new iterator and restore state. it2 = ds.__iter__() - assert isinstance(it2, prefetch.SupportsSlicedStateManagement) - it2.set_shard_states(shard_state) + assert isinstance(it2, dataset.SupportsSlicedStateManagement) + cast(dataset.SupportsSlicedStateManagement, it2).set_shard_states( + shard_state + ) # Check get_shard_states() returns the set shard states correctly. - self.assertEqual(it2.get_shard_states(), expected_shard_state) + self.assertEqual( + cast(dataset.SupportsSlicedStateManagement, it2).get_shard_states(), + expected_shard_state, + ) # Check get_state() internal values. - state = it2.get_state() # pyrefly: ignore[missing-attribute] + state = it2.get_state() self.assertEqual(state["next_index_in_cycle"], 0) self.assertEqual( state["next_index_in_datasets"], expected_next_index_in_datasets @@ -504,11 +514,13 @@ def test_setting_shard_state_with_exhausted_states(self): ] # Create a new iterator and restore state. - assert isinstance(it, prefetch.SupportsSlicedStateManagement) - it.set_shard_states(shard_state) + assert isinstance(it, dataset.SupportsSlicedStateManagement) + cast(dataset.SupportsSlicedStateManagement, it).set_shard_states( + shard_state + ) # Check get_state() internal values. - state = it.get_state() # pyrefly: ignore[missing-attribute] + state = it.get_state() self.assertEqual(state["next_index_in_cycle"], 0) self.assertEqual(state["next_index_in_datasets"], 3) self.assertEqual(state["iterators_in_use_indices"], [2, 0]) diff --git a/grain/_src/python/dataset/transformations/prefetch.py b/grain/_src/python/dataset/transformations/prefetch.py index b2f5fa690..8209bc26d 100644 --- a/grain/_src/python/dataset/transformations/prefetch.py +++ b/grain/_src/python/dataset/transformations/prefetch.py @@ -91,35 +91,6 @@ def _getitem( return stats.record_bytes_consumed(parent[index]) -@typing.runtime_checkable -class SupportsSlicedStateManagement(Protocol): - """Iterators that support setting a sliced state. - - This protocol is used to support elastic resizing of iterators. - """ - - def get_shard_states(self) -> Sequence[Any]: - """Returns the states of all shards managed by this iterator. - - Used for elastic resizing to capture the current progress of each - shard. - """ - ... - - def set_shard_states(self, shard_states: Sequence[Any]): - """Sets the states of all shards managed by this iterator. - - Used for elastic resizing to restore the progress of each shard. - - Args: - shard_states: A sequence of dictionaries, one for each shard. Each dict - must contain 'exhausted' key with value bool indicating if the shard is - exhausted and 'state' key with value Any representing the state of the - shard. - """ - ... - - class PrefetchIterDataset(dataset.IterDataset[T]): """Iterable dataset that uses a thread pool for prefetching."""