fix(router): flush queued spend to Redis before syncing to avoid clobbering memory with stale spend

litellm/router_strategy/budget_limiter.py

@@ -519,33 +519,43 @@ class RouterBudgetLimiting(CustomLogger):
                     DEFAULT_REDIS_SYNC_INTERVAL
                 )  # Still wait DEFAULT_REDIS_SYNC_INTERVAL seconds on error before retrying
 
-    async def _push_in_memory_increments_to_redis(self):
+    async def _push_in_memory_increments_to_redis(self) -> Optional[Dict[str, float]]:
         """
         How this works:
         - async_log_success_event collects all provider spend increments in `redis_increment_operation_queue`
         - This function pushes all increments to Redis in a batched pipeline to optimize performance
 
+        Awaits the pipeline so that once this returns Redis reflects the queued increments, and returns the
+        resulting per-key Redis values so callers can merge them back into memory without clobbering
+        concurrent in-memory increments.
+
         Only runs if Redis is initialized
         """
         try:
             if not self.dual_cache.redis_cache:
-                return  # Redis is not initialized
+                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",
-                self.redis_increment_operation_queue,
+                queue,
             )
-            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,
-                    )
-                )
+            increment_result = await self.dual_cache.redis_cache.async_increment_pipeline(
+                increment_list=queue,
+            )
+            if increment_result is None:
+                return None
 
-            self.redis_increment_operation_queue = []
+            return {op["key"]: float(value) for op, value in zip(queue, increment_result)}
 
         except Exception as e:
             verbose_router_logger.error(f"Error syncing in-memory cache with Redis: {str(e)}")
+            return None
 
     async def _sync_in_memory_spend_with_redis(self):
         """
@@ -565,11 +575,7 @@ class RouterBudgetLimiting(CustomLogger):
             if self.dual_cache.redis_cache is None:
                 return
 
-            # 1. Push all provider spend increments to Redis
-            await self._push_in_memory_increments_to_redis()
-
-            # 2. Fetch all current provider spend from Redis to update in-memory cache
-            cache_keys = []
+            cache_keys: List[str] = []
 
             if self.provider_budget_config is not None:
                 for provider, config in self.provider_budget_config.items():
@@ -589,15 +595,33 @@ class RouterBudgetLimiting(CustomLogger):
                         continue
                     cache_keys.append(f"tag_spend:{tag}:{config.budget_duration}")
 
-            # Batch fetch current spend values from Redis
-            redis_values = await self.dual_cache.redis_cache.async_batch_get_cache(key_list=cache_keys)
+            in_memory_before: List[Optional[float]] = await self.dual_cache.in_memory_cache.async_batch_get_cache(
+                keys=cache_keys
+            )
+            in_memory_before_dict: Dict[str, float] = {
+                key: float(value or 0) for key, value in zip(cache_keys, in_memory_before)
+            }
 
-            # Update in-memory cache with Redis values
-            if isinstance(redis_values, dict):  # Check if redis_values is a dictionary
-                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(f"Updated in-memory cache for {key}: {value}")
+            # 1. Push all provider spend increments to Redis
+            await self._push_in_memory_increments_to_redis()
+
+            # 2. Fetch all current provider spend from Redis to update in-memory cache
+            redis_values: Dict[str, Optional[float]] = await self.dual_cache.redis_cache.async_batch_get_cache(
+                key_list=cache_keys
+            )
+
+            for key in cache_keys:
+                redis_value = redis_values.get(key)
+                if redis_value is None:
+                    continue
+                redis_spend = float(redis_value)
+                before = in_memory_before_dict.get(key, 0.0)
+                current: Optional[float] = await self.dual_cache.in_memory_cache.async_get_cache(key=key)
+                after = float(current or 0)
+                delta = after - before
+                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}")
 
         except Exception as e:
             verbose_router_logger.error(f"Error syncing in-memory cache with Redis: {str(e)}")

tests/test_litellm/router_strategy/test_budget_limiter.py

@@ -0,0 +1,118 @@
+import asyncio
+
+import pytest
+
+from litellm.caching.caching import DualCache
+from litellm.caching.redis_cache import RedisPipelineIncrementOperation
+from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
+from litellm.types.utils import BudgetConfig
+
+
+class FakeRedisCache:
+    """Minimal async Redis stand-in whose increment pipeline yields control before
+    applying, so fire-and-forget scheduling would race a subsequent read."""
+
+    def __init__(self, store):
+        self.store = dict(store)
+
+    async def async_get_cache(self, key, **kwargs):
+        return self.store.get(key)
+
+    async def async_set_cache(self, key, value, **kwargs):
+        self.store[key] = value
+
+    async def async_batch_get_cache(self, key_list, **kwargs):
+        return {key: self.store.get(key) for key in key_list}
+
+    async def async_increment_pipeline(self, increment_list, **kwargs):
+        await asyncio.sleep(0.05)
+        results = []
+        for op in increment_list:
+            self.store[op["key"]] = self.store.get(op["key"], 0.0) + op["increment_value"]
+            results.append(self.store[op["key"]])
+        return results
+
+
+@pytest.fixture
+def disable_budget_sync(monkeypatch):
+    async def noop(*args, **kwargs):
+        return None
+
+    monkeypatch.setattr(
+        "litellm.router_strategy.budget_limiter.RouterBudgetLimiting.periodic_sync_in_memory_spend_with_redis",
+        noop,
+    )
+
+
+async def _make_limiter(fake_redis, provider):
+    limiter = RouterBudgetLimiting(
+        dual_cache=DualCache(redis_cache=fake_redis),
+        provider_budget_config={provider: BudgetConfig(time_period="1d", budget_limit=1000)},
+    )
+    # let the background _init_provider_budget_in_cache tasks settle
+    await asyncio.sleep(0.1)
+    return limiter
+
+
+@pytest.mark.asyncio
+async def test_sync_awaits_pipeline_before_reading_redis(disable_budget_sync):
+    """Regression for #32614: sync must flush queued increments to Redis before
+    reading them back, otherwise in-memory spend is clobbered with stale Redis spend."""
+    provider = "openai"
+    spend_key = f"provider_spend:{provider}:1d"
+
+    fake_redis = FakeRedisCache({spend_key: 100.0})
+    limiter = await _make_limiter(fake_redis, provider)
+
+    # in-memory already reflects the incremented spend (160), the delta (+60) is still queued for Redis
+    await limiter.dual_cache.in_memory_cache.async_set_cache(key=spend_key, value=160.0)
+    fake_redis.store[spend_key] = 100.0
+    limiter.redis_increment_operation_queue = [
+        RedisPipelineIncrementOperation(key=spend_key, increment_value=60.0, ttl=86400)
+    ]
+
+    await limiter._sync_in_memory_spend_with_redis()
+
+    in_memory_spend = await limiter.dual_cache.in_memory_cache.async_get_cache(spend_key)
+    assert float(in_memory_spend) == 160.0
+    assert fake_redis.store[spend_key] == 160.0
+
+
+@pytest.mark.asyncio
+async def test_sync_pulls_in_other_instance_spend(disable_budget_sync):
+    """Redis ahead of memory (another instance spent) should win."""
+    provider = "anthropic"
+    spend_key = f"provider_spend:{provider}:1d"
+
+    fake_redis = FakeRedisCache({spend_key: 100.0})
+    limiter = await _make_limiter(fake_redis, provider)
+
+    await limiter.dual_cache.in_memory_cache.async_set_cache(key=spend_key, value=100.0)
+    fake_redis.store[spend_key] = 250.0
+    limiter.redis_increment_operation_queue = []
+
+    await limiter._sync_in_memory_spend_with_redis()
+
+    in_memory_spend = await limiter.dual_cache.in_memory_cache.async_get_cache(spend_key)
+    assert float(in_memory_spend) == 250.0
+
+
+@pytest.mark.asyncio
+async def test_push_returns_awaited_redis_values(disable_budget_sync):
+    """_push_in_memory_increments_to_redis must await the pipeline and return the
+    resulting per-key Redis values (not schedule a background task)."""
+    provider = "openai"
+    spend_key = f"provider_spend:{provider}:1d"
+
+    fake_redis = FakeRedisCache({spend_key: 10.0})
+    limiter = await _make_limiter(fake_redis, provider)
+
+    limiter.redis_increment_operation_queue = [
+        RedisPipelineIncrementOperation(key=spend_key, increment_value=5.0, ttl=86400)
+    ]
+
+    result = await limiter._push_in_memory_increments_to_redis()
+
+    assert result == {spend_key: 15.0}
+    assert fake_redis.store[spend_key] == 15.0
+    assert limiter.redis_increment_operation_queue == []