From 6cb7a6f6009e2265f303fa5d1a1b5de530d25db2 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 15 Sep 2026 12:55:35 -0500 Subject: [PATCH] fix(router): preserve concurrent spend during Redis synchronization --- litellm/router_strategy/budget_limiter.py | 40 ++++++++++------ .../router_strategy/test_budget_limiter.py | 48 ++++++++++++++++++- 2 files changed, 73 insertions(+), 15 deletions(-) diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 2838c3c6509..6fd2a2534ce 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -388,17 +388,17 @@ class RouterBudgetLimiting(CustomLogger): - Increments the spend in memory cache (so spend instantly updated in memory) - Queues the increment operation to Redis Pipeline (using batched pipeline to optimize performance. Using Redis for multi instance environment of LiteLLM) """ - await self.dual_cache.in_memory_cache.async_increment( - key=spend_key, - value=response_cost, - ttl=ttl, - ) increment_op: Final = RedisPipelineIncrementOperation( key=spend_key, increment_value=response_cost, ttl=ttl, ) async with self._get_redis_increment_queue_lock(): + await self.dual_cache.in_memory_cache.async_increment( + key=spend_key, + value=response_cost, + ttl=ttl, + ) self.redis_increment_operation_queue.append(increment_op) def _get_redis_increment_queue_lock(self) -> asyncio.Lock: @@ -635,18 +635,21 @@ class RouterBudgetLimiting(CustomLogger): # No need to sync if Redis cache is not initialized if self.dual_cache.redis_cache is None: return - await self._flush_increments_then_copy_redis_spend() + async with self._redis_increment_flush_lock: + try: + await self._flush_increments_then_copy_redis_spend() + except asyncio.CancelledError: + await asyncio.shield(self._requeue_detached_increment_operations()) + raise except Exception as e: log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) async def _flush_increments_then_copy_redis_spend(self) -> None: - if not await self._push_in_memory_increments_to_redis(): + redis_cache: Final = self.dual_cache.redis_cache + if redis_cache is None or not await self._flush_queued_increment_operations(redis_cache): return cache_keys: Final = [] - redis_cache: Final = self.dual_cache.redis_cache - if redis_cache is None: - return if self.provider_budget_config is not None: for provider, config in self.provider_budget_config.items(): @@ -668,11 +671,20 @@ class RouterBudgetLimiting(CustomLogger): redis_values: Final = await redis_cache.async_batch_get_cache(key_list=cache_keys) - if isinstance(redis_values, dict): + if not isinstance(redis_values, dict): + return + async with self._get_redis_increment_queue_lock(): for key, value in redis_values.items(): - if value is not None: - await self.dual_cache.in_memory_cache.async_set_cache(key=key, value=float(value)) - verbose_router_logger.debug("Updated in-memory cache for %s: %s", key, value) + if value is None: + continue + pending_spend: Final = sum( + operation["increment_value"] + for operation in self.redis_increment_operation_queue + if operation["key"] == key + ) + updated_spend: Final = float(value) + pending_spend + await self.dual_cache.in_memory_cache.async_set_cache(key=key, value=updated_spend) + verbose_router_logger.debug("Updated in-memory cache for %s: %s", key, updated_spend) def _get_budget_config_for_deployment( self, diff --git a/tests/test_litellm/router_strategy/test_budget_limiter.py b/tests/test_litellm/router_strategy/test_budget_limiter.py index bb0b72f757d..29a646e963f 100644 --- a/tests/test_litellm/router_strategy/test_budget_limiter.py +++ b/tests/test_litellm/router_strategy/test_budget_limiter.py @@ -21,12 +21,16 @@ class _MockRedisCache: pipeline_started: asyncio.Event | None = None, allow_pipeline_to_complete: asyncio.Event | None = None, should_fail_pipeline: bool = False, + read_started: asyncio.Event | None = None, + allow_read_to_complete: asyncio.Event | None = None, ) -> None: self.values = initial_values self.events: list[str] = [] self.pipeline_started = pipeline_started self.allow_pipeline_to_complete = allow_pipeline_to_complete self.should_fail_pipeline = should_fail_pipeline + self.read_started = read_started + self.allow_read_to_complete = allow_read_to_complete async def async_increment_pipeline( self, increment_list: list[RedisPipelineIncrementOperation], **kwargs: object @@ -46,7 +50,12 @@ class _MockRedisCache: async def async_batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, float | None]: self.events.append("batch_get") - return {key: self.values.get(key) for key in key_list} + snapshot = {key: self.values.get(key) for key in key_list} + if self.read_started is not None: + self.read_started.set() + if self.allow_read_to_complete is not None: + await self.allow_read_to_complete.wait() + return snapshot class _MockInMemoryCache: @@ -256,3 +265,40 @@ async def test_should_requeue_increments_when_flush_is_cancelled_and_redis_fails assert redis_cache.values[_SPEND_KEY] == 0.0 assert budget_limiter.redis_increment_operation_queue == [_increment(10.0)] assert budget_limiter._detached_increment_operations is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("pause_during", ["write", "read"]) +async def test_sync_preserves_spend_recorded_during_redis_io(pause_during: str) -> None: + io_started = asyncio.Event() + allow_io_to_complete = asyncio.Event() + redis_cache = _MockRedisCache( + initial_values={_SPEND_KEY: 100.0}, + pipeline_started=io_started if pause_during == "write" else None, + allow_pipeline_to_complete=allow_io_to_complete if pause_during == "write" else None, + read_started=io_started if pause_during == "read" else None, + allow_read_to_complete=allow_io_to_complete if pause_during == "read" else None, + ) + in_memory_cache = _MockInMemoryCache(initial_values={_SPEND_KEY: 160.0}) + budget_limiter = _new_router_budget_limiter( + redis_cache=redis_cache, + in_memory_cache=in_memory_cache, + redis_increment_operation_queue=[_increment(60.0)], + provider_budget_config={"openai": BudgetConfig(time_period="1d", budget_limit=175.0)}, + ) + + sync_task = asyncio.create_task(budget_limiter._sync_in_memory_spend_with_redis()) + await asyncio.wait_for(io_started.wait(), timeout=1) + await budget_limiter._increment_spend_in_current_window(_SPEND_KEY, 20.0, 86400) + allow_io_to_complete.set() + await sync_task + + assert in_memory_cache.values[_SPEND_KEY] == 180.0 + assert redis_cache.values[_SPEND_KEY] == 160.0 + assert budget_limiter.redis_increment_operation_queue == [_increment(20.0)] + + await budget_limiter._sync_in_memory_spend_with_redis() + + assert in_memory_cache.values[_SPEND_KEY] == 180.0 + assert redis_cache.values[_SPEND_KEY] == 180.0 + assert budget_limiter.redis_increment_operation_queue == []