fix(router): preserve concurrent spend during Redis synchronization

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,

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 == []