From a5f47a271ad982052ac7057a28491a7ea15ad8b8 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 13:55:36 -0700 Subject: [PATCH] fix(proxy): re-queue budget window spend increments when the commit fails Budget enforcement trusts a current LiteLLM_BudgetWindowSpend row without reconciling it against LiteLLM_SpendLogs, so an increment dropped after a failed commit let the entity spend past its window limit after the next counter reseed. Failed increments now go back on the in-memory queue, or back to the Redis buffer, and retry on the next scheduler tick like every other spend category. --- litellm/proxy/db/db_spend_update_writer.py | 44 +++--- .../redis_update_buffer.py | 2 + .../test_redis_update_buffer.py | 140 +++++++++--------- .../proxy/db/test_db_spend_update_writer.py | 65 ++++++-- 4 files changed, 151 insertions(+), 100 deletions(-) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 193f41901cd..1ce7a959fec 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -936,6 +936,7 @@ class DBSpendUpdateWriter: "daily_org_spend_update_transactions": daily_org_spend_update_transactions, "daily_end_user_spend_update_transactions": daily_end_user_spend_update_transactions, "daily_agent_spend_update_transactions": daily_agent_spend_update_transactions, + "window_spend_update_transactions": window_spend_update_transactions, } if db_spend_update_transactions is not None: @@ -1008,6 +1009,7 @@ class DBSpendUpdateWriter: prisma_client=prisma_client, window_spend_transactions=window_spend_update_transactions, ) + uncommitted.pop("window_spend_update_transactions", None) except Exception as e: spend_log_error( "Spend tracking - failed to commit spend updates from Redis to DB. " @@ -1129,10 +1131,20 @@ class DBSpendUpdateWriter: await self.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions() ) - await DBSpendUpdateWriter._commit_window_spend_updates( - prisma_client=prisma_client, - window_spend_transactions=window_spend_update_transactions, - ) + try: + await DBSpendUpdateWriter._commit_window_spend_updates( + prisma_client=prisma_client, + window_spend_transactions=window_spend_update_transactions, + ) + except Exception as e: # noqa: BLE001 # the increments go back on the queue; the rest of the flush must run + spend_log_error( + "Spend tracking - failed to commit budget window spend updates. " + "Re-queued %d window increments for retry on next tick. Error: %s", + len(window_spend_update_transactions), + str(e), + exc=e, + ) + await self.window_spend_update_queue.update_queue.put(window_spend_update_transactions) ################## Tool Registry Upserts ################## await self._flush_tool_discovery_queue(prisma_client=prisma_client) @@ -1206,27 +1218,19 @@ class DBSpendUpdateWriter: """ Commit per-budget-window spend increments to LiteLLM_BudgetWindowSpend. - Failures are logged rather than raised: these rows exist so budget - enforcement can stop aggregating LiteLLM_SpendLogs, and the read path - falls back to that aggregate, so a failed commit must not abort the - entity and daily spend commits that share this scheduler tick. + Raises on failure so the caller re-queues the increments: budget + enforcement trusts a current row without reconciling it against + LiteLLM_SpendLogs, so a dropped increment would let the entity spend + past its window limit after the next counter reseed. """ from litellm.proxy.db.budget_window_spend_writer import ( commit_window_spend_updates, ) - try: - await commit_window_spend_updates( - prisma_client=prisma_client, - transactions=window_spend_transactions, - ) - except Exception as e: # noqa: BLE001 # any DB failure here must stay contained to the window rows - spend_log_error( - "Spend tracking - failed to commit budget window spend updates. %d window increments lost. Error: %s", - len(window_spend_transactions), - str(e), - exc=e, - ) + await commit_window_spend_updates( + prisma_client=prisma_client, + transactions=window_spend_transactions, + ) async def _drain_and_commit_daily_tag_spend_from_redis( self, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 8c53c03031b..fa096fa0bf2 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -456,6 +456,7 @@ class RedisUpdateBuffer: daily_end_user_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, daily_agent_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, daily_tag_spend_update_transactions: Mapping[str, BaseDailySpendTransaction] | None = None, + window_spend_update_transactions: Sequence[WindowSpendTransaction] | None = None, ) -> None: """ Re-push transactions that were popped from Redis but not committed to the DB. @@ -477,6 +478,7 @@ class RedisUpdateBuffer: (daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY), (daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY), (daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY), + (window_spend_update_transactions, REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY), ) rpush_list: Final = tuple( diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 78aa10929e8..504654e103a 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -1,4 +1,5 @@ import json +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -22,9 +23,7 @@ def redis_update_buffer(mock_redis_cache): @pytest.mark.asyncio -async def test_store_in_memory_spend_updates_uses_pipeline( - redis_update_buffer, mock_redis_cache -): +async def test_store_in_memory_spend_updates_uses_pipeline(redis_update_buffer, mock_redis_cache): """ Verify store_in_memory_spend_updates_in_redis calls async_rpush_pipeline once with the correct operations and skips empty queues. @@ -33,35 +32,29 @@ async def test_store_in_memory_spend_updates_uses_pipeline( # Create mock queues - only 3 of 6 have data spend_update_queue = AsyncMock() - spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = ( - AsyncMock(return_value={"key_list_transactions": {"key1": 1.0}}) + spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock( + return_value={"key_list_transactions": {"key1": 1.0}} ) daily_spend_queue = AsyncMock() - daily_spend_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={"user_key1": {"spend": 1.0}}) + daily_spend_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock( + return_value={"user_key1": {"spend": 1.0}} ) daily_team_queue = AsyncMock() - daily_team_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={"team_key1": {"spend": 2.0}}) + daily_team_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock( + return_value={"team_key1": {"spend": 2.0}} ) # Empty queues daily_org_queue = AsyncMock() - daily_org_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={}) - ) + daily_org_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={}) daily_end_user_queue = AsyncMock() - daily_end_user_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value=None) - ) + daily_end_user_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value=None) daily_agent_queue = AsyncMock() - daily_agent_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={}) - ) + daily_agent_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={}) await redis_update_buffer.store_in_memory_spend_updates_in_redis( spend_update_queue=spend_update_queue, @@ -82,9 +75,7 @@ async def test_store_in_memory_spend_updates_uses_pipeline( @pytest.mark.asyncio -async def test_store_in_memory_spend_updates_restores_on_rpush_failure( - redis_update_buffer, mock_redis_cache -): +async def test_store_in_memory_spend_updates_restores_on_rpush_failure(redis_update_buffer, mock_redis_cache): """ If async_rpush_pipeline raises, the already-drained transactions must be put back into the in-memory queues so the next scheduler tick retries. @@ -98,9 +89,7 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure( SpendUpdateQueue, ) - mock_redis_cache.async_rpush_pipeline = AsyncMock( - side_effect=ConnectionError("redis went away") - ) + mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis went away")) spend_queue = SpendUpdateQueue() daily_user_queue = DailySpendUpdateQueue() @@ -145,16 +134,12 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure( # After restore, the main spend queue should hold one item per # (entity_type, entity_id) pair with the aggregated cost - restored_spend = ( - await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() - ) + restored_spend = await spend_queue.flush_and_get_aggregated_db_spend_update_transactions() assert restored_spend["key_list_transactions"] == {"key-abc": 1.5} assert restored_spend["team_list_transactions"] == {"team-xyz": 2.5} # Daily user queue should hold the same aggregated dict - restored_daily = ( - await daily_user_queue.flush_and_get_aggregated_daily_spend_update_transactions() - ) + restored_daily = await daily_user_queue.flush_and_get_aggregated_daily_spend_update_transactions() assert restored_daily == { "user1_day_model": { "spend": 1.0, @@ -165,9 +150,7 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure( @pytest.mark.asyncio -async def test_store_in_memory_spend_updates_all_empty_returns_early( - redis_update_buffer, mock_redis_cache -): +async def test_store_in_memory_spend_updates_all_empty_returns_early(redis_update_buffer, mock_redis_cache): """ When all queues are empty, pipeline should never be called. """ @@ -175,13 +158,9 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early( # All queues return empty empty_queue = AsyncMock() - empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock( - return_value={} - ) + empty_queue.flush_and_get_aggregated_db_spend_update_transactions = AsyncMock(return_value={}) empty_daily_queue = AsyncMock() - empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = ( - AsyncMock(return_value={}) - ) + empty_daily_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock(return_value={}) await redis_update_buffer.store_in_memory_spend_updates_in_redis( spend_update_queue=empty_queue, @@ -196,9 +175,7 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early( @pytest.mark.asyncio -async def test_get_all_transactions_from_redis_buffer_pipeline( - redis_update_buffer, mock_redis_cache -): +async def test_get_all_transactions_from_redis_buffer_pipeline(redis_update_buffer, mock_redis_cache): """ Verify get_all_transactions_from_redis_buffer_pipeline correctly parses and aggregates results from async_lpop_pipeline. @@ -287,9 +264,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline( mock_redis_cache.async_lpop_pipeline.assert_called_once() from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY - popped_keys = [ - op["key"] for op in mock_redis_cache.async_lpop_pipeline.call_args.kwargs["lpop_list"] - ] + popped_keys = [op["key"] for op in mock_redis_cache.async_lpop_pipeline.call_args.kwargs["lpop_list"]] assert popped_keys[6] == REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY @@ -302,9 +277,7 @@ async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis(): @pytest.mark.asyncio -async def test_restore_transactions_to_redis_pushes_only_provided( - redis_update_buffer, mock_redis_cache -): +async def test_restore_transactions_to_redis_pushes_only_provided(redis_update_buffer, mock_redis_cache): """ restore_transactions_to_redis re-pushes only the transaction sets it was given, to their matching buffer keys, so uncommitted spend can be retried. @@ -338,9 +311,42 @@ async def test_restore_transactions_to_redis_pushes_only_provided( @pytest.mark.asyncio -async def test_restore_transactions_to_redis_noop_when_empty( - redis_update_buffer, mock_redis_cache -): +async def test_restored_window_spend_transactions_drain_back_unchanged(redis_update_buffer, mock_redis_cache): + """A window commit that fails after the destructive lpop must be re-pushed + in the store path's encoding, so the next drain returns the same increments.""" + from litellm.constants import REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY + from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( + build_window_spend_transaction, + ) + + window_transactions = ( + build_window_spend_transaction( + entity_type="key", + entity_id="hashed-token", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=3.0, + request_id="req-1", + started_at=datetime(2026, 8, 10, 12, 0, tzinfo=timezone.utc), + ), + ) + mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[1]) + + await redis_update_buffer.restore_transactions_to_redis(window_spend_update_transactions=window_transactions) + + rpush_list = mock_redis_cache.async_rpush_pipeline.call_args.kwargs["rpush_list"] + assert [op["key"] for op in rpush_list] == [REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY] + + mock_redis_cache.async_lpop_pipeline = AsyncMock( + return_value=[None, None, None, None, None, None, list(rpush_list[0]["values"])] + ) + drained = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() + + assert drained[6] == window_transactions + + +@pytest.mark.asyncio +async def test_restore_transactions_to_redis_noop_when_empty(redis_update_buffer, mock_redis_cache): """Nothing to restore -> no Redis call.""" mock_redis_cache.async_rpush_pipeline = AsyncMock() await redis_update_buffer.restore_transactions_to_redis() @@ -348,15 +354,11 @@ async def test_restore_transactions_to_redis_noop_when_empty( @pytest.mark.asyncio -async def test_restore_transactions_to_redis_swallows_redis_error( - redis_update_buffer, mock_redis_cache -): +async def test_restore_transactions_to_redis_swallows_redis_error(redis_update_buffer, mock_redis_cache): """A Redis failure during restore must not propagate to the caller's finally block.""" from redis.exceptions import RedisError - mock_redis_cache.async_rpush_pipeline = AsyncMock( - side_effect=RedisError("redis down") - ) + mock_redis_cache.async_rpush_pipeline = AsyncMock(side_effect=RedisError("redis down")) await redis_update_buffer.restore_transactions_to_redis( db_spend_update_transactions={"key_list_transactions": {"key1": 1.0}}, @@ -472,9 +474,7 @@ def test_get_transaction_buffer_redis_cache_parses_string_flag(monkeypatch): @pytest.mark.asyncio -async def test_store_in_memory_spend_updates_pushes_budget_window_spend( - redis_update_buffer, mock_redis_cache -): +async def test_store_in_memory_spend_updates_pushes_budget_window_spend(redis_update_buffer, mock_redis_cache): """The budget window queue has to ride the same rpush as the daily queues, otherwise multi-pod deployments never persist per-window spend.""" from datetime import datetime, timezone @@ -519,15 +519,17 @@ async def test_store_in_memory_spend_updates_pushes_budget_window_spend( assert len(rpush_list) == 1 assert rpush_list[0]["key"] == REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY pushed = json.loads(rpush_list[0]["values"][0]) - assert pushed == [{ - "entity_type": "key", - "entity_id": "hashed-token", - "window_duration": "30d", - "window_start": "2026-08-01T00:00:00.000000", - "spend": 1.25, - "request_ids": ["req-1"], - "started_at": "2026-08-10T12:00:00.000000", - }] + assert pushed == [ + { + "entity_type": "key", + "entity_id": "hashed-token", + "window_duration": "30d", + "window_start": "2026-08-01T00:00:00.000000", + "spend": 1.25, + "request_ids": ["req-1"], + "started_at": "2026-08-10T12:00:00.000000", + } + ] @pytest.mark.asyncio 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 e07d326e8be..d28cf8c9c6a 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 @@ -2569,19 +2569,21 @@ async def test_window_spend_transactions_are_not_committed_without_the_pod_lock( @pytest.mark.asyncio -async def test_failed_window_spend_commit_does_not_abort_the_rest_of_the_flush(): - """Window rows are an optimization over aggregating LiteLLM_SpendLogs, so a - failure must not take the tool registry flush down with it.""" +async def test_failed_window_spend_commit_requeues_the_increments_and_continues_the_flush(): + """Budget enforcement trusts a current window row without reconciling it + against LiteLLM_SpendLogs, so a dropped increment would let the key spend + past its limit after the next reseed. The increments must go back on the + queue, and the tool registry flush must still run.""" db_writer = DBSpendUpdateWriter() - await db_writer.window_spend_update_queue.add_update( - build_window_spend_transaction( - entity_type="key", - entity_id="hashed-token", - window_duration="30d", - window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), - spend=0.5, - ) + transaction = build_window_spend_transaction( + entity_type="key", + entity_id="hashed-token", + window_duration="30d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=0.5, + request_id="req-1", ) + await db_writer.window_spend_update_queue.add_update(transaction) db = _WindowSpendFakeDB() db.query_raw = AsyncMock(side_effect=Exception("connection reset")) db_writer._flush_tool_discovery_queue = AsyncMock() @@ -2593,6 +2595,47 @@ async def test_failed_window_spend_commit_does_not_abort_the_rest_of_the_flush() ) db_writer._flush_tool_discovery_queue.assert_called_once() + requeued = await db_writer.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions() + assert requeued == (transaction,) + + +@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 + the popped increments back exactly like the other spend categories.""" + db_writer = DBSpendUpdateWriter() + window_transactions = ( + build_window_spend_transaction( + entity_type="team", + entity_id="team-1", + window_duration="7d", + window_start=datetime(2026, 8, 1, tzinfo=timezone.utc), + spend=2.0, + request_id="req-1", + ), + ) + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(None, None, None, None, None, None, window_transactions) + ) + mock_redis_update_buffer.restore_transactions_to_redis = AsyncMock() + db_writer.redis_update_buffer = mock_redis_update_buffer + db_writer.pod_lock_manager = AsyncMock() + db_writer.pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + db = _WindowSpendFakeDB() + db.query_raw = AsyncMock(side_effect=Exception("connection reset")) + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=_WindowSpendFakePrisma(db), + n_retry_times=0, + proxy_logging_obj=MagicMock(), + ) + + assert _window_spend_upserts(db) == [] + mock_redis_update_buffer.restore_transactions_to_redis.assert_awaited_once_with( + window_spend_update_transactions=window_transactions + ) + db_writer.pod_lock_manager.release_lock.assert_awaited_once() @pytest.mark.asyncio