[WIP] : Fix resume issues with combined streaming dataset in dataloader

src/litdata/streaming/dataloader.py

@@ -665,14 +665,16 @@ def state_dict(self) -> Dict[str, Any]:
                 "latest_worker_idx": self._latest_worker_idx,
             }
 
-        num_samples_yieled = [0 for _ in range(len(list(self._num_samples_yielded_combined.values())[0]))]
+        # Initialize a list to track the number of samples yielded for each dataset
+        num_samples_yieled = [0 for _ in range(len(self.dataset._datasets))]
+
         for worker_idx in self._num_samples_yielded_combined:
             for dataset_idx, samples_yieled in enumerate(self._num_samples_yielded_combined[worker_idx]):
                 num_samples_yieled[dataset_idx] += samples_yieled
 
         return {
             "dataset": self.dataset.state_dict(self.num_workers, self.batch_size, num_samples_yieled),
-            "current_epoch": self.current_epoch if self.restore else self.current_epoch - 1,
+            "current_epoch": self.current_epoch,
             "latest_worker_idx": self._latest_worker_idx,
             "num_samples_yielded": deepcopy(self._num_samples_yielded_combined),
         }
@@ -701,21 +703,25 @@ def load_state_dict(self, obj: Dict[str, Any]) -> None:
 
         # Inform we are resuming and disable resetting the StreamingDataLoader state.
         # This is toggle back to False when the `__iter__` method of the StreamingDataLoader completes.
-        # self.restore = True
-
         if isinstance(self.dataset, CombinedStreamingDataset):
             self.dataset._set_use_streaming_dataloader(True)
             self.dataset.load_state_dict(obj)
 
-            # Inform that the dataloader is resuming.
-            # TODO: Check if the number of samples yielded is less than the length of the dataset.
-            # Also, len is not available for CombinedStreamingDataset incase of provided weights.
-            self.restore = True
+            total_samples_yielded = sum([sum(samples) for samples in self._num_samples_yielded_combined.values()])
+
+            # Check if we need to restore for the case without weights.
+            if self.dataset._iterate_over_all and total_samples_yielded < len(self.dataset):  # type: ignore
+                self.restore = True
+
+            # Check if we need to restore for the case with weights.
+            # Note: `len` is not available for CombinedStreamingDataset in case of provided weights.
+            # TODO: handle the case with weights.
+            if not self.dataset._iterate_over_all:
+                self.restore = True
 
         elif isinstance(self.dataset, StreamingDataset):
             self.dataset.load_state_dict(obj["dataset"])
 
-            # Inform that the dataloader is resuming.
             if self._num_samples_yielded_streaming < len(self.dataset):
                 self.restore = True
         else:

tests/streaming/test_combined.py

@@ -419,7 +419,7 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
                     "region_of_interest": ANY,
                 },
             },
-            "current_epoch": 0,
+            "current_epoch": 1,
             "latest_worker_idx": 0,
             "num_samples_yielded": {0: [2, 0]},
         },
@@ -456,7 +456,7 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
                     "region_of_interest": ANY,
                 },
             },
-            "current_epoch": 0,
+            "current_epoch": 1,
             "latest_worker_idx": 1,
             "num_samples_yielded": {0: [2, 0], 1: [2, 0]},
         },
@@ -493,7 +493,7 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
                     "region_of_interest": ANY,
                 },
             },
-            "current_epoch": 0,
+            "current_epoch": 1,
             "latest_worker_idx": 2,
             "num_samples_yielded": {0: [2, 0], 1: [2, 0], 2: [2, 0]},
         },
@@ -530,7 +530,7 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
                     "region_of_interest": ANY,
                 },
             },
-            "current_epoch": 0,
+            "current_epoch": 1,
             "latest_worker_idx": 0,
             "num_samples_yielded": {0: [3, 1], 1: [2, 0], 2: [2, 0]},
         },
@@ -567,7 +567,7 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
                     "region_of_interest": ANY,
                 },
             },
-            "current_epoch": 0,
+            "current_epoch": 1,
             "latest_worker_idx": 1,
             "num_samples_yielded": {0: [3, 1], 1: [3, 1], 2: [2, 0]},
         },
@@ -604,7 +604,7 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
                     "region_of_interest": ANY,
                 },
             },
-            "current_epoch": 0,
+            "current_epoch": 1,
             "latest_worker_idx": 2,
             "num_samples_yielded": {0: [3, 1], 1: [3, 1], 2: [3, 1]},
         },
@@ -641,7 +641,7 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
                     "region_of_interest": ANY,
                 },
             },
-            "current_epoch": 0,
+            "current_epoch": 1,
             "latest_worker_idx": 0,
             "num_samples_yielded": {0: [4, 1], 1: [3, 1], 2: [3, 1]},
         },
@@ -681,7 +681,7 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
                     "region_of_interest": ANY,
                 },
             },
-            "current_epoch": 1,
+            "current_epoch": 2,
             "latest_worker_idx": 0,
             "num_samples_yielded": {0: [2, 0]},
         },
@@ -718,7 +718,7 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
                     "region_of_interest": ANY,
                 },
             },
-            "current_epoch": 1,
+            "current_epoch": 2,
             "latest_worker_idx": 1,
             "num_samples_yielded": {0: [2, 0], 1: [2, 0]},
         },
@@ -755,7 +755,7 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
                     "region_of_interest": ANY,
                 },
             },
-            "current_epoch": 1,
+            "current_epoch": 2,
             "latest_worker_idx": 2,
             "num_samples_yielded": {0: [2, 0], 1: [2, 0], 2: [2, 0]},
         },
@@ -792,7 +792,7 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
                     "region_of_interest": ANY,
                 },
             },
-            "current_epoch": 1,
+            "current_epoch": 2,
             "latest_worker_idx": 0,
             "num_samples_yielded": {0: [3, 1], 1: [2, 0], 2: [2, 0]},
         },
@@ -829,7 +829,7 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
                     "region_of_interest": ANY,
                 },
             },
-            "current_epoch": 1,
+            "current_epoch": 2,
             "latest_worker_idx": 1,
             "num_samples_yielded": {0: [3, 1], 1: [3, 1], 2: [2, 0]},
         },
@@ -866,7 +866,7 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
                     "region_of_interest": ANY,
                 },
             },
-            "current_epoch": 1,
+            "current_epoch": 2,
             "latest_worker_idx": 2,
             "num_samples_yielded": {0: [3, 1], 1: [3, 1], 2: [3, 1]},
         },
@@ -903,7 +903,7 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
                     "region_of_interest": ANY,
                 },
             },
-            "current_epoch": 1,
+            "current_epoch": 2,
             "latest_worker_idx": 0,
             "num_samples_yielded": {0: [4, 1], 1: [3, 1], 2: [3, 1]},
         },
@@ -920,6 +920,6 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
         states_23.append(dataloader.state_dict())
 
     assert sum(not torch.equal(b1, b2) for b1, b2 in zip(batches_2[2:], batches_23)) == 0
-    assert states_23[0]["current_epoch"] == 1
+    assert states_23[0]["current_epoch"] == 2
 
     assert not dataloader.restore

tests/streaming/test_dataloader.py

@@ -330,3 +330,57 @@ def test_resume_dataloader_with_new_dataset(tmpdir):
     dataloader.load_state_dict(dataloader_state)
     for _ in dataloader:
         assert dataloader.current_epoch == 2, "Current epoch should be 2"
+
+
+@pytest.mark.timeout(120)
+def test_combined_dataset_dataloader_states(tmpdir):
+    datasets = [str(tmpdir.join(f"dataset_{i}")) for i in range(2)]
+    for dataset in datasets:
+        cache = Cache(input_dir=dataset, chunk_bytes="64MB")
+        for i in range(50):
+            cache[i] = i
+        cache.done()
+        cache.merge()
+
+    dataset_1 = StreamingDataset(datasets[0], shuffle=True)
+    dataset_2 = StreamingDataset(datasets[1], shuffle=True)
+    combined_dataset = CombinedStreamingDataset(datasets=[dataset_1, dataset_2])
+
+    # Test dataloader without explicit num workers
+    dataloader = StreamingDataLoader(combined_dataset, batch_size=4)
+    assert not dataloader.restore
+    dataloader.load_state_dict(dataloader.state_dict())
+    assert dataloader.restore
+    batch = next(iter(dataloader))
+    assert len(batch) == 4, "Batch size should be 4"
+    assert len(dataloader) == 25, "Dataloader length should be 25 (50+50 items / batch size 4)"
+
+    # Test dataloader with num workers
+    dataloader = StreamingDataLoader(combined_dataset, batch_size=4, num_workers=2)
+    assert len(dataloader) == 25, "Dataloader length should be 25 (50+50 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())
+    assert dataloader.restore
+
+    # 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"
+    assert not dataloader.restore
+
+    # 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"
+
+    # TODO: Add more conditions to check the state of the dataloader