Set limits for fetcher.done

src/lightning/pytorch/CHANGELOG.md

@@ -173,6 +173,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - The `FSDPStrategy.load_optimizer_state_dict` and `FSDPStrategy.load_model_state_dict` are a no-op now ([#18358](https://github.com/Lightning-AI/lightning/pull/18358))
 
 
+- The `Trainer.num_val_batches`, `Trainer.num_test_batches` and `Trainer.num_sanity_val_batches` now return a list of sizes per dataloader instead of a single integer ([#18441](https://github.com/Lightning-AI/lightning/pull/18441))
+
 - The `*_step(dataloader_iter)` flavor now no longer takes the `batch_idx` in the signature ([#18390](https://github.com/Lightning-AI/lightning/pull/18390))
 - Calling `next(dataloader_iter)` now returns a triplet `(batch, batch_idx, dataloader_idx)` ([#18390](https://github.com/Lightning-AI/lightning/pull/18390))
 - Calling `next(combined_loader)` now returns a triplet `(batch, batch_idx, dataloader_idx)` ([#18390](https://github.com/Lightning-AI/lightning/pull/18390))

src/lightning/pytorch/loops/evaluation_loop.py

@@ -38,7 +38,7 @@
 )
 from lightning.pytorch.trainer.connectors.logger_connector.result import _OUT_DICT, _ResultCollection
 from lightning.pytorch.trainer.states import RunningStage, TrainerFn
-from lightning.pytorch.utilities.combined_loader import _Sequential, CombinedLoader
+from lightning.pytorch.utilities.combined_loader import CombinedLoader
 from lightning.pytorch.utilities.data import has_len_all_ranks
 from lightning.pytorch.utilities.exceptions import SIGTERMException
 from lightning.pytorch.utilities.model_helpers import is_overridden
@@ -63,8 +63,7 @@ def __init__(
         self.verbose = verbose
         self.inference_mode = inference_mode
         self.batch_progress = _BatchProgress()  # across dataloaders
-        #  list in "sequential" mode, number otherwise
-        self._max_batches: Union[int, float, List[Union[int, float]]] = []
+        self._max_batches: List[Union[int, float]] = []
 
         self._results = _ResultCollection(training=False)
         self._logged_outputs: List[_OUT_DICT] = []
@@ -85,24 +84,17 @@ def num_dataloaders(self) -> int:
         return len(combined_loader.flattened)
 
     @property
-    def max_batches(self) -> Union[int, float, List[Union[int, float]]]:
-        """In "sequential" mode, the max number of batches to run per dataloader.
-
-        Otherwise, the max batches to run.
-
-        """
+    def max_batches(self) -> List[Union[int, float]]:
+        """The max number of batches to run per dataloader."""
         max_batches = self._max_batches
         if not self.trainer.sanity_checking:
             return max_batches
-        sanity_val_steps = self.trainer.num_sanity_val_steps
-        if isinstance(max_batches, list):
-            return [min(sanity_val_steps, batches) for batches in max_batches]
-        return min(sanity_val_steps, max_batches)
+        return [min(self.trainer.num_sanity_val_steps, batches) for batches in max_batches]
 
     @property
     def skip(self) -> bool:
         """Returns whether the evaluation should be skipped."""
-        return sum(self.max_batches) == 0 if isinstance(self.max_batches, list) else self.max_batches == 0
+        return sum(self.max_batches) == 0
 
     @property
     def _should_reload_val_dl(self) -> bool:
@@ -195,17 +187,13 @@ def setup_data(self) -> None:
         if trainer.datamodule is not None:
             allow_zero_length |= trainer.datamodule.allow_zero_length_dataloader_with_multiple_devices
 
-        if self._is_sequential:
-            self._max_batches = []
-            for dl in combined_loader.flattened:
-                # determine number of batches
-                length = len(dl) if has_len_all_ranks(dl, trainer.strategy, allow_zero_length) else float("inf")
-                limit_batches = getattr(trainer, f"limit_{stage.dataloader_prefix}_batches")
-                num_batches = _parse_num_batches(stage, length, limit_batches)
-                self._max_batches.append(num_batches)
-        else:
-            has_len_all_ranks_ = has_len_all_ranks(combined_loader, trainer.strategy, allow_zero_length)
-            self._max_batches = len(combined_loader) if has_len_all_ranks_ else float("inf")
+        self._max_batches = []
+        for dl in combined_loader.flattened:
+            # determine number of batches
+            length = len(dl) if has_len_all_ranks(dl, trainer.strategy, allow_zero_length) else float("inf")
+            limit_batches = getattr(trainer, f"limit_{stage.dataloader_prefix}_batches")
+            num_batches = _parse_num_batches(stage, length, limit_batches)
+            self._max_batches.append(num_batches)
 
         # this depends on the data used, so reset it too
         self._seen_batches_per_dataloader = defaultdict(int)
@@ -237,13 +225,10 @@ def reset(self) -> None:
                 # some users want validation shuffling based on the training progress
                 _set_sampler_epoch(dl, trainer.fit_loop.epoch_progress.current.processed)
 
+        # set the per-dataloader limits
+        combined_loader.limits = self.max_batches
         data_fetcher.setup(combined_loader)
         iter(data_fetcher)  # creates the iterator inside the fetcher
-        if isinstance(combined_loader._iterator, _Sequential):
-            # set the per-dataloader limits
-            max_batches = self.max_batches
-            assert isinstance(max_batches, list)
-            combined_loader._iterator.limits = max_batches
 
         # add the previous `fetched` value to properly track `is_last_batch` with no prefetching
         data_fetcher.fetched += self.batch_progress.current.ready

src/lightning/pytorch/loops/fetchers.py

@@ -43,12 +43,10 @@ def combined_loader(self) -> CombinedLoader:
 
     def setup(self, combined_loader: CombinedLoader) -> None:
         self._combined_loader = combined_loader
-        self.length = sized_len(combined_loader)
-        self.done = self.length == 0
 
     def __iter__(self) -> "_DataFetcher":
-        self.reset()
         self.iterator = iter(self.combined_loader)
+        self.reset()
         return self
 
     def __next__(self) -> _ITERATOR_RETURN:
@@ -68,7 +66,10 @@ def __next__(self) -> _ITERATOR_RETURN:
 
     def reset(self) -> None:
         self.fetched = 0
-        self.done = False
+        # teardown calls `reset()`, and if it happens early, `combined_loader` can still be None
+        if self._combined_loader is not None:
+            self.length = sized_len(self.combined_loader)
+            self.done = self.length == 0
 
     def teardown(self) -> None:
         self.reset()
@@ -189,6 +190,9 @@ def length(self) -> Optional[int]:
 
     def __next__(self) -> _ITERATOR_RETURN:
         fetcher = self.data_fetcher
+        if fetcher.done:
+            # The iterator may still have items, but the length is reached (determined by the limits)
+            raise StopIteration
         batch, batch_idx, dataloader_idx = super(_DataLoaderIterDataFetcher, fetcher).__next__()
         # save the state so the loops can access it
         fetcher._batch = batch

src/lightning/pytorch/loops/fit_loop.py

@@ -17,7 +17,8 @@
 import torch
 
 import lightning.pytorch as pl
-from lightning.fabric.utilities.data import _set_sampler_epoch
+from lightning.fabric.utilities.data import _set_sampler_epoch, sized_len
+from lightning.fabric.utilities.warnings import PossibleUserWarning
 from lightning.pytorch.loops import _Loop
 from lightning.pytorch.loops.fetchers import _DataFetcher
 from lightning.pytorch.loops.progress import _Progress
@@ -39,7 +40,6 @@
 from lightning.pytorch.utilities.exceptions import MisconfigurationException, SIGTERMException
 from lightning.pytorch.utilities.model_helpers import is_overridden
 from lightning.pytorch.utilities.rank_zero import rank_zero_debug, rank_zero_info, rank_zero_warn
-from lightning.pytorch.utilities.warnings import PossibleUserWarning
 
 log = logging.getLogger(__name__)
 
@@ -244,13 +244,29 @@ def setup_data(self) -> None:
         if trainer.datamodule is not None:
             allow_zero_length |= trainer.datamodule.allow_zero_length_dataloader_with_multiple_devices
 
+        limits = []
+        for dl in combined_loader.flattened:
+            # determine number of batches
+            length = len(dl) if has_len_all_ranks(dl, trainer.strategy, allow_zero_length) else float("inf")
+            num_batches = _parse_num_batches(stage, length, trainer.limit_train_batches)
+            limits.append(num_batches)
+
+        combined_loader.limits = limits
+
+        training = trainer.training
+        trainer.training = True
+        self._data_fetcher = _select_data_fetcher(trainer)
+        trainer.training = training
+
+        self._data_fetcher.setup(combined_loader)
+        iter(self._data_fetcher)  # creates the iterator inside the fetcher
+        max_batches = sized_len(combined_loader)
+        self.max_batches = max_batches if max_batches is not None else float("inf")
         has_len_all_ranks_ = has_len_all_ranks(combined_loader, trainer.strategy, allow_zero_length)
-        self.max_batches = len(combined_loader) if has_len_all_ranks_ else float("inf")
+
         if self.max_batches == 0:
             return
 
-        self.max_batches = _parse_num_batches(stage, self.max_batches, trainer.limit_train_batches)
-
         # store epoch of dataloader reset for reload_dataloaders_every_n_epochs
         self._last_train_dl_reload_epoch = trainer.current_epoch
 
@@ -308,8 +324,6 @@ def on_run_start(self) -> None:
             self.epoch_loop.val_loop.setup_data()
             trainer.training = True
 
-        self._data_fetcher = _select_data_fetcher(trainer)
-
         call._call_callback_hooks(trainer, "on_train_start")
         call._call_lightning_module_hook(trainer, "on_train_start")
         call._call_strategy_hook(trainer, "on_train_start")
@@ -344,7 +358,6 @@ def advance(self) -> None:
                 f" The available modes are: {[m for m in _SUPPORTED_MODES if m != 'sequential']}"
             )
         assert (data_fetcher := self._data_fetcher) is not None
-        data_fetcher.setup(combined_loader)
         with self.trainer.profiler.profile("run_training_epoch"):
             self.epoch_loop.run(data_fetcher)
 

src/lightning/pytorch/loops/prediction_loop.py

@@ -35,7 +35,7 @@
     _request_dataloader,
 )
 from lightning.pytorch.trainer.states import RunningStage, TrainerFn
-from lightning.pytorch.utilities.combined_loader import _Sequential, CombinedLoader
+from lightning.pytorch.utilities.combined_loader import CombinedLoader
 from lightning.pytorch.utilities.data import has_len_all_ranks
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
 from lightning.pytorch.utilities.types import _PREDICT_OUTPUT
@@ -169,11 +169,12 @@ def reset(self) -> None:
         assert combined_loader is not None
         if combined_loader._mode != "sequential":
             raise ValueError('`trainer.predict()` only supports the `CombinedLoader(mode="sequential")` mode.')
+
+        # set the per-dataloader limits
+        combined_loader.limits = self.max_batches
         data_fetcher.setup(combined_loader)
         iter(data_fetcher)  # creates the iterator inside the fetcher
-        assert isinstance(combined_loader._iterator, _Sequential)
-        # set the per-dataloader limits
-        combined_loader._iterator.limits = self.max_batches
+
         # add the previous `fetched` value to properly track `is_last_batch` with no prefetching
         data_fetcher.fetched += self.batch_progress.current.ready
         data_fetcher._start_profiler = self._on_before_fetch

src/lightning/pytorch/loops/training_epoch_loop.py

@@ -134,7 +134,7 @@ def run(self, data_fetcher: _DataFetcher) -> None:
         while not self.done:
             try:
                 self.advance(data_fetcher)
-                self.on_advance_end()
+                self.on_advance_end(data_fetcher)
                 self._restarting = False
             except StopIteration:
                 break
@@ -164,7 +164,9 @@ def reset(self) -> None:
             self.val_loop.batch_progress.total.reset()
 
     def on_run_start(self, data_fetcher: _DataFetcher) -> None:
-        iter(data_fetcher)  # creates the iterator inside the fetcher
+        if self.trainer.current_epoch > 0:  # `iter()` was called once in `FitLoop.setup_data()` already
+            iter(data_fetcher)  # creates the iterator inside the fetcher
+
         # add the previous `fetched` value to properly track `is_last_batch` with no prefetching
         data_fetcher.fetched += self.batch_progress.current.ready
         data_fetcher._start_profiler = self._on_before_fetch
@@ -183,7 +185,7 @@ def advance(self, data_fetcher: _DataFetcher) -> None:
             StopIteration: When the epoch is canceled by the user returning -1
 
         """
-        if self.restarting and self._should_check_val_fx():
+        if self.restarting and self._should_check_val_fx(data_fetcher):
             # skip training and run validation in `on_advance_end`
             return
         # we are going to train first so the val loop does not need to restart
@@ -200,6 +202,7 @@ def advance(self, data_fetcher: _DataFetcher) -> None:
             # TODO: we should instead use the batch_idx returned by the fetcher, however, that will require saving the
             # fetcher state so that the batch_idx is correct after restarting
             batch_idx = self.batch_idx + 1
+        # Note: `is_last_batch` is not yet determined if data fetcher is a `_DataLoaderIterDataFetcher`
         self.batch_progress.is_last_batch = data_fetcher.done
 
         trainer = self.trainer
@@ -249,6 +252,8 @@ def advance(self, data_fetcher: _DataFetcher) -> None:
             # update the hook kwargs now that the step method might have consumed the iterator
             batch = data_fetcher._batch
             batch_idx = data_fetcher._batch_idx
+            # update `is_last_batch` again after dataloader_iter was fetched in `training_step()`
+            self.batch_progress.is_last_batch = data_fetcher.done
 
         call._call_callback_hooks(trainer, "on_train_batch_end", batch_output, batch, batch_idx)
         call._call_lightning_module_hook(trainer, "on_train_batch_end", batch_output, batch, batch_idx)
@@ -261,11 +266,11 @@ def advance(self, data_fetcher: _DataFetcher) -> None:
         # -----------------------------------------
         trainer._logger_connector.update_train_step_metrics()
 
-    def on_advance_end(self) -> None:
+    def on_advance_end(self, data_fetcher: _DataFetcher) -> None:
         # -----------------------------------------
         # VALIDATE IF NEEDED
         # -----------------------------------------
-        should_check_val = self._should_check_val_fx()
+        should_check_val = self._should_check_val_fx(data_fetcher)
         if should_check_val:
             # this needs to be set so the correct `trainer._active_loop` is picked
             self.trainer.validating = True
@@ -392,15 +397,15 @@ def _should_check_val_epoch(self) -> bool:
             or (self.trainer.current_epoch + 1) % self.trainer.check_val_every_n_epoch == 0
         )
 
-    def _should_check_val_fx(self) -> bool:
+    def _should_check_val_fx(self, data_fetcher: _DataFetcher) -> bool:
         """Decide if we should run validation."""
         if not self._should_check_val_epoch():
             return False
 
         # val_check_batch is inf for iterable datasets with no length defined
         is_infinite_dataset = self.trainer.val_check_batch == float("inf")
         is_last_batch = self.batch_progress.is_last_batch
-        if is_last_batch and is_infinite_dataset:
+        if is_last_batch and (is_infinite_dataset or isinstance(data_fetcher, _DataLoaderIterDataFetcher)):
             return True
 
         if self.trainer.should_stop and self.trainer.fit_loop._can_stop_early:

src/lightning/pytorch/trainer/trainer.py

@@ -1532,18 +1532,15 @@ def num_training_batches(self) -> Union[int, float]:
         return self.fit_loop.max_batches
 
     @property
-    def num_sanity_val_batches(self) -> Union[int, float, List[Union[int, float]]]:
+    def num_sanity_val_batches(self) -> List[Union[int, float]]:
         """The number of validation batches that will be used during the sanity-checking part of
         ``trainer.fit()``."""
         max_batches = self.fit_loop.epoch_loop.val_loop.max_batches
-        # re-compute the `min` in case this is called outside of the sanity-checking stage
-        sanity_val_steps = self.num_sanity_val_steps
-        if isinstance(max_batches, list):
-            return [min(sanity_val_steps, batches) for batches in max_batches]
-        return min(sanity_val_steps, max_batches)
+        # re-compute the `min` in case this is called outside the sanity-checking stage
+        return [min(self.num_sanity_val_steps, batches) for batches in max_batches]
 
     @property
-    def num_val_batches(self) -> Union[int, float, List[Union[int, float]]]:
+    def num_val_batches(self) -> List[Union[int, float]]:
         """The number of validation batches that will be used during ``trainer.fit()`` or
         ``trainer.validate()``."""
         if self.state.fn == TrainerFn.VALIDATING:
@@ -1553,7 +1550,7 @@ def num_val_batches(self) -> Union[int, float, List[Union[int, float]]]:
         return self.fit_loop.epoch_loop.val_loop._max_batches
 
     @property
-    def num_test_batches(self) -> Union[int, float, List[Union[int, float]]]:
+    def num_test_batches(self) -> List[Union[int, float]]:
         """The number of test batches that will be used during ``trainer.test()``."""
         return self.test_loop.max_batches