scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix(rate-limiting): refund earlier admissions when a later one raises or is cancelled

Browse source 
_atomic_check_and_increment's refund loop only ran on a normal
rejection return, not when a later key's own admission raised (a
transient Redis error) or the coroutine was cancelled mid-call.
Everything admitted earlier in that batch stayed permanently charged --
for concurrency, a leaked reservation the caller never releases,
incorrectly throttling that tag for up to the 1-hour safety TTL. A
try/finally now refunds every earlier admission before the exception
propagates, covering both raised exceptions and cancellation uniformly.
This commit is contained in:
Deepanshu 2026-08-17 15:34:39 -04:00
parent 43654422f5
commit c6f940312c
2 changed files with 64 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

37
litellm/proxy/hooks/tag_rate_limiter.py
View file

@ -815,6 +815,14 @@ class _PROXY_TagRateLimiter( # pyright: ignore[reportUnusedClass] # only refer
a clean rejection into an unhandled exception. A skipped refund
self-heals via the key's TTL -- see `_ttl_for`.
A later key's own admission raising (a transient Redis error, or
this coroutine being cancelled mid-call, e.g. the caller
disconnecting) is treated the same as a normal rejection for refund
purposes: every earlier admission in this batch is refunded via the
`finally` block below before the exception propagates, so a
mid-batch infra failure can't leave a permanently-charged counter or
a leaked concurrency reservation behind for the rest of that key's TTL.
Returns (failing_index, values). On success, failing_index is None
and values holds each key's new post-increment value, same order as
`checks`. On rejection, failing_index is the 0-based index of the
@ -830,20 +838,35 @@ class _PROXY_TagRateLimiter( # pyright: ignore[reportUnusedClass] # only refer
# this can't be expressed as a one-shot comprehension.
admitted_values: Final = [] # mutable-ok: sequential async accumulator, discardable on early rejection; see comment above
for index, (cache, key, limit, increment, ttl) in enumerate(checks):
admitted, value = await self._check_and_increment_one(cache, key, limit, increment, ttl)
admitted = False
try:
admitted, value = await self._check_and_increment_one(cache, key, limit, increment, ttl)
finally:
# Runs on a normal rejection (admitted stays False) and on
# any exception/cancellation from the awaited call above
# (admitted never gets assigned, so it's still the False set
# just before the try) -- either way, everything admitted so
# far in this batch must be refunded before this key's own
# outcome is used.
if not admitted:
await self._refund_admitted(checks, up_to_index=index)
if admitted:
admitted_values.append(value) # mutable-ok: see accumulator comment above
continue
for refund_index in range(index):
refund_cache, refund_key, _limit, refund_increment, _ttl = checks[refund_index]
try:
await self._decrement_floor_zero(refund_cache, refund_key, -refund_increment)
except Exception as e: # noqa: BLE001 - one failed refund must not block refunding the rest
verbose_proxy_logger.warning("tag_rate_limiter: failed to refund %s on rollback: %s", refund_key, e)
return index, (value,)
return None, tuple(admitted_values)
async def _refund_admitted(
self, checks: Sequence[tuple[InternalUsageCache, str, float, float, int]], up_to_index: int
) -> None:
for refund_index in range(up_to_index):
refund_cache, refund_key, _limit, refund_increment, _ttl = checks[refund_index]
try:
await self._decrement_floor_zero(refund_cache, refund_key, -refund_increment)
except Exception as e: # noqa: BLE001 - one failed refund must not block refunding the rest
verbose_proxy_logger.warning("tag_rate_limiter: failed to refund %s on rollback: %s", refund_key, e)
async def async_filter_deployments(
self,
model: str,

34
tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py
View file

@ -1954,6 +1954,40 @@ async def test_refund_failure_on_one_key_does_not_block_others_or_raise(time_con
assert (float(other_value) if other_value is not None else 0.0) == 0.0
@pytest.mark.asyncio
async def test_exception_mid_batch_refunds_every_earlier_admission_before_propagating(time_controller):
"""
Regression test: a later key's own admission raising (a transient Redis
error, or this coroutine being cancelled mid-call) used to skip the
refund loop entirely, since it only ran on a normal rejection return.
An earlier admission in the same batch would then stay permanently
charged -- for concurrency, a leaked reservation the caller never gets
to release, incorrectly throttling that tag until the 1-hour safety TTL
expires.
"""
admitted_key = "{tag_rl:test:exception-refund:a}:requests"
raising_key = "{tag_rl:test:exception-refund:b}:requests"
class _FlakyLimiter(_PROXY_TagRateLimiter):
async def _check_and_increment_one(self, cache, key: str, limit: float, increment: float, ttl: int):
if key == raising_key:
raise RuntimeError("simulated transient redis failure")
return await super()._check_and_increment_one(cache, key, limit, increment, ttl)
flaky = _FlakyLimiter(internal_usage_cache=DualCache(), time_provider=time_controller.now)
with pytest.raises(RuntimeError):
await flaky._atomic_check_and_increment(
[
(flaky.internal_usage_cache, admitted_key, 10.0, 1.0, 60),
(flaky.internal_usage_cache, raising_key, 10.0, 1.0, 60),
]
)
admitted_value = await flaky.internal_usage_cache.async_get_cache(key=admitted_key, litellm_parent_otel_span=None)
assert (float(admitted_value) if admitted_value is not None else 0.0) == 0.0
# ---------------------------------------------------------------------------
# scope_by_key_hash -- opt-in per-calling-key bucket separation
# ---------------------------------------------------------------------------