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