Relax on_train_batch_* hook check with dataloader_iter to a warning (#16062)

Co-authored-by: Carlos Mocholí <[email protected]>

Author: erhoo82

src/pytorch_lightning/CHANGELOG.md

@@ -73,6 +73,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - `MLFlowLogger` now logs hyperparameters and metrics in batched API calls ([#15915](https://github.com/Lightning-AI/lightning/pull/15915))
 
 
+- Overriding the `on_train_batch_{start,end}` hooks in conjunction with taking a `dataloader_iter` in the `training_step` no longer errors out and instead shows a warning ([#16062](https://github.com/Lightning-AI/lightning/pull/16062))
+
+
 ### Deprecated
 
 - Deprecated `description`, `env_prefix` and `env_parse` parameters in `LightningCLI.__init__` in favour of giving them through `parser_kwargs` ([#15651](https://github.com/Lightning-AI/lightning/pull/15651))

src/pytorch_lightning/trainer/configuration_validator.py

@@ -169,20 +169,14 @@ def __verify_manual_optimization_support(trainer: "pl.Trainer", model: "pl.Light
 
 def __check_training_step_requires_dataloader_iter(model: "pl.LightningModule") -> None:
     """Check if the current `training_step` is requesting `dataloader_iter`."""
-    training_step_fx = model.training_step
-    if is_param_in_hook_signature(training_step_fx, "dataloader_iter", explicit=True):
-
-        if is_overridden("on_train_batch_start", model):
-            raise MisconfigurationException(
-                "The model hook `on_train_batch_start` is not compatible with "
-                "taking a `dataloader_iter` argument in your `training_step`."
-            )
-
-        if is_overridden("on_train_batch_end", model):
-            raise MisconfigurationException(
-                "The model hook `on_train_batch_end` is not compatible with "
-                "taking a `dataloader_iter` argument in your `training_step`."
-            )
+    if is_param_in_hook_signature(model.training_step, "dataloader_iter", explicit=True):
+        for hook in ("on_train_batch_start", "on_train_batch_end"):
+            if is_overridden(hook, model):
+                rank_zero_warn(
+                    f"The `batch_idx` argument in `{type(model).__name__}.{hook}` hook may"
+                    " not match with the actual batch index when using a `dataloader_iter`"
+                    " argument in your `training_step`."
+                )
 
         if model.truncated_bptt_steps > 0:
             raise MisconfigurationException(

tests/tests_pytorch/utilities/test_fetching.py

@@ -429,9 +429,9 @@ class InvalidModel(AsyncBoringModel):
         def on_train_batch_start(self, batch, batch_idx):
             pass
 
-    trainer = Trainer(max_epochs=1, default_root_dir=tmpdir)
+    trainer = Trainer(fast_dev_run=1, default_root_dir=tmpdir)
     m = InvalidModel()
-    with pytest.raises(MisconfigurationException, match="The model hook `on_train_batch_start` is not compatible with"):
+    with pytest.warns(match="InvalidModel.on_train_batch_start` hook may not match"):
         trainer.fit(m)
 
 
@@ -443,9 +443,9 @@ class InvalidModel(AsyncBoringModel):
         def on_train_batch_end(self, outputs, batch, batch_idx):
             pass
 
-    trainer = Trainer(max_epochs=1, default_root_dir=tmpdir)
+    trainer = Trainer(fast_dev_run=1, default_root_dir=tmpdir)
     m = InvalidModel()
-    with pytest.raises(MisconfigurationException, match="The model hook `on_train_batch_end` is not compatible with"):
+    with pytest.warns(match="InvalidModel.on_train_batch_end` hook may not match"):
         trainer.fit(m)