Deprecate dataloader_idx from on_train_batch_start/end (#9816)

* deprecate hooks

* dep todo

* explicit

* Apply suggestions from code review

* Apply suggestions from code review

* code review

* base

Author: rohitgr7

pytorch_lightning/accelerators/accelerator.py

@@ -487,6 +487,7 @@ def on_train_end(self) -> None:
         """Called when train ends."""
         return self.training_type_plugin.on_train_end()
 
-    def on_train_batch_start(self, batch: Any, batch_idx: int, dataloader_idx: int) -> None:
+    # TODO: Update this in v1.7 (deprecation: #9816)
+    def on_train_batch_start(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None:
         """Called in the training loop before anything happens for that batch."""
-        return self.training_type_plugin.on_train_batch_start(batch, batch_idx, dataloader_idx)
+        return self.training_type_plugin.on_train_batch_start(batch, batch_idx)

pytorch_lightning/callbacks/base.py

@@ -97,7 +97,12 @@ def on_sanity_check_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningMod
         pass
 
     def on_train_batch_start(
-        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", batch: Any, batch_idx: int, dataloader_idx: int
+        self,
+        trainer: "pl.Trainer",
+        pl_module: "pl.LightningModule",
+        batch: Any,
+        batch_idx: int,
+        unused: Optional[int] = 0,
     ) -> None:
         """Called when the train batch begins."""
         pass
@@ -109,7 +114,7 @@ def on_train_batch_end(
         outputs: STEP_OUTPUT,
         batch: Any,
         batch_idx: int,
-        dataloader_idx: int,
+        unused: Optional[int] = 0,
     ) -> None:
         """Called when the train batch ends."""
         pass

pytorch_lightning/callbacks/gpu_stats_monitor.py

@@ -135,7 +135,7 @@ def on_train_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningMo
 
     @rank_zero_only
     def on_train_batch_start(
-        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", batch: Any, batch_idx: int, dataloader_idx: int
+        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", batch: Any, batch_idx: int
     ) -> None:
         if self._log_stats.intra_step_time:
             self._snap_intra_step_time = time.time()
@@ -161,7 +161,6 @@ def on_train_batch_end(
         outputs: STEP_OUTPUT,
         batch: Any,
         batch_idx: int,
-        dataloader_idx: int,
     ) -> None:
         if self._log_stats.inter_step_time:
             self._snap_inter_step_time = time.time()

pytorch_lightning/callbacks/model_checkpoint.py

@@ -279,7 +279,6 @@ def on_train_batch_end(
         outputs: STEP_OUTPUT,
         batch: Any,
         batch_idx: int,
-        dataloader_idx: int,
     ) -> None:
         """Save checkpoint on train batch end if we meet the criteria for `every_n_train_steps`"""
         if self._should_skip_saving_checkpoint(trainer):
@@ -304,9 +303,7 @@ def on_train_batch_end(
 
         self.save_checkpoint(trainer)
 
-    def on_train_epoch_end(
-        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", unused: Optional = None
-    ) -> None:
+    def on_train_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
         """Save a checkpoint at the end of the training epoch."""
         # as we advance one step at end of training, we use `global_step - 1` to avoid saving duplicates
         trainer.fit_loop.global_step -= 1

pytorch_lightning/callbacks/progress/base.py

@@ -35,8 +35,8 @@ def __init__(self):
             def disable(self):
                 self.enable = False
 
-            def on_train_batch_end(self, trainer, pl_module, outputs):
-                super().on_train_batch_end(trainer, pl_module, outputs)  # don't forget this :)
+            def on_train_batch_end(self, trainer, pl_module, outputs, batch_idx):
+                super().on_train_batch_end(trainer, pl_module, outputs, batch_idx)  # don't forget this :)
                 percent = (self.train_batch_idx / self.total_train_batches) * 100
                 sys.stdout.flush()
                 sys.stdout.write(f'{percent:.01f} percent complete \r')
@@ -161,7 +161,7 @@ def on_train_start(self, trainer, pl_module):
     def on_train_epoch_start(self, trainer, pl_module):
         self._train_batch_idx = trainer.fit_loop.epoch_loop.batch_progress.current.completed
 
-    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
         self._train_batch_idx += 1
 
     def on_validation_start(self, trainer, pl_module):

pytorch_lightning/callbacks/progress/rich_progress.py

@@ -369,8 +369,8 @@ def on_predict_epoch_start(self, trainer, pl_module):
         super().on_predict_epoch_start(trainer, pl_module)
         self.predict_progress_bar_id = self._add_task(self.total_predict_batches, self.predict_description)
 
-    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
-        super().on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx)
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
+        super().on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx)
         self._update(self.main_progress_bar_id)
 
     def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):

pytorch_lightning/callbacks/progress/tqdm_progress.py

@@ -231,8 +231,8 @@ def on_train_epoch_start(self, trainer, pl_module):
         reset(self.main_progress_bar, total=total_batches, current=self.train_batch_idx)
         self.main_progress_bar.set_description(f"Epoch {trainer.current_epoch}")
 
-    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
-        super().on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx)
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
+        super().on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx)
         total_batches = self.total_train_batches + self.total_val_batches
         total_batches = convert_inf(total_batches)
         if self._should_update(self.train_batch_idx, total_batches):

pytorch_lightning/core/hooks.py

@@ -79,25 +79,25 @@ def on_pretrain_routine_end(self) -> None:
         - training_start
         """
 
-    def on_train_batch_start(self, batch: Any, batch_idx: int, dataloader_idx: int) -> None:
+    def on_train_batch_start(self, batch: Any, batch_idx: int, unused: Optional[int] = 0) -> None:
         """Called in the training loop before anything happens for that batch.
 
         If you return -1 here, you will skip training for the rest of the current epoch.
 
         Args:
             batch: The batched data as it is returned by the training DataLoader.
             batch_idx: the index of the batch
-            dataloader_idx: the index of the dataloader
+            unused: Deprecated argument. Will be removed in v1.7.
         """
 
-    def on_train_batch_end(self, outputs: STEP_OUTPUT, batch: Any, batch_idx: int, dataloader_idx: int) -> None:
+    def on_train_batch_end(self, outputs: STEP_OUTPUT, batch: Any, batch_idx: int, unused: Optional[int] = 0) -> None:
         """Called in the training loop after the batch.
 
         Args:
             outputs: The outputs of training_step_end(training_step(x))
             batch: The batched data as it is returned by the training DataLoader.
             batch_idx: the index of the batch
-            dataloader_idx: the index of the dataloader
+            unused: Deprecated argument. Will be removed in v1.7.
         """
 
     def on_validation_batch_start(self, batch: Any, batch_idx: int, dataloader_idx: int) -> None:

pytorch_lightning/loops/batch/training_batch_loop.py

@@ -24,6 +24,7 @@
 from pytorch_lightning.loops.utilities import _get_active_optimizers
 from pytorch_lightning.trainer.supporters import TensorRunningAccum
 from pytorch_lightning.utilities import AttributeDict
+from pytorch_lightning.utilities.signature_utils import is_param_in_hook_signature
 from pytorch_lightning.utilities.warnings import WarningCache
 
 _OUTPUTS_TYPE = List[Union[_OPTIMIZER_LOOP_OUTPUTS_TYPE, _MANUAL_LOOP_OUTPUTS_TYPE]]
@@ -76,7 +77,14 @@ def run(self, batch: Any, batch_idx: int) -> AttributeDict:
             return AttributeDict(signal=-1)
 
         # hook
-        response = self.trainer.call_hook("on_train_batch_start", batch, batch_idx, 0)
+        # TODO: Update this in v1.7 (deprecation: #9816)
+        model_fx = self.trainer.lightning_module.on_train_batch_start
+        extra_kwargs = (
+            {"dataloader_idx": 0}
+            if callable(model_fx) and is_param_in_hook_signature(model_fx, "dataloader_idx", explicit=True)
+            else {}
+        )
+        response = self.trainer.call_hook("on_train_batch_start", batch, batch_idx, **extra_kwargs)
         if response == -1:
             return AttributeDict(signal=-1)
 

pytorch_lightning/loops/epoch/training_epoch_loop.py

@@ -27,6 +27,7 @@
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.fetching import AbstractDataFetcher
 from pytorch_lightning.utilities.model_helpers import is_overridden
+from pytorch_lightning.utilities.signature_utils import is_param_in_hook_signature
 
 _OUTPUTS_TYPE = List[_BATCH_OUTPUTS_TYPE]
 
@@ -170,7 +171,15 @@ def advance(self, *args: Any, **kwargs: Any) -> None:
             automatic=self.trainer.lightning_module.trainer.lightning_module.automatic_optimization,
             num_optimizers=len(self.trainer.optimizers),
         )
-        self.trainer.call_hook("on_train_batch_end", batch_end_outputs, batch, self.batch_idx, 0)
+
+        # TODO: Update this in v1.7 (deprecation: #9816)
+        model_fx = self.trainer.lightning_module.on_train_batch_end
+        extra_kwargs = (
+            {"dataloader_idx": 0}
+            if callable(model_fx) and is_param_in_hook_signature(model_fx, "dataloader_idx", explicit=True)
+            else {}
+        )
+        self.trainer.call_hook("on_train_batch_end", batch_end_outputs, batch, batch_idx, **extra_kwargs)
         self.trainer.call_hook("on_batch_end")
         self.trainer.logger_connector.on_batch_end()