284 lines
8.4 KiB
Python
284 lines
8.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Logging configuration for oMLX.
|
|
|
|
This module provides centralized logging configuration with support for:
|
|
- Standard logging with configurable levels
|
|
- Structured JSON logging (optional)
|
|
- Request context tracking
|
|
- File logging with daily rotation
|
|
- Consistent formatting across all modules
|
|
"""
|
|
|
|
import logging
|
|
import sys
|
|
from contextvars import ContextVar
|
|
from logging.handlers import TimedRotatingFileHandler
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
# Context variable for request ID tracking
|
|
_request_id: ContextVar[Optional[str]] = ContextVar("request_id", default=None)
|
|
|
|
|
|
class RequestContextFilter(logging.Filter):
|
|
"""
|
|
Add request_id to log records.
|
|
|
|
This filter adds the current request ID (if set) to all log records,
|
|
enabling request-level log correlation.
|
|
"""
|
|
|
|
def filter(self, record: logging.LogRecord) -> bool:
|
|
"""Add request_id attribute to log record."""
|
|
record.request_id = _request_id.get() or "-"
|
|
return True
|
|
|
|
|
|
class AdminStatsAccessFilter(logging.Filter):
|
|
"""Suppress repetitive uvicorn access logs for admin polling endpoints."""
|
|
|
|
def filter(self, record: logging.LogRecord) -> bool:
|
|
msg = record.getMessage()
|
|
if "/admin/api/stats" in msg:
|
|
return False
|
|
if "/admin/api/login" in msg:
|
|
return False
|
|
if "/admin/api/hf/tasks" in msg:
|
|
return False
|
|
if "/admin/api/oq/tasks" in msg:
|
|
return False
|
|
return True
|
|
|
|
|
|
class ColoredFormatter(logging.Formatter):
|
|
"""
|
|
Colored log formatter for terminal output.
|
|
|
|
Uses ANSI color codes to highlight different log levels.
|
|
"""
|
|
|
|
COLORS = {
|
|
5: "\033[90m", # Gray (TRACE)
|
|
logging.DEBUG: "\033[36m", # Cyan
|
|
logging.INFO: "\033[32m", # Green
|
|
logging.WARNING: "\033[33m", # Yellow
|
|
logging.ERROR: "\033[31m", # Red
|
|
logging.CRITICAL: "\033[35m", # Magenta
|
|
}
|
|
RESET = "\033[0m"
|
|
|
|
def format(self, record: logging.LogRecord) -> str:
|
|
"""Format the log record with colors."""
|
|
color = self.COLORS.get(record.levelno, "")
|
|
record.levelname = f"{color}{record.levelname}{self.RESET}"
|
|
return super().format(record)
|
|
|
|
|
|
def get_request_id() -> Optional[str]:
|
|
"""Get the current request ID from context."""
|
|
return _request_id.get()
|
|
|
|
|
|
def set_request_id(request_id: Optional[str]) -> None:
|
|
"""Set the current request ID in context."""
|
|
_request_id.set(request_id)
|
|
|
|
|
|
def configure_logging(
|
|
level: str = "INFO",
|
|
format_style: str = "standard",
|
|
include_request_id: bool = True,
|
|
colored: bool = True,
|
|
) -> None:
|
|
"""
|
|
Configure application logging.
|
|
|
|
Args:
|
|
level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL).
|
|
format_style: "standard" for plain text, "json" for structured JSON.
|
|
include_request_id: Whether to include request_id in log format.
|
|
colored: Whether to use colored output (only for standard format).
|
|
"""
|
|
# Determine log level
|
|
level_name = level.upper()
|
|
log_level = 5 if level_name == "TRACE" else getattr(logging, level_name, logging.INFO)
|
|
|
|
# Build format string
|
|
if include_request_id:
|
|
format_str = "%(asctime)s - %(name)s - %(levelname)s - [%(request_id)s] - %(message)s"
|
|
else:
|
|
format_str = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
|
|
|
# Create handler
|
|
handler = logging.StreamHandler(sys.stderr)
|
|
handler.setLevel(log_level)
|
|
|
|
# Create formatter
|
|
if format_style == "json":
|
|
formatter = JsonFormatter(format_str)
|
|
elif colored and sys.stderr.isatty():
|
|
formatter = ColoredFormatter(format_str)
|
|
else:
|
|
formatter = logging.Formatter(format_str)
|
|
|
|
handler.setFormatter(formatter)
|
|
|
|
# Add request context filter if needed
|
|
if include_request_id:
|
|
handler.addFilter(RequestContextFilter())
|
|
|
|
# Configure root logger
|
|
root_logger = logging.getLogger()
|
|
root_logger.setLevel(log_level)
|
|
|
|
# Remove existing handlers to avoid duplicates
|
|
root_logger.handlers.clear()
|
|
root_logger.addHandler(handler)
|
|
|
|
# Set specific loggers
|
|
logging.getLogger("omlx").setLevel(log_level)
|
|
logging.getLogger("uvicorn").setLevel(log_level)
|
|
logging.getLogger("uvicorn.access").addFilter(AdminStatsAccessFilter())
|
|
|
|
# Suppress noisy third-party loggers unless trace level
|
|
third_party_level = log_level if log_level <= 5 else logging.INFO
|
|
logging.getLogger("httpx").setLevel(third_party_level)
|
|
logging.getLogger("httpcore").setLevel(third_party_level)
|
|
|
|
|
|
class JsonFormatter(logging.Formatter):
|
|
"""
|
|
JSON log formatter for structured logging.
|
|
|
|
Outputs logs as JSON objects for easy parsing by log aggregators.
|
|
"""
|
|
|
|
def format(self, record: logging.LogRecord) -> str:
|
|
"""Format the log record as JSON."""
|
|
import json
|
|
import time
|
|
|
|
log_data = {
|
|
"timestamp": self.formatTime(record, self.datefmt),
|
|
"level": record.levelname,
|
|
"logger": record.name,
|
|
"message": record.getMessage(),
|
|
}
|
|
|
|
# Add request_id if available
|
|
request_id = getattr(record, "request_id", None)
|
|
if request_id and request_id != "-":
|
|
log_data["request_id"] = request_id
|
|
|
|
# Add exception info if present
|
|
if record.exc_info:
|
|
log_data["exception"] = self.formatException(record.exc_info)
|
|
|
|
# Add extra fields
|
|
for key in ["request_id", "model", "tokens", "latency_ms"]:
|
|
if hasattr(record, key) and key not in log_data:
|
|
log_data[key] = getattr(record, key)
|
|
|
|
return json.dumps(log_data)
|
|
|
|
|
|
def get_logger(
|
|
name: str,
|
|
request_id: Optional[str] = None,
|
|
) -> logging.Logger:
|
|
"""
|
|
Get a logger with optional request context.
|
|
|
|
Args:
|
|
name: Logger name (usually __name__).
|
|
request_id: Optional request ID to set in context.
|
|
|
|
Returns:
|
|
Logger instance.
|
|
"""
|
|
if request_id:
|
|
set_request_id(request_id)
|
|
|
|
return logging.getLogger(name)
|
|
|
|
|
|
class RequestLogContext:
|
|
"""
|
|
Context manager for request-scoped logging.
|
|
|
|
Usage:
|
|
with RequestLogContext(request_id="abc123"):
|
|
logger.info("Processing request")
|
|
"""
|
|
|
|
def __init__(self, request_id: str):
|
|
self.request_id = request_id
|
|
self.previous_id: Optional[str] = None
|
|
|
|
def __enter__(self) -> "RequestLogContext":
|
|
self.previous_id = _request_id.get()
|
|
_request_id.set(self.request_id)
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb) -> None:
|
|
_request_id.set(self.previous_id)
|
|
|
|
|
|
def configure_file_logging(
|
|
log_dir: Path,
|
|
level: str = "INFO",
|
|
include_request_id: bool = True,
|
|
retention_days: int = 7,
|
|
) -> None:
|
|
"""
|
|
Configure file logging with daily rotation.
|
|
|
|
Adds a file handler to the root logger that writes to {log_dir}/server.log
|
|
with automatic daily rotation. Old log files are automatically deleted
|
|
after retention_days.
|
|
|
|
Args:
|
|
log_dir: Directory to store log files.
|
|
level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL).
|
|
include_request_id: Whether to include request_id in log format.
|
|
retention_days: Number of days to retain old logs.
|
|
"""
|
|
# Ensure log directory exists
|
|
log_dir = Path(log_dir)
|
|
log_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
level_name = level.upper()
|
|
log_level = 5 if level_name == "TRACE" else getattr(logging, level_name, logging.INFO)
|
|
|
|
# Build format string (no colors for file)
|
|
if include_request_id:
|
|
format_str = "%(asctime)s - %(name)s - %(levelname)s - [%(request_id)s] - %(message)s"
|
|
else:
|
|
format_str = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
|
|
|
# Create file handler with daily rotation
|
|
# File: server.log, rotated files: server.log.YYYY-MM-DD
|
|
log_file = log_dir / "server.log"
|
|
|
|
file_handler = TimedRotatingFileHandler(
|
|
filename=log_file,
|
|
when="midnight",
|
|
interval=1,
|
|
backupCount=retention_days,
|
|
encoding="utf-8",
|
|
)
|
|
file_handler.suffix = "%Y-%m-%d" # Results in server.log.2024-01-15
|
|
file_handler.setLevel(log_level)
|
|
|
|
formatter = logging.Formatter(format_str)
|
|
file_handler.setFormatter(formatter)
|
|
|
|
# Add request context filter
|
|
if include_request_id:
|
|
file_handler.addFilter(RequestContextFilter())
|
|
|
|
# Add to root logger
|
|
root_logger = logging.getLogger()
|
|
root_logger.addHandler(file_handler)
|