[RFC] Retool for iterator only

test/nodes/test_adapters.py

@@ -42,6 +42,7 @@ def test_iterable(self):
         n = 20
         node = IterableWrapper(range(n))
         for epoch in range(2):
+            node.reset()
             result = list(node)
             self.assertEqual(len(result), n)
             for i, j in enumerate(result):
@@ -61,8 +62,9 @@ def test_generator(self):
 
     def test_iterable_dataset(self):
         n = 20
-        node = IterableWrapper(DummyIterableDataset(n))
+        node = IterableWrapper(DummyIterableDataset(n, name="test"))
         for epoch in range(2):
+            node.reset()
             result = list(node)
             self.assertEqual(len(result), n)
             for i, row in enumerate(result):
@@ -84,6 +86,7 @@ def test_default_sampler(self):
         n = 20
         node = MapStyleWrapper(DummyMapDataset(n), sampler=range(n))
         for epoch in range(2):
+            node.reset()
             result = list(node)
             self.assertEqual(len(result), n)
             for i, row in enumerate(result):
@@ -97,6 +100,7 @@ def test_random_sampler(self):
         node = MapStyleWrapper(ds, sampler=RandomSampler(ds))
         results = []
         for epoch in range(2):
+            node.reset()
             result = list(node)
             results.append(result)
             self.assertEqual(len(result), n)
@@ -116,6 +120,7 @@ def test_dict(self):
         sampler = list(d.keys())
         node = MapStyleWrapper(d, sampler=sampler)
         for epoch in range(2):
+            node.reset()
             result = list(node)
             self.assertEqual(len(result), n)
             for i, row in enumerate(result):
@@ -145,9 +150,10 @@ def test_sampler_wrapper(self):
 
         results = []
         for epoch in range(2):
+            node.reset()
+            self.assertEqual(node.epoch, epoch)
             result = list(node)
             results.append(result)
-            self.assertEqual(node._epoch, epoch)
             self.assertEqual(len(result), n)
             self.assertEqual(set(result), set(range(n)))
 
@@ -167,6 +173,7 @@ def test_distributed_sampler(self):
         node = SamplerWrapper(sampler=sampler)
 
         for epoch in range(4):
+            node.reset()
             result = list(node)
             self.assertEqual(result, exp[epoch])
 

test/nodes/test_base_node.py

@@ -6,26 +6,10 @@
 
 from torch.testing._internal.common_utils import TestCase
 from torchdata.nodes.adapters import IterableWrapper
-from torchdata.nodes.base_node import BaseNodeIterator
 
 from .utils import run_test_save_load_state
 
 
 class TestBaseNode(TestCase):
-    def test_started_finished(self) -> None:
-        x = IterableWrapper(range(10))
-        for _ in range(3):  # test multi-epoch
-            it = iter(x)
-            self.assertIsInstance(it, BaseNodeIterator)
-            self.assertFalse(it.started())
-            self.assertFalse(it.finished())
-
-            for _ in it:
-                self.assertTrue(it.started())
-                self.assertFalse(it.finished())
-
-            self.assertTrue(it.started())
-            self.assertTrue(it.finished())
-
     def test_save_load_state(self):
         run_test_save_load_state(self, IterableWrapper(range(10)), 5)

test/nodes/test_map.py

@@ -73,6 +73,7 @@ def _test_map(self, in_order, method) -> None:
 
         results: List[List[dict]] = [[], []]
         for epoch in range(2):
+            node.reset()
             for batch in node:
                 results[epoch].extend(batch)
 
@@ -119,7 +120,6 @@ def test_save_load_state_thread(self, midpoint: int, in_order: bool, snapshot_fr
         method = "thread"
         batch_size = 6
         n = 80
-        multiprocessing_context = None if IS_WINDOWS else "forkserver"
         src = MockSource(num_samples=n)
         node = Batcher(src, batch_size=batch_size, drop_last=False)
         node = ParallelMapper(
@@ -128,7 +128,6 @@ def test_save_load_state_thread(self, midpoint: int, in_order: bool, snapshot_fr
             num_workers=4,
             in_order=in_order,
             method=method,
-            multiprocessing_context=multiprocessing_context,
             snapshot_frequency=snapshot_frequency,
         )
         node = Prefetcher(node, prefetch_factor=2)

test/nodes/test_pin_memory.py

@@ -33,6 +33,7 @@ def test_pin_memory(self) -> None:
 
         # 2 epochs
         for epoch in range(2):
+            root.reset()
             results = list(root)
             self.assertEqual(len(results), 3, epoch)
             for i in range(3):

test/nodes/test_prefetch.py

@@ -24,6 +24,7 @@ def test_prefetcher(self) -> None:
 
         # Test multi epoch shutdown and restart
         for _ in range(2):
+            root.reset()
             results = list(root)
             self.assertEqual(len(results), 3)
             for i in range(3):

test/nodes/test_snapshot_store.py

@@ -5,12 +5,8 @@
 # LICENSE file in the root directory of this source tree.
 
 from torch.testing._internal.common_utils import TestCase
-from torchdata.nodes.adapters import IterableWrapper
-from torchdata.nodes.base_node import BaseNodeIterator
 from torchdata.nodes.snapshot_store import DequeSnapshotStore
 
-from .utils import run_test_save_load_state
-
 
 class TestDequeSnapshotStore(TestCase):
     def test_snapshot_store(self) -> None:

test/nodes/utils.py

@@ -50,22 +50,33 @@ def __call__(self, x):
 
 class IterInitError(BaseNode[int]):
     def __init__(self, msg: str = "Iter Init Error") -> None:
+        super().__init__()
         self.msg = msg
 
-    def iterator(self, initial_state: Optional[Dict[str, Any]]) -> Iterator[int]:
+    def reset(self, initial_state: Optional[Dict[str, Any]] = None):
+        super().reset(initial_state)
         raise ValueError(self.msg)
 
+    def next(self):
+        raise ValueError("next() should not be called")
+
     def get_state(self) -> Dict[str, Any]:
         return {}
 
 
 class DummyIterableDataset(torch.utils.data.IterableDataset):
-    def __init__(self, num_samples: int) -> None:
+    def __init__(self, num_samples: int, name: str) -> None:
         self.num_samples = num_samples
+        self.name = name
 
     def __iter__(self) -> Iterator[dict]:
         for i in range(self.num_samples):
-            yield {"step": i, "test_tensor": torch.tensor([i]), "test_str": f"str_{i}"}
+            yield {
+                "name": self.name,
+                "step": i,
+                "test_tensor": torch.tensor([i]),
+                "test_str": f"str_{i}",
+            }
 
 
 class DummyMapDataset(torch.utils.data.Dataset):
@@ -104,8 +115,6 @@ def run_test_save_load_state(test, node: BaseNode, midpoint: int):
     for val in it:
         results_1.append(val)
 
-    assert len(results_1) == len(results)
-
     ##############################
     # Test restoring from midpoint
     x.load_state_dict(state_dict)
@@ -118,7 +127,7 @@ def run_test_save_load_state(test, node: BaseNode, midpoint: int):
 
     ##############################
     # Test restoring from midpoint of epoch 1
-    x.load_state_dict(state_dict_1, restart_on_stop_iteration=True)
+    x.load_state_dict(state_dict_1)
     results_after_2 = list(x)
     test.assertEqual(results_after_2, results_1[midpoint:])
 

torchdata/nodes/__init__.py

@@ -19,6 +19,7 @@
     "Batcher",
     "DataLoader",
     "IterableWrapper",
+    "Loader",
     "MapStyleWrapper",
     "Mapper",
     "ParallelMapper",

torchdata/nodes/adapters.py

@@ -34,33 +34,37 @@ class IterableWrapper(BaseNode[T]):
     ITERABLE_KEY = "iterable"
 
     def __init__(self, iterable: Iterable[T]):
+        super().__init__()
         self.iterable = iterable
-        self._num_yielded = 0
+        self._it: Optional[Iterator[T]] = None
 
-    def iterator(self, initial_state: Optional[Dict[str, Any]]) -> Iterator[T]:
+    def reset(self, initial_state: Optional[Dict[str, Any]] = None):
         self._num_yielded = 0
+        self._it = None
+        super().reset(initial_state)
         if initial_state is not None:
             self._num_yielded = initial_state[self.NUM_YIELDED_KEY]
             if isinstance(self.iterable, Stateful):
                 self.iterable.load_state_dict(initial_state[self.ITERABLE_KEY])
-                it = iter(self.iterable)
+                self._it = iter(self.iterable)
             else:
-                it = iter(self.iterable)
+                self._it = iter(self.iterable)
                 # Naively fast-forwarding
                 for i in range(self._num_yielded):
                     try:
-                        next(it)
+                        next(self._it)
                     except StopIteration:
                         raise ValueError(
                             f"Tried to fast-forward {self._num_yielded} items during init but "
                             f"hit StopIteration after {i} items, this is likely a bug or malformed state_dict"
                         )
         else:
-            it = iter(self.iterable)
+            self._it = iter(self.iterable)
 
-        for item in it:
-            self._num_yielded += 1
-            yield item
+    def next(self) -> T:
+        item = next(self._it)  # type: ignore [arg-type, union-attr]
+        self._num_yielded += 1
+        return item
 
     def get_state(self) -> Dict[str, Any]:
         state_dict: Dict[str, Any] = {self.NUM_YIELDED_KEY: self._num_yielded}
@@ -92,68 +96,71 @@ class SamplerWrapper(BaseNode[T]):
     :param epoch_updater: Optional[Callable[[int], int]] = None - callback to update epoch at start of new iteration. It's called at the beginning of each iterator request, except the first one.
     """
 
+    NEXT_EPOCH_KEY = "_next_epoch"
     NUM_YIELDED_KEY = "_num_yielded"
-    SAMPLER_KEY = "sampler"
     EPOCH_KEY = "_epoch"
-    STARTED_KEY = "_started"
-
-    @classmethod
-    def _default_epoch_updater(cls, epoch: int) -> int:
-        return epoch + 1
+    SAMPLER_KEY = "_sampler"
 
     def __init__(
         self,
         sampler: Sampler[T],
         initial_epoch: int = 0,
         epoch_updater: Optional[Callable[[int], int]] = None,
     ):
+        super().__init__()
         self.sampler = sampler
-        self.epoch_updater = epoch_updater or self._default_epoch_updater
+        self.epoch = initial_epoch
         self._num_yielded = 0
-        self._epoch = initial_epoch
         self._started = False
+        self.epoch_updater = epoch_updater or self._default_epoch_updater
+        self._it: Optional[Iterator[T]] = None
 
-    def iterator(self, initial_state: Optional[Dict[str, Any]]) -> Iterator[T]:
-        it: Iterator[T]
-        self._num_yielded = 0
+    def reset(self, initial_state: Optional[Dict[str, Any]] = None):
+        super().reset(initial_state)
         if initial_state is not None:
             self._num_yielded = initial_state[self.NUM_YIELDED_KEY]
-            self._epoch = initial_state[self.EPOCH_KEY]
-            self._started = initial_state[self.STARTED_KEY]
-
+            self.epoch = initial_state[self.EPOCH_KEY]
             if isinstance(self.sampler, Stateful):
                 self.sampler.load_state_dict(initial_state[self.SAMPLER_KEY])
-                it = iter(self.sampler)
+                self._it = iter(self.sampler)  # type: ignore [assignment]
             else:
                 if hasattr(self.sampler, "set_epoch"):
-                    self.sampler.set_epoch(self._epoch)
-                it = iter(self.sampler)
+                    print("Setting epoch", self.epoch)
+                    self.sampler.set_epoch(self.epoch)
+                self._it = iter(self.sampler)
                 for i in range(self._num_yielded):
                     try:
-                        next(it)
+                        next(self._it)  # type: ignore [arg-type]
                     except StopIteration:
                         raise ValueError(
                             f"Tried to fast-forward {self._num_yielded} items during init but "
                             f"hit StopIteration after {i} items, this is likely a bug or malformed state_dict"
                         )
         else:
-            if self._started:  # don't call first time
-                self._epoch = self.epoch_updater(self._epoch)
+            self._num_yielded = 0
+            if self._started:
+                # Don't update epoch unless iterator has started
+                self.epoch = self.epoch_updater(self.epoch)
             if hasattr(self.sampler, "set_epoch"):
-                self.sampler.set_epoch(self._epoch)
-            it = iter(self.sampler)
+                self.sampler.set_epoch(self.epoch)
+            self._it = iter(self.sampler)
+        self._started = False
 
+    def next(self) -> T:
         self._started = True
-        for item in it:
-            self._num_yielded += 1
-            yield item
+        item = next(self._it)  # type: ignore [arg-type, union-attr]
+        self._num_yielded += 1
+        return item
 
     def get_state(self) -> Dict[str, Any]:
         state_dict: Dict[str, Any] = {
             self.NUM_YIELDED_KEY: self._num_yielded,
-            self.EPOCH_KEY: self._epoch,
-            self.STARTED_KEY: self._started,
+            self.EPOCH_KEY: self.epoch,
         }
         if isinstance(self.sampler, Stateful):
             state_dict[self.SAMPLER_KEY] = self.sampler.state_dict()
         return state_dict
+
+    @classmethod
+    def _default_epoch_updater(cls, epoch: int) -> int:
+        return epoch + 1