diff --git a/litellm/integrations/clickhouse/clickhouse_batch_logger.py b/litellm/integrations/clickhouse/clickhouse_batch_logger.py new file mode 100644 index 00000000000..3dd3dbf2962 --- /dev/null +++ b/litellm/integrations/clickhouse/clickhouse_batch_logger.py @@ -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 diff --git a/litellm/integrations/clickhouse/clickhouse_writer.py b/litellm/integrations/clickhouse/clickhouse_writer.py deleted file mode 100644 index 735cc0dc6af..00000000000 --- a/litellm/integrations/clickhouse/clickhouse_writer.py +++ /dev/null @@ -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, - ) diff --git a/tests/test_litellm/integrations/clickhouse/test_clickhouse_writer.py b/tests/test_litellm/integrations/clickhouse/test_clickhouse_writer.py deleted file mode 100644 index 3690b30c3c6..00000000000 --- a/tests/test_litellm/integrations/clickhouse/test_clickhouse_writer.py +++ /dev/null @@ -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