From 92e6ea4350d5647610d2d48200c32834b0c18089 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 12 Aug 2026 10:55:05 -0700 Subject: [PATCH] fix(proxy): retry spend-log writes on any error before dropping the batch update_spend_logs pops the batch off spend_log_transactions before writing and deliberately never puts it back, so an attempt not made here is billing data gone for good. The inner retry only caught DB_CONNECTION_ERROR_TYPES, which is httpx transport errors, while the failures concurrency actually produces (deadlock detected, serialization failure, statement and lock timeouts, pool exhaustion) arrive as prisma errors and got zero attempts before the whole batch, up to MAX_LOGS_PER_INTERVAL rows, was discarded Retrying any error is safe here: the write is idempotent (create_many with skip_duplicates), and a row Postgres rejects on its data never reaches this handler because _create_spend_logs_with_poison_isolation drops it and returns, so a poisoned batch cannot spin. Retries stay bounded by n_retry_times and the batch is still surfaced once they are spent --- basedpyright-code-budget.json | 2 +- litellm/proxy/utils.py | 22 +++-- .../test_proxy_update_spend.py | 84 +++++++++++++++++++ 3 files changed, 100 insertions(+), 8 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 6e3cbdff9d0..1931a4fb4c7 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -117,7 +117,7 @@ "limit": 118 }, "reportUnnecessaryComparison": { - "limit": 699 + "limit": 698 }, "reportUnnecessaryContains": { "limit": 5 diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index cce8379ab25..e8d0add8609 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -25,7 +25,6 @@ from litellm.constants import ( SPEND_LOG_WRITE_BATCH_MAX_BYTES, ) from litellm.proxy._types import ( - DB_CONNECTION_ERROR_TYPES, DB_RETRY_SAFE_ERROR_TYPES, CommonProxyErrors, ProxyErrorTypes, @@ -5561,18 +5560,27 @@ class ProxyUpdateSpend: "%s logs processed. Remaining in queue: %s", len(logs_to_process), remaining_count ) break - except DB_CONNECTION_ERROR_TYPES as e: - if i is None: - i = 0 + except Exception as e: + # The batch was popped off the queue before this write, and the + # handler below deliberately does not put it back, so whatever is + # not retried here is billing data lost for good. Retry on any + # error rather than only DB_CONNECTION_ERROR_TYPES (transport): + # the errors concurrency actually produces - deadlock detected, + # serialization failure, statement/lock timeout, pool exhaustion - + # reach us as prisma errors, and got zero attempts. Retrying them + # is safe because the write is idempotent (create_many with + # skip_duplicates), and a row Postgres rejects on its data never + # gets here: _create_spend_logs_with_poison_isolation drops it and + # returns, so a poisoned batch cannot spin. + if i >= n_retry_times: + raise verbose_proxy_logger.warning( - "Spend tracking - DB connection error writing spend logs, retry %d/%d. logs_count=%d, error=%s", + "Spend tracking - error writing spend logs, retry %d/%d. logs_count=%d, error=%s", i + 1, n_retry_times, len(logs_to_process), str(e), ) - if i >= n_retry_times: - raise await asyncio.sleep(2**i) except Exception as e: # Logs already removed from queue at start - don't put them back diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py index 9d6f53841ed..76a82bd7a6f 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py @@ -290,6 +290,90 @@ async def test_update_spend_logs_failure_raises_after_retries( ) +@pytest.mark.asyncio +async def test_update_spend_logs_retries_a_postgres_level_error_and_persists_the_batch( + mock_prisma_client: Any, + make_spend_log_row: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A deadlock / serialization failure must not cost the batch. + + The batch is popped off ``spend_log_transactions`` before the write and the + failure handler deliberately does not put it back, so an attempt not made + here is billing data lost for good. Only transport errors + (``DB_CONNECTION_ERROR_TYPES``) used to be retried, while the errors + concurrency actually produces - deadlock detected, serialization failure, + lock timeout, pool exhaustion - arrive as prisma errors and got zero + attempts. Retrying is safe because the write is idempotent. + """ + from prisma.errors import TransactionError + + async def _fake_sleep(_: float) -> None: + return None + + monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep) + + attempts: List[List[str]] = [] + written: List[str] = [] + + async def _create_many(*, data: Any, skip_duplicates: bool) -> None: + ids = [row["request_id"] for row in data] + attempts.append(ids) + if len(attempts) == 1: + raise TransactionError("deadlock detected") + written.extend(ids) + + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock(side_effect=_create_many) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=3, + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + logs_to_process=[make_spend_log_row(request_id="r1"), make_spend_log_row(request_id="r2")], + ) + + assert written == ["r1", "r2"], f"the batch was dropped instead of retried; attempts={attempts}" + assert len(attempts) == 2 + + +@pytest.mark.asyncio +async def test_update_spend_logs_stops_retrying_a_postgres_level_error_at_the_limit( + mock_prisma_client: Any, + make_spend_log_row: Any, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Retries stay bounded by ``n_retry_times``: a persistently failing write must + still surface rather than spin the flush job forever.""" + from prisma.errors import TransactionError + + async def _fake_sleep(_: float) -> None: + return None + + monkeypatch.setattr(utils_mod.asyncio, "sleep", _fake_sleep) + + mock_prisma_client.db.litellm_spendlogs.create_many = AsyncMock( + side_effect=TransactionError("deadlock detected") + ) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + with pytest.raises(TransactionError): + await ProxyUpdateSpend.update_spend_logs( + n_retry_times=2, + prisma_client=mock_prisma_client, + db_writer_client=None, + proxy_logging_obj=proxy_logging, + logs_to_process=[make_spend_log_row(request_id="r1")], + ) + + assert mock_prisma_client.db.litellm_spendlogs.create_many.await_count == 3, ( + "expected the initial write plus n_retry_times retries" + ) + + def _data_error(message: str) -> Any: from prisma.errors import DataError