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