scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 114

Merge pull request #34887 from RayJueWang/litellm_fix_spend_deadlock_retry

Browse source 
fix(proxy): retry spend updates on Postgres deadlock instead of dropping them
This commit is contained in:
Mateo Wang 2026-08-19 15:53:45 -07:00 • committed by GitHub
commit da9d406e8d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 387 additions and 91 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

149
litellm/proxy/db/db_spend_update_writer.py
View file

@ -1204,6 +1204,29 @@ class DBSpendUpdateWriter:
except Exception as e:
verbose_proxy_logger.debug("_flush_tool_discovery_queue error (non-blocking): %s", e)
@staticmethod
async def _handle_spend_update_failure(
e: Exception,
attempt: int,
n_retry_times: int,
start_time: float,
proxy_logging_obj: ProxyLogging,
) -> None:
"""Retry a failed spend-update transaction on connection errors or deadlocks, else re-raise."""
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.utils import _raise_failed_update_spend_exception
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(
self,
prisma_client: PrismaClient,
@ -1215,10 +1238,7 @@ class DBSpendUpdateWriter:
Commits all the spend `UPDATE` transactions to the Database
"""
from litellm.proxy.utils import (
ProxyUpdateSpend,
_raise_failed_update_spend_exception,
)
from litellm.proxy.utils import ProxyUpdateSpend
### UPDATE USER TABLE ###
user_list_transactions: Final = db_spend_update_transactions["user_list_transactions"]
@ -1238,18 +1258,13 @@ class DBSpendUpdateWriter:
data={"spend": {"increment": response_cost}},
)
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 self._handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
### UPDATE END-USER TABLE ###
@ -1281,18 +1296,13 @@ class DBSpendUpdateWriter:
},
)
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 self._handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
### UPDATE TEAM TABLE ###
@ -1314,18 +1324,13 @@ class DBSpendUpdateWriter:
data={"spend": {"increment": response_cost}},
)
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 self._handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
### UPDATE TEAM Membership TABLE with spend ###
@ -1361,18 +1366,13 @@ class DBSpendUpdateWriter:
)
# Transaction succeeded, break out of retry loop
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 self._handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
# Invalidate cache for updated team memberships
@ -1403,25 +1403,13 @@ class DBSpendUpdateWriter:
data={"spend": {"increment": response_cost}},
)
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(
# Sleep a random amount to avoid retrying and deadlocking again: when two transactions deadlock they are
# cancelled basically at the same time, so if they wait the same time they will also retry at the same time
# and thus they are more likely to deadlock again.
# Instead, we sleep a random amount so that they retry at slightly different times, lowering the chance of
# repeated deadlocks, and therefore of exceeding the retry limit.
random.uniform(2**i, 2 ** (i + 1))
)
except Exception as e:
_raise_failed_update_spend_exception(
e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
await self._handle_spend_update_failure(
e=e,
attempt=i,
n_retry_times=n_retry_times,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
### UPDATE TAG TABLE ###
@ -1470,8 +1458,6 @@ class DBSpendUpdateWriter:
prisma_client: Prisma client instance
proxy_logging_obj: Proxy logging object
"""
from litellm.proxy.utils import _raise_failed_update_spend_exception
verbose_proxy_logger.debug("%s Spend transactions: %s", entity_name, transactions)
if transactions is not None and len(transactions.keys()) > 0:
for i in range(n_retry_times + 1):
@ -1493,17 +1479,13 @@ class DBSpendUpdateWriter:
data={"spend": {"increment": response_cost}},
)
break
except DB_RETRY_SAFE_ERROR_TYPES as e:
if i >= n_retry_times:
_raise_failed_update_spend_exception(
e=e,
start_time=start_time,
proxy_logging_obj=proxy_logging_obj,
)
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,
)
# fmt: off
@ -1672,7 +1654,16 @@ class DBSpendUpdateWriter:
break
except DB_RETRY_SAFE_ERROR_TYPES as e:
except Exception as e:
from litellm.proxy.db.exception_handler import (
PrismaDBExceptionHandler,
)
is_retryable = isinstance(
e, DB_RETRY_SAFE_ERROR_TYPES
) or PrismaDBExceptionHandler.is_deadlock_error(e)
if not is_retryable:
raise
if i >= n_retry_times:
_raise_failed_update_spend_exception(
e=e,

16
litellm/proxy/db/exception_handler.py
View file

@ -166,6 +166,22 @@ class PrismaDBExceptionHandler:
return True
return False
@staticmethod
def is_deadlock_error(e: Exception) -> bool:
"""True iff ``e`` is a Postgres deadlock (P2034 / 40P01) surfaced through prisma."""
import prisma
if not isinstance(e, prisma.errors.PrismaError):
return False
if getattr(e, "code", None) == "P2034":
return True
error_message = str(e).lower()
return (
"deadlock detected" in error_message
or "40p01" in error_message
or "write conflict or a deadlock" in error_message
)
@staticmethod
def is_prisma_engine_internal_error(e: Exception) -> bool:
"""True iff ``e`` is a non-``PrismaError`` exception raised from inside

16
litellm/proxy/utils.py
View file

@ -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(

184
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -2533,3 +2533,187 @@ async def test_daily_transaction_internal_call_keeps_spend_but_not_request_count
assert internal["autorouter_savings_spend"] == 0.0
assert user_sent["api_requests"] == 1
assert user_sent["successful_requests"] == 1
def _deadlock_error():
from prisma.errors import RawQueryError
return RawQueryError(
data={"user_facing_error": {"error_code": "P2034", "meta": {"table": "LiteLLM_VerificationToken"}}}
)
def _empty_spend_transactions(**overrides):
base = {
"user_list_transactions": {},
"end_user_list_transactions": {},
"key_list_transactions": {},
"team_list_transactions": {},
"team_member_list_transactions": {},
"org_list_transactions": {},
"tag_list_transactions": {},
"agent_list_transactions": {},
}
return {**base, **overrides}
def _good_tx(mock_batcher):
tx = AsyncMock()
tx.__aenter__ = AsyncMock(return_value=tx)
tx.__aexit__ = AsyncMock(return_value=False)
tx.batch_ = MagicMock(
return_value=AsyncMock(
__aenter__=AsyncMock(return_value=mock_batcher),
__aexit__=AsyncMock(return_value=False),
)
)
return tx
def _failing_tx(error):
tx = MagicMock()
tx.__aenter__ = AsyncMock(side_effect=error)
tx.__aexit__ = AsyncMock(return_value=False)
return tx
@pytest.mark.asyncio
async def test_commit_spend_updates_retries_deadlock_then_commits(monkeypatch):
"""Regression: a deadlock on the key-spend UPDATE is retried and commits the increment exactly once."""
slept = []
monkeypatch.setattr(
"litellm.proxy.db.db_spend_update_writer.asyncio.sleep",
AsyncMock(side_effect=lambda s: slept.append(s)),
)
mock_batcher = MagicMock()
mock_prisma_client = MagicMock()
mock_prisma_client.db.tx = MagicMock(side_effect=[_failing_tx(_deadlock_error()), _good_tx(mock_batcher)])
proxy_logging = MagicMock()
proxy_logging.failure_handler = AsyncMock()
await DBSpendUpdateWriter()._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=3,
proxy_logging_obj=proxy_logging,
db_spend_update_transactions=_empty_spend_transactions(key_list_transactions={"sk-abc": 0.5}),
)
assert mock_prisma_client.db.tx.call_count == 2
mock_batcher.litellm_verificationtoken.update_many.assert_called_once()
call_kwargs = mock_batcher.litellm_verificationtoken.update_many.call_args[1]
assert call_kwargs["where"] == {"token": "sk-abc"}
assert call_kwargs["data"]["spend"] == {"increment": 0.5}
assert len(slept) == 1
proxy_logging.failure_handler.assert_not_called()
@pytest.mark.asyncio
async def test_commit_spend_updates_raises_after_exhausting_deadlock_retries(monkeypatch):
"""A deadlock that never clears must surface after the retry budget is spent, not loop or swallow."""
monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", AsyncMock(return_value=None))
mock_prisma_client = MagicMock()
mock_prisma_client.db.tx = MagicMock(side_effect=lambda *a, **k: _failing_tx(_deadlock_error()))
proxy_logging = MagicMock()
proxy_logging.failure_handler = AsyncMock()
from prisma.errors import RawQueryError
with pytest.raises(RawQueryError):
await DBSpendUpdateWriter()._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=2,
proxy_logging_obj=proxy_logging,
db_spend_update_transactions=_empty_spend_transactions(key_list_transactions={"sk-abc": 0.5}),
)
assert mock_prisma_client.db.tx.call_count == 3
@pytest.mark.asyncio
async def test_commit_spend_updates_does_not_retry_non_deadlock_data_error(monkeypatch):
"""A non-retryable data-layer error raises on the first attempt, never retried against the increment."""
monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", AsyncMock(return_value=None))
from prisma.errors import UniqueViolationError
non_deadlock = UniqueViolationError(
data={"user_facing_error": {"error_code": "P2002", "meta": {"table": "LiteLLM_VerificationToken"}}}
)
mock_prisma_client = MagicMock()
mock_prisma_client.db.tx = MagicMock(side_effect=lambda *a, **k: _failing_tx(non_deadlock))
proxy_logging = MagicMock()
proxy_logging.failure_handler = AsyncMock()
with pytest.raises(UniqueViolationError):
await DBSpendUpdateWriter()._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=3,
proxy_logging_obj=proxy_logging,
db_spend_update_transactions=_empty_spend_transactions(key_list_transactions={"sk-abc": 0.5}),
)
mock_prisma_client.db.tx.assert_called_once()
@pytest.mark.asyncio
async def test_update_daily_spend_retries_deadlock(monkeypatch):
"""The daily-spend upsert path retries a deadlock on the bulk upsert and then drains successfully."""
mock_prisma_client = MagicMock()
mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[_deadlock_error(), None])
proxy_logging = MagicMock()
proxy_logging.failure_handler = AsyncMock()
monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", AsyncMock(return_value=None))
daily_spend_transactions = {"k1": _daily_txn()}
await DBSpendUpdateWriter._update_daily_spend(
n_retry_times=3,
prisma_client=mock_prisma_client,
proxy_logging_obj=proxy_logging,
daily_spend_transactions=daily_spend_transactions,
entity_type="user",
entity_id_field="user_id",
)
assert mock_prisma_client.db.execute_raw.call_count == 2
assert daily_spend_transactions == {}
proxy_logging.failure_handler.assert_not_called()
@pytest.mark.parametrize(
"transactions_key, sample_key",
[
("user_list_transactions", "user-1"),
("team_list_transactions", "team-1"),
("team_member_list_transactions", "team_id::team-1::user_id::user-1"),
("org_list_transactions", "org-1"),
("tag_list_transactions", "tag-1"),
("agent_list_transactions", "agent-1"),
],
)
@pytest.mark.asyncio
async def test_commit_spend_updates_retries_deadlock_on_every_entity_path(monkeypatch, transactions_key, sample_key):
"""Every per-entity spend path, not just keys, retries a deadlock instead of dropping the increment."""
monkeypatch.setattr("litellm.proxy.db.db_spend_update_writer.asyncio.sleep", AsyncMock(return_value=None))
mock_batcher = MagicMock()
mock_prisma_client = MagicMock()
mock_prisma_client.db.tx = MagicMock(side_effect=[_failing_tx(_deadlock_error()), _good_tx(mock_batcher)])
proxy_logging = MagicMock()
proxy_logging.failure_handler = AsyncMock()
proxy_logging.call_details = {}
await DBSpendUpdateWriter()._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=3,
proxy_logging_obj=proxy_logging,
db_spend_update_transactions=_empty_spend_transactions(**{transactions_key: {sample_key: 0.5}}),
)
assert mock_prisma_client.db.tx.call_count == 2
proxy_logging.failure_handler.assert_not_called()

33
tests/test_litellm/proxy/db/test_exception_handler.py
View file

@ -549,3 +549,36 @@ def test_handle_db_exception_surfaces_a_permanent_fault_even_when_degraded_mode_
with pytest.raises(BinaryNotFoundError):
PrismaDBExceptionHandler.handle_db_exception(BinaryNotFoundError("query engine binary not found"))
@pytest.mark.parametrize(
"error",
[
RawQueryError(data={"user_facing_error": {"error_code": "P2034", "meta": {"table": "t"}}}),
PrismaError("Transaction failed due to a write conflict or a deadlock. Please retry your transaction"),
RawQueryError(data={"user_facing_error": {"message": "deadlock detected", "meta": {"table": "t"}}}),
RawQueryError(
data={"user_facing_error": {"message": "ERROR: 40P01: deadlock detected", "meta": {"table": "t"}}}
),
],
)
def test_is_deadlock_error_matches_postgres_deadlock(error):
"""A Postgres deadlock surfaced through prisma (P2034 or 40P01 / "deadlock detected" text) is recognized."""
assert PrismaDBExceptionHandler.is_deadlock_error(error) is True
@pytest.mark.parametrize(
"error",
[
UniqueViolationError(data={"user_facing_error": {"error_code": "P2002", "meta": {"table": "t"}}}),
RecordNotFoundError(data={"user_facing_error": {"meta": {"table": "t"}}}),
PrismaError("validation failed on query"),
PrismaError("can't reach database server"),
httpx.ConnectError("connection refused"),
RuntimeError("deadlock detected"),
ValueError("40P01"),
],
)
def test_is_deadlock_error_excludes_non_deadlocks(error):
"""Non-deadlock prisma errors, connectivity failures, and non-prisma exceptions are not treated as deadlocks."""
assert PrismaDBExceptionHandler.is_deadlock_error(error) is False

80
tests/test_litellm/proxy/utils/prisma_and_spend/test_proxy_update_spend.py
View file

@ -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