Self-review of the recent Trainer changes (#14916)

Author: carmocca

pyproject.toml

@@ -58,7 +58,6 @@ warn_no_return = "False"
 module = [
     "pytorch_lightning.callbacks.progress.rich_progress",
     "pytorch_lightning.trainer.trainer",
-    "pytorch_lightning.trainer.connectors.checkpoint_connector",
     "lightning_app.api.http_methods",
     "lightning_app.api.request_types",
     "lightning_app.cli.app-template.app",

src/pytorch_lightning/trainer/teardown.py renamed to src/pytorch_lightning/trainer/call.py

@@ -14,13 +14,14 @@
 import traceback
 from typing import Any, Callable
 
+import pytorch_lightning as pl
 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
 
 
-def call_and_handle_interrupt(trainer: Any, trainer_fn: Callable, *args: Any, **kwargs: Any) -> Any:
+def _call_and_handle_interrupt(trainer: "pl.Trainer", trainer_fn: Callable, *args: Any, **kwargs: Any) -> Any:
     r"""
     Error handling, intended to be used only for main trainer function entry points (fit, validate, test, predict)
     as all errors should funnel through them

src/pytorch_lightning/trainer/connectors/checkpoint_connector.py

@@ -13,11 +13,9 @@
 # limitations under the License.
 
 import logging
-import operator
 import os
 import re
 from copy import deepcopy
-from functools import partial
 from typing import Any, Dict, Optional
 
 import torch
@@ -172,7 +170,7 @@ def _set_ckpt_path(
                     " or last checkpoint available. No checkpoint will be loaded."
                 )
                 return None
-            ckpt_path = max(candidates_ts.keys(), key=partial(operator.getitem, candidates_ts))
+            ckpt_path = max(candidates_ts, key=candidates_ts.get)  # type: ignore[arg-type]
 
         elif ckpt_path == "hpc":
             if not self._hpc_resume_path:

src/pytorch_lightning/trainer/setup.py

@@ -11,11 +11,11 @@
 # 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.
-
 """Houses the methods used to set up the Trainer."""
 
-from typing import Any, Optional, Union
+from typing import Optional, Union
 
+import pytorch_lightning as pl
 from lightning_lite.utilities.warnings import PossibleUserWarning
 from pytorch_lightning.accelerators import (
     CUDAAccelerator,
@@ -38,8 +38,8 @@
 from pytorch_lightning.utilities.rank_zero import rank_zero_info, rank_zero_warn
 
 
-def init_debugging_flags(
-    trainer: Any,
+def _init_debugging_flags(
+    trainer: "pl.Trainer",
     limit_train_batches: Optional[Union[int, float]],
     limit_val_batches: Optional[Union[int, float]],
     limit_test_batches: Optional[Union[int, float]],
@@ -128,7 +128,7 @@ def _determine_batch_limits(batches: Optional[Union[int, float]], name: str) ->
     )
 
 
-def init_profiler(trainer: Any, profiler: Optional[Union[Profiler, str]]) -> None:
+def _init_profiler(trainer: "pl.Trainer", profiler: Optional[Union[Profiler, str]]) -> None:
     if isinstance(profiler, str):
         PROFILERS = {
             "simple": SimpleProfiler,
@@ -147,8 +147,7 @@ def init_profiler(trainer: Any, profiler: Optional[Union[Profiler, str]]) -> Non
     trainer.profiler = profiler or PassThroughProfiler()
 
 
-def log_device_info(trainer: Any) -> None:
-
+def _log_device_info(trainer: "pl.Trainer") -> None:
     if CUDAAccelerator.is_available():
         gpu_available = True
         gpu_type = " (cuda)"

src/pytorch_lightning/trainer/trainer.py

@@ -66,7 +66,7 @@
 )
 from pytorch_lightning.profilers import Profiler
 from pytorch_lightning.strategies import ParallelStrategy, Strategy
-from pytorch_lightning.trainer import setup, teardown
+from pytorch_lightning.trainer import call, setup
 from pytorch_lightning.trainer.configuration_validator import verify_loop_configurations
 from pytorch_lightning.trainer.connectors.accelerator_connector import _LITERAL_WARN, AcceleratorConnector
 from pytorch_lightning.trainer.connectors.callback_connector import CallbackConnector
@@ -393,7 +393,6 @@ def __init__(
         Trainer._log_api_event("init")
         log.detail(f"{self.__class__.__name__}: Initializing trainer with parameters: {locals()}")
         self.state = TrainerState()
-        self.num_sanity_val_steps: int
 
         # init connectors
         self._data_connector = DataConnector(self, multiple_trainloader_mode)
@@ -498,15 +497,16 @@ def __init__(
         self.tuner.on_trainer_init(auto_lr_find, auto_scale_batch_size)
 
         # configure profiler
-        setup.init_profiler(self, profiler)
+        setup._init_profiler(self, profiler)
 
         # init logger flags
         self._loggers: List[Logger]
         self._logger_connector.on_trainer_init(logger, log_every_n_steps, move_metrics_to_cpu)
 
         # init debugging flags
         self.val_check_interval: Union[int, float]
-        setup.init_debugging_flags(
+        self.num_sanity_val_steps: Union[float, int]
+        setup._init_debugging_flags(
             self,
             limit_train_batches,
             limit_val_batches,
@@ -522,7 +522,7 @@ def __init__(
         self._call_callback_hooks("on_init_end")
 
     def _setup_on_init(self) -> None:
-        setup.log_device_info(self)
+        setup._log_device_info(self)
 
         self.should_stop = False
         self.state = TrainerState()
@@ -568,7 +568,7 @@ def fit(
         if not isinstance(model, pl.LightningModule):
             raise TypeError(f"`Trainer.fit()` requires a `LightningModule`, got: {model.__class__.__qualname__}")
         self.strategy._lightning_module = model
-        teardown.call_and_handle_interrupt(
+        call._call_and_handle_interrupt(
             self, self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
         )
 
@@ -648,7 +648,7 @@ def validate(
         if model is not None and not isinstance(model, pl.LightningModule):
             raise TypeError(f"`Trainer.validate()` requires a `LightningModule`, got: {model.__class__.__qualname__}")
         self.strategy._lightning_module = model or self.lightning_module
-        return teardown.call_and_handle_interrupt(
+        return call._call_and_handle_interrupt(
             self, self._validate_impl, model, dataloaders, ckpt_path, verbose, datamodule
         )
 
@@ -740,7 +740,7 @@ def test(
         if model is not None and not isinstance(model, pl.LightningModule):
             raise TypeError(f"`Trainer.test()` requires a `LightningModule`, got: {model.__class__.__qualname__}")
         self.strategy._lightning_module = model or self.lightning_module
-        return teardown.call_and_handle_interrupt(
+        return call._call_and_handle_interrupt(
             self, self._test_impl, model, dataloaders, ckpt_path, verbose, datamodule
         )
 
@@ -831,7 +831,7 @@ def predict(
         if model is not None and not isinstance(model, pl.LightningModule):
             raise TypeError(f"`Trainer.predict()` requires a `LightningModule`, got: {model.__class__.__qualname__}")
         self.strategy._lightning_module = model or self.lightning_module
-        return teardown.call_and_handle_interrupt(
+        return call._call_and_handle_interrupt(
             self, self._predict_impl, model, dataloaders, datamodule, return_predictions, ckpt_path
         )