From 708ff0b91090eb7acd02e619bc8b20ebabf3f0f8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:20:59 -0700 Subject: [PATCH] fix(proxy): retry end-user spend updates on Postgres deadlock instead of dropping them --- litellm/proxy/db/db_spend_update_writer.py | 6 ++ litellm/proxy/utils.py | 16 ++-- .../test_proxy_update_spend.py | 80 ++++++++++++++++++- 3 files changed, 90 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 19a82d556d8..65a271d4029 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -1219,6 +1219,12 @@ class DBSpendUpdateWriter: is_retryable = isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) or PrismaDBExceptionHandler.is_deadlock_error(e) if not is_retryable or attempt >= n_retry_times: _raise_failed_update_spend_exception(e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj) + verbose_proxy_logger.warning( + "Retrying spend update after retryable DB error (attempt %s/%s): %s", + attempt + 1, + n_retry_times, + e, + ) await asyncio.sleep(random.uniform(2**attempt, 2 ** (attempt + 1))) async def _commit_spend_updates_to_db( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a743526e975..41187af2bd8 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -30,7 +30,6 @@ from litellm.constants import ( SPEND_LOG_WRITE_BATCH_MAX_BYTES, ) from litellm.proxy._types import ( - DB_RETRY_SAFE_ERROR_TYPES, CommonProxyErrors, ProxyErrorTypes, ProxyException, @@ -5960,15 +5959,14 @@ class ProxyUpdateSpend: ) break - except DB_RETRY_SAFE_ERROR_TYPES as e: - if i >= n_retry_times: # If we've reached the maximum number of retries - _raise_failed_update_spend_exception( - e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj - ) - # Optionally, sleep for a bit before retrying - await asyncio.sleep(2**i) # Exponential backoff except Exception as e: - _raise_failed_update_spend_exception(e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj) + await DBSpendUpdateWriter._handle_spend_update_failure( + e=e, + attempt=i, + n_retry_times=n_retry_times, + start_time=start_time, + proxy_logging_obj=proxy_logging_obj, + ) @staticmethod async def update_spend_logs( 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 dd21bbc9e8a..7057a112c83 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 @@ -83,8 +83,8 @@ async def test_update_end_user_spend_retries_on_connect_error( mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch ) -> None: """``DB_RETRY_SAFE_ERROR_TYPES`` (ConnectError, statements provably never - sent) retries with backoff; once retries are exhausted the original - exception bubbles up via ``_raise_failed_update_spend_exception``. + sent) retries with jittered backoff; once retries are exhausted the + original exception bubbles up via ``_raise_failed_update_spend_exception``. """ import httpx import litellm.proxy.utils as utils_mod @@ -107,7 +107,8 @@ async def test_update_end_user_spend_retries_on_connect_error( proxy_logging_obj=proxy_logging, end_user_list_transactions={"u": 1.0}, ) - assert sleeps == [1.0] + assert len(sleeps) == 1 + assert 1.0 <= sleeps[0] <= 2.0 @pytest.mark.asyncio @@ -149,6 +150,79 @@ async def test_update_end_user_spend_non_connection_error_raises_immediately( ) +def _end_user_deadlock_error() -> Exception: + from prisma.errors import RawQueryError + + return RawQueryError(data={"user_facing_error": {"error_code": "P2034", "meta": {"table": "LiteLLM_EndUserTable"}}}) + + +def _failing_tx(error: Exception) -> Any: + tx = MagicMock() + tx.__aenter__ = AsyncMock(side_effect=error) + tx.__aexit__ = AsyncMock(return_value=False) + return tx + + +@pytest.mark.asyncio +async def test_update_end_user_spend_retries_on_deadlock_then_commits( + mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + """Regression for #27989: a Postgres deadlock (P2034/40P01) on the end-user + spend batch is retried with jittered backoff and the increments land, + instead of raising immediately and dropping the flushed spend.""" + sleeps: list[float] = [] + + async def _fake_sleep(seconds: float) -> None: + sleeps.append(seconds) + + monkeypatch.setattr(asyncio, "sleep", _fake_sleep) + + batcher = MagicMock() + batcher.litellm_endusertable.upsert = MagicMock() + transaction = MagicMock() + transaction.batch_ = lambda: _AsyncCM(batcher) + mock_prisma_client.db.tx = MagicMock(side_effect=[_failing_tx(_end_user_deadlock_error()), _AsyncCM(transaction)]) + + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + await ProxyUpdateSpend.update_end_user_spend( + n_retry_times=3, + prisma_client=mock_prisma_client, + proxy_logging_obj=proxy_logging, + end_user_list_transactions={"end-user-1": 0.25}, + ) + + assert mock_prisma_client.db.tx.call_count == 2 + batcher.litellm_endusertable.upsert.assert_called_once() + assert batcher.litellm_endusertable.upsert.call_args.kwargs["where"] == {"user_id": "end-user-1"} + assert len(sleeps) == 1 + assert 1.0 <= sleeps[0] <= 2.0 + proxy_logging.failure_handler.assert_not_called() + + +@pytest.mark.asyncio +async def test_update_end_user_spend_raises_after_exhausting_deadlock_retries( + mock_prisma_client: Any, monkeypatch: pytest.MonkeyPatch +) -> None: + from prisma.errors import RawQueryError + + monkeypatch.setattr(asyncio, "sleep", AsyncMock(return_value=None)) + mock_prisma_client.db.tx = MagicMock(side_effect=lambda timeout: _failing_tx(_end_user_deadlock_error())) + proxy_logging = MagicMock() + proxy_logging.failure_handler = AsyncMock() + + with pytest.raises(RawQueryError): + await ProxyUpdateSpend.update_end_user_spend( + n_retry_times=2, + prisma_client=mock_prisma_client, + proxy_logging_obj=proxy_logging, + end_user_list_transactions={"end-user-1": 0.25}, + ) + + assert mock_prisma_client.db.tx.call_count == 3 + + @pytest.mark.asyncio async def test_update_spend_logs_writes_batches_via_create_many( mock_prisma_client: Any, make_spend_log_row: Any