Internal

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",

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)

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

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])

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."""