Update shuffle

Differential Revision: D80646907

Pull Request resolved: #857

Author: moto-meta

src/spdl/source/utils.py

@@ -206,23 +206,26 @@ def __init__(
     ) -> None:
         self.src = src
         self._epoch = epoch
-        self._shuffle_last = shuffle_last
+        self._shuffle_first: bool = not shuffle_last
 
     def _shuffle(self) -> None:
         t0 = time.monotonic()
         self.src.shuffle(seed=self._epoch)
         if (elapsed := time.monotonic() - t0) > 3:
             _LG.warning("Shuffling took %.2f sec.", elapsed)
-
-    def __iter__(self) -> Iterator[T]:
-        if not self._shuffle_last:
-            self._shuffle()
-
-        yield from self.src
         self._epoch += 1
 
-        if self._shuffle_last:
+    def __iter__(self) -> Iterator[T]:
+        if self._shuffle_first:
             self._shuffle()
+            yield from self.src
+        else:
+            try:
+                yield from self.src
+            finally:
+                # in case the iteration is stopped in the middle.
+                # shuffle is called when the iterator is deleted.
+                self._shuffle()
 
     def __len__(self) -> int:
         if isinstance(self.src, Sized):

tests/spdl_unittest/dataloader/source_utils_test.py

@@ -0,0 +1,93 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+
+from collections.abc import Iterator
+
+from spdl.source.utils import embed_shuffle
+
+
+class IterableWithShuffle_:
+    def __init__(self, n: int) -> None:
+        self.vals = list(range(n))
+        self._seed: int | None = None
+
+    def __iter__(self) -> Iterator[int]:
+        yield from self.vals
+
+    def shuffle(self, seed: int) -> None:
+        # rotate
+        self._seed = seed
+        self.vals = self.vals[1:] + self.vals[:1]
+
+
+def test_embed_shuffle():
+    """Iterable created by embed_shuffle calls shuffle automatically"""
+
+    foo = IterableWithShuffle_(3)
+    assert foo._seed is None
+    iterable = embed_shuffle(foo)
+    assert list(iterable) == [1, 2, 0]
+    assert foo._seed == 0
+    assert list(iterable) == [2, 0, 1]
+    assert foo._seed == 1
+    assert list(iterable) == [0, 1, 2]
+    assert foo._seed == 2
+
+
+def test_embed_shuffle_halt():
+    """The value is shuffled with different seed even after an iteration is halted."""
+
+    foo = IterableWithShuffle_(5)
+    iterable = embed_shuffle(foo)
+
+    iterator = iter(iterable)
+    assert foo._seed is None
+    assert next(iterator) == 1
+    assert foo._seed == 0
+    assert next(iterator) == 2
+    del iterator
+
+    iterator = iter(iterable)
+    assert next(iterator) == 2
+    assert foo._seed == 1
+    assert next(iterator) == 3
+    del iterator
+
+
+def test_embed_shuffle_shuffle_after():
+    """Iterable created by embed_shuffle calls shuffle automatically after iteration"""
+
+    foo = IterableWithShuffle_(3)
+    iterable = embed_shuffle(foo, shuffle_last=True)
+    assert foo._seed is None
+    assert list(iterable) == [0, 1, 2]
+    assert foo._seed == 0
+    assert list(iterable) == [1, 2, 0]
+    assert foo._seed == 1
+    assert list(iterable) == [2, 0, 1]
+    assert foo._seed == 2
+
+
+def test_embed_shuffle_shuffle_after_halt():
+    """The value is shuffled with different seed even after an iteration is halted."""
+
+    foo = IterableWithShuffle_(5)
+    iterable = embed_shuffle(foo, shuffle_last=True)
+
+    iterator = iter(iterable)
+    assert next(iterator) == 0
+    assert next(iterator) == 1
+    assert foo._seed is None
+    del iterator
+    assert foo._seed == 0
+
+    iterator = iter(iterable)
+    assert next(iterator) == 1
+    assert next(iterator) == 2
+    del iterator
+    assert foo._seed == 1