fix(proxy): retry end-user spend updates on Postgres deadlock instead of dropping them

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(

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(

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