diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index b5e5b93023f..34689a0616a 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -529,34 +529,41 @@ class RouterBudgetLimiting(CustomLogger): resulting per-key Redis values so callers can merge them back into memory without clobbering concurrent in-memory increments. + If the pipeline fails, the snapshot is restored to the queue so the increments are retried on the + next push rather than lost. Returns None in that case. + Only runs if Redis is initialized """ + if not self.dual_cache.redis_cache: + return None # Redis is not initialized + + if len(self.redis_increment_operation_queue) == 0: + return None + + queue = self.redis_increment_operation_queue + self.redis_increment_operation_queue = [] + + verbose_router_logger.debug( + "Pushing Redis Increment Pipeline for queue: %s", + queue, + ) try: - if not self.dual_cache.redis_cache: - return None # Redis is not initialized - - if len(self.redis_increment_operation_queue) == 0: - return None - - queue = self.redis_increment_operation_queue - self.redis_increment_operation_queue = [] - - verbose_router_logger.debug( - "Pushing Redis Increment Pipeline for queue: %s", - queue, - ) increment_result = await self.dual_cache.redis_cache.async_increment_pipeline( increment_list=queue, ) - if increment_result is None: - return None - - return {op["key"]: float(value) for op, value in zip(queue, increment_result)} - except Exception as e: + # Put the snapshot back (ahead of anything queued during the await) so the + # increments are retried on the next push instead of being silently lost; + # otherwise other instances would never see this spend in Redis + self.redis_increment_operation_queue = queue + self.redis_increment_operation_queue verbose_router_logger.error(f"Error syncing in-memory cache with Redis: {str(e)}") return None + if increment_result is None: + return None + + return {op["key"]: float(value) for op, value in zip(queue, increment_result)} + async def _sync_in_memory_spend_with_redis(self): """ Ensures in-memory cache is updated with latest Redis values for all provider spends. @@ -619,6 +626,11 @@ class RouterBudgetLimiting(CustomLogger): current: Optional[float] = await self.dual_cache.in_memory_cache.async_get_cache(key=key) after = float(current or 0) delta = after - before + # When after <= redis_spend, Redis reflects the push and is the authoritative + # cross-instance total; add only the local increments that landed mid-sync (delta). + # When after > redis_spend, Redis is missing local spend (push failed or was + # skipped), so the in-memory total is already authoritative; applying delta on + # top of the stale Redis value would drop pre-snapshot local spend merged = redis_spend + delta if after <= redis_spend else after await self.dual_cache.in_memory_cache.async_set_cache(key=key, value=merged) verbose_router_logger.debug(f"Updated in-memory cache for {key}: {merged}") diff --git a/tests/test_litellm/router_strategy/test_budget_limiter.py b/tests/test_litellm/router_strategy/test_budget_limiter.py index dc338b6b937..333437e2370 100644 --- a/tests/test_litellm/router_strategy/test_budget_limiter.py +++ b/tests/test_litellm/router_strategy/test_budget_limiter.py @@ -33,6 +33,20 @@ class FakeRedisCache: return results +class FlakyRedisCache(FakeRedisCache): + """Fails the increment pipeline while `failures` > 0, then behaves like FakeRedisCache.""" + + def __init__(self, store, failures=0): + super().__init__(store) + self.failures = failures + + async def async_increment_pipeline(self, increment_list, **kwargs): + if self.failures > 0: + self.failures -= 1 + raise ConnectionError("redis unreachable") + return await super().async_increment_pipeline(increment_list, **kwargs) + + @pytest.fixture def disable_budget_sync(monkeypatch): async def noop(*args, **kwargs): @@ -116,3 +130,26 @@ async def test_push_returns_awaited_redis_values(disable_budget_sync): assert result == {spend_key: 15.0} assert fake_redis.store[spend_key] == 15.0 assert limiter.redis_increment_operation_queue == [] + + +@pytest.mark.asyncio +async def test_push_requeues_increments_when_pipeline_fails(disable_budget_sync): + """A failed Redis pipeline must restore the snapshot to the queue so the increments + are retried on the next push instead of being silently lost.""" + provider = "openai" + spend_key = f"provider_spend:{provider}:1d" + + fake_redis = FlakyRedisCache({spend_key: 10.0}) + limiter = await _make_limiter(fake_redis, provider) + + op = RedisPipelineIncrementOperation(key=spend_key, increment_value=5.0, ttl=86400) + limiter.redis_increment_operation_queue = [op] + fake_redis.failures = 1 + + assert await limiter._push_in_memory_increments_to_redis() is None + assert limiter.redis_increment_operation_queue == [op] + assert fake_redis.store[spend_key] == 10.0 + + assert await limiter._push_in_memory_increments_to_redis() == {spend_key: 15.0} + assert limiter.redis_increment_operation_queue == [] + assert fake_redis.store[spend_key] == 15.0