mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(model_max_budget): enforce cap against max of local and shared Redis spend
Redis-first read alone regressed the single-replica path: when Redis was attached but had not yet seen the spend key, admission treated the miss as no spend and stopped enforcing the budget. Read both the pod-local counter and the shared Redis counter and enforce against the larger, so the cap holds whether another replica's flush or this pod's own increment is ahead, and Redis failures degrade to the local value.
This commit is contained in:
parent
507081e706
commit
91e974e199
2 changed files with 91 additions and 17 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue