fix(router-budget): lock increment queue and stabilize tests

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):

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},
+    ]