122 lines
3.7 KiB
Python
122 lines
3.7 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Logging handlers and middleware for structured logging.
|
|
|
|
This module provides FastAPI middleware and structlog processors for:
|
|
- Request/response logging with timing
|
|
- Sensitive data filtering in logs
|
|
- Structured logging configuration
|
|
- Error handling with detailed context
|
|
|
|
Key Components:
|
|
- LoggingMiddleware: FastAPI middleware for request/response logging
|
|
- filter_sensitive_data: Structlog processor for data sanitization
|
|
- get_logger: Factory function for structured loggers
|
|
"""
|
|
|
|
import time
|
|
|
|
import structlog
|
|
|
|
logger = structlog.get_logger(__name__)
|
|
|
|
# Paths and suffixes excluded from request logging.
|
|
_EXCLUDED_PATHS = frozenset(
|
|
{
|
|
"/api/train/status",
|
|
"/api/train/metrics",
|
|
"/api/train/hardware",
|
|
"/api/system",
|
|
}
|
|
)
|
|
_EXCLUDED_SUFFIXES = (".png", ".jpg", ".jpeg", ".ico", ".woff", ".woff2", ".ttf")
|
|
|
|
|
|
class LoggingMiddleware:
|
|
"""Pure ASGI middleware for request logging.
|
|
|
|
Unlike Starlette's BaseHTTPMiddleware, this wraps the ``send``
|
|
callable directly -- no anyio memory channel, no extra coroutine
|
|
context switches per response body chunk. This matters for SSE
|
|
streaming where hundreds of small chunks are sent per second.
|
|
"""
|
|
|
|
def __init__(self, app):
|
|
self.app = app
|
|
|
|
async def __call__(self, scope, receive, send):
|
|
if scope["type"] != "http":
|
|
await self.app(scope, receive, send)
|
|
return
|
|
|
|
start_time = time.time()
|
|
status_code = None
|
|
|
|
async def send_wrapper(message):
|
|
nonlocal status_code
|
|
if message["type"] == "http.response.start":
|
|
status_code = message.get("status", 0)
|
|
await send(message)
|
|
|
|
try:
|
|
await self.app(scope, receive, send_wrapper)
|
|
except Exception as e:
|
|
path = scope.get("path", "")
|
|
method = scope.get("method", "")
|
|
logger.error(
|
|
"request_failed",
|
|
path = path,
|
|
method = method,
|
|
error = str(e),
|
|
exc_info = True,
|
|
)
|
|
raise
|
|
|
|
# Log after response completes (same exclusion logic as before)
|
|
process_time = (time.time() - start_time) * 1000
|
|
path = scope.get("path", "")
|
|
is_excluded = (
|
|
path in _EXCLUDED_PATHS
|
|
or path.startswith("/assets/")
|
|
or path.endswith(_EXCLUDED_SUFFIXES)
|
|
)
|
|
if not is_excluded:
|
|
method = scope.get("method", "")
|
|
logger.info(
|
|
"request_completed",
|
|
method = method,
|
|
path = path,
|
|
status_code = status_code,
|
|
process_time_ms = round(process_time, 2),
|
|
)
|
|
|
|
|
|
def filter_sensitive_data(logger, method_name, event_dict):
|
|
"""Structlog processor to filter out base64 data from logs."""
|
|
|
|
def filter_value(value):
|
|
if (
|
|
isinstance(value, str)
|
|
and len(value) > 100
|
|
and ("," in value or "/" in value)
|
|
):
|
|
# Likely base64 data, truncate it
|
|
return value[:20] + "..."
|
|
elif isinstance(value, dict):
|
|
return {k: filter_value(v) for k, v in value.items()}
|
|
elif isinstance(value, list):
|
|
return [filter_value(item) for item in value]
|
|
return value
|
|
|
|
return {k: filter_value(v) for k, v in event_dict.items()}
|
|
|
|
|
|
def get_logger(name: str) -> structlog.BoundLogger:
|
|
"""Get a logger instance for a specific module.
|
|
Args:
|
|
name: Usually __name__ of the module
|
|
Returns:
|
|
A bound structured logger
|
|
"""
|
|
return structlog.get_logger(name)
|