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 74

fix(proxy): requeue daily spend rows when the commit fails without the Redis buffer

Browse source 
With the Redis transaction buffer off, each daily spend queue (user, team, org,
end user, agent) was drained into a dict and handed to the bulk upsert. When
the upsert raised after its retries, the drained dict was discarded and the
exception escaped update_spend, so those rows never reached the daily rollup
tables and /user/daily/activity stayed short forever while /spend/logs had
every request.

Each daily queue now flushes through one helper that puts the uncommitted
remainder back on the queue for the next tick and moves on to the next table,
the same shape the window-spend step already used.
This commit is contained in:
mateo-berri 2026-09-18 13:50:56 -07:00
parent 88e150bb59
commit 72a14246a6
2 changed files with 118 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

99
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 ReadOnly, TypedDict
@ -143,6 +143,20 @@ 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: ...
def _timed_request_duration_ms(
payload: dict | SpendLogsPayload,
request_status: Literal["success", "failure"],
@ -1288,6 +1302,34 @@ class DBSpendUpdateWriter:
cronjob_id=DB_SPEND_UPDATE_JOB_NAME,
)
async def _flush_daily_spend_queue(
self,
queue: DailySpendUpdateQueue,
entity_type: Literal["user", "team", "org", "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 # the uncommitted rows go back on the queue; the other tables must still flush
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,
@ -1316,74 +1358,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 ##################

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

@ -2809,6 +2809,61 @@ 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) -> None:
super().__init__()
self.failing_table = failing_table
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 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]]
@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_window_spend_commit_from_redis_is_restored_to_redis():
"""The Redis drain is destructive, so a failed window commit has to push