From 288d4035293d8a4ed6ccff9a55376b2550e83c3c Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 22 Apr 2026 14:10:40 -0700 Subject: [PATCH] [Fix] Preserve in-memory spend updates when Redis rpush fails store_in_memory_spend_updates_in_redis drained the in-memory queues into local variables before the rpush pipeline. If rpush raised (cloud Redis hiccup, timeout, connection blip), those already-drained transactions were garbage-collected with the scheduler job, silently losing all spend aggregated during that tick. Wrap the rpush in try/except. On failure, re-enqueue the aggregated transactions into their respective in-memory queues so the next scheduler tick retries. Add a unit test that seeds real queues, simulates an rpush failure, and asserts the transactions land back in-memory. --- .../redis_update_buffer.py | 101 +++++++++++++++++- .../test_redis_update_buffer.py | 83 ++++++++++++++ 2 files changed, 181 insertions(+), 3 deletions(-) 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 b8537c2be9..3a008d265f 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -29,6 +29,8 @@ from litellm.proxy._types import ( DailyTeamSpendTransaction, DailyUserSpendTransaction, DBSpendUpdateTransactions, + Litellm_EntityType, + SpendUpdateQueueItem, ) from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( @@ -259,9 +261,36 @@ class RedisUpdateBuffer: if len(rpush_list) == 0: return - result_lengths = await self.redis_cache.async_rpush_pipeline( - rpush_list=rpush_list, - ) + try: + result_lengths = await self.redis_cache.async_rpush_pipeline( + rpush_list=rpush_list, + ) + except Exception as e: + # The in-memory queues were already drained above. If we let the + # exception propagate without restoring, the aggregated spend is + # permanently lost. Re-enqueue so the next scheduler tick retries. + verbose_proxy_logger.error( + "Spend tracking - failed to push aggregated spend updates to Redis. " + "Restoring %d transaction sets to in-memory queues for retry on next tick. " + "Error: %s", + len(rpush_list), + str(e), + ) + await self._restore_spend_updates_to_in_memory_queues( + db_spend_update_transactions=db_spend_update_transactions, + daily_spend_update_transactions=daily_spend_update_transactions, + daily_team_spend_update_transactions=daily_team_spend_update_transactions, + 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, + spend_update_queue=spend_update_queue, + daily_spend_update_queue=daily_spend_update_queue, + daily_team_spend_update_queue=daily_team_spend_update_queue, + daily_org_spend_update_queue=daily_org_spend_update_queue, + daily_end_user_spend_update_queue=daily_end_user_spend_update_queue, + daily_agent_spend_update_queue=daily_agent_spend_update_queue, + ) + return # Emit gauge events for each queue for i, queue_size in enumerate(result_lengths): @@ -271,6 +300,72 @@ class RedisUpdateBuffer: service=service_types[i], ) + @staticmethod + async def _restore_spend_updates_to_in_memory_queues( + db_spend_update_transactions: Optional[DBSpendUpdateTransactions], + daily_spend_update_transactions: Optional[Dict[str, DailyUserSpendTransaction]], + daily_team_spend_update_transactions: Optional[ + Dict[str, DailyTeamSpendTransaction] + ], + daily_org_spend_update_transactions: Optional[ + Dict[str, DailyOrganizationSpendTransaction] + ], + daily_end_user_spend_update_transactions: Optional[ + Dict[str, DailyEndUserSpendTransaction] + ], + daily_agent_spend_update_transactions: Optional[ + Dict[str, DailyAgentSpendTransaction] + ], + spend_update_queue: SpendUpdateQueue, + daily_spend_update_queue: DailySpendUpdateQueue, + daily_team_spend_update_queue: DailySpendUpdateQueue, + daily_org_spend_update_queue: DailySpendUpdateQueue, + daily_end_user_spend_update_queue: DailySpendUpdateQueue, + daily_agent_spend_update_queue: DailySpendUpdateQueue, + ) -> None: + """ + Put drained-but-unpushed transactions back into in-memory queues. + + Called when the Redis rpush pipeline raises. Without this, all spend + data aggregated during the current scheduler tick is permanently lost + because the source queues were already drained before the rpush. + """ + entity_type_field_pairs = [ + (Litellm_EntityType.USER, "user_list_transactions"), + (Litellm_EntityType.END_USER, "end_user_list_transactions"), + (Litellm_EntityType.KEY, "key_list_transactions"), + (Litellm_EntityType.TEAM, "team_list_transactions"), + (Litellm_EntityType.TEAM_MEMBER, "team_member_list_transactions"), + (Litellm_EntityType.ORGANIZATION, "org_list_transactions"), + (Litellm_EntityType.TAG, "tag_list_transactions"), + (Litellm_EntityType.AGENT, "agent_list_transactions"), + ] + if db_spend_update_transactions is not None: + for entity_type, field in entity_type_field_pairs: + entities = db_spend_update_transactions.get(field) or {} # type: ignore[call-overload] + for entity_id, cost in entities.items(): + await spend_update_queue.add_update( + SpendUpdateQueueItem( + entity_type=entity_type, + entity_id=entity_id, + response_cost=cost, + ) + ) + + daily_pairs = [ + (daily_spend_update_transactions, daily_spend_update_queue), + (daily_team_spend_update_transactions, daily_team_spend_update_queue), + (daily_org_spend_update_transactions, daily_org_spend_update_queue), + ( + daily_end_user_spend_update_transactions, + daily_end_user_spend_update_queue, + ), + (daily_agent_spend_update_transactions, daily_agent_spend_update_queue), + ] + for daily_txns, daily_queue in daily_pairs: + if daily_txns: + await daily_queue.update_queue.put(daily_txns) + @staticmethod def _number_of_transactions_to_store_in_redis( db_spend_update_transactions: DBSpendUpdateTransactions, 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 44602125ff..0587e3bce1 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 @@ -87,6 +87,89 @@ async def test_store_in_memory_spend_updates_uses_pipeline( assert len(rpush_list) == 3 +@pytest.mark.asyncio +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. + Without this, any transient Redis hiccup silently loses spend data. + """ + from litellm.proxy._types import Litellm_EntityType + from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( + DailySpendUpdateQueue, + ) + from litellm.proxy.db.db_transaction_queue.spend_update_queue import ( + SpendUpdateQueue, + ) + + mock_redis_cache.async_rpush_pipeline = AsyncMock( + side_effect=ConnectionError("redis went away") + ) + + spend_queue = SpendUpdateQueue() + daily_user_queue = DailySpendUpdateQueue() + daily_team_queue = DailySpendUpdateQueue() + daily_org_queue = DailySpendUpdateQueue() + daily_end_user_queue = DailySpendUpdateQueue() + daily_agent_queue = DailySpendUpdateQueue() + + # Seed real queues with data so flush_and_get_aggregated returns it + await spend_queue.add_update( + { + "entity_type": Litellm_EntityType.KEY, + "entity_id": "key-abc", + "response_cost": 1.5, + } + ) + await spend_queue.add_update( + { + "entity_type": Litellm_EntityType.TEAM, + "entity_id": "team-xyz", + "response_cost": 2.5, + } + ) + await daily_user_queue.add_update( + { + "user1_day_model": { + "spend": 1.0, + "prompt_tokens": 10, + "completion_tokens": 20, + } + } + ) + + await redis_update_buffer.store_in_memory_spend_updates_in_redis( + spend_update_queue=spend_queue, + daily_spend_update_queue=daily_user_queue, + daily_team_spend_update_queue=daily_team_queue, + daily_org_spend_update_queue=daily_org_queue, + daily_end_user_spend_update_queue=daily_end_user_queue, + daily_agent_spend_update_queue=daily_agent_queue, + ) + + # 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() + ) + 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() + ) + assert restored_daily == { + "user1_day_model": { + "spend": 1.0, + "prompt_tokens": 10, + "completion_tokens": 20, + } + } + + @pytest.mark.asyncio async def test_store_in_memory_spend_updates_all_empty_returns_early( redis_update_buffer, mock_redis_cache