Fix end of epoch StatefulDataLoader restart

test/stateful_dataloader/test_state_dict.py

@@ -1314,7 +1314,7 @@ def test(self):
                 dataset=dataset,
                 num_workers=num_workers,
                 collate_fn=identity,
-                multiprocessing_context="forkserver" if IS_MACOS and num_workers else None,
+                multiprocessing_context=("forkserver" if IS_MACOS and num_workers else None),
             )
             it = iter(dl)
             # Fetch at least one batch from each worker
@@ -1325,7 +1325,10 @@ def test(self):
             if num_workers > 0:
                 for i in range(num_workers):
                     # Ensure worker state is stored only once if the dataset is also the iterator
-                    self.assertEqual(state_dict["_snapshot"]["_worker_snapshots"][f"worker_{i}"]["dataset_state"], None)
+                    self.assertEqual(
+                        state_dict["_snapshot"]["_worker_snapshots"][f"worker_{i}"]["dataset_state"],
+                        None,
+                    )
                     self.assertTrue(
                         state_dict["_snapshot"]["_worker_snapshots"][f"worker_{i}"]["fetcher_state"][
                             "dataset_iter_state"
@@ -1441,6 +1444,154 @@ def test_fast_state_dict_request_skip_steps(self) -> None:
         self._run_test(17, 19)
 
 
+class TestMultiEpochSDL_shard0(TestCase):
+    def get_map_dl(self, data_size, num_workers, batch_size, shuffle):
+        dataset = DummyMapDataset(data_size, shuffle=False)
+        return StatefulDataLoader(
+            dataset=dataset,
+            num_workers=num_workers,
+            batch_size=batch_size,
+            shuffle=shuffle,
+            multiprocessing_context=("forkserver" if IS_MACOS and num_workers else None),
+        )
+
+    def _run(self, data_size, num_workers, batch_size, shuffle):
+        dl1 = self.get_map_dl(
+            data_size=data_size,
+            num_workers=num_workers,
+            batch_size=batch_size,
+            shuffle=shuffle,
+        )
+        # Run through the dataloader for 2 epochs and count the number of items yielded
+        num_items_yielded = 0
+        dl1_items = []
+        for _ in range(2):
+            for batch in dl1:
+                dl1_items.append(batch)
+                num_items_yielded += 1
+        # Save the state dict
+        state_dict = dl1.state_dict()
+        # Create a new StatefulDataLoader instance and load the state dict
+        new_dl1 = self.get_map_dl(
+            data_size=data_size,
+            num_workers=num_workers,
+            batch_size=batch_size,
+            shuffle=shuffle,
+        )
+        new_dl1.load_state_dict(state_dict)
+        # Run through the new dataloader for another 2 epochs and count the number of items yielded
+        additional_num_items_yielded = 0
+        for i in range(2):
+            epoch_num_items_yielded = 0
+            for batch in new_dl1:
+                dl1_items.append(batch)
+                epoch_num_items_yielded += 1
+            additional_num_items_yielded += epoch_num_items_yielded
+        # Check that the total number of items yielded is correct
+        self.assertEqual(num_items_yielded + additional_num_items_yielded, data_size * 4)
+
+        # now run a second dataloder for 4 epochs and check if the order is same.
+        dl2 = self.get_map_dl(
+            data_size=data_size,
+            num_workers=num_workers,
+            batch_size=batch_size,
+            shuffle=shuffle,
+        )
+        dl2_items = []
+        for _ in range(4):
+            for batch in dl2:
+                dl2_items.append(batch)
+
+        self.assertEqual(dl1_items, dl2_items)
+
+    def test_main_process(self):
+        self._run(100, 0, 1, False)
+
+    def test_multiprocess(self):
+        self._run(100, 2, 1, False)
+
+    def test_main_process_shuffle(self):
+        self._run(100, 0, 1, True)
+
+    def test_multiprocess_shuffle(self):
+        self._run(100, 2, 1, True)
+
+
+class TestEndOfEpochBehavior_shard0(TestCase):
+    def get_map_dl(self, data_size, num_workers, batch_size, shuffle):
+        dataset = DummyMapDataset(data_size, shuffle=False)
+        return StatefulDataLoader(
+            dataset=dataset,
+            num_workers=num_workers,
+            batch_size=batch_size,
+            shuffle=shuffle,
+            multiprocessing_context=("forkserver" if IS_MACOS and num_workers else None),
+        )
+
+    def _count_items_yielded(self, data_loader: StatefulDataLoader) -> int:
+        num_items_yielded = 0
+        for batch in data_loader:
+            num_items_yielded += 1
+        return num_items_yielded
+
+    def _run(self, data_size, num_workers, batch_size, shuffle):
+        dl = self.get_map_dl(
+            data_size=data_size,
+            num_workers=num_workers,
+            batch_size=batch_size,
+            shuffle=shuffle,
+        )
+        # Run through the dataloader for 1 epoch and count the number of items yielded
+        num_items_yielded = 0
+
+        for batch in dl:
+            num_items_yielded += 1
+            sd_in = dl.state_dict()
+        sd_out = dl.state_dict()
+
+        self.assertEqual(num_items_yielded, data_size)
+
+        # Create a new StatefulDataLoader instance and load the state dict saved before the end of epoch
+        dl_sd_in = self.get_map_dl(
+            data_size=data_size,
+            num_workers=num_workers,
+            batch_size=batch_size,
+            shuffle=shuffle,
+        )
+        dl_sd_in.load_state_dict(sd_in)
+
+        # Run through the new dataloader for 1 epoch and count the number of items yielded
+        # num_items_yielded should be 0 since the state dict was saved before the end of epoch
+        num_items_yielded = self._count_items_yielded(dl_sd_in)
+        self.assertEqual(num_items_yielded, 0)
+
+        # Create a new StatefulDataLoader instance and load the state dict saved after the end of epoch
+        dl_sd_out = self.get_map_dl(
+            data_size=data_size,
+            num_workers=num_workers,
+            batch_size=batch_size,
+            shuffle=shuffle,
+        )
+        dl_sd_out.load_state_dict(sd_out)
+
+        # Run through the new dataloader for 1 epoch and count the number of items yielded
+        # num_items_yielded should be data_size since the state dict was saved after the end of epoch
+        num_items_yielded = self._count_items_yielded(dl_sd_out)
+        self.assertEqual(num_items_yielded, data_size)
+
+    def test_main_process(self):
+        self._run(100, 0, 1, False)
+
+    def test_multiprocess(self):
+        self._run(100, 2, 1, False)
+
+    def test_main_process_shuffle(self):
+        self._run(100, 0, 1, True)
+
+    def test_multiprocess_shuffle(self):
+        self._run(100, 2, 1, True)
+
+
 class TestMultiEpochState_shard0(TestCase):
     def get_iterable_dl(self, pw, num_workers):
         data_size = [25, 50, 100, 75]