StreamingDataLoader: Resolve fault tolerance with the CombinedStreamingDataset and multiple workers (#19326)

Author: tchaton

.github/workflows/ci-tests-data.yml

@@ -87,7 +87,7 @@ jobs:
       #     ls -lh $PYPI_CACHE_DIR
 
       - name: Install package & dependencies
-        timeout-minutes: 20
+        timeout-minutes: 30
         run: |
           pip install -e ".[data-dev]" -U --prefer-binary -f ${TORCH_URL}
           pip list

src/lightning/data/streaming/combined.py

@@ -17,6 +17,10 @@
 from torch.utils.data import IterableDataset
 
 from lightning.data.streaming.dataset import StreamingDataset
+from lightning.data.utilities.env import _WorkerEnv
+
+__NUM_SAMPLES_YIELDED_KEY__ = "__NUM_SAMPLES_YIELDED__"
+__SAMPLES_KEY__ = "__SAMPLES__"
 
 
 class CombinedStreamingDataset(IterableDataset):
@@ -31,6 +35,8 @@ class CombinedStreamingDataset(IterableDataset):
     def __init__(
         self, datasets: List[StreamingDataset], seed: int = 42, weights: Optional[Sequence[float]] = None
     ) -> None:
+        self._check_datasets(datasets)
+
         self._seed = seed
         self._datasets = datasets
         self._weights = weights
@@ -43,53 +49,128 @@ def __init__(
             self._weights = [w / sum(weights) for w in weights]
 
         self._iterator: Optional[_CombinedDatasetIterator] = None
+        self._use_streaming_dataloader = False
+        self._num_samples_yielded: Optional[List[int]] = None
+        self._current_epoch = 0
+
+    def set_epoch(self, current_epoch: int) -> None:
+        """Set the current epoch to the datasets on epoch starts.
+
+        When using the StreamingDataLoader, this is done automatically
+
+        """
+        self._current_epoch = current_epoch
+        for dataset in self._datasets:
+            dataset.set_epoch(current_epoch)
+
+    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.")
+
+    def _set_use_streaming_dataloader(self, use_streaming_dataloader: bool) -> None:
+        # Used to prevent returning num_samples_yielded when using PyTorch DataLoader
+        self._use_streaming_dataloader = use_streaming_dataloader
 
     def __len__(self) -> int:
         assert self._weights
         return int(min([1 / w * len(d) for w, d in zip(self._weights, self._datasets) if w > 0]))
 
     def __iter__(self) -> Iterator[Any]:
         assert self._weights
-        self._iterator = _CombinedDatasetIterator(self._datasets, self._seed, self._weights)
+
+        worker_env = _WorkerEnv.detect()
+
+        num_samples_yielded = None
+
+        if self._num_samples_yielded is not None and worker_env.rank in self._num_samples_yielded:
+            num_samples_yielded = self._num_samples_yielded[worker_env.rank]
+
+        self._iterator = _CombinedDatasetIterator(
+            self._datasets,
+            self._seed,
+            self._weights,
+            self._use_streaming_dataloader,
+            num_samples_yielded,
+        )
         return self._iterator
 
-    def state_dict(self, num_workers: int, batch_size: int) -> Dict[str, Any]:
+    def state_dict(
+        self, num_workers: int, batch_size: int, num_samples_yielded: Optional[List[int]] = None
+    ) -> Dict[str, Any]:
         if self._iterator is None:
-            return {}
+            if num_samples_yielded is None:
+                return {}
+            return _state_dict(self._datasets, num_samples_yielded, num_workers, batch_size)
         return self._iterator.state_dict(num_workers, batch_size)
 
     def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
-        if len(state_dict) != len(self._datasets):
+        if not state_dict:
+            return
+
+        if len(state_dict["dataset"]) != len(self._datasets):
             raise RuntimeError(f"The provided state doesn't match the current number of datasets: {self._datasets}.")
 
         for dataset_idx, dataset in enumerate(self._datasets):
-            if str(dataset_idx) not in state_dict:
+            if str(dataset_idx) not in state_dict["dataset"]:
                 raise RuntimeError(f"The provided state doesn't contain the index {dataset_idx}.")
 
-            dataset.load_state_dict(state_dict[str(dataset_idx)])
+            dataset.load_state_dict(state_dict["dataset"][str(dataset_idx)])
+
+        # Used to iterate over the sampler to avoid sampling the same samples
+        if self._use_streaming_dataloader:
+            self._num_samples_yielded = state_dict["num_samples_yielded"]
 
 
 class _CombinedDatasetIterator(Iterator):
-    def __init__(self, datasets: List[StreamingDataset], seed: int, weights: Sequence[float]) -> None:
+    def __init__(
+        self,
+        datasets: List[StreamingDataset],
+        seed: int,
+        weights: Sequence[float],
+        use_streaming_dataloader: bool,
+        num_samples_yielded: Optional[Any] = None,
+    ) -> None:
         self._datasets = datasets
         self._dataset_iters = [iter(dataset) for dataset in datasets]
         self._dataset_indexes = list(range(len(datasets)))
         self._num_samples_yielded = [0 for _ in range(len(datasets))]
         self._weights = weights
         self._rng = random.Random(seed)
 
+        if num_samples_yielded is not None:
+            self._num_samples_yielded = num_samples_yielded
+            for _ in range(sum(num_samples_yielded)):
+                self._rng.choices(self._dataset_indexes, weights=self._weights, k=1)
+
+        self._use_streaming_dataloader = use_streaming_dataloader
+
     def __next__(self) -> Any:
         # randomly select a dataset index
         (dataset_index,) = self._rng.choices(self._dataset_indexes, weights=self._weights, k=1)
 
         # keep track the sample was fetched
         self._num_samples_yielded[dataset_index] += 1
 
+        sample = next(self._dataset_iters[dataset_index])
+
         # return a new sample
-        return next(self._dataset_iters[dataset_index])
+        if self._use_streaming_dataloader:
+            return {
+                __SAMPLES_KEY__: sample,
+                __NUM_SAMPLES_YIELDED_KEY__: self._num_samples_yielded,
+            }
+        return sample
 
     def state_dict(self, num_workers: int = 0, batch_size: int = 1) -> Dict[str, Any]:
-        return {
-            str(dataset_idx): dataset.state_dict(self._num_samples_yielded[dataset_idx], num_workers, batch_size)
-            for dataset_idx, dataset in enumerate(self._datasets)
-        }
+        return _state_dict(self._datasets, self._num_samples_yielded, num_workers, batch_size)
+
+
+def _state_dict(
+    datasets: List[StreamingDataset], num_samples_yielded: List[int], num_workers: int = 0, batch_size: int = 1
+) -> Dict[str, Any]:
+    return {
+        str(dataset_idx): dataset.state_dict(
+            num_samples_yielded=num_samples_yielded[dataset_idx], num_workers=num_workers, batch_size=batch_size
+        )
+        for dataset_idx, dataset in enumerate(datasets)
+    }

src/lightning/data/streaming/dataloader.py

@@ -15,7 +15,9 @@
 import inspect
 import logging
 import os
+from copy import deepcopy
 from importlib import reload
+from itertools import cycle
 from typing import Any, Callable, Dict, List, Optional, Union
 
 import torch
@@ -32,7 +34,11 @@
 from torch.utils.data.sampler import BatchSampler, Sampler
 
 from lightning.data.streaming import Cache
-from lightning.data.streaming.combined import CombinedStreamingDataset
+from lightning.data.streaming.combined import (
+    __NUM_SAMPLES_YIELDED_KEY__,
+    __SAMPLES_KEY__,
+    CombinedStreamingDataset,
+)
 from lightning.data.streaming.constants import _DEFAULT_CHUNK_BYTES, _TORCH_GREATER_EQUAL_2_1_0, _VIZ_TRACKER_AVAILABLE
 from lightning.data.streaming.dataset import StreamingDataset
 from lightning.data.streaming.sampler import CacheBatchSampler
@@ -341,6 +347,35 @@ def _get_iterator(self) -> "_BaseDataLoaderIter":
         return _MultiProcessingDataLoaderIterPatch(self)
 
 
+class _StreamingMultiProcessingDataLoaderIter(_MultiProcessingDataLoaderIter):
+    def __init__(self, loader: DataLoader) -> None:
+        self._loader = loader
+        self._indexes = (
+            list(range(self._loader._latest_worker_idx, self._loader.num_workers))
+            if self._loader._latest_worker_idx > 0
+            else []
+        )
+        super().__init__(loader)
+
+    def _try_put_index(self) -> None:
+        # Used to restart on the right DataLoader worker
+        if self._loader.restore and self._indexes:
+            assert self._tasks_outstanding < self._prefetch_factor * self._num_workers
+
+            try:
+                index = self._next_index()
+            except StopIteration:
+                return
+            worker_queue_idx = self._indexes.pop(0)
+
+            self._index_queues[worker_queue_idx].put((self._send_idx, index))
+            self._task_info[self._send_idx] = (worker_queue_idx,)
+            self._tasks_outstanding += 1
+            self._send_idx += 1
+        else:
+            super()._try_put_index()
+
+
 class StreamingDataLoader(DataLoader):
     """The `StreamingDataLoader` keeps track of the number of samples fetched in order to enable resumability of the
     dataset."""
@@ -355,27 +390,82 @@ def __init__(
         num_workers: int = 0,
         **kwargs: Any,
     ) -> None:  # pyright: ignore
+        if not isinstance(dataset, (StreamingDataset, CombinedStreamingDataset)):
+            raise RuntimeError(
+                "The provided dataset should be either an instance of StreamingDataset or CombinedStreamingDataset."
+                f" Found {dataset}."
+            )
+
+        self.current_epoch = 0
         self.batch_size = batch_size
         self.num_workers = num_workers
-        self.num_samples_yielded = 0
+        self._num_samples_yielded_streaming = 0
+        self._num_samples_yielded_combined: Dict[int, List[Any]] = {}
+        self.rng_state: Optional[Any] = None
+        self._worker_idx = cycle(list(range(self.num_workers if self.num_workers > 0 else 1)))
+        self._worker_idx_iter: Optional[Any] = None
+        self._latest_worker_idx = 0
+        self.restore = False
         super().__init__(dataset, *args, batch_size=batch_size, num_workers=num_workers, **kwargs)  # type: ignore
 
     def __iter__(self) -> Any:
+        if not self.restore:
+            self._latest_worker_idx = 0
+            self._worker_idx = cycle(list(range(self.num_workers if self.num_workers > 0 else 1)))
+            self._worker_idx_iter = iter(self._worker_idx)
+            self.current_epoch += 1
+            self._num_samples_yielded_combined = {}
+            self._num_samples_yielded_streaming = 0
+
+        self.dataset.set_epoch(self.current_epoch)
+
         if isinstance(self.dataset, StreamingDataset):
             assert self.batch_size
-            self.num_samples_yielded = 0
             for batch in super().__iter__():
-                self.num_samples_yielded += self.batch_size
+                self._latest_worker_idx = next(self._worker_idx_iter)  # type: ignore
+                self._num_samples_yielded_streaming += self.batch_size
                 yield batch
         else:
-            yield from super().__iter__()
+            self.dataset._set_use_streaming_dataloader(True)
+            assert self.batch_size
+            # TODO: Inject a custom collate function to avoid collating the __NUM_SAMPLES_YIELDED__ key
+            for batch in super().__iter__():
+                self._latest_worker_idx = next(self._worker_idx_iter)  # type: ignore
+                if isinstance(batch, dict) and __NUM_SAMPLES_YIELDED_KEY__ in batch:
+                    self._num_samples_yielded_combined[self._latest_worker_idx] = [
+                        sample[-1].item() if self.batch_size > 1 else sample.item()
+                        for sample in batch[__NUM_SAMPLES_YIELDED_KEY__]
+                    ]
+
+                    yield batch[__SAMPLES_KEY__]
+                else:
+                    yield batch
+
+        self.restore = False
 
     def state_dict(self) -> Dict[str, Any]:
         if isinstance(self.dataset, StreamingDataset):
             assert self.batch_size
-            num_samples = self.num_samples_yielded
-            return self.dataset.state_dict(num_samples, self.num_workers, self.batch_size)
-        return self.dataset.state_dict(self.num_workers, self.batch_size)
+            return {
+                "dataset": self.dataset.state_dict(
+                    self._num_samples_yielded_streaming, self.num_workers, self.batch_size
+                ),
+                "current_epoch": self.current_epoch,
+                "num_samples_yielded": self._num_samples_yielded_streaming,
+                "latest_worker_idx": self._latest_worker_idx,
+            }
+
+        num_samples_yieled = [0 for _ in range(len(list(self._num_samples_yielded_combined.values())[0]))]
+        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,
+            "latest_worker_idx": self._latest_worker_idx,
+            "num_samples_yielded": deepcopy(self._num_samples_yielded_combined),
+        }
 
     def load_state_dict(self, obj: Dict[str, Any]) -> None:
         """Load a dict containing training state (called from non-worker process).
@@ -386,7 +476,34 @@ def load_state_dict(self, obj: Dict[str, Any]) -> None:
             obj (Any): The state.
 
         """
-        if isinstance(self.dataset, (StreamingDataset, CombinedStreamingDataset)):
+        self.current_epoch = obj["current_epoch"]
+
+        if isinstance(self.dataset, StreamingDataset):
+            self._num_samples_yielded_streaming = obj["num_samples_yielded"]
+        else:
+            self._num_samples_yielded_combined = obj["num_samples_yielded"]
+
+        # Used to restart on the next DataLoader worker from the previous run.
+        self._latest_worker_idx = obj["latest_worker_idx"] + 1
+        self._worker_idx_iter = iter(self._worker_idx)
+        for _ in range(self._latest_worker_idx):
+            next(self._worker_idx_iter)
+
+        # 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)
+        elif isinstance(self.dataset, StreamingDataset):
+            self.dataset.load_state_dict(obj["dataset"])
         else:
             raise RuntimeError("The provided dataset should be a `StreamingDataset` or a `CombinedStreamingDataset`.")
+
+    def _get_iterator(self) -> "_BaseDataLoaderIter":
+        """Overriden to ensure the `Cache.done()` method is triggered on iteration done."""
+        if self.num_workers == 0:
+            return _SingleProcessDataLoaderIter(self)
+        self.check_worker_number_rationality()
+        return _StreamingMultiProcessingDataLoaderIter(self)