From 72a14246a6ffef7724492b8c0da2798c30a32148 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 13:50:56 -0700 Subject: [PATCH] fix(proxy): requeue daily spend rows when the commit fails without the Redis buffer 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. --- litellm/proxy/db/db_spend_update_writer.py | 99 ++++++++++++------- .../proxy/db/test_db_spend_update_writer.py | 55 +++++++++++ 2 files changed, 118 insertions(+), 36 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index c13b852484e..a4f1713d61d 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -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 ################## diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index b3f5a60877d..3feb117edb7 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -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