From c57edc06157fff54b9e7e1bc1dcf82b5e5515faf Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Tue, 10 Feb 2026 13:41:47 -0600 Subject: [PATCH] fix(router-budget): lock increment queue and stabilize tests --- litellm/router_strategy/budget_limiter.py | 33 +++++-- .../router_strategy/test_budget_limiter.py | 87 ++++++++++++++++--- 2 files changed, 101 insertions(+), 19 deletions(-) diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index c8601c7e1ae..fdd4a45e0e1 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -52,6 +52,7 @@ class RouterBudgetLimiting(CustomLogger): ): self.dual_cache = dual_cache self.redis_increment_operation_queue: List[RedisPipelineIncrementOperation] = [] + self._redis_increment_queue_lock = asyncio.Lock() asyncio.create_task(self.periodic_sync_in_memory_spend_with_redis()) self.provider_budget_config: Optional[ GenericBudgetConfigType @@ -348,7 +349,17 @@ class RouterBudgetLimiting(CustomLogger): increment_value=response_cost, ttl=ttl, ) - self.redis_increment_operation_queue.append(increment_op) + async with self._get_redis_increment_queue_lock(): + self.redis_increment_operation_queue.append(increment_op) + + def _get_redis_increment_queue_lock(self) -> asyncio.Lock: + queue_lock: Optional[asyncio.Lock] = getattr( + self, "_redis_increment_queue_lock", None + ) + if queue_lock is None: + queue_lock = asyncio.Lock() + self._redis_increment_queue_lock = queue_lock + return queue_lock async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """Original method now uses helper functions""" @@ -493,8 +504,9 @@ class RouterBudgetLimiting(CustomLogger): # Snapshot pending increments and clear queue before await, so new writes are queued # for the next sync cycle while this batch is flushed to Redis. - increment_operations_to_flush = self.redis_increment_operation_queue - self.redis_increment_operation_queue = [] + async with self._get_redis_increment_queue_lock(): + increment_operations_to_flush = self.redis_increment_operation_queue + self.redis_increment_operation_queue = [] verbose_router_logger.debug( "Pushing Redis Increment Pipeline for queue: %s", @@ -505,14 +517,17 @@ class RouterBudgetLimiting(CustomLogger): increment_list=increment_operations_to_flush, ) - except Exception as e: + except Exception: if len(increment_operations_to_flush) > 0: # Retry these increments on a future sync cycle. - self.redis_increment_operation_queue = ( - increment_operations_to_flush + self.redis_increment_operation_queue - ) - verbose_router_logger.error( - f"Error syncing in-memory cache with Redis: {str(e)}" + async with self._get_redis_increment_queue_lock(): + self.redis_increment_operation_queue = ( + increment_operations_to_flush + + self.redis_increment_operation_queue + ) + verbose_router_logger.exception( + "Error pushing queued Redis increment operations to Redis", + exc_info=True, ) async def _sync_in_memory_spend_with_redis(self): diff --git a/tests/test_litellm/router_strategy/test_budget_limiter.py b/tests/test_litellm/router_strategy/test_budget_limiter.py index c60ece45efb..f79b2f92bed 100644 --- a/tests/test_litellm/router_strategy/test_budget_limiter.py +++ b/tests/test_litellm/router_strategy/test_budget_limiter.py @@ -1,5 +1,6 @@ import asyncio from types import SimpleNamespace +from typing import Optional import pytest @@ -8,13 +9,27 @@ from litellm.types.utils import BudgetConfig class _MockRedisCache: - def __init__(self, initial_values): + def __init__( + self, + initial_values, + pipeline_started: Optional[asyncio.Event] = None, + allow_pipeline_to_complete: Optional[asyncio.Event] = None, + should_fail_pipeline: bool = False, + ): self.values = initial_values self.events = [] + self.pipeline_started = pipeline_started + self.allow_pipeline_to_complete = allow_pipeline_to_complete + self.should_fail_pipeline = should_fail_pipeline async def async_increment_pipeline(self, increment_list, **kwargs): self.events.append("increment_pipeline:start") - await asyncio.sleep(0.05) + if self.pipeline_started is not None: + self.pipeline_started.set() + if self.allow_pipeline_to_complete is not None: + await self.allow_pipeline_to_complete.wait() + if self.should_fail_pipeline: + raise RuntimeError("redis down") for op in increment_list: key = op["key"] current = float(self.values.get(key, 0.0) or 0.0) @@ -30,6 +45,11 @@ class _MockInMemoryCache: def __init__(self, initial_values): self.values = initial_values + async def async_increment(self, key, value, ttl, **kwargs): + current = float(self.values.get(key, 0.0) or 0.0) + self.values[key] = current + float(value) + return self.values[key] + async def async_set_cache(self, key, value, **kwargs): self.values[key] = float(value) @@ -37,7 +57,13 @@ class _MockInMemoryCache: @pytest.mark.asyncio async def test_should_await_redis_pipeline_before_sync_reads(): spend_key = "provider_spend:openai:1d" - redis_cache = _MockRedisCache(initial_values={spend_key: 100.0}) + pipeline_started = asyncio.Event() + allow_pipeline_to_complete = asyncio.Event() + redis_cache = _MockRedisCache( + initial_values={spend_key: 100.0}, + pipeline_started=pipeline_started, + allow_pipeline_to_complete=allow_pipeline_to_complete, + ) in_memory_cache = _MockInMemoryCache(initial_values={spend_key: 160.0}) budget_limiter = RouterBudgetLimiting.__new__(RouterBudgetLimiting) @@ -57,8 +83,13 @@ async def test_should_await_redis_pipeline_before_sync_reads(): "ttl": 86400, } ] + budget_limiter._redis_increment_queue_lock = asyncio.Lock() - await budget_limiter._sync_in_memory_spend_with_redis() + sync_task = asyncio.create_task(budget_limiter._sync_in_memory_spend_with_redis()) + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + assert "batch_get" not in redis_cache.events + allow_pipeline_to_complete.set() + await sync_task assert redis_cache.values[spend_key] == 160.0 assert in_memory_cache.values[spend_key] == 160.0 @@ -73,22 +104,58 @@ async def test_should_await_redis_pipeline_before_sync_reads(): @pytest.mark.asyncio async def test_should_requeue_increments_when_redis_pipeline_fails(): spend_key = "provider_spend:openai:1d" - - class _FailingRedisCache: - async def async_increment_pipeline(self, increment_list, **kwargs): - raise RuntimeError("redis down") - + redis_cache = _MockRedisCache( + initial_values={}, + should_fail_pipeline=True, + ) budget_limiter = RouterBudgetLimiting.__new__(RouterBudgetLimiting) budget_limiter.dual_cache = SimpleNamespace( - redis_cache=_FailingRedisCache(), + redis_cache=redis_cache, in_memory_cache=SimpleNamespace(), ) budget_limiter.redis_increment_operation_queue = [ {"key": spend_key, "increment_value": 10.0, "ttl": 86400} ] + budget_limiter._redis_increment_queue_lock = asyncio.Lock() await budget_limiter._push_in_memory_increments_to_redis() assert budget_limiter.redis_increment_operation_queue == [ {"key": spend_key, "increment_value": 10.0, "ttl": 86400} ] + + +@pytest.mark.asyncio +async def test_should_keep_new_increments_when_pipeline_flush_fails(): + spend_key = "provider_spend:openai:1d" + pipeline_started = asyncio.Event() + allow_pipeline_to_complete = asyncio.Event() + redis_cache = _MockRedisCache( + initial_values={}, + pipeline_started=pipeline_started, + allow_pipeline_to_complete=allow_pipeline_to_complete, + should_fail_pipeline=True, + ) + in_memory_cache = _MockInMemoryCache(initial_values={spend_key: 0.0}) + budget_limiter = RouterBudgetLimiting.__new__(RouterBudgetLimiting) + budget_limiter.dual_cache = SimpleNamespace( + redis_cache=redis_cache, + in_memory_cache=in_memory_cache, + ) + budget_limiter.redis_increment_operation_queue = [ + {"key": spend_key, "increment_value": 10.0, "ttl": 86400} + ] + budget_limiter._redis_increment_queue_lock = asyncio.Lock() + + push_task = asyncio.create_task(budget_limiter._push_in_memory_increments_to_redis()) + await asyncio.wait_for(pipeline_started.wait(), timeout=1) + await budget_limiter._increment_spend_in_current_window( + spend_key=spend_key, response_cost=20.0, ttl=86400 + ) + allow_pipeline_to_complete.set() + await push_task + + assert budget_limiter.redis_increment_operation_queue == [ + {"key": spend_key, "increment_value": 10.0, "ttl": 86400}, + {"key": spend_key, "increment_value": 20.0, "ttl": 86400}, + ]