Add experimental option to use the rich logging handler (#1810)

Author: oelachqar

src/oumi/utils/logging.py

@@ -78,7 +78,7 @@ def configure_logger(
 
     device_rank = _detect_rank()
 
-    formatter = logging.Formatter(
+    default_formatter = logging.Formatter(
         "[%(asctime)s][%(name)s]"
         f"[rank{device_rank}]"
         "[pid:%(process)d][%(threadName)s]"
@@ -87,8 +87,12 @@ def configure_logger(
 
     # Add a console handler to the logger for only global leader.
     if device_rank == 0:
-        console_handler = logging.StreamHandler(sys.stdout)
-        console_handler.setFormatter(formatter)
+        if _should_use_rich_logging():
+            console_handler = _configure_rich_handler(device_rank, level)
+        else:
+            console_handler = logging.StreamHandler(sys.stdout)
+            console_handler.setFormatter(default_formatter)
+
         console_handler.setLevel(level.upper())
         logger.addHandler(console_handler)
 
@@ -98,13 +102,76 @@ def configure_logger(
         log_dir.mkdir(parents=True, exist_ok=True)
 
         file_handler = logging.FileHandler(log_dir / f"rank_{device_rank:04d}.log")
-        file_handler.setFormatter(formatter)
+        file_handler.setFormatter(default_formatter)
         file_handler.setLevel(level.upper())
         logger.addHandler(file_handler)
 
     logger.propagate = False
 
 
+def _should_use_rich_logging() -> bool:
+    """Determines if rich logging should be used based on environment variables.
+
+    Note: Rich logging is experimental, and may be removed in the future.
+        Currently it is disabled by default.
+    """
+    # Check if explicitly disabled
+    if os.environ.get("OUMI_ENABLE_RICH_LOGGING", "").lower() in (
+        "1",
+        "yes",
+        "on",
+        "true",
+        "y",
+    ):
+        return sys.stdout.isatty()  # is in a terminal
+
+    return False
+
+
+def _configure_rich_handler(
+    device_rank: int,
+    level: str,
+) -> logging.Handler:
+    """Configures a rich logging handler."""
+    try:
+        from rich.console import Console
+        from rich.logging import RichHandler
+        from rich.traceback import install
+    except ImportError:
+        raise ImportError(
+            "Rich logging is not installed. Please install it with `pip install rich`."
+        )
+
+    use_detailed_logging = level.upper() == "DEBUG"
+
+    if use_detailed_logging:
+        # Add extra logging for debugging
+        install(show_locals=True, suppress=[])
+
+    console = Console()
+    console_handler = RichHandler(
+        console=console,
+        show_time=True,
+        show_level=True,
+        show_path=True,
+        enable_link_path=True,
+        markup=False,
+        rich_tracebacks=use_detailed_logging,
+        tracebacks_show_locals=use_detailed_logging,
+        locals_max_length=20,
+        locals_max_string=80,
+    )
+
+    if use_detailed_logging:
+        rich_formatter = logging.Formatter(
+            f"[rank-{device_rank}][pid-%(process)d][%(threadName)s] %(message)s"
+        )
+    else:
+        rich_formatter = logging.Formatter(f"[rank-{device_rank}] %(message)s")
+    console_handler.setFormatter(rich_formatter)
+    return console_handler
+
+
 def update_logger_level(name: str, level: str = "info") -> None:
     """Updates the log level of the logger.