Logging: Switch to loguru

Loguru is a flexible logger that allows for easier hooking and imports
into Rich with no problems. Also makes progress bars stick to the
bottom of the terminal window.

Signed-off-by: kingbri <bdashore3@proton.me>

common/auth.py

@@ -6,12 +6,9 @@ import secrets
 import yaml
 from fastapi import Header, HTTPException
 from pydantic import BaseModel
+from loguru import logger
 from typing import Optional
 
-from common.logger import init_logger
-
-logger = init_logger(__name__)
-
 
 class AuthKeys(BaseModel):
     """

common/config.py

@@ -1,11 +1,9 @@
 import yaml
 import pathlib
+from loguru import logger
 
-from common.logger import init_logger
 from common.utils import unwrap
 
-logger = init_logger(__name__)
-
 GLOBAL_CONFIG: dict = {}
 
 

common/gen_logging.py

@@ -2,12 +2,9 @@
 Functions for logging generation events.
 """
 from pydantic import BaseModel
+from loguru import logger
 from typing import Dict, Optional
 
-from common.logger import init_logger
-
-logger = init_logger(__name__)
-
 
 class LogPreferences(BaseModel):
     """Logging preference config."""

common/logger.py

@@ -1,71 +1,104 @@
 """
-Logging utility.
-https://github.com/PygmalionAI/aphrodite-engine/blob/main/aphrodite/common/logger.py
+Internal logging utility.
 """
 
 import logging
-import sys
-import colorlog
+from loguru import logger
+from rich.console import Console
+from rich.progress import (
+    Progress,
+    TextColumn,
+    BarColumn,
+    TimeRemainingColumn,
+    TaskProgressColumn,
+    MofNCompleteColumn,
+)
 
-_FORMAT = "%(log_color)s%(levelname)s: %(message)s"
-_DATE_FORMAT = "%m-%d %H:%M:%S"
+from common.utils import unwrap
+
+RICH_CONSOLE = Console()
 
 
-class ColoredFormatter(colorlog.ColoredFormatter):
-    """Adds logging prefix to newlines to align multi-line messages."""
+def get_loading_progress_bar():
+    """Gets a pre-made progress bar for loading tasks."""
 
-    def __init__(self, fmt, datefmt=None, log_colors=None, reset=True, style="%"):
-        super().__init__(
-            fmt, datefmt=datefmt, log_colors=log_colors, reset=reset, style=style
+    return Progress(
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        TaskProgressColumn(),
+        MofNCompleteColumn(),
+        TimeRemainingColumn(),
+        console=RICH_CONSOLE,
+    )
+
+
+def _log_formatter(record: dict) -> str:
+    """Log message formatter."""
+
+    color_map = {
+        "TRACE": "dim blue",
+        "DEBUG": "cyan",
+        "INFO": "green",
+        "SUCCESS": "bold green",
+        "WARNING": "yellow",
+        "ERROR": "red",
+        "CRITICAL": "bold white on red",
+    }
+    level = record.get("level")
+    level_color = color_map.get(level.name, "cyan")
+
+    message = unwrap(record.get("message"), "")
+    lines = message.splitlines()
+
+    # Replace once loguru allows for turning off str.format
+    message = message.replace("{", "{{").replace("}", "}}")
+
+    fmt = ""
+    if len(lines) > 1:
+        fmt = "\n".join(
+            [
+                f"[{level_color}]{level.name + ':' :<10}[/{level_color}]{line}"
+                for line in lines
+            ]
+        )
+    else:
+        fmt = f"[{level_color}]{level.name + ':' :<10}[/{level_color}]{message}"
+
+    return fmt
+
+
+# Uvicorn log handler
+# Uvicorn log portions inspired from https://github.com/encode/uvicorn/discussions/2027#discussioncomment-6432362
+class UvicornLoggingHandler(logging.Handler):
+    def emit(self, record: logging.LogRecord) -> None:
+        logger.opt(exception=record.exc_info).log(
+            record.levelname, self.format(record).rstrip()
         )
 
-    def format(self, record):
-        msg = super().format(record)
-        if record.message != "":
-            parts = msg.split(record.message)
-            msg = msg.replace("\n", "\r\n" + parts[0])
-        return msg
 
-
-_root_logger = logging.getLogger("aphrodite")
-_default_handler = None
-
-
-def _setup_logger():
-    _root_logger.setLevel(logging.DEBUG)
-    global _default_handler
-    if _default_handler is None:
-        _default_handler = logging.StreamHandler(sys.stdout)
-        _default_handler.flush = sys.stdout.flush  # type: ignore
-        _default_handler.setLevel(logging.INFO)
-        _root_logger.addHandler(_default_handler)
-    fmt = ColoredFormatter(
-        _FORMAT,
-        datefmt=_DATE_FORMAT,
-        log_colors={
-            "DEBUG": "cyan",
-            "INFO": "green",
-            "WARNING": "yellow",
-            "ERROR": "red",
-            "CRITICAL": "red,bg_white",
+# Uvicorn config for logging. Passed into run when creating all loggers in server
+UVICORN_LOG_CONFIG = {
+    "version": 1,
+    "disable_existing_loggers": False,
+    "handlers": {
+        "uvicorn": {
+            "class": (
+                f"{UvicornLoggingHandler.__module__}.{UvicornLoggingHandler.__qualname__}",
+            )
         },
-        reset=True,
+    },
+    "root": {"handlers": ["uvicorn"], "propagate": False, "level": "TRACE"},
+}
+
+
+def setup_logger():
+    """Bootstrap the logger."""
+
+    logger.remove()
+
+    logger.add(
+        RICH_CONSOLE.print,
+        level="DEBUG",
+        format=_log_formatter,
+        colorize=True,
     )
-    _default_handler.setFormatter(fmt)
-    # Setting this will avoid the message
-    # being propagated to the parent logger.
-    _root_logger.propagate = False
-
-
-# The logger is initialized when the module is imported.
-# This is thread-safe as the module is only imported once,
-# guaranteed by the Python GIL.
-_setup_logger()
-
-
-def init_logger(name: str):
-    logger = logging.getLogger(name)
-    logger.setLevel(logging.DEBUG)
-    logger.addHandler(_default_handler)
-    logger.propagate = False
-    return logger

common/sampling.py

@@ -1,17 +1,14 @@
 """Common functions for sampling parameters"""
 
 import pathlib
-from typing import Dict, List, Optional, Union
-from pydantic import AliasChoices, BaseModel, Field
 import yaml
+from loguru import logger
+from pydantic import AliasChoices, BaseModel, Field
+from typing import Dict, List, Optional, Union
 
-from common.logger import init_logger
 from common.utils import unwrap, prune_dict
 
 
-logger = init_logger(__name__)
-
-
 # Common class for sampler params
 class BaseSamplerRequest(BaseModel):
     """Common class for sampler params that are used in APIs"""

common/utils.py

@@ -1,21 +1,10 @@
 """Common utility functions"""
 
 import traceback
+from loguru import logger
 from pydantic import BaseModel
-from rich.progress import (
-    Progress,
-    TextColumn,
-    BarColumn,
-    TimeRemainingColumn,
-    TaskProgressColumn,
-    MofNCompleteColumn,
-)
 from typing import Optional
 
-from common.logger import init_logger
-
-logger = init_logger(__name__)
-
 
 def load_progress(module, modules):
     """Wrapper callback for load progress."""
@@ -66,18 +55,6 @@ def get_sse_packet(json_data: str):
     return f"data: {json_data}\n\n"
 
 
-def get_loading_progress_bar():
-    """Gets a pre-made progress bar for loading tasks."""
-
-    return Progress(
-        TextColumn("[progress.description]{task.description}"),
-        BarColumn(),
-        TaskProgressColumn(),
-        MofNCompleteColumn(),
-        TimeRemainingColumn(),
-    )
-
-
 def unwrap(wrapped, default=None):
     """Unwrap function for Optionals."""
     if wrapped is None: