update

Author: thomas

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

@@ -90,6 +90,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.