Bugfix: inconsistent streaming dataloader state (specific to StreamingDataset)

src/litdata/streaming/combined.py

@@ -134,6 +134,11 @@ def set_drop_last(self, drop_last: bool) -> None:
         for dataset in self._datasets:
             dataset.set_drop_last(drop_last)
 
+    def reset_state_dict(self) -> None:
+        """Reset the state of the dataset."""
+        for dataset in self._datasets:
+            dataset.reset_state_dict()
+
     def _check_datasets(self, datasets: List[StreamingDataset]) -> None:
         if any(not isinstance(d, StreamingDataset) for d in datasets):
             raise RuntimeError("The provided datasets should be instances of the StreamingDataset.")

src/litdata/streaming/dataloader.py

@@ -615,6 +615,7 @@ def __iter__(self) -> Any:
             self.current_epoch += 1
             self._num_samples_yielded_combined = {}
             self._num_samples_yielded_streaming = 0
+            self.dataset.reset_state_dict()
 
         self.dataset.set_epoch(self.current_epoch)
 

src/litdata/streaming/dataset.py

@@ -338,12 +338,14 @@ def __next__(self) -> Any:
         # Prevent to create more batch on a given process
         if self.global_index >= self.stop_length:
             self.current_epoch += 1
+            self.reset_state_dict()
             raise StopIteration
 
         # Lazily re-populate the interval to reduce memory usage.
         if len(self.current_indexes) == 0:
             if self.chunk_index == self.num_chunks:
                 self.current_epoch += 1
+                self.reset_state_dict()
                 raise StopIteration
 
             # reset index
@@ -388,7 +390,7 @@ def state_dict(self, num_samples_yielded: int, num_workers: int, batch_size: int
 
         return {
             "num_samples_yielded": num_samples_yielded,
-            "num_workers": num_workers,
+            "num_workers": num_workers or 1,
             "batch_size": batch_size,
             "current_epoch": self.current_epoch,
             "input_dir_path": self.input_dir.path,
@@ -407,6 +409,9 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
             # the state is restored within the workers
             self._state_dict = state_dict
 
+    def reset_state_dict(self) -> None:
+        self._state_dict = None
+
     def _validate_state_dict(self) -> None:
         assert self._state_dict
         assert self.worker_env

tests/streaming/test_combined.py

@@ -17,6 +17,9 @@ class TestCombinedStreamingDataset(CombinedStreamingDataset):
     def _check_datasets(self, datasets) -> None:
         pass
 
+    def reset_state_dict(self):
+        pass
+
 
 def test_combined_dataset_num_samples_yield():
     dataset = TestCombinedStreamingDataset(

tests/streaming/test_dataloader.py

@@ -56,6 +56,9 @@ class TestCombinedStreamingDataset(CombinedStreamingDataset):
     def _check_datasets(self, datasets) -> None:
         pass
 
+    def reset_state_dict(self):
+        pass
+
 
 def test_streaming_dataloader():
     dataset = TestCombinedStreamingDataset(
@@ -202,3 +205,46 @@ def test_dataloader_no_workers(tmpdir):
     assert len(dataset) == 1000
     assert len(dataloader) == 1000
     assert len(dataset) == 1000
+
+
+@pytest.mark.timeout(120)
+def test_dataloader_with_loading_states(tmpdir):
+    cache = Cache(input_dir=str(tmpdir), chunk_bytes="64MB")
+    for i in range(100):
+        cache[i] = i
+    cache.done()
+    cache.merge()
+
+    dataset = StreamingDataset(str(tmpdir), shuffle=True)
+
+    # Test dataloader without explicit num workers
+    dataloader = StreamingDataLoader(dataset, batch_size=4)
+    dataloader.load_state_dict(dataloader.state_dict())
+    batch = next(iter(dataloader))
+    assert len(batch) == 4, "Batch size should be 4"
+    assert len(dataloader) == 25, "Dataloader length should be 25 (100 items / batch size 4)"
+
+    # Test dataloader with num workers
+    dataloader = StreamingDataLoader(dataset, batch_size=4, num_workers=2)
+    assert len(dataloader) == 25, "Dataloader length should be 25 (100 items / batch size 4)"
+
+    # Verify dataloader state after partial iteration
+    for batch_idx, batch in enumerate(dataloader):
+        assert dataloader.current_epoch == 1, "Current epoch should be 1"
+        if batch_idx == 10:
+            break
+    dataloader.load_state_dict(dataloader.state_dict())
+
+    # Verify remaining batches in the first epoch
+    count = 0
+    for _ in dataloader:
+        assert dataloader.current_epoch == 1, "Current epoch should be 1"
+        count += 1
+    assert count == 15, "There should be atleast 15 batches remaining in the first epoch"
+
+    # Verify batches in the second epoch
+    count = 0
+    for _ in dataloader:
+        assert dataloader.current_epoch == 2, "Current epoch should be 2"
+        count += 1
+    assert count >= 25, "There should be at least 25 batches in the second epoch"