From 98b496433086b5a56f87e610c0349607b7553f55 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 24 Feb 2026 12:31:45 -0800 Subject: [PATCH 1/2] perf(proxy): pipeline Redis RPUSH/LPOP in spend update cycle Replace 14 sequential Redis round-trips (7 RPUSH + 7 LPOP) per spend update cycle with 2 pipelined calls (1 RPUSH pipeline + 1 LPOP pipeline). This reduces connection pool contention at scale (50+ pods). - Add RedisPipelineRpushOperation and RedisPipelineLpopOperation TypedDicts - Add async_rpush_pipeline() and async_lpop_pipeline() to RedisCache - Refactor store_in_memory_spend_updates_in_redis() to use pipeline - Add get_all_transactions_from_redis_buffer_pipeline() for batched drain - Update _commit_spend_updates_to_db_with_redis() to use pipeline drain - Existing individual methods preserved for backward compatibility --- litellm/caching/redis_cache.py | 187 ++++++++++++++- litellm/proxy/db/db_spend_update_writer.py | 31 +-- .../redis_update_buffer.py | 149 ++++++++---- litellm/types/caching.py | 18 ++ .../test_litellm/caching/test_redis_cache.py | 213 ++++++++++++++++++ .../test_redis_update_buffer.py | 194 ++++++++++++++++ .../proxy/db/test_db_spend_update_writer.py | 43 ++++ 7 files changed, 773 insertions(+), 62 deletions(-) create mode 100644 tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index dcc2df5f91..9b93d890ce 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -22,7 +22,11 @@ from litellm._logging import print_verbose, verbose_logger from litellm.constants import DEFAULT_REDIS_MAJOR_VERSION from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs from litellm.litellm_core_utils.coroutine_checker import coroutine_checker -from litellm.types.caching import RedisPipelineIncrementOperation +from litellm.types.caching import ( + RedisPipelineIncrementOperation, + RedisPipelineLpopOperation, + RedisPipelineRpushOperation, +) from litellm.types.services import ServiceTypes from .base_cache import BaseCache @@ -1320,6 +1324,75 @@ class RedisCache(BaseCache): ) raise e + async def _pipeline_rpush_helper( + self, + pipe: pipeline, + rpush_list: List[RedisPipelineRpushOperation], + ) -> List[int]: + """Helper function for pipeline rpush operations""" + for rpush_op in rpush_list: + pipe.rpush(rpush_op["key"], *rpush_op["values"]) + results = await pipe.execute() + # Preserve positional correspondence — raise on per-command errors + for r in results: + if isinstance(r, Exception): + raise r + return results + + async def async_rpush_pipeline( + self, + rpush_list: List[RedisPipelineRpushOperation], + ) -> List[int]: + """ + Use Redis Pipelines for bulk RPUSH operations + + Args: + rpush_list: List of RedisPipelineRpushOperation dicts containing: + - key: str + - values: List[Any] + + Returns: + List[int]: List lengths after each push + """ + if len(rpush_list) == 0: + return [] + + _redis_client: Any = self.init_async_client() + start_time = time.time() + + try: + async with _redis_client.pipeline(transaction=False) as pipe: + results = await self._pipeline_rpush_helper(pipe, rpush_list) + + ## LOGGING ## + end_time = time.time() + _duration = end_time - start_time + asyncio.create_task( + self.service_logger_obj.async_service_success_hook( + service=ServiceTypes.REDIS, + duration=_duration, + call_type=f"async_rpush_pipeline <- {_get_call_stack_info()}", + ) + ) + return results + except Exception as e: + ## LOGGING ## + end_time = time.time() + _duration = end_time - start_time + asyncio.create_task( + self.service_logger_obj.async_service_failure_hook( + service=ServiceTypes.REDIS, + duration=_duration, + error=e, + call_type=f"async_rpush_pipeline <- {_get_call_stack_info()}", + ) + ) + verbose_logger.error( + "LiteLLM Redis Caching: async_rpush_pipeline() - Got exception from REDIS %s", + str(e), + ) + raise e + async def handle_lpop_count_for_older_redis_versions( self, pipe: pipeline, key: str, count: int ) -> List[bytes]: @@ -1400,3 +1473,115 @@ class RedisCache(BaseCache): f"LiteLLM Redis Cache LPOP: - Got exception from REDIS : {str(e)}" ) raise e + + async def _pipeline_lpop_helper( + self, + pipe: pipeline, + lpop_list: List[RedisPipelineLpopOperation], + ) -> List[Optional[List[str]]]: + """Helper function for pipeline lpop operations. + + For Redis >= 7, queues one LPOP(key, count) per operation. + For Redis < 7, queues `count` individual LPOP(key) commands per operation. + """ + major_version = self._parse_redis_major_version() + + if major_version >= 7: + for lpop_op in lpop_list: + pipe.lpop(lpop_op["key"], lpop_op["count"]) + raw_results = await pipe.execute() + else: + # For Redis < 7, LPOP doesn't support count param. + # Issue `count` individual LPOP commands per key, all in one pipeline. + counts: List[int] = [] + for lpop_op in lpop_list: + count = lpop_op["count"] or 1 + counts.append(count) + for _ in range(count): + pipe.lpop(lpop_op["key"]) + flat_results = await pipe.execute() + + # Re-group the flat results back into per-key lists + raw_results = [] + offset = 0 + for count in counts: + key_results = [ + r for r in flat_results[offset : offset + count] if r is not None + ] + raw_results.append(key_results if key_results else None) + offset += count + + # Decode bytes -> str for each result set + decoded_results: List[Optional[List[str]]] = [] + for r in raw_results: + if r is None: + decoded_results.append(None) + elif isinstance(r, list): + try: + decoded_results.append( + [ + item.decode("utf-8") if isinstance(item, bytes) else item + for item in r + if item is not None + ] + or None + ) + except Exception: + decoded_results.append(r) # type: ignore + else: + decoded_results.append(None) + return decoded_results + + async def async_lpop_pipeline( + self, + lpop_list: List[RedisPipelineLpopOperation], + ) -> List[Optional[List[str]]]: + """ + Use Redis Pipelines for bulk LPOP operations + + Args: + lpop_list: List of RedisPipelineLpopOperation dicts containing: + - key: str + - count: Optional[int] + + Returns: + List[Optional[List[str]]]: Decoded results per key, None if key was empty + """ + if len(lpop_list) == 0: + return [] + + _redis_client: Any = self.init_async_client() + start_time = time.time() + + try: + async with _redis_client.pipeline(transaction=False) as pipe: + results = await self._pipeline_lpop_helper(pipe, lpop_list) + + ## LOGGING ## + end_time = time.time() + _duration = end_time - start_time + asyncio.create_task( + self.service_logger_obj.async_service_success_hook( + service=ServiceTypes.REDIS, + duration=_duration, + call_type=f"async_lpop_pipeline <- {_get_call_stack_info()}", + ) + ) + return results + except Exception as e: + ## LOGGING ## + end_time = time.time() + _duration = end_time - start_time + asyncio.create_task( + self.service_logger_obj.async_service_failure_hook( + service=ServiceTypes.REDIS, + duration=_duration, + error=e, + call_type=f"async_lpop_pipeline <- {_get_call_stack_info()}", + ) + ) + verbose_logger.error( + "LiteLLM Redis Caching: async_lpop_pipeline() - Got exception from REDIS %s", + str(e), + ) + raise e diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 03628fda47..b9cd794e78 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -509,9 +509,16 @@ class DBSpendUpdateWriter: verbose_proxy_logger.debug("acquired lock for spend updates") try: - db_spend_update_transactions = ( - await self.redis_update_buffer.get_all_update_transactions_from_redis_buffer() - ) + ( + db_spend_update_transactions, + daily_spend_update_transactions, + daily_team_spend_update_transactions, + daily_org_spend_update_transactions, + daily_end_user_spend_update_transactions, + daily_agent_spend_update_transactions, + daily_tag_spend_update_transactions, + ) = await self.redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() + if db_spend_update_transactions is not None: await self._commit_spend_updates_to_db( prisma_client=prisma_client, @@ -520,9 +527,6 @@ class DBSpendUpdateWriter: db_spend_update_transactions=db_spend_update_transactions, ) - daily_spend_update_transactions = ( - await self.redis_update_buffer.get_all_daily_spend_update_transactions_from_redis_buffer() - ) if daily_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_user_spend( n_retry_times=n_retry_times, @@ -530,9 +534,6 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_spend_update_transactions, ) - daily_team_spend_update_transactions = ( - await self.redis_update_buffer.get_all_daily_team_spend_update_transactions_from_redis_buffer() - ) if daily_team_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_team_spend( n_retry_times=n_retry_times, @@ -541,9 +542,6 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_team_spend_update_transactions, ) - daily_org_spend_update_transactions = ( - await self.redis_update_buffer.get_all_daily_org_spend_update_transactions_from_redis_buffer() - ) if daily_org_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_org_spend( n_retry_times=n_retry_times, @@ -552,9 +550,6 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_org_spend_update_transactions, ) - daily_tag_spend_update_transactions = ( - await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() - ) if daily_tag_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_tag_spend( n_retry_times=n_retry_times, @@ -562,9 +557,6 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_tag_spend_update_transactions, ) - daily_end_user_spend_update_transactions = ( - await self.redis_update_buffer.get_all_daily_end_user_spend_update_transactions_from_redis_buffer() - ) if daily_end_user_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_end_user_spend( n_retry_times=n_retry_times, @@ -572,9 +564,6 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, daily_spend_transactions=daily_end_user_spend_update_transactions, ) - daily_agent_spend_update_transactions = ( - await self.redis_update_buffer.get_all_daily_agent_spend_update_transactions_from_redis_buffer() - ) if daily_agent_spend_update_transactions is not None: await DBSpendUpdateWriter.update_daily_agent_spend( n_retry_times=n_retry_times, 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 37b42e26bc..ef8dbb2c7b 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -6,7 +6,7 @@ This is to prevent deadlocks and improve reliability import asyncio import json -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache @@ -36,6 +36,7 @@ from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( ) from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue from litellm.secret_managers.main import str_to_bool +from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation from litellm.types.services import ServiceTypes if TYPE_CHECKING: @@ -195,47 +196,44 @@ class RedisUpdateBuffer: "ALL DAILY SPEND UPDATE TRANSACTIONS: %s", daily_spend_update_transactions ) - await self._store_transactions_in_redis( - transactions=db_spend_update_transactions, - redis_key=REDIS_UPDATE_BUFFER_KEY, - service_type=ServiceTypes.REDIS_SPEND_UPDATE_QUEUE, + # Build a list of rpush operations, skipping empty/None transaction sets + _queue_configs: List[Tuple[Any, str, ServiceTypes]] = [ + (db_spend_update_transactions, REDIS_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_SPEND_UPDATE_QUEUE), + (daily_spend_update_transactions, REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE), + (daily_team_spend_update_transactions, REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE), + (daily_org_spend_update_transactions, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE), + (daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE), + (daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE), + (daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE), + ] + + rpush_list: List[RedisPipelineRpushOperation] = [] + service_types: List[ServiceTypes] = [] + for transactions, redis_key, service_type in _queue_configs: + if transactions is None or len(transactions) == 0: + continue + rpush_list.append( + RedisPipelineRpushOperation( + key=redis_key, + values=[safe_dumps(transactions)], + ) + ) + service_types.append(service_type) + + if len(rpush_list) == 0: + return + + result_lengths = await self.redis_cache.async_rpush_pipeline( + rpush_list=rpush_list, ) - await self._store_transactions_in_redis( - transactions=daily_spend_update_transactions, - redis_key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, - service_type=ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE, - ) - - await self._store_transactions_in_redis( - transactions=daily_team_spend_update_transactions, - redis_key=REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, - service_type=ServiceTypes.REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE, - ) - - await self._store_transactions_in_redis( - transactions=daily_org_spend_update_transactions, - redis_key=REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, - service_type=ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE, - ) - - await self._store_transactions_in_redis( - transactions=daily_end_user_spend_update_transactions, - redis_key=REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, - service_type=ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE, - ) - - await self._store_transactions_in_redis( - transactions=daily_agent_spend_update_transactions, - redis_key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, - service_type=ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE, - ) - - await self._store_transactions_in_redis( - transactions=daily_tag_spend_update_transactions, - redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, - service_type=ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE, - ) + # Emit gauge events for each queue + for i, queue_size in enumerate(result_lengths): + if i < len(service_types): + await self._emit_new_item_added_to_redis_buffer_event( + queue_size=queue_size, + service=service_types[i], + ) @staticmethod def _number_of_transactions_to_store_in_redis( @@ -317,6 +315,77 @@ class RedisUpdateBuffer: return combined_transaction + async def get_all_transactions_from_redis_buffer_pipeline( + self, + ) -> Tuple[ + Optional[DBSpendUpdateTransactions], + Optional[Dict[str, DailyUserSpendTransaction]], + Optional[Dict[str, DailyTeamSpendTransaction]], + Optional[Dict[str, DailyOrganizationSpendTransaction]], + Optional[Dict[str, DailyEndUserSpendTransaction]], + Optional[Dict[str, DailyAgentSpendTransaction]], + Optional[Dict[str, DailyTagSpendTransaction]], + ]: + """ + Drains all 7 Redis buffer queues in a single pipeline round-trip. + + Returns a 7-tuple of parsed results in this order: + 0: DBSpendUpdateTransactions + 1: daily user spend + 2: daily team spend + 3: daily org spend + 4: daily end-user spend + 5: daily agent spend + 6: daily tag spend + """ + if self.redis_cache is None: + return None, None, None, None, None, None, None + + lpop_list: List[RedisPipelineLpopOperation] = [ + RedisPipelineLpopOperation(key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), + RedisPipelineLpopOperation(key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), + RedisPipelineLpopOperation(key=REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), + RedisPipelineLpopOperation(key=REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), + RedisPipelineLpopOperation(key=REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), + RedisPipelineLpopOperation(key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), + RedisPipelineLpopOperation(key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT), + ] + + raw_results = await self.redis_cache.async_lpop_pipeline(lpop_list=lpop_list) + + # Pad with None if pipeline returned fewer results than expected + while len(raw_results) < 7: + raw_results.append(None) + + # Slot 0: DBSpendUpdateTransactions + db_spend: Optional[DBSpendUpdateTransactions] = None + if raw_results[0] is not None: + parsed = self._parse_list_of_transactions(raw_results[0]) + if len(parsed) > 0: + db_spend = self._combine_list_of_transactions(parsed) + + # Slots 1-6: daily spend categories + daily_results: List[Optional[Dict[str, Any]]] = [] + for slot in range(1, 7): + if raw_results[slot] is None: + daily_results.append(None) + else: + list_of_daily = [json.loads(t) for t in raw_results[slot]] # type: ignore + aggregated = DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( + list_of_daily + ) + daily_results.append(aggregated) + + return ( + db_spend, + cast(Optional[Dict[str, DailyUserSpendTransaction]], daily_results[0]), + cast(Optional[Dict[str, DailyTeamSpendTransaction]], daily_results[1]), + cast(Optional[Dict[str, DailyOrganizationSpendTransaction]], daily_results[2]), + cast(Optional[Dict[str, DailyEndUserSpendTransaction]], daily_results[3]), + cast(Optional[Dict[str, DailyAgentSpendTransaction]], daily_results[4]), + cast(Optional[Dict[str, DailyTagSpendTransaction]], daily_results[5]), + ) + async def get_all_daily_spend_update_transactions_from_redis_buffer( self, ) -> Optional[Dict[str, DailyUserSpendTransaction]]: diff --git a/litellm/types/caching.py b/litellm/types/caching.py index ad2ffeaf9b..7126ba3e9b 100644 --- a/litellm/types/caching.py +++ b/litellm/types/caching.py @@ -52,6 +52,24 @@ class RedisPipelineSetOperation(TypedDict): ttl: Optional[int] +class RedisPipelineRpushOperation(TypedDict): + """ + TypedDict for 1 Redis Pipeline RPUSH Operation + """ + + key: str + values: List[Any] + + +class RedisPipelineLpopOperation(TypedDict): + """ + TypedDict for 1 Redis Pipeline LPOP Operation + """ + + key: str + count: Optional[int] + + DynamicCacheControl = TypedDict( "DynamicCacheControl", { diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index f3e7953b8a..31eeb85488 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -122,6 +122,219 @@ async def test_handle_lpop_count_for_older_redis_versions(monkeypatch): assert mock_pipeline.execute.call_count == 2 +@pytest.mark.asyncio +async def test_async_rpush_pipeline_executes_all_operations(monkeypatch, redis_no_ping): + """Verify that multiple rpush ops are batched into a single pipeline execute""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + + mock_redis_instance = AsyncMock() + mock_pipeline = MagicMock() + mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) + mock_pipeline.__aexit__ = AsyncMock(return_value=None) + mock_pipeline.rpush = MagicMock() + mock_pipeline.execute = AsyncMock(return_value=[3, 5, 1]) + mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) + + from litellm.types.caching import RedisPipelineRpushOperation + + rpush_list = [ + RedisPipelineRpushOperation(key="key1", values=["a", "b"]), + RedisPipelineRpushOperation(key="key2", values=["c"]), + RedisPipelineRpushOperation(key="key3", values=["d", "e", "f"]), + ] + + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + result = await redis_cache.async_rpush_pipeline(rpush_list=rpush_list) + + assert result == [3, 5, 1] + assert mock_pipeline.rpush.call_count == 3 + mock_pipeline.rpush.assert_any_call("key1", "a", "b") + mock_pipeline.rpush.assert_any_call("key2", "c") + mock_pipeline.rpush.assert_any_call("key3", "d", "e", "f") + mock_pipeline.execute.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_rpush_pipeline_empty_list_returns_empty(monkeypatch, redis_no_ping): + """Empty rpush_list should return empty list without touching Redis""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + + mock_redis_instance = AsyncMock() + + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + result = await redis_cache.async_rpush_pipeline(rpush_list=[]) + + assert result == [] + mock_redis_instance.pipeline.assert_not_called() + + +@pytest.mark.asyncio +async def test_async_rpush_pipeline_raises_on_redis_error(monkeypatch, redis_no_ping): + """Pipeline errors should propagate""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + + mock_redis_instance = AsyncMock() + mock_pipeline = MagicMock() + mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) + mock_pipeline.__aexit__ = AsyncMock(return_value=None) + mock_pipeline.rpush = MagicMock() + mock_pipeline.execute = AsyncMock(side_effect=ConnectionError("Redis down")) + mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) + + from litellm.types.caching import RedisPipelineRpushOperation + + rpush_list = [RedisPipelineRpushOperation(key="key1", values=["a"])] + + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + with pytest.raises(ConnectionError, match="Redis down"): + await redis_cache.async_rpush_pipeline(rpush_list=rpush_list) + + +@pytest.mark.asyncio +async def test_async_lpop_pipeline_single_round_trip(monkeypatch, redis_no_ping): + """Verify that multiple lpop ops are batched into a single pipeline execute""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + redis_cache.redis_version = "7.0.0" + + mock_redis_instance = AsyncMock() + mock_pipeline = MagicMock() + mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) + mock_pipeline.__aexit__ = AsyncMock(return_value=None) + mock_pipeline.lpop = MagicMock() + mock_pipeline.execute = AsyncMock(return_value=[ + [b"val1", b"val2"], # key1 results + None, # key2 empty + [b"val3"], # key3 results + ]) + mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) + + from litellm.types.caching import RedisPipelineLpopOperation + + lpop_list = [ + RedisPipelineLpopOperation(key="key1", count=10), + RedisPipelineLpopOperation(key="key2", count=10), + RedisPipelineLpopOperation(key="key3", count=5), + ] + + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + results = await redis_cache.async_lpop_pipeline(lpop_list=lpop_list) + + assert len(results) == 3 + assert results[0] == ["val1", "val2"] + assert results[1] is None + assert results[2] == ["val3"] + mock_pipeline.execute.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_lpop_pipeline_redis_lt7_regroups_flat_results(monkeypatch, redis_no_ping): + """Verify Redis < 7 fallback issues individual LPOPs and regroups correctly""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + redis_cache.redis_version = "6.2.0" + + mock_redis_instance = AsyncMock() + mock_pipeline = MagicMock() + mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) + mock_pipeline.__aexit__ = AsyncMock(return_value=None) + mock_pipeline.lpop = MagicMock() + + # With count=3 for key1 and count=2 for key2, we get 5 individual LPOP commands + # Simulate: key1 has 2 values then None, key2 has 1 value then None + mock_pipeline.execute = AsyncMock(return_value=[ + b"val1", b"val2", None, # 3 LPOPs for key1 + b"val3", None, # 2 LPOPs for key2 + ]) + mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) + + from litellm.types.caching import RedisPipelineLpopOperation + + lpop_list = [ + RedisPipelineLpopOperation(key="key1", count=3), + RedisPipelineLpopOperation(key="key2", count=2), + ] + + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + results = await redis_cache.async_lpop_pipeline(lpop_list=lpop_list) + + assert len(results) == 2 + assert results[0] == ["val1", "val2"] # 2 values, None filtered out + assert results[1] == ["val3"] # 1 value, None filtered out + # All 5 individual LPOPs should be queued, but only 1 execute() call + assert mock_pipeline.lpop.call_count == 5 + mock_pipeline.execute.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_rpush_pipeline_raises_on_per_command_error(monkeypatch, redis_no_ping): + """Verify that per-command errors in pipeline results are raised, not silently dropped""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + + mock_redis_instance = AsyncMock() + mock_pipeline = MagicMock() + mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) + mock_pipeline.__aexit__ = AsyncMock(return_value=None) + mock_pipeline.rpush = MagicMock() + # Simulate: first RPUSH succeeds, second returns a per-command error + mock_pipeline.execute = AsyncMock(return_value=[3, Exception("WRONGTYPE")]) + mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) + + from litellm.types.caching import RedisPipelineRpushOperation + + rpush_list = [ + RedisPipelineRpushOperation(key="key1", values=["a"]), + RedisPipelineRpushOperation(key="key2", values=["b"]), + ] + + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + with pytest.raises(Exception, match="WRONGTYPE"): + await redis_cache.async_rpush_pipeline(rpush_list=rpush_list) + + +@pytest.mark.asyncio +async def test_async_lpop_pipeline_empty_list(monkeypatch, redis_no_ping): + """Empty lpop_list should return empty list without touching Redis""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + + mock_redis_instance = AsyncMock() + + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + result = await redis_cache.async_lpop_pipeline(lpop_list=[]) + + assert result == [] + mock_redis_instance.pipeline.assert_not_called() + + +@pytest.mark.asyncio +async def test_async_lpop_pipeline_propagates_redis_exception(monkeypatch, redis_no_ping): + """Pipeline errors should propagate""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + redis_cache.redis_version = "7.0.0" + + mock_redis_instance = AsyncMock() + mock_pipeline = MagicMock() + mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) + mock_pipeline.__aexit__ = AsyncMock(return_value=None) + mock_pipeline.lpop = MagicMock() + mock_pipeline.execute = AsyncMock(side_effect=ConnectionError("Redis down")) + mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) + + from litellm.types.caching import RedisPipelineLpopOperation + + lpop_list = [RedisPipelineLpopOperation(key="key1", count=10)] + + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + with pytest.raises(ConnectionError, match="Redis down"): + await redis_cache.async_lpop_pipeline(lpop_list=lpop_list) + + @pytest.mark.asyncio @pytest.mark.parametrize( "redis_version", 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 new file mode 100644 index 0000000000..2a380370c3 --- /dev/null +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -0,0 +1,194 @@ +import json +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path + +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer +from litellm.types.caching import RedisPipelineRpushOperation + + +@pytest.fixture +def mock_redis_cache(): + """Create a mock RedisCache instance""" + mock = AsyncMock() + return mock + + +@pytest.fixture +def redis_update_buffer(mock_redis_cache): + """Create a RedisUpdateBuffer with a mock RedisCache""" + return RedisUpdateBuffer(redis_cache=mock_redis_cache) + + +@pytest.mark.asyncio +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. + """ + mock_redis_cache.async_rpush_pipeline = AsyncMock(return_value=[3, 5, 2]) + + # Create mock queues - only 3 of 7 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}} + ) + + daily_spend_queue = AsyncMock() + 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}} + ) + + # Empty queues + daily_org_queue = AsyncMock() + 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_agent_queue = AsyncMock() + daily_agent_queue.flush_and_get_aggregated_daily_spend_update_transactions = AsyncMock( + return_value={} + ) + + daily_tag_queue = AsyncMock() + daily_tag_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, + daily_spend_update_queue=daily_spend_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, + daily_tag_spend_update_queue=daily_tag_queue, + ) + + # Should be called exactly once (pipeline) + mock_redis_cache.async_rpush_pipeline.assert_called_once() + + # Verify only 3 operations were included (empty ones skipped) + call_args = mock_redis_cache.async_rpush_pipeline.call_args + rpush_list = call_args.kwargs["rpush_list"] + assert len(rpush_list) == 3 + + +@pytest.mark.asyncio +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. + """ + mock_redis_cache.async_rpush_pipeline = AsyncMock() + + # All queues return empty + empty_queue = AsyncMock() + 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={} + ) + + await redis_update_buffer.store_in_memory_spend_updates_in_redis( + spend_update_queue=empty_queue, + daily_spend_update_queue=empty_daily_queue, + daily_team_spend_update_queue=empty_daily_queue, + daily_org_spend_update_queue=empty_daily_queue, + daily_end_user_spend_update_queue=empty_daily_queue, + daily_agent_spend_update_queue=empty_daily_queue, + daily_tag_spend_update_queue=empty_daily_queue, + ) + + mock_redis_cache.async_rpush_pipeline.assert_not_called() + + +@pytest.mark.asyncio +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. + """ + # Simulate pipeline results: slot 0 = spend updates, slots 1-6 = daily categories + db_spend_json = json.dumps( + { + "key_list_transactions": {"key1": 1.0, "key2": 2.0}, + "user_list_transactions": {"user1": 0.5}, + "end_user_list_transactions": {}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + } + ) + daily_user_json = json.dumps({"user_key1": {"spend": 1.0, "api_requests": 1}}) + daily_team_json = json.dumps({"team_key1": {"spend": 2.0, "api_requests": 2}}) + + mock_redis_cache.async_lpop_pipeline = AsyncMock( + return_value=[ + [db_spend_json], # slot 0: db spend updates + [daily_user_json], # slot 1: daily user + [daily_team_json], # slot 2: daily team + None, # slot 3: daily org (empty) + None, # slot 4: daily end-user (empty) + None, # slot 5: daily agent (empty) + None, # slot 6: daily tag (empty) + ] + ) + + result = await redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline() + + assert len(result) == 7 + db_spend, daily_user, daily_team, daily_org, daily_end_user, daily_agent, daily_tag = result + + # Verify db spend was parsed correctly + assert db_spend is not None + assert db_spend["key_list_transactions"]["key1"] == 1.0 + assert db_spend["key_list_transactions"]["key2"] == 2.0 + assert db_spend["user_list_transactions"]["user1"] == 0.5 + + # Verify daily user was parsed + assert daily_user is not None + assert daily_user["user_key1"]["spend"] == 1.0 + + # Verify daily team was parsed + assert daily_team is not None + assert daily_team["team_key1"]["spend"] == 2.0 + + # Verify empty slots + assert daily_org is None + assert daily_end_user is None + assert daily_agent is None + assert daily_tag is None + + # Verify pipeline was called once with correct keys + mock_redis_cache.async_lpop_pipeline.assert_called_once() + + +@pytest.mark.asyncio +async def test_get_all_transactions_from_redis_buffer_pipeline_no_redis(): + """When redis_cache is None, should return all Nones""" + buffer = RedisUpdateBuffer(redis_cache=None) + result = await buffer.get_all_transactions_from_redis_buffer_pipeline() + assert result == (None, None, None, None, None, None, None) 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 0fa0d4cf10..2a0428e2f0 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 @@ -1076,3 +1076,46 @@ async def test_commit_key_spend_updates_includes_last_active(): last_active = call_kwargs["data"]["last_active"] assert isinstance(last_active, datetime) assert before_call <= last_active <= after_call + + +@pytest.mark.asyncio +async def test_commit_spend_updates_uses_pipeline(): + """ + Verify that _commit_spend_updates_to_db_with_redis uses + get_all_transactions_from_redis_buffer_pipeline instead of 7 individual calls. + """ + db_writer = DBSpendUpdateWriter() + + mock_redis_update_buffer = AsyncMock() + mock_redis_update_buffer.store_in_memory_spend_updates_in_redis = AsyncMock() + # Return all-None tuple (no data to commit) + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline = AsyncMock( + return_value=(None, None, None, None, None, None, None) + ) + db_writer.redis_update_buffer = mock_redis_update_buffer + + mock_pod_lock_manager = AsyncMock() + mock_pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + mock_pod_lock_manager.release_lock = AsyncMock() + db_writer.pod_lock_manager = mock_pod_lock_manager + + mock_prisma_client = MagicMock() + mock_proxy_logging = MagicMock() + + await db_writer._commit_spend_updates_to_db_with_redis( + prisma_client=mock_prisma_client, + n_retry_times=1, + proxy_logging_obj=mock_proxy_logging, + ) + + # Pipeline method should be called once + mock_redis_update_buffer.get_all_transactions_from_redis_buffer_pipeline.assert_called_once() + + # Individual methods should NOT be called + mock_redis_update_buffer.get_all_update_transactions_from_redis_buffer.assert_not_called() + mock_redis_update_buffer.get_all_daily_spend_update_transactions_from_redis_buffer.assert_not_called() + mock_redis_update_buffer.get_all_daily_team_spend_update_transactions_from_redis_buffer.assert_not_called() + mock_redis_update_buffer.get_all_daily_org_spend_update_transactions_from_redis_buffer.assert_not_called() + mock_redis_update_buffer.get_all_daily_end_user_spend_update_transactions_from_redis_buffer.assert_not_called() + mock_redis_update_buffer.get_all_daily_agent_spend_update_transactions_from_redis_buffer.assert_not_called() + mock_redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer.assert_not_called() From 68a30a39e68d631959920f7c658baaf8bea2df06 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Tue, 24 Feb 2026 14:22:57 -0800 Subject: [PATCH 2/2] fix(proxy): add LPOP pipeline error checking and fix org spend ServiceType - Add per-command error check in _pipeline_lpop_helper to match _pipeline_rpush_helper, preventing silent data loss on WRONGTYPE errors - Fix pre-existing bug: org spend queue metric was using REDIS_DAILY_SPEND_UPDATE_QUEUE instead of REDIS_DAILY_ORG_SPEND_UPDATE_QUEUE - Add test for per-command LPOP pipeline error propagation --- litellm/caching/redis_cache.py | 5 ++++ .../redis_update_buffer.py | 2 +- .../test_litellm/caching/test_redis_cache.py | 30 +++++++++++++++++++ 3 files changed, 36 insertions(+), 1 deletion(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 9b93d890ce..fa9b94bc2a 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1511,6 +1511,11 @@ class RedisCache(BaseCache): raw_results.append(key_results if key_results else None) offset += count + # Raise on per-command errors (matches _pipeline_rpush_helper behavior) + for r in raw_results: + if isinstance(r, Exception): + raise r + # Decode bytes -> str for each result set decoded_results: List[Optional[List[str]]] = [] for r in raw_results: 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 ef8dbb2c7b..504fc5aba0 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -201,7 +201,7 @@ class RedisUpdateBuffer: (db_spend_update_transactions, REDIS_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_SPEND_UPDATE_QUEUE), (daily_spend_update_transactions, REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE), (daily_team_spend_update_transactions, REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE), - (daily_org_spend_update_transactions, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE), + (daily_org_spend_update_transactions, REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_ORG_SPEND_UPDATE_QUEUE), (daily_end_user_spend_update_transactions, REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE), (daily_agent_spend_update_transactions, REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE), (daily_tag_spend_update_transactions, REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE), diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index 31eeb85488..8260651182 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -296,6 +296,36 @@ async def test_async_rpush_pipeline_raises_on_per_command_error(monkeypatch, red await redis_cache.async_rpush_pipeline(rpush_list=rpush_list) +@pytest.mark.asyncio +async def test_async_lpop_pipeline_raises_on_per_command_error(monkeypatch, redis_no_ping): + """Verify that per-command errors in LPOP pipeline results are raised, not silently dropped""" + monkeypatch.setenv("REDIS_HOST", "https://my-test-host") + redis_cache = RedisCache() + redis_cache.redis_version = "7.0.0" + + mock_redis_instance = AsyncMock() + mock_pipeline = MagicMock() + mock_pipeline.__aenter__ = AsyncMock(return_value=mock_pipeline) + mock_pipeline.__aexit__ = AsyncMock(return_value=None) + mock_pipeline.lpop = MagicMock() + # Simulate: first LPOP succeeds, second returns a per-command error + mock_pipeline.execute = AsyncMock( + return_value=[[b"val1"], Exception("WRONGTYPE")] + ) + mock_redis_instance.pipeline = MagicMock(return_value=mock_pipeline) + + from litellm.types.caching import RedisPipelineLpopOperation + + lpop_list = [ + RedisPipelineLpopOperation(key="key1", count=10), + RedisPipelineLpopOperation(key="key2", count=10), + ] + + with patch.object(redis_cache, "init_async_client", return_value=mock_redis_instance): + with pytest.raises(Exception, match="WRONGTYPE"): + await redis_cache.async_lpop_pipeline(lpop_list=lpop_list) + + @pytest.mark.asyncio async def test_async_lpop_pipeline_empty_list(monkeypatch, redis_no_ping): """Empty lpop_list should return empty list without touching Redis"""