261 lines
9.0 KiB
Python
261 lines
9.0 KiB
Python
"""Concise model connection and latency diagnostics for production logs."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import atexit
|
|
import logging
|
|
import queue
|
|
import threading
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
from urllib.parse import urlsplit
|
|
from uuid import UUID
|
|
|
|
from langchain_core.callbacks import BaseCallbackHandler
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class _TimingLogDispatcher:
|
|
"""Keep slow log handlers off both the event loop and its tool executor.
|
|
|
|
Only already-sanitized, small lifecycle records enter this bounded queue.
|
|
A daemon thread invokes the existing handlers; producers never wait for
|
|
their I/O. Full queues drop diagnostics and report the count when the sink
|
|
recovers, instead of retaining unbounded records or stalling generation.
|
|
"""
|
|
|
|
def __init__(self, *, capacity: int = 4096) -> None:
|
|
self._queue: queue.Queue[logging.LogRecord] = queue.Queue(maxsize=capacity)
|
|
self._lock = threading.Lock()
|
|
self._thread: threading.Thread | None = None
|
|
self._closed = False
|
|
self._dropped = 0
|
|
|
|
def emit(self, level: int, message: str, *args: Any) -> None:
|
|
if not logger.isEnabledFor(level):
|
|
return
|
|
# Preserve the event's wall-clock timestamp even if delivery is delayed.
|
|
# Formatting, filters and handler I/O run only on the consumer thread.
|
|
record = logger.makeRecord(logger.name, level, __file__, 0, message, args, None)
|
|
with self._lock:
|
|
if self._closed:
|
|
return
|
|
try:
|
|
self._queue.put_nowait(record)
|
|
except queue.Full:
|
|
self._dropped += 1
|
|
return
|
|
if self._thread is None:
|
|
self._thread = threading.Thread(target=self._run, name="deerflow-model-timing-log", daemon=True)
|
|
self._thread.start()
|
|
|
|
def _run(self) -> None:
|
|
while True:
|
|
try:
|
|
record = self._queue.get(timeout=0.1)
|
|
except queue.Empty:
|
|
with self._lock:
|
|
if self._closed:
|
|
return
|
|
continue
|
|
try:
|
|
logger.handle(record)
|
|
except Exception:
|
|
# Diagnostic sinks must not fail a model call or kill the worker.
|
|
pass
|
|
finally:
|
|
self._queue.task_done()
|
|
with self._lock:
|
|
dropped, self._dropped = self._dropped, 0
|
|
if dropped:
|
|
try:
|
|
logger.warning("LLM_TIMING_LOG_DROPPED count=%d reason=queue_full", dropped)
|
|
except Exception:
|
|
pass
|
|
|
|
def close(self, timeout: float = 0.5) -> None:
|
|
"""Best-effort drain without making process exit wait on a log sink."""
|
|
with self._lock:
|
|
self._closed = True
|
|
thread = self._thread
|
|
if thread is not None and thread is not threading.current_thread():
|
|
thread.join(timeout=timeout)
|
|
|
|
|
|
_timing_logs = _TimingLogDispatcher()
|
|
atexit.register(_timing_logs.close)
|
|
|
|
|
|
def _safe_base_url(value: str | None) -> str:
|
|
"""Log only the upstream origin, never credentials, query, or URL path."""
|
|
if not value:
|
|
return "-"
|
|
try:
|
|
parsed = urlsplit(value)
|
|
except (TypeError, ValueError):
|
|
return "-"
|
|
if parsed.scheme not in {"http", "https"} or not parsed.hostname:
|
|
return "-"
|
|
host = parsed.hostname
|
|
try:
|
|
port = parsed.port
|
|
except ValueError:
|
|
return "-"
|
|
return f"{parsed.scheme}://{host}:{port}" if port is not None else f"{parsed.scheme}://{host}"
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class _CallTiming:
|
|
started: float
|
|
first_chunk: float | None = None
|
|
|
|
|
|
class ModelTimingCallback(BaseCallbackHandler):
|
|
"""Optionally record timing inline and deliver logs on a dedicated queue."""
|
|
|
|
# LangChain otherwise sends each synchronous callback, including every
|
|
# token's no-op after the first, through the shared tool thread pool.
|
|
# When enabled, inline work is limited to timestamps, short locked state
|
|
# changes and a non-blocking enqueue, never synchronous handler I/O.
|
|
run_inline = False
|
|
|
|
def __init__(self, *, model_name: str, base_url: str | None) -> None:
|
|
from deerflow.config.runtime_performance_config import get_runtime_performance_config
|
|
|
|
self.run_inline = get_runtime_performance_config().model_timing_nonblocking_enabled
|
|
self._model_name = model_name
|
|
self._base_url = _safe_base_url(base_url)
|
|
self._calls: dict[str, _CallTiming] = {}
|
|
self._lock = threading.Lock()
|
|
|
|
def _log(self, level: int, message: str, *args: Any) -> None:
|
|
if self.run_inline:
|
|
_timing_logs.emit(level, message, *args)
|
|
else:
|
|
# The default-off path retains LangChain's prior callback dispatch
|
|
# and synchronous logging behavior for deployment compatibility.
|
|
logger.log(level, message, *args)
|
|
|
|
def on_chat_model_start(
|
|
self,
|
|
serialized: dict[str, Any],
|
|
messages: list[list[Any]],
|
|
*,
|
|
run_id: UUID,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
rid = str(run_id)
|
|
now = time.monotonic()
|
|
with self._lock:
|
|
self._calls[rid] = _CallTiming(started=now)
|
|
self._log(
|
|
logging.INFO,
|
|
"LLM_CONNECT_START model=%s base_url=%s run_id=%s message_batches=%d message_count=%d",
|
|
self._model_name,
|
|
self._base_url,
|
|
rid,
|
|
len(messages),
|
|
sum(len(batch) for batch in messages),
|
|
)
|
|
|
|
def on_llm_new_token(
|
|
self,
|
|
token: str,
|
|
*,
|
|
run_id: UUID,
|
|
chunk: Any | None = None,
|
|
**kwargs: Any,
|
|
) -> None:
|
|
"""Treat the first callback as connection success, including reasoning-only chunks."""
|
|
rid = str(run_id)
|
|
now = time.monotonic()
|
|
with self._lock:
|
|
timing = self._calls.get(rid)
|
|
if timing is None or timing.first_chunk is not None:
|
|
return
|
|
timing.first_chunk = now
|
|
first_chunk_ms = int((now - timing.started) * 1000)
|
|
self._log(
|
|
logging.INFO,
|
|
"LLM_CONNECT_OK model=%s base_url=%s run_id=%s first_chunk_ms=%d",
|
|
self._model_name,
|
|
self._base_url,
|
|
rid,
|
|
first_chunk_ms,
|
|
)
|
|
|
|
def on_llm_end(self, response: Any, *, run_id: UUID, **kwargs: Any) -> None:
|
|
self._finish(run_id=run_id, error=None)
|
|
|
|
def on_llm_error(self, error: BaseException, *, run_id: UUID, **kwargs: Any) -> None:
|
|
self._finish(run_id=run_id, error=error)
|
|
|
|
def _finish(self, *, run_id: UUID, error: BaseException | None) -> None:
|
|
rid = str(run_id)
|
|
now = time.monotonic()
|
|
cancelled = isinstance(error, (asyncio.CancelledError, GeneratorExit))
|
|
with self._lock:
|
|
timing = self._calls.pop(rid, None)
|
|
|
|
if timing is None:
|
|
if cancelled:
|
|
self._log(
|
|
logging.INFO,
|
|
"LLM_REQUEST_CANCELLED model=%s base_url=%s run_id=%s status=cancelled error_type=%s",
|
|
self._model_name,
|
|
self._base_url,
|
|
rid,
|
|
type(error).__name__,
|
|
)
|
|
return
|
|
self._log(
|
|
logging.ERROR,
|
|
"LLM_REQUEST_FAILED model=%s base_url=%s run_id=%s error_type=%s",
|
|
self._model_name,
|
|
self._base_url,
|
|
rid,
|
|
type(error).__name__ if error else "UnknownLifecycle",
|
|
)
|
|
return
|
|
|
|
total_ms = int((now - timing.started) * 1000)
|
|
first_chunk_ms = int((timing.first_chunk - timing.started) * 1000) if timing.first_chunk else None
|
|
if error is None:
|
|
self._log(
|
|
logging.INFO,
|
|
"LLM_REQUEST_DONE model=%s base_url=%s run_id=%s status=success first_chunk_ms=%s total_ms=%d",
|
|
self._model_name,
|
|
self._base_url,
|
|
rid,
|
|
first_chunk_ms if first_chunk_ms is not None else "none",
|
|
total_ms,
|
|
)
|
|
return
|
|
|
|
if cancelled:
|
|
self._log(
|
|
logging.INFO,
|
|
"LLM_REQUEST_CANCELLED model=%s base_url=%s run_id=%s status=cancelled first_chunk_ms=%s total_ms=%d error_type=%s",
|
|
self._model_name,
|
|
self._base_url,
|
|
rid,
|
|
first_chunk_ms if first_chunk_ms is not None else "none",
|
|
total_ms,
|
|
type(error).__name__,
|
|
)
|
|
return
|
|
|
|
self._log(
|
|
logging.ERROR,
|
|
"LLM_REQUEST_FAILED model=%s base_url=%s run_id=%s status=error first_chunk_ms=%s total_ms=%d error_type=%s",
|
|
self._model_name,
|
|
self._base_url,
|
|
rid,
|
|
first_chunk_ms if first_chunk_ms is not None else "none",
|
|
total_ms,
|
|
type(error).__name__,
|
|
)
|