feat: integrate logging w/ rerun

moai/conf/engine/modules/rerun.yaml

@@ -11,3 +11,6 @@ world_coordinates: RUF
 add_floor: true
 root: /
 memory_limit: 75% #NOTE: can also be absolute, `4GB` for example
+log:
+  handler: moai/logs # null
+  level: ${hydra:hydra_logging.root.level}

moai/engine/modules/rerun.py

@@ -1,10 +1,13 @@
+import logging
 import typing
 
 import colour
 import numpy as np
 import toolz
+from omegaconf.omegaconf import DictConfig
 
 from moai.utils.color.colormap import random_color
+from moai.utils.funcs import get
 
 try:
     import rerun as rr
@@ -16,6 +19,15 @@
 __all__ = ["Rerun"]
 
 
+class RerunFilter(logging.Filter):
+    def __init__(self) -> None:
+        super().__init__("rerun")
+
+    def filter(self, record):
+        record.msg = record.msg.removeprefix(":moai: ")
+        return record
+
+
 class Rerun:
     r"""
     This logger should support all the different types of rr logging.
@@ -39,6 +51,7 @@ def __init__(
         add_floor: bool = True,
         root: str = "/",
         memory_limit: str = "75%",
+        log: DictConfig = {},
     ) -> None:
         # NOTE: https://github.com/iterative/dvc/issues/9731
         rr.init(name)
@@ -62,6 +75,15 @@ def __init__(
         if add_floor:
             self._create_floor(root)
         # NOTE: add more global gizmos (e.g. tr-axis)
+        if handler_name := get(log, "handler"):
+            handler = rr.LoggingHandler(handler_name)
+            handler.addFilter(RerunFilter())
+            logging.getLogger().addHandler(handler)
+            if level := get(log, "level"):
+                if isinstance(level, str):
+                    level = level.upper()
+                level = logging.getLevelName(level)
+                logging.getLogger().setLevel(level)
 
     def _create_scalar_plots(self, root: str, plots) -> None:
         for plot in plots:

moai/engine/progressbar.py

@@ -1,4 +1,5 @@
 import rich.progress
+from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import RichProgressBar
 from pytorch_lightning.callbacks.progress.rich_progress import RichProgressBarTheme
 
@@ -39,7 +40,7 @@ def __init__(self) -> None:
     #         CustomTimeColumn(style=self.theme.time),
     #         ProcessingSpeedColumn(style=self.theme.processing_speed),
     #     ]
-    def configure_columns(self, trainer: "pl.Trainer") -> list:
+    def configure_columns(self, trainer: Trainer) -> list:
         original = super().configure_columns(trainer)
         moai_column = rich.progress.TextColumn(":moai:")
         spinner_column = rich.progress.SpinnerColumn(spinner_name="dots5")