mirror of
https://github.com/harvard-edge/cs249r_book.git
synced 2026-05-02 18:50:17 -05:00
e3cc9f7af3
Remove redundant ml_ prefix from ml_workflow chapter files and update all Quarto config references. Consolidate custom scripts into native binder subcommands and archive obsolete tooling.
248 lines
8.0 KiB
Python
248 lines
8.0 KiB
Python
"""
|
|
Centralized logging configuration for MLSysBook tools.
|
|
|
|
This module provides a standardized logging setup with support for structured
|
|
logging, different output formats, and configurable log levels.
|
|
"""
|
|
|
|
import logging
|
|
import logging.handlers
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import Optional, Dict, Any
|
|
from datetime import datetime
|
|
|
|
from rich.console import Console
|
|
from rich.logging import RichHandler
|
|
from rich.traceback import install
|
|
|
|
from .config import get_config
|
|
|
|
|
|
class StructuredFormatter(logging.Formatter):
|
|
"""Custom formatter that adds structured information to log records."""
|
|
|
|
def __init__(self, include_context: bool = True) -> None:
|
|
"""Initialize the formatter.
|
|
|
|
Args:
|
|
include_context: Whether to include additional context in log records
|
|
"""
|
|
super().__init__()
|
|
self.include_context = include_context
|
|
|
|
def format(self, record: logging.LogRecord) -> str:
|
|
"""Format a log record with structured information."""
|
|
# Add timestamp
|
|
record.timestamp = datetime.utcnow().isoformat()
|
|
|
|
# Add context information if available
|
|
if self.include_context and hasattr(record, 'context'):
|
|
record.context_str = f" | Context: {record.context}"
|
|
else:
|
|
record.context_str = ""
|
|
|
|
# Create the formatted message
|
|
if record.levelno >= logging.ERROR and record.exc_info:
|
|
# For errors with exceptions, include the full traceback
|
|
formatted = (
|
|
f"[{record.timestamp}] {record.levelname:8} | "
|
|
f"{record.name} | {record.getMessage()}{record.context_str}\n"
|
|
f"Exception: {self.formatException(record.exc_info)}"
|
|
)
|
|
else:
|
|
formatted = (
|
|
f"[{record.timestamp}] {record.levelname:8} | "
|
|
f"{record.name} | {record.getMessage()}{record.context_str}"
|
|
)
|
|
|
|
return formatted
|
|
|
|
|
|
class ProgressAwareHandler(RichHandler):
|
|
"""Rich handler that works well with progress bars and other rich output."""
|
|
|
|
def __init__(self, *args, **kwargs) -> None:
|
|
# Create a separate console for logging to avoid conflicts
|
|
console = Console(stderr=True, force_terminal=True)
|
|
kwargs['console'] = console
|
|
kwargs['show_time'] = False # We'll handle time in our formatter
|
|
kwargs['show_path'] = False
|
|
super().__init__(*args, **kwargs)
|
|
|
|
|
|
def setup_logging(
|
|
name: Optional[str] = None,
|
|
level: Optional[str] = None,
|
|
log_to_file: Optional[bool] = None,
|
|
log_file_path: Optional[Path] = None,
|
|
enable_rich_tracebacks: bool = True,
|
|
context: Optional[Dict[str, Any]] = None
|
|
) -> logging.Logger:
|
|
"""Set up logging configuration for MLSysBook tools.
|
|
|
|
Args:
|
|
name: Logger name (defaults to calling module)
|
|
level: Log level (defaults to config value)
|
|
log_to_file: Whether to log to file (defaults to config value)
|
|
log_file_path: Path for log file (defaults to config value)
|
|
enable_rich_tracebacks: Whether to enable rich exception formatting
|
|
context: Additional context to include in all log messages
|
|
|
|
Returns:
|
|
Configured logger instance
|
|
"""
|
|
config = get_config()
|
|
|
|
# Use provided values or fall back to config
|
|
level = level or config.log_level
|
|
log_to_file = log_to_file if log_to_file is not None else config.log_to_file
|
|
log_file_path = log_file_path or config.log_file_path
|
|
|
|
# Set up rich tracebacks if enabled
|
|
if enable_rich_tracebacks:
|
|
install(show_locals=True)
|
|
|
|
# Create logger
|
|
logger = logging.getLogger(name or __name__)
|
|
logger.setLevel(getattr(logging, level.upper()))
|
|
|
|
# Clear any existing handlers
|
|
logger.handlers.clear()
|
|
|
|
# Add console handler with rich formatting
|
|
console_handler = ProgressAwareHandler()
|
|
console_handler.setLevel(getattr(logging, level.upper()))
|
|
console_formatter = logging.Formatter("%(name)s | %(message)s")
|
|
console_handler.setFormatter(console_formatter)
|
|
logger.addHandler(console_handler)
|
|
|
|
# Add file handler if requested
|
|
if log_to_file and log_file_path:
|
|
# Ensure log directory exists
|
|
log_file_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Create rotating file handler
|
|
file_handler = logging.handlers.RotatingFileHandler(
|
|
log_file_path,
|
|
maxBytes=10 * 1024 * 1024, # 10MB
|
|
backupCount=5,
|
|
encoding='utf-8'
|
|
)
|
|
file_handler.setLevel(logging.DEBUG) # Always debug level for files
|
|
file_formatter = StructuredFormatter(include_context=True)
|
|
file_handler.setFormatter(file_formatter)
|
|
logger.addHandler(file_handler)
|
|
|
|
# Add context if provided
|
|
if context:
|
|
logger = LoggerAdapter(logger, context)
|
|
|
|
return logger
|
|
|
|
|
|
def get_logger(
|
|
name: Optional[str] = None,
|
|
context: Optional[Dict[str, Any]] = None
|
|
) -> logging.Logger:
|
|
"""Get a logger instance with optional context.
|
|
|
|
Args:
|
|
name: Logger name (defaults to calling module)
|
|
context: Additional context to include in log messages
|
|
|
|
Returns:
|
|
Logger instance
|
|
"""
|
|
logger = logging.getLogger(name or __name__)
|
|
|
|
# If logger doesn't have handlers, set it up
|
|
if not logger.handlers:
|
|
logger = setup_logging(name)
|
|
|
|
# Add context if provided
|
|
if context:
|
|
logger = LoggerAdapter(logger, context)
|
|
|
|
return logger
|
|
|
|
|
|
class LoggerAdapter(logging.LoggerAdapter):
|
|
"""Adapter that adds context information to log records."""
|
|
|
|
def __init__(self, logger: logging.Logger, extra: Dict[str, Any]) -> None:
|
|
"""Initialize the adapter with context information.
|
|
|
|
Args:
|
|
logger: Base logger instance
|
|
extra: Context information to add to log records
|
|
"""
|
|
super().__init__(logger, extra)
|
|
|
|
def process(self, msg: str, kwargs: Dict[str, Any]) -> tuple:
|
|
"""Process a log message to add context information."""
|
|
if 'extra' not in kwargs:
|
|
kwargs['extra'] = {}
|
|
|
|
# Add our context to the extra information
|
|
kwargs['extra'].update(self.extra)
|
|
kwargs['extra']['context'] = self.extra
|
|
|
|
return msg, kwargs
|
|
|
|
|
|
class ProgressLogger:
|
|
"""Logger that works well with progress bars and other rich output."""
|
|
|
|
def __init__(self, name: Optional[str] = None) -> None:
|
|
"""Initialize the progress logger.
|
|
|
|
Args:
|
|
name: Logger name
|
|
"""
|
|
self.logger = get_logger(name)
|
|
self._progress_active = False
|
|
|
|
def start_progress(self) -> None:
|
|
"""Indicate that a progress operation is starting."""
|
|
self._progress_active = True
|
|
|
|
def end_progress(self) -> None:
|
|
"""Indicate that a progress operation has ended."""
|
|
self._progress_active = False
|
|
|
|
def log(self, level: str, message: str, **kwargs) -> None:
|
|
"""Log a message with progress awareness.
|
|
|
|
Args:
|
|
level: Log level (debug, info, warning, error, critical)
|
|
message: Log message
|
|
**kwargs: Additional keyword arguments for logging
|
|
"""
|
|
if self._progress_active:
|
|
# For progress operations, use a simpler format
|
|
print(f"[{level.upper()}] {message}", file=sys.stderr)
|
|
else:
|
|
# Use normal logging
|
|
getattr(self.logger, level.lower())(message, **kwargs)
|
|
|
|
def debug(self, message: str, **kwargs) -> None:
|
|
"""Log a debug message."""
|
|
self.log('debug', message, **kwargs)
|
|
|
|
def info(self, message: str, **kwargs) -> None:
|
|
"""Log an info message."""
|
|
self.log('info', message, **kwargs)
|
|
|
|
def warning(self, message: str, **kwargs) -> None:
|
|
"""Log a warning message."""
|
|
self.log('warning', message, **kwargs)
|
|
|
|
def error(self, message: str, **kwargs) -> None:
|
|
"""Log an error message."""
|
|
self.log('error', message, **kwargs)
|
|
|
|
def critical(self, message: str, **kwargs) -> None:
|
|
"""Log a critical message."""
|
|
self.log('critical', message, **kwargs)
|