Revert "Add BatchSizeFinder callback (#11089)"

This reverts commit d1a3a3e.

Author: rohitgr7

docs/source-pytorch/api_references.rst

@@ -30,7 +30,6 @@ callbacks
     BackboneFinetuning
     BaseFinetuning
     BasePredictionWriter
-    BatchSizeFinder
     Callback
     DeviceStatsMonitor
     EarlyStopping

docs/source-pytorch/extensions/callbacks.rst

@@ -87,7 +87,6 @@ Lightning has a few built-in callbacks.
     BackboneFinetuning
     BaseFinetuning
     BasePredictionWriter
-    BatchSizeFinder
     Callback
     DeviceStatsMonitor
     EarlyStopping

src/pytorch_lightning/CHANGELOG.md

@@ -9,11 +9,6 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
-- Added `BatchSizeFinder` callback ([#11089](https://github.com/PyTorchLightning/pytorch-lightning/pull/11089))
-
-
-- Tuner now supports a new `method` argument which will determine when to run the `BatchSizeFinder`: one of `fit`, `validate`, `test` or `predict` ([#11089](https://github.com/PyTorchLightning/pytorch-lightning/pull/11089))
-
 
 - Added prefix to log message in `seed_everything` with rank info ([#13290](https://github.com/Lightning-AI/lightning/issues/13290))
 

src/pytorch_lightning/callbacks/__init__.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from pytorch_lightning.callbacks.batch_size_finder import BatchSizeFinder
 from pytorch_lightning.callbacks.callback import Callback
 from pytorch_lightning.callbacks.checkpoint import Checkpoint
 from pytorch_lightning.callbacks.device_stats_monitor import DeviceStatsMonitor
@@ -33,8 +32,6 @@
 __all__ = [
     "BackboneFinetuning",
     "BaseFinetuning",
-    "BasePredictionWriter",
-    "BatchSizeFinder",
     "Callback",
     "Checkpoint",
     "DeviceStatsMonitor",
@@ -45,6 +42,7 @@
     "ModelCheckpoint",
     "ModelPruning",
     "ModelSummary",
+    "BasePredictionWriter",
     "ProgressBarBase",
     "QuantizationAwareTraining",
     "RichModelSummary",

src/pytorch_lightning/callbacks/batch_size_finder.py

@@ -28,7 +28,6 @@
     RichProgressBar,
     TQDMProgressBar,
 )
-from pytorch_lightning.callbacks.batch_size_finder import BatchSizeFinder
 from pytorch_lightning.callbacks.rich_model_summary import RichModelSummary
 from pytorch_lightning.callbacks.timer import Timer
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
@@ -230,30 +229,19 @@ def _attach_model_callbacks(self) -> None:
 
     @staticmethod
     def _reorder_callbacks(callbacks: List[Callback]) -> List[Callback]:
-        """Moves all the tuner specific callbacks at the beginning of the list and all the `ModelCheckpoint`
-        callbacks to the end of the list. The sequential order within the group of checkpoint callbacks is
-        preserved, as well as the order of all other callbacks.
+        """Moves all Checkpoint callbacks to the end of the list. The sequential order within the group of
+        checkpoint callbacks is preserved, as well as the order of all other callbacks.
 
         Args:
             callbacks: A list of callbacks.
 
         Return:
-            A new list in which the first elements are tuner specific callbacks and last elements are ModelCheckpoints
-            if there were any present in the input.
+            A new list in which the last elements are Checkpoint if there were any present in the
+            input.
         """
-        tuner_callbacks: List[Callback] = []
-        other_callbacks: List[Callback] = []
-        checkpoint_callbacks: List[Callback] = []
-
-        for cb in callbacks:
-            if isinstance(cb, BatchSizeFinder):
-                tuner_callbacks.append(cb)
-            elif isinstance(cb, Checkpoint):
-                checkpoint_callbacks.append(cb)
-            else:
-                other_callbacks.append(cb)
-
-        return tuner_callbacks + other_callbacks + checkpoint_callbacks
+        checkpoints: List[Callback] = [c for c in callbacks if isinstance(c, Checkpoint)]
+        not_checkpoints = [c for c in callbacks if not isinstance(c, Checkpoint)]
+        return not_checkpoints + checkpoints
 
 
 def _configure_external_callbacks() -> List[Callback]:

src/pytorch_lightning/trainer/connectors/callback_connector.py

@@ -16,7 +16,6 @@
 
 from lightning_lite.utilities.distributed import distributed_available
 from pytorch_lightning.trainer.states import TrainerStatus
-from pytorch_lightning.utilities.exceptions import _TunerExitException
 from pytorch_lightning.utilities.rank_zero import rank_zero_warn
 
 
@@ -35,13 +34,6 @@ def call_and_handle_interrupt(trainer: Any, trainer_fn: Callable, *args: Any, **
             return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
         else:
             return trainer_fn(*args, **kwargs)
-
-    except _TunerExitException:
-        trainer._call_teardown_hook()
-        trainer._teardown()
-        trainer.state.status = TrainerStatus.FINISHED
-        trainer.state.stage = None
-
     # TODO: Unify both exceptions below, where `KeyboardError` doesn't re-raise
     except KeyboardInterrupt as exception:
         rank_zero_warn("Detected KeyboardInterrupt, attempting graceful shutdown...")

src/pytorch_lightning/trainer/teardown.py

@@ -41,7 +41,6 @@
 from torch import Tensor
 from torch.optim import Optimizer
 from torch.utils.data import DataLoader
-from typing_extensions import Literal
 
 import pytorch_lightning as pl
 from lightning_lite.utilities.cloud_io import get_filesystem
@@ -890,11 +889,9 @@ def tune(
         model: "pl.LightningModule",
         train_dataloaders: Optional[Union[TRAIN_DATALOADERS, LightningDataModule]] = None,
         val_dataloaders: Optional[EVAL_DATALOADERS] = None,
-        dataloaders: Optional[EVAL_DATALOADERS] = None,
         datamodule: Optional[LightningDataModule] = None,
         scale_batch_size_kwargs: Optional[Dict[str, Any]] = None,
         lr_find_kwargs: Optional[Dict[str, Any]] = None,
-        method: Literal["fit", "validate", "test", "predict"] = "fit",
     ) -> _TunerResult:
         r"""
         Runs routines to tune hyperparameters before training.
@@ -908,34 +905,44 @@ def tune(
 
             val_dataloaders: A :class:`torch.utils.data.DataLoader` or a sequence of them specifying validation samples.
 
-            dataloaders: A :class:`torch.utils.data.DataLoader` or a sequence of them specifying val/test/predict
-                samples used for running tuner on validation/testing/prediction.
-
             datamodule: An instance of :class:`~pytorch_lightning.core.datamodule.LightningDataModule`.
 
             scale_batch_size_kwargs: Arguments for :func:`~pytorch_lightning.tuner.batch_size_scaling.scale_batch_size`
 
             lr_find_kwargs: Arguments for :func:`~pytorch_lightning.tuner.lr_finder.lr_find`
-
-            method: Method to run tuner on. It can be any of ``("fit", "validate", "test", "predict")``.
         """
         if not isinstance(model, pl.LightningModule):
             raise TypeError(f"`Trainer.tune()` requires a `LightningModule`, got: {model.__class__.__qualname__}")
 
         Trainer._log_api_event("tune")
 
+        self.state.fn = TrainerFn.TUNING
+        self.state.status = TrainerStatus.RUNNING
+        self.tuning = True
+
+        # if a datamodule comes in as the second arg, then fix it for the user
+        if isinstance(train_dataloaders, LightningDataModule):
+            datamodule = train_dataloaders
+            train_dataloaders = None
+        # If you supply a datamodule you can't supply train_dataloader or val_dataloaders
+        if (train_dataloaders is not None or val_dataloaders is not None) and datamodule is not None:
+            raise MisconfigurationException(
+                "You cannot pass `train_dataloader` or `val_dataloaders` to `trainer.tune(datamodule=...)`"
+            )
+
+        # links data to the trainer
+        self._data_connector.attach_data(
+            model, train_dataloaders=train_dataloaders, val_dataloaders=val_dataloaders, datamodule=datamodule
+        )
+
         with isolate_rng():
             result = self.tuner._tune(
-                model,
-                train_dataloaders,
-                val_dataloaders,
-                dataloaders,
-                datamodule,
-                scale_batch_size_kwargs=scale_batch_size_kwargs,
-                lr_find_kwargs=lr_find_kwargs,
-                method=method,
+                model, scale_batch_size_kwargs=scale_batch_size_kwargs, lr_find_kwargs=lr_find_kwargs
             )
 
+        assert self.state.stopped
+        self.tuning = False
+
         return result
 
     def _restore_modules_and_callbacks(self, checkpoint_path: Optional[_PATH] = None) -> None: