123 lines
3.7 KiB
Python
123 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
|
|
|
|
"""Structured logging handlers and middleware.
|
|
|
|
LoggingMiddleware (request/response logging with timing),
|
|
filter_sensitive_data (structlog processor for sanitization), and
|
|
get_logger (factory for structured loggers).
|
|
"""
|
|
|
|
import re
|
|
import time
|
|
|
|
import structlog
|
|
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
|
|
|
from utils.native_path_leases import redact_native_paths
|
|
|
|
logger = structlog.get_logger(__name__)
|
|
_NATIVE_PATH_LEASE_RE = re.compile(
|
|
r"(?i)(\b(?:native_path_lease|nativePathLease)[\"']?\s*[:=]\s*[\"']?)[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+"
|
|
)
|
|
_EXCLUDED_PATHS = {
|
|
"/api/train/status",
|
|
"/api/train/metrics",
|
|
"/api/train/hardware",
|
|
"/api/system",
|
|
}
|
|
_EXCLUDED_SUFFIXES = (
|
|
".png",
|
|
".jpg",
|
|
".jpeg",
|
|
".svg",
|
|
".ico",
|
|
".woff",
|
|
".woff2",
|
|
".ttf",
|
|
)
|
|
|
|
|
|
class LoggingMiddleware:
|
|
"""ASGI request logger that avoids BaseHTTPMiddleware streaming wrappers."""
|
|
|
|
def __init__(self, app: ASGIApp) -> None:
|
|
self.app = app
|
|
|
|
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
|
if scope["type"] != "http":
|
|
await self.app(scope, receive, send)
|
|
return
|
|
|
|
path = scope["path"]
|
|
excluded = (
|
|
path in _EXCLUDED_PATHS
|
|
or path.startswith("/assets/")
|
|
or path.endswith(_EXCLUDED_SUFFIXES)
|
|
)
|
|
start_time = time.perf_counter()
|
|
status_code = 500
|
|
|
|
async def send_wrapper(message: Message) -> None:
|
|
nonlocal status_code
|
|
if message["type"] == "http.response.start":
|
|
status_code = message["status"]
|
|
await send(message)
|
|
|
|
try:
|
|
await self.app(scope, receive, send_wrapper)
|
|
except Exception as exc:
|
|
logger.error(
|
|
"request_failed",
|
|
path = path,
|
|
method = scope["method"],
|
|
status_code = status_code,
|
|
error = str(exc),
|
|
process_time_ms = round((time.perf_counter() - start_time) * 1000, 2),
|
|
exc_info = True,
|
|
)
|
|
raise
|
|
else:
|
|
if not excluded:
|
|
logger.info(
|
|
"request_completed",
|
|
method = scope["method"],
|
|
path = path,
|
|
status_code = status_code,
|
|
process_time_ms = round((time.perf_counter() - start_time) * 1000, 2),
|
|
)
|
|
|
|
|
|
def filter_sensitive_data(logger, method_name, event_dict):
|
|
"""Structlog processor to redact native path leases from logs."""
|
|
|
|
def filter_value(value):
|
|
if isinstance(value, str):
|
|
try:
|
|
value = redact_native_paths(value)
|
|
except Exception:
|
|
pass
|
|
value = _NATIVE_PATH_LEASE_RE.sub(r"\1<redacted native path lease>", value)
|
|
return value
|
|
elif isinstance(value, dict):
|
|
return {
|
|
k: "<redacted native path lease>"
|
|
if str(k).replace("_", "").lower() == "nativepathlease"
|
|
else 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: "<redacted native path lease>"
|
|
if str(k).replace("_", "").lower() == "nativepathlease"
|
|
else filter_value(v)
|
|
for k, v in event_dict.items()
|
|
}
|
|
|
|
|
|
def get_logger(name: str) -> structlog.BoundLogger:
|
|
"""Get a bound structured logger for a module (name is usually __name__)."""
|
|
return structlog.get_logger(name)
|