fix(router): await budget redis pipeline before sync reads

litellm/proxy/hooks/model_max_budget_limiter.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 from typing import Final
 
@@ -28,6 +29,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
     def __init__(self, dual_cache: DualCache):
         self.dual_cache = dual_cache
         self.redis_increment_operation_queue = []
+        self._redis_increment_queue_lock = asyncio.Lock()
         self.deployment_budget_config = None
 
     async def is_key_within_model_budget(

litellm/router_strategy/budget_limiter.py

@@ -100,6 +100,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: GenericBudgetConfigType | None = provider_budget_config
         self.deployment_budget_config: GenericBudgetConfigType | None = None
@@ -393,7 +394,11 @@ 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:
+        return self._redis_increment_queue_lock
 
     async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
         """Original method now uses helper functions"""
@@ -527,25 +532,31 @@ class RouterBudgetLimiting(CustomLogger):
 
         Only runs if Redis is initialized
         """
+        increment_operations_to_flush: list[RedisPipelineIncrementOperation] = []
         try:
             if not self.dual_cache.redis_cache:
                 return  # Redis is not initialized
 
+            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",
-                self.redis_increment_operation_queue,
+                increment_operations_to_flush,
             )
-            if len(self.redis_increment_operation_queue) > 0:
-                asyncio.create_task(
-                    self.dual_cache.redis_cache.async_increment_pipeline(
-                        increment_list=self.redis_increment_operation_queue,
-                    )
+            if len(increment_operations_to_flush) > 0:
+                await self.dual_cache.redis_cache.async_increment_pipeline(
+                    increment_list=increment_operations_to_flush,
                 )
 
-            self.redis_increment_operation_queue = []
-
-        except Exception as e:
-            verbose_router_logger.error("Error syncing in-memory cache with Redis: %s", e)
+        except Exception:
+            if len(increment_operations_to_flush) > 0:
+                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")
 
     async def _sync_in_memory_spend_with_redis(self):
         """

tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py

@@ -14,6 +14,7 @@ from litellm.proxy.hooks.model_max_budget_limiter import (
     _PROXY_VirtualKeyModelMaxBudgetLimiter,
 )
 from litellm.proxy._types import UserAPIKeyAuth
+from litellm.types.caching import RedisPipelineIncrementOperation
 from litellm.types.utils import BudgetConfig as GenericBudgetInfo
 
 
@@ -452,6 +453,23 @@ async def test_async_log_success_event_pushes_redis_increments_when_redis_config
             mock_push.assert_awaited_once()
 
 
+@pytest.mark.asyncio
+async def test_model_budget_limiter_initializes_redis_increment_queue_lock():
+    dual_cache = DualCache()
+    limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache)
+    spend_key = "virtual_key_spend:test-key:gpt-4:1d"
+
+    await limiter._increment_spend_in_current_window(
+        spend_key=spend_key, response_cost=0.01, ttl=86400
+    )
+
+    assert limiter.redis_increment_operation_queue == [
+        RedisPipelineIncrementOperation(
+            key=spend_key, increment_value=0.01, ttl=86400
+        )
+    ]
+
+
 @pytest.mark.asyncio
 async def test_get_fallback_model_within_budget_returns_none_without_fallbacks(
     budget_limiter,

tests/test_litellm/router_strategy/test_budget_limiter.py

@@ -0,0 +1,160 @@
+import asyncio
+from types import SimpleNamespace
+from typing import Optional
+
+import pytest
+
+from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
+from litellm.types.caching import RedisPipelineIncrementOperation
+from litellm.types.utils import BudgetConfig
+
+
+class _MockRedisCache:
+    def __init__(
+        self,
+        initial_values: dict[str, float],
+        pipeline_started: Optional[asyncio.Event] = None,
+        allow_pipeline_to_complete: Optional[asyncio.Event] = None,
+        should_fail_pipeline: bool = False,
+    ) -> 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
+
+    async def async_increment_pipeline(
+        self, increment_list: list[RedisPipelineIncrementOperation], **kwargs: object
+    ) -> None:
+        self.events.append("increment_pipeline:start")
+        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)
+            self.values[key] = current + float(op["increment_value"])
+        self.events.append("increment_pipeline:done")
+
+    async def async_batch_get_cache(self, key_list: list[str], **kwargs: object) -> dict[str, Optional[float]]:
+        self.events.append("batch_get")
+        return {key: self.values.get(key) for key in key_list}
+
+
+class _MockInMemoryCache:
+    def __init__(self, initial_values: dict[str, float]) -> None:
+        self.values = initial_values
+
+    async def async_increment(self, key: str, value: float, ttl: int, **kwargs: object) -> float:
+        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: str, value: float, **kwargs: object) -> None:
+        self.values[key] = float(value)
+
+
+@pytest.mark.asyncio
+async def test_should_await_redis_pipeline_before_sync_reads() -> None:
+    spend_key = "provider_spend:openai:1d"
+    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)
+    budget_limiter.dual_cache = SimpleNamespace(
+        redis_cache=redis_cache,
+        in_memory_cache=in_memory_cache,
+    )
+    budget_limiter.provider_budget_config = {"openai": BudgetConfig(time_period="1d", budget_limit=500.0)}
+    budget_limiter.deployment_budget_config = None
+    budget_limiter.tag_budget_config = None
+    budget_limiter.redis_increment_operation_queue = [
+        RedisPipelineIncrementOperation(
+            key=spend_key,
+            increment_value=60.0,
+            ttl=86400,
+        )
+    ]
+    budget_limiter._redis_increment_queue_lock = asyncio.Lock()
+
+    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
+    assert budget_limiter.redis_increment_operation_queue == []
+    assert redis_cache.events == [
+        "increment_pipeline:start",
+        "increment_pipeline:done",
+        "batch_get",
+    ]
+
+
+@pytest.mark.asyncio
+async def test_should_requeue_increments_when_redis_pipeline_fails() -> None:
+    spend_key = "provider_spend:openai:1d"
+    redis_cache = _MockRedisCache(
+        initial_values={},
+        should_fail_pipeline=True,
+    )
+    budget_limiter = RouterBudgetLimiting.__new__(RouterBudgetLimiting)
+    budget_limiter.dual_cache = SimpleNamespace(
+        redis_cache=redis_cache,
+        in_memory_cache=SimpleNamespace(),
+    )
+    budget_limiter.redis_increment_operation_queue = [
+        RedisPipelineIncrementOperation(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 == [
+        RedisPipelineIncrementOperation(key=spend_key, increment_value=10.0, ttl=86400)
+    ]
+
+
+@pytest.mark.asyncio
+async def test_should_keep_new_increments_when_pipeline_flush_fails() -> None:
+    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 = [
+        RedisPipelineIncrementOperation(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 == [
+        RedisPipelineIncrementOperation(key=spend_key, increment_value=10.0, ttl=86400),
+        RedisPipelineIncrementOperation(key=spend_key, increment_value=20.0, ttl=86400),
+    ]