[WIP] Feat/logger_map

src/lightning/pytorch/trainer/connectors/logger_connector/logger_connector.py

@@ -11,7 +11,7 @@
 # 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 collections.abc import Iterable
+from collections.abc import Iterable, Mapping
 from typing import Any, Optional, Union
 
 from lightning_utilities.core.apply_func import apply_to_collection
@@ -82,6 +82,8 @@ def configure_logger(self, logger: Union[bool, Logger, Iterable[Logger]]) -> Non
                 )
                 logger_ = CSVLogger(save_dir=self.trainer.default_root_dir)  # type: ignore[assignment]
             self.trainer.loggers = [logger_]
+        elif isinstance(logger, Mapping):
+            self.trainer.loggers = logger
         elif isinstance(logger, Iterable):
             self.trainer.loggers = list(logger)
         else:

src/lightning/pytorch/trainer/trainer.py

@@ -23,7 +23,7 @@
 import logging
 import math
 import os
-from collections.abc import Generator, Iterable
+from collections.abc import Generator, Iterable, Mapping
 from contextlib import contextmanager
 from datetime import timedelta
 from typing import Any, Optional, Union
@@ -1643,8 +1643,26 @@ def loggers(self) -> list[Logger]:
         return self._loggers
 
     @loggers.setter
-    def loggers(self, loggers: Optional[list[Logger]]) -> None:
-        self._loggers = loggers if loggers else []
+    def loggers(self, loggers: Optional[Union[list[Logger], Mapping[str, Logger]]]) -> None:
+        self._logger_keys: list[Union[str, int]]
+        if isinstance(loggers, Mapping):
+            self._loggers = list(loggers.values())
+            self._logger_keys = list(loggers.keys())
+        else:
+            self._loggers = loggers if loggers else []
+            self._logger_keys = []
+
+    @property
+    def logger_map(self) -> dict[Union[str, int], Logger]:
+        """A mapping of logger keys to :class:`~lightning.pytorch.loggers.logger.Logger` used.
+
+        .. code-block:: python
+            tb_logger = trainer.logger_map.get("tensorboard", None)
+            if tb_logger:
+                tb_logger.log_hyperparams({"lr": 0.001})
+
+        """
+        return dict(zip(self._logger_keys, self._loggers))
 
     @property
     def callback_metrics(self) -> _OUT_DICT:

tests/tests_pytorch/trainer/properties/test_loggers.py

@@ -23,29 +23,47 @@ def test_trainer_loggers_property():
     """Test for correct initialization of loggers in Trainer."""
     logger1 = CustomLogger()
     logger2 = CustomLogger()
+    logger3 = CustomLogger()
 
     # trainer.loggers should be a copy of the input list
     trainer = Trainer(logger=[logger1, logger2])
 
     assert trainer.loggers == [logger1, logger2]
+    assert trainer.logger_map == {}
 
     # trainer.loggers should create a list of size 1
     trainer = Trainer(logger=logger1)
 
     assert trainer.logger == logger1
     assert trainer.loggers == [logger1]
+    assert trainer.logger_map == {}
+
+    trainer.loggers.append(logger2)
+    assert trainer.loggers == [logger1, logger2]
+    assert trainer.logger_map == {}
 
     # trainer.loggers should be a list of size 1 holding the default logger
     trainer = Trainer(logger=True)
 
     assert trainer.loggers == [trainer.logger]
     assert isinstance(trainer.logger, TensorBoardLogger)
 
+    trainer = Trainer(logger={"log1": logger1, "log2": logger2})
+    assert trainer.logger == logger1
+    assert trainer.loggers == [logger1, logger2]
+    assert isinstance(trainer.logger_map, dict)
+    assert trainer.logger_map == {"log1": logger1, "log2": logger2}
+
+    trainer.loggers.append(logger3)
+    assert trainer.loggers == [logger1, logger2, logger3]
+    assert trainer.logger_map == {"log1": logger1, "log2": logger2}
+
 
 def test_trainer_loggers_setters():
     """Test the behavior of setters for trainer.logger and trainer.loggers."""
     logger1 = CustomLogger()
     logger2 = CustomLogger()
+    logger3 = CustomLogger()
 
     trainer = Trainer()
     assert type(trainer.logger) is TensorBoardLogger
@@ -59,22 +77,40 @@ def test_trainer_loggers_setters():
     trainer.logger = None
     assert trainer.logger is None
     assert trainer.loggers == []
+    assert isinstance(trainer.loggers, list)
+    assert trainer.logger_map == {}
 
     # Test setters for trainer.loggers
     trainer.loggers = [logger1, logger2]
     assert trainer.loggers == [logger1, logger2]
+    assert isinstance(trainer.loggers, list)
+    assert trainer.logger_map == {}
 
     trainer.loggers = [logger1]
     assert trainer.loggers == [logger1]
     assert trainer.logger == logger1
+    assert trainer.logger_map == {}
 
     trainer.loggers = []
     assert trainer.loggers == []
     assert trainer.logger is None
+    assert isinstance(trainer.loggers, list)
+    assert trainer.logger_map == {}
 
     trainer.loggers = None
     assert trainer.loggers == []
     assert trainer.logger is None
+    assert isinstance(trainer.loggers, list)
+    assert trainer.logger_map == {}
+
+    trainer.loggers = {"log1": logger1, "log2": logger2}
+    assert trainer.loggers == [logger1, logger2]
+    assert isinstance(trainer.loggers, list)
+    assert isinstance(trainer.logger_map, dict)
+    assert trainer.logger_map == {"log1": logger1, "log2": logger2}
+
+    trainer.loggers.append(logger3)
+    assert trainer.logger_map == {"log1": logger1, "log2": logger2}
 
 
 @pytest.mark.parametrize(
@@ -94,3 +130,4 @@ def test_no_logger(tmp_path, logger_value):
     assert trainer.logger is None
     assert trainer.loggers == []
     assert trainer.log_dir == str(tmp_path)
+    assert trainer.logger_map == {}