Make StreamingDataLoader shuffle to set shuffle to datasets. (#19481)

Author: tchaton

src/lightning/data/streaming/combined.py

@@ -66,6 +66,11 @@ def set_epoch(self, current_epoch: int) -> None:
         for dataset in self._datasets:
             dataset.set_epoch(current_epoch)
 
+    def set_shuffle(self, shuffle: bool) -> None:
+        """Set the current shuffle to the datasets."""
+        for dataset in self._datasets:
+            dataset.set_shuffle(shuffle)
+
     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/lightning/data/streaming/dataloader.py

@@ -541,6 +541,7 @@ def __init__(
         profile_batches: Union[bool, int] = False,
         profile_dir: Optional[str] = None,
         prefetch_factor: Optional[int] = None,
+        shuffle: Optional[bool] = None,
         **kwargs: Any,
     ) -> None:  # pyright: ignore
         if not isinstance(dataset, (StreamingDataset, CombinedStreamingDataset)):
@@ -549,6 +550,11 @@ def __init__(
                 f" Found {dataset}."
             )
 
+        if shuffle is not None:
+            dataset.set_shuffle(shuffle)
+
+        shuffle = None
+
         if profile_batches and not _VIZ_TRACKER_AVAILABLE:
             raise ModuleNotFoundError("To use profile_batches, viztracer is required. Run `pip install viztracer`")
 

src/lightning/data/streaming/dataset.py

@@ -107,6 +107,9 @@ def __init__(
         self.serializers = serializers
         self._state_dict: Optional[Dict[str, Any]] = None
 
+    def set_shuffle(self, shuffle: bool) -> None:
+        self.shuffle = shuffle
+
     def set_epoch(self, current_epoch: int) -> None:
         """Set the current epoch to the dataset on epoch starts.
 

tests/tests_data/streaming/test_dataloader.py

@@ -12,6 +12,10 @@ def __init__(self, size, step):
         self.size = size
         self.step = step
         self.counter = 0
+        self.shuffle = None
+
+    def set_shuffle(self, shuffle):
+        self.shuffle = shuffle
 
     def __len__(self):
         return self.size
@@ -92,3 +96,14 @@ def test_dataloader_profiling(profile, tmpdir, monkeypatch):
         batches.append(batch)
 
     assert os.path.exists(os.path.join(tmpdir, "result.json"))
+
+
+def test_dataloader_shuffle():
+    dataset = TestCombinedStreamingDataset(
+        [TestStatefulDataset(10, 1), TestStatefulDataset(10, -1)], 42, weights=(0.5, 0.5)
+    )
+    assert dataset._datasets[0].shuffle is None
+    assert dataset._datasets[1].shuffle is None
+    StreamingDataLoader(dataset, batch_size=2, num_workers=1, shuffle=True)
+    assert dataset._datasets[0].shuffle
+    assert dataset._datasets[1].shuffle