scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 73

Merge pull request #41878 from BerriAI/litellm_requeue_daily_spend_without_redis_buffer

Browse source 
fix(proxy): requeue daily spend rows when the commit fails without the Redis buffer
This commit is contained in:
Mateo Wang 2026-09-18 17:27:36 -07:00 • committed by GitHub
commit f6d4766ebe
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 325 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -15,7 +15,7 @@ import traceback
from collections.abc import Mapping, Sequence
from datetime import datetime, timedelta, timezone
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, cast, overload
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, overload
from urllib.parse import quote, unquote
from typing_extensions import LiteralString, ReadOnly, TypedDict
@ -31,6 +31,7 @@ from litellm.constants import (
from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.proxy._types import (
DB_CONNECTION_ERROR_TYPES,
DB_RETRY_SAFE_ERROR_TYPES,
BaseDailySpendTransaction,
DailyAgentSpendTransaction,
@ -65,6 +66,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
WindowSpendTransaction,
WindowSpendUpdateQueue,
)
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING
from litellm.proxy.spend_tracking.compression_savings import (
extract_compression_saved_tokens,
@ -147,6 +149,30 @@ class _SpendTransactionManager(Protocol):
async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> bool | None: ...
_DailySpendTransactionT = TypeVar("_DailySpendTransactionT", bound=BaseDailySpendTransaction)
class _DailySpendCommit(Protocol[_DailySpendTransactionT]):
async def __call__(
self,
*,
n_retry_times: int,
prisma_client: PrismaClient,
proxy_logging_obj: ProxyLogging,
daily_spend_transactions: dict[str, _DailySpendTransactionT],
) -> None: ...
_DATA_REJECTED_SQLSTATE_CLASSES: Final = frozenset({"22", "23"})
def _daily_spend_commit_failure_is_requeue_safe(e: Exception) -> bool:
if isinstance(e, DB_CONNECTION_ERROR_TYPES):
return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES)
sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e)
return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES
def _timed_request_duration_ms(
payload: dict | SpendLogsPayload,
request_status: Literal["success", "failure"],
@ -1372,6 +1398,36 @@ class DBSpendUpdateWriter:
cronjob_id=DB_SPEND_UPDATE_JOB_NAME,
)
async def _flush_daily_spend_queue(
self,
queue: DailySpendUpdateQueue,
entity_type: Literal["user", "team", "org", "tag", "end_user", "agent"],
commit: _DailySpendCommit[_DailySpendTransactionT],
n_retry_times: int,
prisma_client: PrismaClient,
proxy_logging_obj: ProxyLogging,
) -> None:
transactions: Final = await queue.flush_and_get_aggregated_daily_spend_update_transactions()
try:
await commit(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=cast(dict[str, _DailySpendTransactionT], transactions),
)
except Exception as e: # noqa: BLE001 # whatever failed here, the other tables must still flush
if not transactions:
return
spend_log_error(
"Spend tracking - failed to commit daily %s spend updates. "
"Re-queued %d rows for retry on next tick. Error: %s",
entity_type,
len(transactions),
str(e),
exc=e,
)
await queue.add_update(transactions)
async def _commit_spend_updates_to_db_without_redis_buffer(
self,
prisma_client: PrismaClient,
@ -1400,74 +1456,59 @@ class DBSpendUpdateWriter:
################## Daily Spend Update Transactions ##################
# Aggregate all in memory daily spend transactions and commit to db
daily_spend_update_transactions: Final = cast(
dict[str, DailyUserSpendTransaction],
await self.daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(),
)
await DBSpendUpdateWriter.update_daily_user_spend(
await self._flush_daily_spend_queue(
queue=self.daily_spend_update_queue,
entity_type="user",
commit=DBSpendUpdateWriter.update_daily_user_spend,
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_spend_update_transactions,
)
################## Daily Team Spend Update Transactions ##################
# Aggregate all in memory daily team spend transactions and commit to db
daily_team_spend_update_transactions: Final = cast(
dict[str, DailyTeamSpendTransaction],
await self.daily_team_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(),
)
await DBSpendUpdateWriter.update_daily_team_spend(
await self._flush_daily_spend_queue(
queue=self.daily_team_spend_update_queue,
entity_type="team",
commit=DBSpendUpdateWriter.update_daily_team_spend,
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_team_spend_update_transactions,
)
################## Daily Organization Spend Update Transactions ##################
# Aggregate all in memory daily org spend transactions and commit to db
daily_org_spend_update_transactions: Final = cast(
dict[str, DailyOrganizationSpendTransaction],
await self.daily_org_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(),
)
await DBSpendUpdateWriter.update_daily_org_spend(
await self._flush_daily_spend_queue(
queue=self.daily_org_spend_update_queue,
entity_type="org",
commit=DBSpendUpdateWriter.update_daily_org_spend,
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_org_spend_update_transactions,
)
# NOTE: Daily tag spend is committed by a separate scheduler job.
################## Daily End-User Spend Update Transactions ##################
# Aggregate all in memory daily end-user spend transactions and commit to db
daily_end_user_spend_update_transactions: Final = cast(
dict[str, DailyEndUserSpendTransaction],
await self.daily_end_user_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(),
)
await DBSpendUpdateWriter.update_daily_end_user_spend(
await self._flush_daily_spend_queue(
queue=self.daily_end_user_spend_update_queue,
entity_type="end_user",
commit=DBSpendUpdateWriter.update_daily_end_user_spend,
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_end_user_spend_update_transactions,
)
################## Daily Agent Spend Update Transactions ##################
# Aggregate all in memory daily agent spend transactions and commit to db
daily_agent_spend_update_transactions: Final = cast(
dict[str, DailyAgentSpendTransaction],
await self.daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(),
)
await DBSpendUpdateWriter.update_daily_agent_spend(
await self._flush_daily_spend_queue(
queue=self.daily_agent_spend_update_queue,
entity_type="agent",
commit=DBSpendUpdateWriter.update_daily_agent_spend,
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_agent_spend_update_transactions,
)
################## Budget Window Spend Update Transactions ##################
@ -1504,19 +1545,15 @@ class DBSpendUpdateWriter:
Commit only tag spend updates to database.
This is called by a separate scheduler job at a longer interval.
"""
daily_tag_spend_update_transactions: Final = cast(
dict[str, DailyTagSpendTransaction],
await self.daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(),
await self._flush_daily_spend_queue(
queue=self.daily_tag_spend_update_queue,
entity_type="tag",
commit=DBSpendUpdateWriter.update_daily_tag_spend,
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
)
if daily_tag_spend_update_transactions:
await DBSpendUpdateWriter.update_daily_tag_spend(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_tag_spend_update_transactions,
)
async def _commit_daily_tag_spend_to_db_with_redis(
self,
prisma_client: PrismaClient,
@ -2103,13 +2140,25 @@ class DBSpendUpdateWriter:
sql, params = build_bulk_upsert(table=table, batch=merged_batch)
await prisma_client.db.execute_raw(sql, *params)
except Exception as batch_error:
# Log detailed error information for debugging batch upsert failures
# This helps diagnose issues like unique constraint violations
if _daily_spend_commit_failure_is_requeue_safe(batch_error):
spend_log_error(
"Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s",
entity_type,
table.name,
len(transactions_to_process),
str(batch_error),
exc=batch_error,
)
raise
for key in transactions_to_process:
daily_spend_transactions.pop(key, None)
spend_log_error(
"Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s",
"Spend tracking - dropped %d daily %s spend rows: the failed statement may have "
"applied or the database refused the data, so re-sending it is not safe. "
"Table: %s, Error: %s",
len(transactions_to_process),
entity_type,
table.name,
len(transactions_to_process),
str(batch_error),
exc=batch_error,
)

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

@ -1,6 +1,8 @@
from collections.abc import Awaitable, Callable, Iterator
from typing import Any, Final, TypeVar
from pydantic import TypeAdapter, ValidationError
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import (
DB_CONNECTION_ERROR_TYPES,
@ -17,6 +19,8 @@ _TRANSIENT_DB_UNAVAILABLE_MESSAGE: Final = (
"Service Unavailable, the authentication database is temporarily unreachable. Please retry shortly."
)
_DATABASE_ERROR_META: Final = TypeAdapter(dict[str, object])
def _exception_chain(e: BaseException) -> Iterator[BaseException]:
current = e # rebind-ok: advances one link per iteration of the bounded walk
@ -221,6 +225,20 @@ class PrismaDBExceptionHandler:
or "write conflict or a deadlock" in error_message
)
@staticmethod
def postgres_sqlstate(e: Exception) -> str | None:
"""The SQLSTATE Postgres attached to a failed statement, as prisma surfaces it, or None."""
import prisma
if not isinstance(e, _exception_types(prisma.errors.DataError)):
return None
try:
meta: Final = _DATABASE_ERROR_META.validate_python(getattr(e, "meta", None))
except ValidationError:
return None
code: Final = meta.get("code")
return code if isinstance(code, str) else None
@staticmethod
def is_read_only_transaction_error(e: Exception) -> bool:
"""True iff ``e`` is Postgres SQLSTATE 25006 surfaced through prisma: the

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

@ -12,7 +12,9 @@ from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, call, patch
import httpx
import pytest
from prisma.errors import RawQueryError
from redis.exceptions import DataError
import litellm
@ -1776,6 +1778,33 @@ async def test_update_daily_spend_keeps_failed_transactions_for_retry():
assert daily_spend_transactions == expected
@pytest.mark.asyncio
async def test_update_daily_spend_drops_the_batch_whose_failure_cannot_be_resent():
"""A reply lost after the statement was sent may already have applied, so the batch is
taken out of the caller's dict before the error propagates: whichever requeue the caller
runs afterwards, the Redis restore included, cannot send it a second time."""
def lose_the_reply() -> int:
raise httpx.ReadTimeout("no reply")
prisma_client = _RecordingPrisma(execute_raw=lose_the_reply)
daily_spend_transactions = {"user-key": _daily_txn(user_id="user-1")}
proxy_logging_obj = MagicMock()
proxy_logging_obj.failure_handler = AsyncMock()
with pytest.raises(httpx.ReadTimeout):
await DBSpendUpdateWriter._update_daily_spend(
n_retry_times=0,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_spend_transactions,
entity_type="user",
entity_id_field="user_id",
)
assert daily_spend_transactions == {}
@pytest.mark.asyncio
async def test_commit_key_spend_updates_includes_last_active():
"""
@ -2952,6 +2981,162 @@ async def test_failed_window_spend_commit_requeues_the_increments_and_continues_
assert requeued == (transaction,)
class _DailySpendFakeDB(_WindowSpendFakeDB):
"""Records the daily rollup upserts it is handed and fails the ones aimed at one table."""
def __init__(self, failing_table: str | None, failure: Exception | None = None) -> None:
super().__init__()
self.failing_table = failing_table
self.failure = failure
self.execute_raw_calls: list[Statement] = []
async def execute_raw(self, query: str, *args: object) -> int:
if self.failing_table is not None and self.failing_table in query:
raise self.failure if self.failure is not None else Exception("connection reset")
self.execute_raw_calls.append((query, args))
return len(args)
def _daily_upserts(db: _DailySpendFakeDB, table: str) -> list[Statement]:
return [statement for statement in db.execute_raw_calls if table in statement[0]]
def _postgres_rejection(sqlstate: str) -> RawQueryError:
return RawQueryError(
data={"user_facing_error": {"error_code": "P2010", "meta": {"code": sqlstate, "message": "db error"}}}
)
@pytest.mark.parametrize(
("failure", "lands_on_the_next_tick"),
[
pytest.param(httpx.ReadTimeout("no reply"), False, id="reply lost after the statement was sent"),
pytest.param(httpx.ConnectError("refused"), True, id="statement never reached the database"),
pytest.param(_postgres_rejection("22021"), False, id="postgres refused the data itself"),
pytest.param(_postgres_rejection("23502"), False, id="postgres refused a constraint violation"),
pytest.param(_postgres_rejection("42P01"), True, id="table missing"),
pytest.param(_postgres_rejection("57014"), True, id="statement cancelled"),
],
)
@pytest.mark.asyncio
async def test_failed_daily_spend_commit_is_requeued_only_when_the_rows_are_provably_uncommitted(
failure: Exception, lands_on_the_next_tick: bool
):
"""A lost reply means the statement may already have applied, and re-sending it stacks a
second increment into the same transaction (LIT-4823); a row Postgres refuses would fail
every tick forever. Both are dropped loudly. Every other failure left nothing committed,
so its rows go back on the queue and land on the next tick."""
db_writer = DBSpendUpdateWriter()
await db_writer.daily_spend_update_queue.add_update({"user-key": _daily_txn(user_id="user-1")})
db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=failure)
db_writer._flush_tool_discovery_queue = AsyncMock()
proxy_logging_obj = MagicMock()
proxy_logging_obj.failure_handler = AsyncMock()
await db_writer._commit_spend_updates_to_db_without_redis_buffer(
prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj
)
db.failing_table = None
await db_writer._commit_spend_updates_to_db_without_redis_buffer(
prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj
)
assert len(_daily_upserts(db, "LiteLLM_DailyUserSpend")) == (1 if lands_on_the_next_tick else 0)
assert db_writer.daily_spend_update_queue.update_queue.empty()
@pytest.mark.asyncio
async def test_failed_daily_spend_commit_drops_only_the_batch_that_was_sent():
"""A tick holding more than one batch of 100 rows sends them one statement at a time, and
a reply lost on one statement says nothing about the batches after it: only the batch that
was on the wire is dropped, the ones never sent go back on the queue and land next tick."""
db_writer = DBSpendUpdateWriter()
await db_writer.daily_spend_update_queue.add_update(
{f"user-{i:03d}": _daily_txn(user_id=f"user-{i:03d}") for i in range(150)}
)
db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=httpx.ReadTimeout("no reply"))
db_writer._flush_tool_discovery_queue = AsyncMock()
proxy_logging_obj = MagicMock()
proxy_logging_obj.failure_handler = AsyncMock()
await db_writer._commit_spend_updates_to_db_without_redis_buffer(
prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj
)
db.failing_table = None
await db_writer._commit_spend_updates_to_db_without_redis_buffer(
prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj
)
(upsert,) = _daily_upserts(db, "LiteLLM_DailyUserSpend")
assert _row_values(upsert, "user_id") == [f"user-{i:03d}" for i in range(100, 150)]
assert db_writer.daily_spend_update_queue.update_queue.empty()
@pytest.mark.asyncio
async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other_tables():
"""With the Redis buffer off, a daily batch that failed to commit was discarded along
with the tick's exception, so the Usage page stayed short of LiteLLM_SpendLogs for good.
The uncommitted rows must go back on their queue and land on the next tick, and the
other daily tables must still be flushed on the failing tick."""
db_writer = DBSpendUpdateWriter()
await db_writer.daily_spend_update_queue.add_update({"user-key": _daily_txn(user_id="user-1")})
team_txn = {key: value for key, value in _daily_txn().items() if key != "user_id"} | {"team_id": "team-1"}
await db_writer.daily_team_spend_update_queue.add_update({"team-key": team_txn})
db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend")
db_writer._flush_tool_discovery_queue = AsyncMock()
proxy_logging_obj = MagicMock()
proxy_logging_obj.failure_handler = AsyncMock()
await db_writer._commit_spend_updates_to_db_without_redis_buffer(
prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj
)
assert _daily_upserts(db, "LiteLLM_DailyUserSpend") == []
(team_upsert,) = _daily_upserts(db, "LiteLLM_DailyTeamSpend")
assert _row_values(team_upsert, "team_id") == ["team-1"]
db_writer._flush_tool_discovery_queue.assert_called_once()
db.failing_table = None
await db_writer._commit_spend_updates_to_db_without_redis_buffer(
prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj
)
(user_upsert,) = _daily_upserts(db, "LiteLLM_DailyUserSpend")
assert _row_values(user_upsert, "user_id") == ["user-1"]
assert _row_values(user_upsert, "spend") == [0.1]
assert len(_daily_upserts(db, "LiteLLM_DailyTeamSpend")) == 1
assert db_writer.daily_spend_update_queue.update_queue.empty()
@pytest.mark.asyncio
async def test_failed_daily_tag_spend_commit_requeues_the_rows():
"""The tag rollup drains on its own scheduler job with the same no-Redis drop:
a failed LiteLLM_DailyTagSpend commit has to put the rows back for the next tick."""
db_writer = DBSpendUpdateWriter()
tag_txn = {key: value for key, value in _daily_txn().items() if key != "user_id"} | {"tag": "tag-1"}
await db_writer.daily_tag_spend_update_queue.add_update({"tag-key": tag_txn})
db = _DailySpendFakeDB(failing_table="LiteLLM_DailyTagSpend")
proxy_logging_obj = MagicMock()
proxy_logging_obj.failure_handler = AsyncMock()
await db_writer._commit_daily_tag_spend_to_db(
prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj
)
assert _daily_upserts(db, "LiteLLM_DailyTagSpend") == []
assert not db_writer.daily_tag_spend_update_queue.update_queue.empty()
db.failing_table = None
await db_writer._commit_daily_tag_spend_to_db(
prisma_client=_WindowSpendFakePrisma(db), n_retry_times=0, proxy_logging_obj=proxy_logging_obj
)
(tag_upsert,) = _daily_upserts(db, "LiteLLM_DailyTagSpend")
assert _row_values(tag_upsert, "tag") == ["tag-1"]
assert _row_values(tag_upsert, "spend") == [0.1]
assert db_writer.daily_tag_spend_update_queue.update_queue.empty()
@pytest.mark.asyncio
async def test_failed_window_spend_commit_from_redis_is_restored_to_redis():
"""The Redis drain is destructive, so a failed window commit has to push

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

@ -665,6 +665,28 @@ def test_is_deadlock_error_excludes_non_deadlocks(error):
assert PrismaDBExceptionHandler.is_deadlock_error(error) is False
@pytest.mark.parametrize(
("error", "sqlstate"),
[
(
RawQueryError(
data={"user_facing_error": {"error_code": "P2010", "meta": {"code": "22021", "message": "m"}}}
),
"22021",
),
(RawQueryError(data={"user_facing_error": {"error_code": "P2010", "meta": {"message": "m"}}}), None),
(RawQueryError(data={"user_facing_error": {"error_code": "P2010", "meta": {"code": 42, "message": "m"}}}), None),
(prisma_errors.DataError(data={"user_facing_error": {"meta": None}}), None),
(PrismaError("db error"), None),
(httpx.ReadTimeout("no reply"), None),
],
)
def test_postgres_sqlstate_reads_the_code_prisma_attached_to_the_failed_statement(error: Exception, sqlstate: str | None):
"""Only a prisma data error carrying Postgres's own error code yields a SQLSTATE; a
codeless or malformed payload, an engine-level error, and a transport error yield None."""
assert PrismaDBExceptionHandler.postgres_sqlstate(error) == sqlstate
READ_ONLY_CONNECTOR_ERROR: Final = (
"Error occurred during query execution:\nConnectorError(ConnectorError { user_facing_error: None, "
'kind: QueryError(PostgresError { code: "25006", message: "cannot execute UPDATE in a read-only transaction", '