Feat: Support for colored log output (#30)

anandhu-eng · web-flow · commit 86d62dc8d71b · 2025-12-02T14:06:16.000-06:00
* Support for colored log output
* make setup_logging() idempotent with force=True parameter
diff --git a/requirements/base.txt b/requirements/base.txt
@@ -26,6 +26,9 @@ datasets==4.1.1
 sentencepiece==0.2.1
 protobuf==6.33.0
 
+# Color support for cross-platform terminals
+colorama==0.4.6
+
 # Tokens will encode to torch.Tensor objects, but we only tokenize, not encode.
 # We don't necessarily need to install torch, but if needed, uncomment this:
 # --extra-index-url https://download.pytorch.org/whl/cpu
diff --git a/src/inference_endpoint/utils/logging.py b/src/inference_endpoint/utils/logging.py
@@ -20,8 +20,74 @@
 """
 
 import logging
+import os
 import sys
 
+from colorama import Fore, Style
+from colorama import init as _colorama_init
+
+# Initialize colorama
+_colorama_init(autoreset=True)
+
+# Map levelname -> color
+_LEVEL_COLORS = {
+    "INFO": Fore.GREEN,
+    "WARNING": Fore.YELLOW,
+    "ERROR": Fore.RED,
+    "CRITICAL": Fore.RED,
+}
+
+
+class ColoredFormatter(logging.Formatter):
+    """Formatter that applies colors to log level names.
+
+    Applies colorama colors (green for INFO, yellow for WARNING, red for ERROR/CRITICAL)
+    to the levelname field only, leaving the rest of the log message unmodified.
+    """
+
+    def __init__(
+        self,
+        fmt: str | None = None,
+        datefmt: str | None = None,
+        style: str = "%",
+        use_color: bool = False,
+    ):
+        """Initialize the formatter.
+
+        Args:
+            fmt: Log format string.
+            datefmt: Date format string.
+            style: Format style (% or {).
+            use_color: Whether to apply colors to levelname. Defaults to False.
+                      Enable by setting FORCE_COLOR_LOGGING environment variable.
+        """
+        super().__init__(fmt=fmt, datefmt=datefmt, style=style)
+        self.use_color = use_color
+
+    def format(self, record: logging.LogRecord) -> str:
+        """Format the log record with optional colors.
+
+        Args:
+            record: The log record to format.
+
+        Returns:
+            Formatted log message with colors applied if enabled.
+        """
+        # If coloring disabled, delegate
+        if not self.use_color:
+            return super().format(record)
+
+        orig = record.levelname
+        color = _LEVEL_COLORS.get(orig)
+        if color:
+            try:
+                record.levelname = f"{color}{orig}{Style.RESET_ALL}"
+                return super().format(record)
+            finally:
+                record.levelname = orig
+
+        return super().format(record)
+
 
 def setup_logging(level: str | None = None, format_string: str | None = None) -> None:
     """
@@ -38,20 +104,25 @@ def setup_logging(level: str | None = None, format_string: str | None = None) ->
     # Default format
     if format_string is None:
         format_string = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    # Disable colors by default to avoid potential formatting overhead during benchmarking.
+    # Colors can be explicitly enabled via FORCE_COLOR_LOGGING environment variable.
+    use_color = os.getenv("FORCE_COLOR_LOGGING") is not None
+
+    handler = logging.StreamHandler(sys.stdout)
+    handler.setFormatter(ColoredFormatter(fmt=format_string, use_color=use_color))
 
-    # Configure root logger
     logging.basicConfig(
-        level=getattr(logging, level.upper()),
-        format=format_string,
-        handlers=[logging.StreamHandler(sys.stdout)],
+        level=getattr(logging, level.upper()), handlers=[handler], force=True
     )
 
     # Set specific logger levels
     logging.getLogger("asyncio").setLevel(logging.WARNING)
     logging.getLogger("urllib3").setLevel(logging.WARNING)
 
     logger = logging.getLogger(__name__)
-    logger.debug(f"Logging configured with level: {level}")
+    logger.debug(
+        f"Logging configured with level: {level}, colors={'on' if use_color else 'off'}"
+    )
 
 
 def get_logger(name: str) -> logging.Logger:
diff --git a/tests/unit/test_logging.py b/tests/unit/test_logging.py
@@ -0,0 +1,230 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+
+"""Unit tests for logging module."""
+
+import logging
+
+import pytest
+from colorama import Fore, Style
+from inference_endpoint.utils.logging import ColoredFormatter, setup_logging
+
+
+class TestColoredFormatter:
+    """Tests for the ColoredFormatter class."""
+
+    def test_colored_formatter_with_color_enabled(self):
+        """Test that ColoredFormatter applies colors when use_color=True."""
+        formatter = ColoredFormatter(fmt="%(levelname)s - %(message)s", use_color=True)
+        record = logging.LogRecord(
+            name="test",
+            level=logging.INFO,
+            pathname="test.py",
+            lineno=1,
+            msg="test message",
+            args=(),
+            exc_info=None,
+        )
+
+        output = formatter.format(record)
+
+        # Check that color code is present in output
+        assert Fore.GREEN in output
+        assert Style.RESET_ALL in output
+        assert "test message" in output
+
+    def test_colored_formatter_with_color_disabled(self):
+        """Test that ColoredFormatter does not apply colors when use_color=False."""
+        formatter = ColoredFormatter(fmt="%(levelname)s - %(message)s", use_color=False)
+        record = logging.LogRecord(
+            name="test",
+            level=logging.INFO,
+            pathname="test.py",
+            lineno=1,
+            msg="test message",
+            args=(),
+            exc_info=None,
+        )
+
+        output = formatter.format(record)
+
+        # Check that no color codes are present
+        assert Fore.GREEN not in output
+        assert Style.RESET_ALL not in output
+        assert "INFO" in output
+        assert "test message" in output
+
+    def test_colored_formatter_levelname_restoration(self):
+        """Test that ColoredFormatter restores original levelname after formatting."""
+        formatter = ColoredFormatter(fmt="%(levelname)s", use_color=True)
+        record = logging.LogRecord(
+            name="test",
+            level=logging.WARNING,
+            pathname="test.py",
+            lineno=1,
+            msg="test",
+            args=(),
+            exc_info=None,
+        )
+
+        original_levelname = record.levelname
+        formatter.format(record)
+
+        # Verify levelname is restored
+        assert record.levelname == original_levelname
+
+    @pytest.mark.parametrize(
+        "level, expected_color",
+        [
+            (logging.INFO, Fore.GREEN),
+            (logging.WARNING, Fore.YELLOW),
+            (logging.ERROR, Fore.RED),
+            (logging.CRITICAL, Fore.RED),
+        ],
+    )
+    def test_colored_formatter_level_colors(self, level, expected_color):
+        """Test that correct colors are applied for each log level."""
+        formatter = ColoredFormatter(fmt="%(levelname)s", use_color=True)
+        record = logging.LogRecord(
+            name="test",
+            level=level,
+            pathname="test.py",
+            lineno=1,
+            msg="test",
+            args=(),
+            exc_info=None,
+        )
+
+        output = formatter.format(record)
+
+        assert expected_color in output
+        assert Style.RESET_ALL in output
+
+    def test_colored_formatter_debug_level_no_color(self):
+        """Test that DEBUG level (unmapped) does not crash or apply colors."""
+        formatter = ColoredFormatter(fmt="%(levelname)s - %(message)s", use_color=True)
+        record = logging.LogRecord(
+            name="test",
+            level=logging.DEBUG,
+            pathname="test.py",
+            lineno=1,
+            msg="debug message",
+            args=(),
+            exc_info=None,
+        )
+
+        output = formatter.format(record)
+
+        # DEBUG is not in the color map, so no color should be applied
+        assert "DEBUG" in output
+        assert "debug message" in output
+
+        # Verify NO colors were applied (since DEBUG is unmapped)
+        assert Fore.GREEN not in output
+        assert Fore.YELLOW not in output
+        assert Fore.RED not in output
+        assert Style.RESET_ALL not in output
+
+
+class TestSetupLogging:
+    """Tests for the setup_logging function."""
+
+    def test_setup_logging_configures_root_logger(self, monkeypatch):
+        """Test that setup_logging configures the root logger with a handler."""
+        # Ensure FORCE_COLOR_LOGGING is not set to ensure colors are disabled by default
+        monkeypatch.delenv("FORCE_COLOR_LOGGING", raising=False)
+
+        # Clear existing handlers
+        root_logger = logging.getLogger()
+        root_logger.handlers.clear()
+
+        setup_logging()
+
+        # Verify root logger has at least one handler
+        assert len(root_logger.handlers) > 0
+
+    def test_setup_logging_colors_enabled_with_force_color_env(self, monkeypatch):
+        """Test that FORCE_COLOR_LOGGING env var enables colors by testing output directly."""
+        monkeypatch.setenv("FORCE_COLOR_LOGGING", "1")
+
+        # Clear existing handlers
+        root_logger = logging.getLogger()
+        root_logger.handlers.clear()
+
+        setup_logging()
+
+        # Create a log record and format it
+        record = logging.LogRecord(
+            name="test",
+            level=logging.INFO,
+            pathname="test.py",
+            lineno=1,
+            msg="test",
+            args=(),
+            exc_info=None,
+        )
+
+        # Get the first handler's formatter
+        handler = root_logger.handlers[0]
+        formatted = handler.formatter.format(record)
+
+        # Verify color codes are present in output when FORCE_COLOR_LOGGING is set
+        assert Fore.GREEN in formatted
+
+    def test_setup_logging_colors_disabled_by_default(self, monkeypatch):
+        """Test that colors are disabled by default when FORCE_COLOR_LOGGING is not set."""
+        monkeypatch.delenv("FORCE_COLOR_LOGGING", raising=False)
+
+        # Clear existing handlers
+        root_logger = logging.getLogger()
+        root_logger.handlers.clear()
+
+        setup_logging()
+
+        # Create a log record and format it
+        record = logging.LogRecord(
+            name="test",
+            level=logging.INFO,
+            pathname="test.py",
+            lineno=1,
+            msg="test",
+            args=(),
+            exc_info=None,
+        )
+
+        # Get the first handler's formatter
+        handler = root_logger.handlers[0]
+        formatted = handler.formatter.format(record)
+
+        # Verify no color codes in output when FORCE_COLOR_LOGGING is not set
+        assert Fore.GREEN not in formatted
+
+    def test_setup_logging_asyncio_logger_level(self, monkeypatch):
+        """Test that asyncio logger is set to WARNING level after calling setup_logging()."""
+        monkeypatch.delenv("FORCE_COLOR_LOGGING", raising=False)
+
+        setup_logging()
+
+        asyncio_logger = logging.getLogger("asyncio")
+        assert asyncio_logger.level == logging.WARNING
+
+    def test_setup_logging_urllib3_logger_level(self, monkeypatch):
+        """Test that urllib3 logger is set to WARNING level after calling setup_logging()."""
+        monkeypatch.delenv("FORCE_COLOR_LOGGING", raising=False)
+
+        setup_logging()
+
+        urllib3_logger = logging.getLogger("urllib3")
+        assert urllib3_logger.level == logging.WARNING
