diff --git a/grain/_src/python/dataset/transformations/interleave.py b/grain/_src/python/dataset/transformations/interleave.py index d1830f616..d38b44343 100644 --- a/grain/_src/python/dataset/transformations/interleave.py +++ b/grain/_src/python/dataset/transformations/interleave.py @@ -100,6 +100,9 @@ def __next__(self) -> T: self._started = True timer = stats.Timer() _ = self._stats # eagerly initialize stats + # Snapshots and restores can materialize children without starting them. + # Start all active children before waiting for the selected child. + self.start_prefetch() while True: if iterator_to_use := self._iterators_in_use[self._next_index_in_cycle]: try: diff --git a/grain/_src/python/dataset/transformations/interleave_test.py b/grain/_src/python/dataset/transformations/interleave_test.py index 218569a98..0af869f1e 100644 --- a/grain/_src/python/dataset/transformations/interleave_test.py +++ b/grain/_src/python/dataset/transformations/interleave_test.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import threading from typing import cast from absl.testing import absltest from absl.testing import flagsaver @@ -97,6 +98,74 @@ def __iter__(self) -> dataset.DatasetIterator: return _IteratorIdDatasetIterator(self._parent.__iter__()) +class InterleavePrefetchStartupTest(parameterized.TestCase): + """Snapshot and restore must preserve concurrent lane startup and order.""" + + @parameterized.product( + debug=(False, True), + resume=("snapshot", "zero_buffer", "live", "fresh"), + ) + def test_starts_all_lanes_before_waiting(self, debug, resume): + started = [threading.Event() for _ in range(3)] + for event in started: + event.set() + + def read(value): + started[value // 100].set() + for event in started: + self.assertTrue( + event.wait(timeout=5), + "interleave waited for one lane before starting all lanes", + ) + return value + + def make_iterator(): + shards = [ + dataset.MapDataset.range(lane * 100, lane * 100 + 20) + .to_iter_dataset( + options.ReadOptions(num_threads=0, prefetch_buffer_size=0) + ) + .map(read) + for lane in range(3) + ] + ds = interleave.InterleaveIterDataset( + shards, + cycle_length=3, + num_make_iter_threads=3, + make_iter_buffer_size=3, + iter_buffer_size=2, + ) + if resume == "zero_buffer": + ds = prefetch.ThreadPrefetchIterDataset(ds, prefetch_buffer_size=0) + return iter(ds) + + expected = [lane * 100 + offset for offset in range(20) for lane in range(3)] + with flagsaver.flagsaver(grain_py_debug_mode=debug): + iterator = make_iterator() + try: + offset = 0 + if resume in ("live", "fresh"): + offset = 5 + for _ in range(offset): + next(iterator) + state = iterator.get_state() + if resume == "fresh": + iterator.close() + iterator = make_iterator() + iterator.set_state(state) + else: + iterator.get_state() + for event in started: + event.clear() + self.assertEqual( + [next(iterator) for _ in range(12)], expected[offset : offset + 12] + ) + finally: + for event in started: + event.set() + iterator.close() + + @absltest.skipThisClass("Base class") class _InterleaveIterDatasetTestBase(parameterized.TestCase):