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:
Devin AI 2026-07-16 16:14:06 +00:00
parent 507081e706
commit 91e974e199
2 changed files with 91 additions and 17 deletions

View file

@ -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,

View file

@ -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,