Refactor data fetcher selection in loops

src/lightning/pytorch/loops/evaluation_loop.py

@@ -217,7 +217,8 @@ def reset(self) -> None:
         if fn != TrainerFn.FITTING:
             self.batch_progress.reset_on_run()
 
-        data_fetcher = _select_data_fetcher(trainer)
+        assert trainer.state.stage is not None
+        data_fetcher = _select_data_fetcher(trainer, trainer.state.stage)
         combined_loader = self._combined_loader
         assert combined_loader is not None
 

src/lightning/pytorch/loops/fit_loop.py

@@ -253,11 +253,7 @@ def setup_data(self) -> None:
 
         combined_loader.limits = limits
 
-        training = trainer.training
-        trainer.training = True
-        self._data_fetcher = _select_data_fetcher(trainer)
-        trainer.training = training
-
+        self._data_fetcher = _select_data_fetcher(trainer, RunningStage.TRAINING)
         self._data_fetcher.setup(combined_loader)
         iter(self._data_fetcher)  # creates the iterator inside the fetcher
         max_batches = sized_len(combined_loader)

src/lightning/pytorch/loops/prediction_loop.py

@@ -165,7 +165,8 @@ def reset(self) -> None:
         """Resets the internal state of the loop for a new run."""
         self.batch_progress.reset_on_run()
 
-        data_fetcher = _select_data_fetcher(self.trainer)
+        assert self.trainer.state.stage is not None
+        data_fetcher = _select_data_fetcher(self.trainer, self.trainer.state.stage)
         combined_loader = self._combined_loader
         assert combined_loader is not None
         if combined_loader._mode != "sequential":

src/lightning/pytorch/loops/utilities.py

@@ -131,15 +131,15 @@ def _reset_progress(loop: _Loop) -> None:
             _reset_progress(v)
 
 
-def _select_data_fetcher(trainer: "pl.Trainer") -> _DataFetcher:
+def _select_data_fetcher(trainer: "pl.Trainer", stage: RunningStage) -> _DataFetcher:
     lightning_module = trainer.lightning_module
-    if trainer.testing:
+    if stage == RunningStage.TESTING:
         step_fx_name = "test_step"
-    elif trainer.training:
+    elif stage == RunningStage.TRAINING:
         step_fx_name = "training_step"
-    elif trainer.validating or trainer.sanity_checking:
+    elif stage in (RunningStage.VALIDATING, RunningStage.SANITY_CHECKING):
         step_fx_name = "validation_step"
-    elif trainer.predicting:
+    elif stage == RunningStage.PREDICTING:
         step_fx_name = "predict_step"
     else:
         raise RuntimeError(f"DataFetcher is unsupported for {trainer.state.stage}")