update to address reviewers' comments

Author: ninginthecloud

CHANGELOG.md

@@ -157,7 +157,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 - Deprecated the `TestTubeLogger` ([#9065](https://github.com/PyTorchLightning/pytorch-lightning/pull/9065))
 
-- Deprecated `on_{train/val/test/predict}_trainer()` from `DataHooks` [#9098](https://github.com/PyTorchLightning/pytorch-lightning/pull/9098)
+- Deprecated `on_{train/val/test/predict}_trainer()` from `LightningModule` and `LightningDataModule` [#9098](https://github.com/PyTorchLightning/pytorch-lightning/pull/9098)
 
 -
 

pytorch_lightning/core/hooks.py

@@ -689,7 +689,7 @@ def on_train_dataloader(self) -> None:
         """
         rank_zero_deprecation(
             "Method `on_train_dataloader` in DataHooks is deprecated and will be removed in v1.7.0."
-            " Please use `train_dataloader()` directly. "
+            " Please use `train_dataloader()` directly."
         )
 
     def on_val_dataloader(self) -> None:
@@ -701,7 +701,7 @@ def on_val_dataloader(self) -> None:
         """
         rank_zero_deprecation(
             "Method `on_val_dataloader` in DataHooks is deprecated and will be removed in v1.7.0."
-            " Please use ``val_dataloader()`` directly. "
+            " Please use `val_dataloader()` directly."
         )
 
     def on_test_dataloader(self) -> None:
@@ -713,7 +713,7 @@ def on_test_dataloader(self) -> None:
         """
         rank_zero_deprecation(
             "Method `on_test_dataloader` in DataHooks is deprecated and will be removed in v1.7.0."
-            " Please use ``test_dataloader()`` directly. "
+            " Please use `test_dataloader()` directly."
         )
 
     def on_predict_dataloader(self) -> None:
@@ -725,7 +725,7 @@ def on_predict_dataloader(self) -> None:
         """
         rank_zero_deprecation(
             "Method `on_predict_dataloader` in DataHooks is deprecated and will be removed in v1.7.0."
-            " Please use ``predict_dataloader()`` directly. "
+            " Please use `predict_dataloader()` directly."
         )
 
     def transfer_batch_to_device(self, batch: Any, device: torch.device, dataloader_idx: int) -> Any:

pytorch_lightning/trainer/configuration_validator.py

@@ -120,6 +120,7 @@ def __verify_eval_loop_configuration(self, model: "pl.LightningModule", stage: s
             rank_zero_warn(f"you passed in a {loader_name} but have no {step_name}. Skipping {stage} loop")
         if has_step and not has_loader:
             rank_zero_warn(f"you defined a {step_name} but have no {loader_name}. Skipping {stage} loop")
+
         # ----------------------------------------------
         # verify model does not have
         # - on_val_dataloader

pytorch_lightning/trainer/connectors/logger_connector/fx_validator.py

@@ -78,8 +78,11 @@ class FxValidator:
         validation_epoch_end=dict(on_step=(False,), on_epoch=(True,)),
         test_epoch_end=dict(on_step=(False,), on_epoch=(True,)),
         configure_optimizers=None,
+        on_train_dataloader=None,
         train_dataloader=None,
+        on_val_dataloader=None,
         val_dataloader=None,
+        on_test_dataloader=None,
         test_dataloader=None,
         prepare_data=None,
         configure_callbacks=None,

tests/models/test_hooks.py

@@ -505,6 +505,7 @@ def training_step(self, batch, batch_idx):
         dict(name="Callback.on_pretrain_routine_end", args=(trainer, model)),
         dict(name="on_pretrain_routine_end"),
         dict(name="Callback.on_sanity_check_start", args=(trainer, model)),
+        dict(name="on_val_dataloader"),
         dict(name="val_dataloader"),
         dict(name="train", args=(False,)),
         dict(name="on_validation_model_eval"),
@@ -519,6 +520,7 @@ def training_step(self, batch, batch_idx):
         dict(name="Callback.on_sanity_check_end", args=(trainer, model)),
         # duplicate `train` because `_run_train` calls it again in case validation wasn't run
         dict(name="train", args=(True,)),
+        dict(name="on_train_dataloader"),
         dict(name="train_dataloader"),
         dict(name="Callback.on_train_start", args=(trainer, model)),
         dict(name="on_train_start"),
@@ -631,8 +633,10 @@ def test_trainer_model_hook_system_fit_no_val_and_resume(tmpdir):
         dict(name="Callback.on_pretrain_routine_end", args=(trainer, model)),
         dict(name="on_pretrain_routine_end"),
         dict(name="train", args=(True,)),
+        dict(name="on_train_dataloader"),
         dict(name="train_dataloader"),
         # even though no validation runs, we initialize the val dataloader for properties like `num_val_batches`
+        dict(name="on_val_dataloader"),
         dict(name="val_dataloader"),
         dict(name="Callback.on_train_start", args=(trainer, model)),
         dict(name="on_train_start"),
@@ -704,6 +708,7 @@ def test_trainer_model_hook_system_eval(tmpdir, batches, verb, noun, dataloader,
         dict(name="setup", kwargs=dict(stage=verb)),
         dict(name="configure_sharded_model"),
         dict(name="Callback.on_configure_sharded_model", args=(trainer, model)),
+        dict(name=f"on_{dataloader}_dataloader"),
         dict(name=f"{dataloader}_dataloader"),
         *(hooks if batches else []),
         dict(name="Callback.teardown", args=(trainer, model), kwargs=dict(stage=verb)),
@@ -735,6 +740,7 @@ def test_trainer_model_hook_system_predict(tmpdir):
         dict(name="setup", kwargs=dict(stage="predict")),
         dict(name="configure_sharded_model"),
         dict(name="Callback.on_configure_sharded_model", args=(trainer, model)),
+        dict(name="on_predict_dataloader"),
         dict(name="predict_dataloader"),
         dict(name="train", args=(False,)),
         dict(name="on_predict_model_eval"),

tests/trainer/logging_/test_logger_connector.py

@@ -218,7 +218,9 @@ def test_fx_validator_integration(tmpdir):
         "on_fit_start": "You can't",
         "on_pretrain_routine_start": "You can't",
         "on_pretrain_routine_end": "You can't",
+        "on_train_dataloader": "You can't",
         "train_dataloader": "You can't",
+        "on_val_dataloader": "You can't",
         "val_dataloader": "You can't",
         "on_validation_end": "You can't",
         "on_train_end": "You can't",
@@ -252,6 +254,7 @@ def test_fx_validator_integration(tmpdir):
         {
             # `lightning_module` ref is now present from the `fit` call
             "on_before_accelerator_backend_setup": "You can't",
+            "on_test_dataloader": "You can't",
             "test_dataloader": "You can't",
             "on_test_model_eval": "You can't",
             "on_test_end": "You can't",
@@ -262,6 +265,7 @@ def test_fx_validator_integration(tmpdir):
     not_supported.update({k: "ResultCollection` is not registered yet" for k in not_supported})
     not_supported.update(
         {
+            "on_test_dataloader": "You can't",
             "predict_dataloader": "ResultCollection` is not registered yet",
             "on_predict_model_eval": "ResultCollection` is not registered yet",
             "on_predict_start": "ResultCollection` is not registered yet",

tests/trainer/test_dataloaders.py

@@ -27,7 +27,6 @@
 from pytorch_lightning.callbacks import ModelCheckpoint
 from pytorch_lightning.utilities.data import has_iterable_dataset, has_len
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
-from pytorch_lightning.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
 from tests.base import EvalModelTemplate
 from tests.helpers.boring_model import BoringModel, RandomDataset, RandomIterableDataset, RandomIterableDatasetWithLen
 from tests.helpers.runif import RunIf
@@ -1471,28 +1470,35 @@ def __iter__(self):
         def __next__(self):
             return next(self._iter)
 
+    class DataLoaderFunc:
+        def __init__(self, loader):
+            self.loader = loader
+
+        def __call__(self):
+            return self.loader
+
     class TestModel(BoringModel):
         def __init__(self):
             super().__init__()
-            self.train_dataloader_called = False
+            self.on_train_dataloader_called = False
             self.on_train_batch_start_called = False
-            self.val_dataloader_called = False
+            self.on_val_dataloader_called = False
             self.on_val_batch_start_called = False
 
-        def train_dataloader(self) -> TRAIN_DATALOADERS:
-            loader = super().train_dataloader()
-            self.train_dataloader_called = True
-            return DataLoaderWrapper(loader)
-
-        def val_dataloader(self) -> EVAL_DATALOADERS:
-            loader = super().val_dataloader()
-            self.val_dataloader_called = True
-            return DataLoaderWrapper(loader)
+        def on_train_dataloader(self) -> None:
+            loader = self.train_dataloader()
+            self.train_dataloader = DataLoaderFunc(DataLoaderWrapper(loader))
+            self.on_train_dataloader_called = True
 
         def on_train_batch_start(self, batch, batch_idx: int, dataloader_idx: int) -> None:
             assert isinstance(self.trainer.train_dataloader.loaders, DataLoaderWrapper)
             self.on_train_batch_start_called = True
 
+        def on_val_dataloader(self) -> None:
+            loader = self.val_dataloader()
+            self.val_dataloader = DataLoaderFunc(DataLoaderWrapper(loader))
+            self.on_val_dataloader_called = True
+
         def on_validation_batch_start(self, batch, batch_idx: int, dataloader_idx: int) -> None:
             assert isinstance(self.trainer.val_dataloaders[0], DataLoaderWrapper)
             self.on_val_batch_start_called = True
@@ -1501,7 +1507,7 @@ def on_validation_batch_start(self, batch, batch_idx: int, dataloader_idx: int)
     model = TestModel()
     trainer.fit(model)
     trainer.test(model)
-    assert model.train_dataloader_called
-    assert model.val_dataloader_called
+    assert model.on_train_dataloader_called
     assert model.on_train_batch_start_called
+    assert model.on_val_dataloader_called
     assert model.on_val_batch_start_called