WIP Add support for self.log from LightningModule in Lite

requirements/fabric/test.txt

@@ -4,3 +4,4 @@ pytest==7.2.0
 pytest-cov==4.0.0
 pre-commit==2.20.0
 click==8.1.3
+tensorboard>=2.9.1, <2.12.0

src/lightning_fabric/CHANGELOG.md

@@ -29,6 +29,13 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added support for managing callbacks via `Fabric(callbacks=...)` and emitting events through `Fabric.call()` ([#16074](https://github.com/Lightning-AI/lightning/issues/16074))
 
 
+- Added Logger support ([#16121](https://github.com/Lightning-AI/lightning/issues/16121))
+  * Added `Fabric(loggers=...)` to support different Logger frameworks in Fabric
+  * Added `Fabric.log` for logging scalars using multiple loggers
+  * Added `Fabric.log_dict` for logging a dictionary of multiple metrics at once
+  * Added `Fabric.loggers` and `Fabric.logger` attributes to access the individual logger instances
+
+
 ### Changed
 
 - The `LightningLite.run()` method is no longer abstract ([#14992](https://github.com/Lightning-AI/lightning/issues/14992))

src/lightning_fabric/fabric.py

@@ -27,6 +27,8 @@
 from torch.optim import Optimizer
 from torch.utils.data import BatchSampler, DataLoader, DistributedSampler, RandomSampler
 
+from lightning_fabric.loggers import Logger
+
 from lightning_fabric.plugins import Precision  # avoid circular imports: # isort: split
 from lightning_fabric.accelerators.accelerator import Accelerator
 from lightning_fabric.connector import _Connector, _PLUGIN_INPUT, _PRECISION_INPUT
@@ -74,6 +76,10 @@ class Fabric:
         precision: Double precision (``64``), full precision (``32``), half precision (``16``),
             or bfloat16 precision (``"bf16"``).
         plugins: One or several custom plugins
+        callbacks: A single callback or a list of callbacks. A callback can contain any arbitrary methods that
+            can be invoked through :meth:`lightning_fabric.fabric.Fabric.call` by the user.
+        loggers: A single logger or a list of loggers. See :meth:`lightning_fabric.fabric.Fabric.log` for more
+            information.
     """
 
     def __init__(
@@ -85,6 +91,7 @@ def __init__(
         precision: _PRECISION_INPUT = 32,
         plugins: Optional[Union[_PLUGIN_INPUT, List[_PLUGIN_INPUT]]] = None,
         callbacks: Optional[Union[List[Any], Any]] = None,
+        loggers: Optional[Union[Logger, List[Logger]]] = None,
     ) -> None:
         self._connector = _Connector(
             accelerator=accelerator,
@@ -99,6 +106,8 @@ def __init__(
         self._precision: Precision = self._strategy.precision
         callbacks = callbacks if callbacks is not None else []
         self._callbacks = callbacks if isinstance(callbacks, list) else [callbacks]
+        loggers = loggers if loggers is not None else []
+        self._loggers = loggers if isinstance(loggers, list) else [loggers]
         self._models_setup: int = 0
 
         self._prepare_run_method()
@@ -148,6 +157,16 @@ def is_global_zero(self) -> bool:
         """Whether this rank is rank zero."""
         return self._strategy.is_global_zero
 
+    @property
+    def loggers(self) -> List[Logger]:
+        """Returns all loggers passed to Fabric."""
+        return self._loggers
+
+    @property
+    def logger(self) -> Logger:
+        """Returns the first logger in the list passed to Fabric, which is considered the main logger."""
+        return self._loggers[0]
+
     def run(self, *args: Any, **kwargs: Any) -> Any:
         """All the code inside this run method gets accelerated by Fabric.
 
@@ -574,6 +593,28 @@ def on_train_epoch_end(self, results):
             # method(self, *args, y=1)
             # method(self, *args, **kwargs)
 
+    def log(self, name: str, value: Any, step: Optional[int] = None) -> None:
+        """Log a scalar to all loggers that were added to Fabric.
+
+        Args:
+            name: The name of the metric to log.
+            value: The metric value to collect.
+            step: Optional step number. Most Logger implementations auto-increment this value by one with every
+                log call.
+        """
+        self.log_dict(metrics={name: value}, step=step)
+
+    def log_dict(self, metrics: Dict[str, Any], step: Optional[int] = None) -> None:
+        """Log multiple scalars at once to all loggers that were added to Fabric.
+
+        Args:
+            metrics: A dictionary where the key is the name of the metric and the value the scalar to be logged.
+            step: Optional step number. Most Logger implementations auto-increment this value by one with every
+                log call.
+        """
+        for logger in self._loggers:
+            logger.log_metrics(metrics=metrics, step=step)
+
     @staticmethod
     def seed_everything(seed: Optional[int] = None, workers: Optional[bool] = None) -> int:
         """Helper function to seed everything without explicitly importing Lightning.

src/lightning_fabric/loggers/__init__.py

@@ -0,0 +1,14 @@
+# Copyright The PyTorch Lightning team.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 lightning_fabric.loggers.logger import Logger  # noqa: F401
+from lightning_fabric.loggers.tensorboard import TensorBoardLogger  # noqa: F401

src/lightning_fabric/loggers/logger.py

@@ -0,0 +1,136 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""Abstract base class used to build new loggers."""
+
+from abc import ABC, abstractmethod
+from argparse import Namespace
+from functools import wraps
+from typing import Any, Callable, Dict, Optional, Union
+
+from torch import Tensor
+from torch.nn import Module
+
+from lightning_fabric.utilities.rank_zero import rank_zero_only
+
+
+class Logger(ABC):
+    """Base class for experiment loggers."""
+
+    @property
+    @abstractmethod
+    def name(self) -> Optional[str]:
+        """Return the experiment name."""
+
+    @property
+    @abstractmethod
+    def version(self) -> Optional[Union[int, str]]:
+        """Return the experiment version."""
+
+    @property
+    def root_dir(self) -> Optional[str]:
+        """Return the root directory where all versions of an experiment get saved, or `None` if the logger does
+        not save data locally."""
+        return None
+
+    @property
+    def log_dir(self) -> Optional[str]:
+        """Return directory the current version of the experiment gets saved, or `None` if the logger does not save
+        data locally."""
+        return None
+
+    @property
+    def group_separator(self) -> str:
+        """Return the default separator used by the logger to group the data into subfolders."""
+        return "/"
+
+    @abstractmethod
+    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
+        """Records metrics. This method logs metrics as soon as it received them.
+
+        Args:
+            metrics: Dictionary with metric names as keys and measured quantities as values
+            step: Step number at which the metrics should be recorded
+        """
+        pass
+
+    @abstractmethod
+    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace], *args: Any, **kwargs: Any) -> None:
+        """Record hyperparameters.
+
+        Args:
+            params: :class:`~argparse.Namespace` or `Dict` containing the hyperparameters
+            args: Optional positional arguments, depends on the specific logger being used
+            kwargs: Optional keyword arguments, depends on the specific logger being used
+        """
+
+    def log_graph(self, model: Module, input_array: Optional[Tensor] = None) -> None:
+        """Record model graph.
+
+        Args:
+            model: the model with an implementation of ``forward``.
+            input_array: input passes to `model.forward`
+        """
+        pass
+
+    def save(self) -> None:
+        """Save log data."""
+
+    def finalize(self, status: str) -> None:
+        """Do any processing that is necessary to finalize an experiment.
+
+        Args:
+            status: Status that the experiment finished with (e.g. success, failed, aborted)
+        """
+        self.save()
+
+
+def rank_zero_experiment(fn: Callable) -> Callable:
+    """Returns the real experiment on rank 0 and otherwise the _DummyExperiment."""
+
+    @wraps(fn)
+    def experiment(self) -> Union[Any, _DummyExperiment]:  # type: ignore[no-untyped-def]
+        """
+        Note:
+            ``self`` is a custom logger instance. The loggers typically wrap an ``experiment`` method
+            with a ``@rank_zero_experiment`` decorator.
+
+            ``Union[Any, _DummyExperiment]`` is used because the wrapped hooks have several return
+            types that are specific to the custom logger. The return type here can be considered as
+            ``Union[return type of logger.experiment, _DummyExperiment]``.
+        """
+
+        @rank_zero_only
+        def get_experiment() -> Callable:
+            return fn(self)
+
+        return get_experiment() or _DummyExperiment()
+
+    return experiment
+
+
+class _DummyExperiment:
+    """Dummy experiment."""
+
+    def nop(self, *args: Any, **kw: Any) -> None:
+        pass
+
+    def __getattr__(self, _: Any) -> Callable:
+        return self.nop
+
+    def __getitem__(self, idx: int) -> "_DummyExperiment":
+        # enables self.logger.experiment[0].add_image(...)
+        return self
+
+    def __setitem__(self, *args: Any, **kwargs: Any) -> None:
+        pass