From 3d7afa19fb4f3b3b554e06000e5b3a93a2a5f293 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 9 Jul 2026 14:13:53 +0000 Subject: [PATCH] fix(router): flush queued spend to Redis before syncing to avoid clobbering memory with stale spend --- litellm/router_strategy/budget_limiter.py | 70 +++++++---- .../router_strategy/test_budget_limiter.py | 118 ++++++++++++++++++ 2 files changed, 165 insertions(+), 23 deletions(-) create mode 100644 tests/test_litellm/router_strategy/test_budget_limiter.py diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 067f38ab11c..b5e5b93023f 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/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)}") diff --git a/tests/test_litellm/router_strategy/test_budget_limiter.py b/tests/test_litellm/router_strategy/test_budget_limiter.py new file mode 100644 index 00000000000..dc338b6b937 --- /dev/null +++ b/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 == []