mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
feat(agent-tracing): add ClickHouseBatchLogger built on CustomBatchLogger
This commit is contained in:
Ishaan Jaff
parent
c533ebaeca
commit
a82445aa53
3 changed files with 104 additions and 135 deletions
104
litellm/integrations/clickhouse/clickhouse_batch_logger.py
Normal file
104
litellm/integrations/clickhouse/clickhouse_batch_logger.py
Normal file
|
|
@ -0,0 +1,104 @@
|
|||
"""
|
||||
Shared base for everything LiteLLM writes to ClickHouse.
|
||||
|
||||
Built on `CustomBatchLogger`: rows accumulate in `log_queue` and are flushed as one
|
||||
gzip JSONEachRow insert, either every `CLICKHOUSE_FLUSH_INTERVAL_SECONDS` or as soon as
|
||||
`batch_size` rows are queued. Subclasses only pick the table and build rows:
|
||||
|
||||
- `ClickHouseSpendLogger` -> spend_logs (LiteLLM requests, via the `clickhouse` callback)
|
||||
- `ClickHouseTraceLogger` -> otel_traces (agent spans, via POST /v1/traces)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from typing import Any, ClassVar
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import (
|
||||
CLICKHOUSE_BATCH_SIZE,
|
||||
CLICKHOUSE_FLUSH_INTERVAL_SECONDS,
|
||||
CLICKHOUSE_MAX_BUFFERED_ROWS,
|
||||
CLICKHOUSE_MAX_RETRIES,
|
||||
)
|
||||
from litellm.integrations.clickhouse.clickhouse_client import ClickHouseClient
|
||||
from litellm.integrations.custom_batch_logger import CustomBatchLogger
|
||||
|
||||
|
||||
def clickhouse_client_from_env() -> ClickHouseClient:
|
||||
return ClickHouseClient(
|
||||
url=os.getenv("CLICKHOUSE_URL", ""),
|
||||
user=os.getenv("CLICKHOUSE_USER", "default"),
|
||||
password=os.getenv("CLICKHOUSE_PASSWORD", ""),
|
||||
database=os.getenv("CLICKHOUSE_DATABASE", "litellm"),
|
||||
)
|
||||
|
||||
|
||||
class ClickHouseBatchLogger(CustomBatchLogger):
|
||||
table: ClassVar[str]
|
||||
|
||||
def __init__(self, client: ClickHouseClient | None = None, **kwargs: Any) -> None:
|
||||
self.client = client or clickhouse_client_from_env()
|
||||
self.rows_written = 0
|
||||
self.rows_dropped = 0
|
||||
self._failed_attempts = 0
|
||||
super().__init__(
|
||||
flush_lock=asyncio.Lock(),
|
||||
batch_size=CLICKHOUSE_BATCH_SIZE,
|
||||
flush_interval=CLICKHOUSE_FLUSH_INTERVAL_SECONDS, # type: ignore[arg-type]
|
||||
**kwargs,
|
||||
)
|
||||
try:
|
||||
asyncio.get_running_loop().create_task(self.periodic_flush())
|
||||
except RuntimeError: # no loop yet (e.g. sync config load); proxy startup calls start()
|
||||
pass
|
||||
|
||||
def start(self) -> None:
|
||||
asyncio.get_running_loop().create_task(self.periodic_flush())
|
||||
|
||||
def is_full(self) -> bool:
|
||||
"""Backpressure signal: producers should reject (429) instead of enqueueing."""
|
||||
return len(self.log_queue) >= CLICKHOUSE_MAX_BUFFERED_ROWS
|
||||
|
||||
def enqueue(self, rows: list[dict[str, Any]]) -> None:
|
||||
"""Never awaits ClickHouse. Kicks off an early flush once a full batch is queued."""
|
||||
self.log_queue.extend(rows)
|
||||
if len(self.log_queue) >= self.batch_size:
|
||||
asyncio.get_running_loop().create_task(self.flush_queue())
|
||||
|
||||
async def flush_queue(self) -> None:
|
||||
# Swap the queue under the lock so rows enqueued during the insert are kept.
|
||||
if self.flush_lock is None:
|
||||
return
|
||||
async with self.flush_lock:
|
||||
while self.log_queue:
|
||||
batch = self.log_queue[: self.batch_size]
|
||||
self.log_queue = self.log_queue[len(batch) :]
|
||||
if not await self._insert(batch):
|
||||
break
|
||||
|
||||
async def async_send_batch(self, *args: Any, **kwargs: Any) -> None:
|
||||
await self.flush_queue()
|
||||
|
||||
async def _insert(self, batch: list[dict[str, Any]]) -> bool:
|
||||
try:
|
||||
await self.client.insert_json_each_row(self.table, batch)
|
||||
self.rows_written += len(batch)
|
||||
self._failed_attempts = 0
|
||||
return True
|
||||
except Exception as e:
|
||||
self._failed_attempts += 1
|
||||
if self._failed_attempts >= CLICKHOUSE_MAX_RETRIES:
|
||||
self.rows_dropped += len(batch)
|
||||
self._failed_attempts = 0
|
||||
verbose_logger.error(
|
||||
"ClickHouse: dropped %s rows for %s after %s attempts: %s",
|
||||
len(batch),
|
||||
self.table,
|
||||
CLICKHOUSE_MAX_RETRIES,
|
||||
e,
|
||||
)
|
||||
else:
|
||||
# put it back; the next periodic flush retries it
|
||||
self.log_queue = batch + self.log_queue
|
||||
verbose_logger.warning("ClickHouse: insert into %s failed, will retry: %s", self.table, e)
|
||||
return False
|
||||
|
|
@ -1,74 +0,0 @@
|
|||
"""
|
||||
Bounded, batched writer: many producers enqueue rows, one background task flushes.
|
||||
|
||||
One instance per worker process. Producers never await ClickHouse — they either
|
||||
enqueue or are told the buffer is full (the OTLP endpoint turns that into a 429).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections import defaultdict
|
||||
from typing import Any
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import (
|
||||
CLICKHOUSE_BATCH_SIZE,
|
||||
CLICKHOUSE_FLUSH_INTERVAL_SECONDS,
|
||||
CLICKHOUSE_MAX_BUFFERED_ROWS,
|
||||
CLICKHOUSE_MAX_RETRIES,
|
||||
)
|
||||
from litellm.integrations.clickhouse.clickhouse_client import ClickHouseClient
|
||||
|
||||
|
||||
class ClickHouseWriter:
|
||||
def __init__(self, client: ClickHouseClient):
|
||||
self.client = client
|
||||
self.buffers: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
self.rows_written: dict[str, int] = defaultdict(int)
|
||||
self.rows_dropped: dict[str, int] = defaultdict(int)
|
||||
self._flush_lock = asyncio.Lock()
|
||||
self._task: asyncio.Task | None = None
|
||||
|
||||
def is_full(self, table: str) -> bool:
|
||||
return len(self.buffers[table]) >= CLICKHOUSE_MAX_BUFFERED_ROWS
|
||||
|
||||
def enqueue(self, table: str, rows: list[dict[str, Any]]) -> None:
|
||||
self.buffers[table].extend(rows)
|
||||
|
||||
def start(self) -> None:
|
||||
if self._task is None:
|
||||
self._task = asyncio.create_task(self._run_forever())
|
||||
|
||||
async def _run_forever(self) -> None:
|
||||
while True:
|
||||
await asyncio.sleep(CLICKHOUSE_FLUSH_INTERVAL_SECONDS)
|
||||
await self.flush()
|
||||
|
||||
async def flush(self) -> None:
|
||||
async with self._flush_lock:
|
||||
for table in list(self.buffers):
|
||||
while self.buffers[table]:
|
||||
batch = self.buffers[table][:CLICKHOUSE_BATCH_SIZE]
|
||||
del self.buffers[table][: len(batch)]
|
||||
await self._insert_with_retry(table, batch)
|
||||
|
||||
async def _insert_with_retry(self, table: str, batch: list[dict[str, Any]]) -> None:
|
||||
for attempt in range(CLICKHOUSE_MAX_RETRIES):
|
||||
try:
|
||||
await self.client.insert_json_each_row(table, batch)
|
||||
self.rows_written[table] += len(batch)
|
||||
return
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
"ClickHouseWriter: insert into %s failed (attempt %s): %s",
|
||||
table,
|
||||
attempt + 1,
|
||||
e,
|
||||
)
|
||||
await asyncio.sleep(2**attempt)
|
||||
self.rows_dropped[table] += len(batch)
|
||||
verbose_logger.error(
|
||||
"ClickHouseWriter: dropped %s rows for %s after %s retries",
|
||||
len(batch),
|
||||
table,
|
||||
CLICKHOUSE_MAX_RETRIES,
|
||||
)
|
||||
|
|
@ -1,61 +0,0 @@
|
|||
"""
|
||||
Tests for the batched ClickHouse writer used by agent tracing.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../.."))
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.integrations.clickhouse import clickhouse_writer as writer_module
|
||||
from litellm.integrations.clickhouse.clickhouse_writer import ClickHouseWriter
|
||||
|
||||
|
||||
def _writer(insert: AsyncMock) -> ClickHouseWriter:
|
||||
client = MagicMock()
|
||||
client.insert_json_each_row = insert
|
||||
return ClickHouseWriter(client=client)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flush_splits_into_batches():
|
||||
insert = AsyncMock()
|
||||
writer = _writer(insert)
|
||||
writer.enqueue("otel_traces", [{"i": i} for i in range(5)])
|
||||
|
||||
with patch.object(writer_module, "CLICKHOUSE_BATCH_SIZE", 2):
|
||||
await writer.flush()
|
||||
|
||||
assert [len(c.args[1]) for c in insert.await_args_list] == [2, 2, 1]
|
||||
assert writer.buffers["otel_traces"] == []
|
||||
assert writer.rows_written["otel_traces"] == 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_full_signals_backpressure():
|
||||
writer = _writer(AsyncMock())
|
||||
with patch.object(writer_module, "CLICKHOUSE_MAX_BUFFERED_ROWS", 3):
|
||||
writer.enqueue("otel_traces", [{}, {}])
|
||||
assert writer.is_full("otel_traces") is False
|
||||
writer.enqueue("otel_traces", [{}])
|
||||
assert writer.is_full("otel_traces") is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failed_batch_is_retried_then_dropped():
|
||||
insert = AsyncMock(side_effect=RuntimeError("clickhouse down"))
|
||||
writer = _writer(insert)
|
||||
writer.enqueue("spend_logs", [{"request_id": "a"}, {"request_id": "b"}])
|
||||
|
||||
with (
|
||||
patch.object(writer_module, "CLICKHOUSE_MAX_RETRIES", 2),
|
||||
patch.object(writer_module.asyncio, "sleep", AsyncMock()),
|
||||
):
|
||||
await writer.flush()
|
||||
|
||||
assert insert.await_count == 2
|
||||
assert writer.rows_dropped["spend_logs"] == 2
|
||||
assert writer.rows_written["spend_logs"] == 0
|
||||
Loading…
Add table
Reference in a new issue