Fix logger typing by using fabric's base logger

Author: carmocca

src/pytorch_lightning/core/module.py

@@ -33,6 +33,7 @@
 
 import lightning_fabric as lf
 import pytorch_lightning as pl
+from lightning_fabric.loggers import Logger
 from lightning_fabric.utilities.apply_func import convert_to_tensors
 from lightning_fabric.utilities.cloud_io import get_filesystem
 from lightning_fabric.utilities.device_dtype_mixin import _DeviceDtypeModuleMixin
@@ -45,7 +46,6 @@
 from pytorch_lightning.core.mixins import HyperparametersMixin
 from pytorch_lightning.core.optimizer import LightningOptimizer
 from pytorch_lightning.core.saving import ModelIO
-from pytorch_lightning.loggers import Logger, TensorBoardLogger
 from pytorch_lightning.trainer.connectors.logger_connector.fx_validator import _FxValidator
 from pytorch_lightning.utilities import GradClipAlgorithmType
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
@@ -291,12 +291,12 @@ def truncated_bptt_steps(self, truncated_bptt_steps: int) -> None:
         self._truncated_bptt_steps = truncated_bptt_steps
 
     @property
-    def logger(self) -> Optional[Union[Logger, TensorBoardLogger]]:
+    def logger(self) -> Optional[Logger]:
         """Reference to the logger object in the Trainer."""
         return self._trainer.logger if self._trainer is not None else None
 
     @property
-    def loggers(self) -> List[Union[Logger, TensorBoardLogger]]:
+    def loggers(self) -> List[Logger]:
         """Reference to the list of loggers in the Trainer."""
         return self.trainer.loggers if self._trainer else []
 

src/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py

@@ -17,9 +17,10 @@
 from torch import Tensor
 
 import pytorch_lightning as pl
+from lightning_fabric.loggers import Logger
 from lightning_fabric.plugins.environments import SLURMEnvironment
 from lightning_fabric.utilities import move_data_to_device
-from pytorch_lightning.loggers import Logger, TensorBoardLogger
+from pytorch_lightning.loggers import TensorBoardLogger
 from pytorch_lightning.trainer.connectors.logger_connector.result import _METRICS, _OUT_DICT, _PBAR_DICT
 from pytorch_lightning.utilities.metrics import metrics_to_scalars
 
@@ -37,7 +38,7 @@ def __init__(self, trainer: "pl.Trainer") -> None:
 
     def on_trainer_init(
         self,
-        logger: Union[bool, Logger, TensorBoardLogger, Iterable[Union[Logger, TensorBoardLogger]]],
+        logger: Union[bool, Logger, Iterable[Logger]],
         log_every_n_steps: int,
         move_metrics_to_cpu: bool,
     ) -> None:
@@ -51,9 +52,7 @@ def should_update_logs(self) -> bool:
         should_log = (self.trainer.fit_loop.epoch_loop._batches_that_stepped + 1) % self.trainer.log_every_n_steps == 0
         return should_log or self.trainer.should_stop
 
-    def configure_logger(
-        self, logger: Union[bool, Logger, TensorBoardLogger, Iterable[Union[Logger, TensorBoardLogger]]]
-    ) -> None:
+    def configure_logger(self, logger: Union[bool, Logger, Iterable[Logger]]) -> None:
         if not logger:
             # logger is None or logger is False
             self.trainer.loggers = []

src/pytorch_lightning/trainer/trainer.py

@@ -43,6 +43,7 @@
 from typing_extensions import Literal
 
 import pytorch_lightning as pl
+from lightning_fabric.loggers import Logger
 from lightning_fabric.utilities.cloud_io import get_filesystem
 from lightning_fabric.utilities.data import _auto_add_worker_init_fn
 from lightning_fabric.utilities.types import _PATH
@@ -51,7 +52,6 @@
 from pytorch_lightning.callbacks import Callback, Checkpoint, EarlyStopping, ProgressBarBase
 from pytorch_lightning.callbacks.prediction_writer import BasePredictionWriter
 from pytorch_lightning.core.datamodule import LightningDataModule
-from pytorch_lightning.loggers import Logger
 from pytorch_lightning.loggers.tensorboard import TensorBoardLogger
 from pytorch_lightning.loops import PredictionLoop, TrainingEpochLoop
 from pytorch_lightning.loops.dataloader.evaluation_loop import EvaluationLoop
@@ -112,7 +112,7 @@ class Trainer:
     @_defaults_from_env_vars
     def __init__(
         self,
-        logger: Union[Union[Logger, TensorBoardLogger], Iterable[Union[Logger, TensorBoardLogger]], bool] = True,
+        logger: Union[Union[Logger], Iterable[Union[Logger]], bool] = True,
         enable_checkpointing: bool = True,
         callbacks: Optional[Union[List[Callback], Callback]] = None,
         default_root_dir: Optional[_PATH] = None,
@@ -448,7 +448,7 @@ def __init__(
         setup._init_profiler(self, profiler)
 
         # init logger flags
-        self._loggers: List[Union[Logger, TensorBoardLogger]]
+        self._loggers: List[Logger]
         self._logger_connector.on_trainer_init(logger, log_every_n_steps, move_metrics_to_cpu)
 
         # init debugging flags
@@ -2068,11 +2068,11 @@ def logger(self, logger: Optional[Logger]) -> None:
             self.loggers = [logger]
 
     @property
-    def loggers(self) -> List[Union[Logger, TensorBoardLogger]]:
+    def loggers(self) -> List[Union[Logger]]:
         return self._loggers
 
     @loggers.setter
-    def loggers(self, loggers: Optional[List[Union[Logger, TensorBoardLogger]]]) -> None:
+    def loggers(self, loggers: Optional[List[Logger]]) -> None:
         self._loggers = loggers if loggers else []
 
     @property