[Bug] Add SharedCycleIteratorState (#8889)

Author: tchaton

CHANGELOG.md

@@ -4,6 +4,11 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
+
+## [1.4.3] - 2021-08-17
+
+- Fixed infinite loop with CycleIterator and multiple loaders ([#8889](https://github.com/PyTorchLightning/pytorch-lightning/pull/8889))
+
 ## [1.4.2] - 2021-08-10
 
 - Fixed recursive call for `apply_to_collection(include_none=False)` ([#8719](https://github.com/PyTorchLightning/pytorch-lightning/pull/8719))
@@ -27,6 +32,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed plateau scheduler stepping on incomplete epoch ([#8861](https://github.com/PyTorchLightning/pytorch-lightning/pull/8861))
 
 
+- Fixed infinite loop with CycleIterator and multiple loaders ([#8889](https://github.com/PyTorchLightning/pytorch-lightning/pull/8889))
+
+
 ## [1.4.0] - 2021-07-27
 
 ### Added

pytorch_lightning/trainer/supporters.py

@@ -14,8 +14,9 @@
 
 import os
 from collections.abc import Iterable, Iterator, Mapping, Sequence
+from dataclasses import dataclass, field
 from functools import partial
-from typing import Any, Callable, Dict, Generator, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -170,12 +171,35 @@ def to_disk(self) -> None:
                 torch.save(outputs, fp)
 
 
+@dataclass
+class SharedCycleIteratorState:
+
+    mode: str = "max_size_cycle"
+    dataloaders: List[DataLoader] = field(default_factory=lambda: [])
+    has_finished: Dict[int, bool] = field(default_factory=lambda: {})
+    has_reset: bool = False
+
+    def reset(self) -> None:
+        for dataloader in self.dataloaders:
+            self.has_finished[id(dataloader)] = False
+        self.has_reset = True
+
+    @property
+    def done(self) -> bool:
+        if not self.has_reset:
+            raise MisconfigurationException("Please, call reset once all dataloaders have been added.")
+        if len(self.dataloaders) == 1:
+            return False
+        decision_fn = all if self.mode == "max_size_cycle" else any
+        return decision_fn(self.has_finished.values())
+
+
 class CycleIterator:
     """
     Iterator for restarting a dataloader if it runs out of samples
     """
 
-    def __init__(self, loader: Any, length: Optional[int] = None):
+    def __init__(self, loader: Any, length: Optional[int] = None, state: SharedCycleIteratorState = None):
         """
         Args:
             loader: the loader to restart for cyclic (and optionally infinite) sampling
@@ -185,6 +209,15 @@ def __init__(self, loader: Any, length: Optional[int] = None):
         if length is None:
             length = float("inf")
 
+        if not state:
+            state = SharedCycleIteratorState()
+            state.dataloaders.append(loader)
+            state.reset()
+        else:
+            state.dataloaders.append(loader)
+
+        self.state = state
+
         self.length = length
         self.loader = loader
         self._loader_iter = None
@@ -205,21 +238,27 @@ def __next__(self) -> Any:
         """
         Fetches the next batch from internal dataloader and restarts
         it if necessary
-
         Returns:
             Any: the resulting batch
-
         Raises:
             StopIteration: if more then :attr:`length` batches have been returned
         """
         # Note: if self.length is `inf`, then the iterator will never stop
-        if self.counter >= self.__len__():
+        if self.counter >= self.__len__() or self.state.done:
             raise StopIteration
 
         try:
             return next(self._loader_iter)
 
         except StopIteration:
+
+            # inform the shared state this loader has completed
+            self.state.has_finished[id(self.loader)] = True
+
+            # check if iteration should be stopped.
+            if self.state.done:
+                raise StopIteration
+
             self._loader_iter = iter(self.loader)
             return next(self._loader_iter)
 
@@ -468,10 +507,14 @@ def _wrap_loaders_max_size_cycle(self) -> Any:
 
         # multiple loaders
         if isinstance(self.loaders, (Sequence, Mapping)):
+            state = SharedCycleIteratorState()
+
             self.loaders = apply_to_collection(
-                self.loaders, Iterable, CycleIterator, length=length, wrong_dtype=(Sequence, Mapping)
+                self.loaders, Iterable, CycleIterator, length=length, state=state, wrong_dtype=(Sequence, Mapping)
             )
 
+            state.reset()
+
     def __iter__(self) -> Any:
         """
         Create and return an iterator, `CombinedLoaderIterator`, for the combined loader.

tests/trainer/test_supporters.py

@@ -20,7 +20,7 @@
 from torch.utils.data import DataLoader, TensorDataset
 from torch.utils.data.dataset import Dataset, IterableDataset
 from torch.utils.data.distributed import DistributedSampler
-from torch.utils.data.sampler import Sampler
+from torch.utils.data.sampler import Sampler, SequentialSampler
 
 from pytorch_lightning import Trainer
 from pytorch_lightning.trainer.supporters import (
@@ -59,6 +59,7 @@ def test_tensor_running_accum_reset():
 
 def test_cycle_iterator():
     """Test the cycling function of `CycleIterator`"""
+
     iterator = CycleIterator(range(100), 1000)
     assert len(iterator) == 1000
     for idx, item in enumerate(iterator):
@@ -216,6 +217,83 @@ def test_combined_loader_sequence_min_size():
     assert idx == len(combined_loader) - 1
 
 
+class TestIterableDataset(IterableDataset):
+    def __init__(self, size: int = 10):
+        self.size = size
+
+    def __iter__(self):
+        self.sampler = SequentialSampler(range(self.size))
+        self.sampler_iter = iter(self.sampler)
+        return self
+
+    def __next__(self):
+        return next(self.sampler_iter)
+
+
+@pytest.mark.parametrize("mode", ["min_size", "max_size_cycle"])
+@pytest.mark.parametrize("use_multiple_dataloaders", [False, True])
+def test_combined_loader_sequence_iterable_dataset(mode, use_multiple_dataloaders):
+    """Test `CombinedLoader` of mode 'min_size' given sequence loaders"""
+    if use_multiple_dataloaders:
+        loaders = [
+            torch.utils.data.DataLoader(TestIterableDataset(10), batch_size=2),
+            torch.utils.data.DataLoader(TestIterableDataset(20), batch_size=2),
+        ]
+    else:
+        loaders = [
+            torch.utils.data.DataLoader(TestIterableDataset(10), batch_size=2),
+        ]
+
+    combined_loader = CombinedLoader(loaders, mode)
+
+    has_break = False
+
+    for idx, item in enumerate(combined_loader):
+        assert isinstance(item, Sequence)
+        assert len(item) == 2 if use_multiple_dataloaders else 1
+        if not use_multiple_dataloaders and idx == 4:
+            has_break = True
+            break
+
+    if mode == "max_size_cycle":
+        assert combined_loader.loaders[0].state.done == (not has_break)
+    expected = (10 if mode == "max_size_cycle" else 5) if use_multiple_dataloaders else 5
+    assert (expected - 1) == idx, (mode, use_multiple_dataloaders)
+
+
+@pytest.mark.parametrize("lengths", [[4, 6], [5, 5], [6, 4]])
+def test_combined_loader_sequence_with_map_and_iterable(lengths):
+    class MyIterableDataset(IterableDataset):
+        def __init__(self, size: int = 10):
+            self.size = size
+
+        def __iter__(self):
+            self.sampler = SequentialSampler(range(self.size))
+            self.iter_sampler = iter(self.sampler)
+            return self
+
+        def __next__(self):
+            return next(self.iter_sampler)
+
+    class MyMapDataset(Dataset):
+        def __init__(self, size: int = 10):
+            self.size = size
+
+        def __getitem__(self, index):
+            return index
+
+        def __len__(self):
+            return self.size
+
+    x, y = lengths
+    loaders = [DataLoader(MyIterableDataset(x)), DataLoader(MyMapDataset(y))]
+    dataloader = CombinedLoader(loaders, mode="max_size_cycle")
+    counter = 0
+    for _ in dataloader:
+        counter += 1
+    assert counter == max(x, y)
+
+
 def test_combined_loader_sequence_max_size_cycle():
     """Test `CombinedLoader` of mode 'max_size_cycle' given sequence loaders"""
     loaders = [