diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 67f392bb4b1..b7d1caecb8a 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -146,26 +146,39 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): async def _get_shared_model_spend(self, cache_key: str) -> Optional[float]: """ - Read the shared, cross-replica spend for a model budget cache key. + Read the cross-replica spend for a model budget cache key. - Redis-first so multi-replica admission decisions use the total spend - accumulated across all pods instead of a stale pod-local in-memory value - (DualCache.async_get_cache returns an in-memory hit before consulting - Redis). Falls back to the local in-memory value only when Redis is - unavailable. + Spend is tracked in two places: a pod-local in-memory counter (this + replica's own increments) and a shared Redis counter (the running total + flushed across every replica). Admission must not trust the pod-local + value alone, otherwise N replicas each admit off their own partial spend + and the combined spend blows past the cap (issue #33325). + + Return the larger of the two so the cap is enforced against the true + accumulated spend: Redis dominates once other replicas have flushed, + while the local value still guards the window between this pod's own + increment and its next Redis flush. Redis failures (including an open + circuit breaker) degrade to the local value instead of failing the + request. """ - if self.dual_cache.redis_cache is not None: - try: - result = await self.dual_cache.redis_cache.async_get_cache(key=cache_key) - return float(result) if result is not None else None - except Exception as e: # noqa: BLE001 # redis (incl. open circuit breaker) failures are non-fatal; fall back to in-memory - verbose_proxy_logger.warning( - "_PROXY_VirtualKeyModelMaxBudgetLimiter: Redis GET failed, falling back to in-memory: %s", - str(e), - ) + local_spend = await self.dual_cache.async_get_cache(key=cache_key, local_only=True) + local_value = float(local_spend) if local_spend is not None else None - result = await self.dual_cache.async_get_cache(key=cache_key, local_only=True) - return float(result) if result is not None else None + if self.dual_cache.redis_cache is None: + return local_value + + try: + redis_spend = await self.dual_cache.redis_cache.async_get_cache(key=cache_key) + except Exception as e: # noqa: BLE001 # redis (incl. open circuit breaker) failures are non-fatal; fall back to in-memory + verbose_proxy_logger.warning( + "_PROXY_VirtualKeyModelMaxBudgetLimiter: Redis GET failed, falling back to in-memory: %s", + str(e), + ) + return local_value + + redis_value = float(redis_spend) if redis_spend is not None else None + candidates = tuple(value for value in (local_value, redis_value) if value is not None) + return max(candidates) if candidates else None async def _get_end_user_spend_for_model( self, diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py index e6db9b87db7..adfaa9a878f 100644 --- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py +++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py @@ -527,6 +527,67 @@ async def test_get_shared_model_spend_falls_back_to_in_memory_without_redis( assert spend == 42.0 +@pytest.mark.asyncio +async def test_get_shared_model_spend_falls_back_to_local_when_redis_key_missing(): + """ + Redis is attached but has not seen this key yet (e.g. this replica served + the only requests so far and its flush is still in flight). Admission must + fall back to the pod-local value instead of treating a Redis miss as zero + spend, otherwise a single replica stops enforcing its own budget the moment + Redis is wired in. + """ + cache_key = "virtual_key_spend:test-key:gpt-4:1d" + dual_cache = DualCache() + dual_cache.redis_cache = _SharedRedisCache(store={}) # shared store has no entry + await dual_cache.in_memory_cache.async_set_cache(key=cache_key, value=7.0) + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + + assert await limiter._get_shared_model_spend(cache_key=cache_key) == 7.0 + assert dual_cache.redis_cache.get_calls > 0 + + +@pytest.mark.asyncio +async def test_get_shared_model_spend_returns_max_of_local_and_redis(): + """ + Local (this pod) and Redis (cross-replica total) can disagree. Admission + must enforce against the larger value so the cap holds whether the local + increment or another replica's flush is ahead. + """ + cache_key = "virtual_key_spend:test-key:gpt-4:1d" + + # Redis ahead of local (other replicas already flushed a higher total). + dual_cache = DualCache() + dual_cache.redis_cache = _SharedRedisCache(store={cache_key: 90.0}) + await dual_cache.in_memory_cache.async_set_cache(key=cache_key, value=10.0) + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + assert await limiter._get_shared_model_spend(cache_key=cache_key) == 90.0 + + # Local ahead of Redis (this pod incremented but has not flushed yet). + dual_cache = DualCache() + dual_cache.redis_cache = _SharedRedisCache(store={cache_key: 10.0}) + await dual_cache.in_memory_cache.async_set_cache(key=cache_key, value=90.0) + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + assert await limiter._get_shared_model_spend(cache_key=cache_key) == 90.0 + + +@pytest.mark.asyncio +async def test_get_shared_model_spend_falls_back_to_local_when_redis_raises(): + """A Redis failure (including an open circuit breaker) must degrade to the + local value rather than failing the admission check.""" + cache_key = "virtual_key_spend:test-key:gpt-4:1d" + + class _RaisingRedisCache: + async def async_get_cache(self, key, parent_otel_span=None, **kwargs): + raise ConnectionError("redis down") + + dual_cache = DualCache() + dual_cache.redis_cache = _RaisingRedisCache() + await dual_cache.in_memory_cache.async_set_cache(key=cache_key, value=5.0) + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + + assert await limiter._get_shared_model_spend(cache_key=cache_key) == 5.0 + + @pytest.mark.asyncio async def test_get_fallback_model_within_budget_returns_none_without_fallbacks( budget_limiter,