fix(proxy): give RPM and TPM counters independent rate-limit windows to stop false 429s at boundaries

Virtual keys with both tpm_limit and rpm_limit got throttled below their
configured limits, and once throttled, 429s persisted for minutes past the
window. Both RPM and TPM counters shared one window-start key, so the pass
that rolled the window left the sibling counter's stale previous-window value
stranded under the freshly reset window and counted it against the new window
(false 429 with counter = limit + 1 on the first request after every
boundary). Additionally the local in-memory over-limit pre-check ignored
window expiry, extending lockouts by up to a full window_size.

- each rate-limit type now owns its window key (...🪟requests /
  🪟tokens) so tracked windows roll independently
- is_cache_list_over_limit treats a counter from an expired window as 0
- legacy shared-window keys are orphaned (one benign reset at deploy)

Fixes #24677
This commit is contained in:
Sisyphus 2026-08-27 23:27:24 +08:00
parent 75736323e6
commit d806d3b776
2 changed files with 316 additions and 16 deletions

View file

@ -390,6 +390,22 @@ class _AsyncLuaScript(Protocol):
"""A Lua script registered against the async Redis client, called with KEYS and ARGV."""
def __call__(self, *, keys: Sequence[str], args: Sequence[object]) -> Awaitable[list[CacheCounterValue]]: ...
WindowKeyType: TypeAlias = Literal["requests", "tokens"]
def descriptor_window_key(descriptor_key: str, descriptor_value: str, rate_limit_type: WindowKeyType) -> str:
"""
Window-start key for a single descriptor counter.
Each rate-limit type owns its window key. Sharing one window key between
the requests and tokens counters made a window roll on one pair leave the
sibling counter's stale previous-window value stranded under a freshly
reset window, which was then counted against the new window — guaranteed
false 429s on the first request after every window boundary (issue #24677).
With per-type windows, every tracked window rolls independently and no
branch can skip a sibling counter's reset.
"""
return f"{{{descriptor_key}:{descriptor_value}}}:window:{rate_limit_type}"
class RateLimitDescriptorRateLimitObject(TypedDict, total=False):
@ -1071,9 +1087,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
keys_to_fetch: list[str],
cache_values: CacheCounterValues,
key_metadata: dict[str, WindowKeyMetadata],
now_int: int | None = None,
) -> RateLimitResponse:
"""
Check if the cache values are over the limit.
``now_int`` enables window-aware evaluation: a counter snapshot whose
window_start shows its window has already expired is treated as 0.
The local in-memory mirror written by a previous request can hold an
over-limit value from a window that has since rolled; rejecting on it
(before the authoritative pass rolls the window) produced 429s that
outlived their window by up to a full window_size (#24677).
"""
statuses: Final[list[RateLimitStatus]] = []
overall_code = "OK"
@ -1086,6 +1110,19 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
requests_limit = key_metadata[window_key]["requests_limit"]
tokens_limit = key_metadata[window_key]["tokens_limit"]
if now_int is not None and counter_value is not None:
window_start = cache_values[i]
if window_start is not None:
try:
window_expired = (now_int - int(window_start)) >= key_metadata[window_key]["window_size"]
except (TypeError, ValueError):
window_expired = False
if window_expired:
# This counter belongs to a window that has already
# rolled over; it must not reject the request that
# starts the new window.
counter_value = 0
# Determine which limit to use for current_limit and limit_remaining
current_limit: int | None = None
rate_limit_type: Literal["requests", "tokens", "max_parallel_requests"] | None = None
@ -1288,7 +1325,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
if cache_values is not None:
rate_limit_response: Final = self.is_cache_list_over_limit(keys_to_fetch, cache_values, key_metadata)
rate_limit_response: Final = self.is_cache_list_over_limit(
keys_to_fetch, cache_values, key_metadata, now_int=now_int
)
if rate_limit_response["overall_code"] == "OVER_LIMIT":
return rate_limit_response
@ -1304,7 +1343,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# For keys that don't exist yet, set them to 0
if cache_values is None:
cache_values = [ # rebind-ok: missing keys default to a zeroed window snapshot
str(now_int) if key.endswith(":window") else 0 for key in keys_to_fetch
str(now_int) if ":window" in key else 0 for key in keys_to_fetch
]
elif self.batch_rate_limiter_script is not None:
# NORMAL MODE: Increment counters in Redis
@ -1342,7 +1381,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
window_size=self.window_size,
)
windowed_response = self.is_cache_list_over_limit(keys_to_fetch, cache_values, key_metadata)
windowed_response = self.is_cache_list_over_limit(
keys_to_fetch, cache_values, key_metadata, now_int=now_int
)
if windowed_response["overall_code"] == "OVER_LIMIT":
return windowed_response
@ -1384,8 +1425,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
max_parallel_requests_limit = rate_limit.get("max_parallel_requests")
window_size = rate_limit.get("window_size") or self.window_size
window_key = f"{{{descriptor_key}:{descriptor_value}}}:window"
if max_parallel_requests_limit is not None:
gauges.append(
ParallelRequestGauge(
@ -1398,24 +1437,44 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
rate_limit_set = False
requests_window_key: str | None = (
descriptor_window_key(descriptor_key, descriptor_value, "requests")
if requests_limit is not None
else None
)
tokens_window_key: str | None = (
descriptor_window_key(descriptor_key, descriptor_value, "tokens") if tokens_limit is not None else None
)
if requests_limit is not None:
rpm_key = self.create_rate_limit_keys(descriptor_key, descriptor_value, "requests")
keys_to_fetch.extend([window_key, rpm_key])
keys_to_fetch.extend((requests_window_key, rpm_key))
rate_limit_set = True
if tokens_limit is not None:
tpm_key = self.create_rate_limit_keys(descriptor_key, descriptor_value, "tokens")
keys_to_fetch.extend([window_key, tpm_key])
keys_to_fetch.extend((tokens_window_key, tpm_key))
rate_limit_set = True
if not rate_limit_set:
continue
key_metadata[window_key] = {
"requests_limit": (int(requests_limit) if requests_limit is not None else None),
"tokens_limit": int(tokens_limit) if tokens_limit is not None else None,
"window_size": int(window_size),
"descriptor_key": descriptor_key,
}
# Per-type window metadata: each counter is evaluated against its
# own window, so the sibling limit stays None for the other type.
if requests_window_key is not None:
requests_metadata: WindowKeyMetadata = {
"requests_limit": int(requests_limit),
"tokens_limit": None,
"window_size": int(window_size),
"descriptor_key": descriptor_key,
}
key_metadata[requests_window_key] = requests_metadata
if tokens_window_key is not None:
tokens_metadata: WindowKeyMetadata = {
"requests_limit": None,
"tokens_limit": int(tokens_limit),
"window_size": int(window_size),
"descriptor_key": descriptor_key,
}
key_metadata[tokens_window_key] = tokens_metadata
return keys_to_fetch, key_metadata, gauges
def _gauge_status(self, gauge: ParallelRequestGauge, in_flight: int, code: str) -> RateLimitStatus:
@ -1736,7 +1795,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
descriptor.get("rate_limit") or RateLimitDescriptorRateLimitObject()
)
window_size: Final = rate_limit.get("window_size") or self.window_size
window_key: Final = f"{{{descriptor_key}:{descriptor_value}}}:window"
keys: Final[list[str]] = []
args: Final[list[int]] = []
@ -1753,6 +1811,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if limit_value is None or inc_amount < 0:
continue
counter_key = self.create_rate_limit_keys(descriptor_key, descriptor_value, rlt)
# Per-type window key: the requests and tokens counters roll their
# windows independently, so a roll on one can never leave the
# sibling counter's stale value stranded under a fresh window
# (issue #24677).
window_key: Final = descriptor_window_key(descriptor_key, descriptor_value, rlt)
# Counter-key TTL and window_size are conceptually distinct
# ("how long the counter Redis key lives" vs "how long the
# sliding window is"). Kept as separate values so a future
@ -4285,7 +4348,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
key=counter_key,
increment_value=increment,
ttl=self.window_size,
window_key=f"{{{scope_key}:{scope_value}}}:window",
# The guarded increment is tokens-only, so it is pinned to the
# tokens window that reserve_io_tokens observed.
window_key=descriptor_window_key(scope_key, scope_value, "tokens"),
expected_window_start=window_identity[0],
reservation_backend=window_identity[1],
)

View file

@ -255,7 +255,7 @@ async def test_rate_limiter_script_return_values_v3(monkeypatch, time_controller
)
# Verify both counter and window values are stored in cache
window_key = f"{{api_key:{_api_key}}}:window"
window_key = f"{{api_key:{_api_key}}}:window:requests"
counter_key = f"{{api_key:{_api_key}}}:requests"
window_value = await local_cache.async_get_cache(key=window_key)
@ -6171,3 +6171,238 @@ async def test_success_hook_leaves_stash_untouched_for_non_batch_responses():
data={}, user_api_key_dict=user, response=ModelResponse(usage=Usage(total_tokens=5))
)
assert get_request_stash().batch_enqueued_reservation == reservation
############################################################
# Issue #24677 regression tests: false 429s at window boundaries.
#
# Root cause: the requests and tokens counters of one descriptor shared a
# single window-start key. A window roll triggered by one counter (the RPM
# pair in the batch pass) left the sibling counter's stale previous-window
# value stranded under the freshly reset window, so the stale value was
# counted against the new window — guaranteed false 429s on the first
# request after every boundary, and window-blind local pre-checks kept
# rejecting for up to a full extra window after a genuine limit hit.
#
# The requests below are spaced like steady production traffic (last write
# well before the boundary) so the previous window's counters are still
# cached when the window rolls, matching the conditions under which the
# bug manifested.
############################################################
@pytest.mark.asyncio
async def test_tpm_counter_resets_at_boundary_with_both_limits_reservation_disabled(
monkeypatch, time_controller
):
"""
Reservation disabled + BOTH rpm and tpm limits: crossing the window
boundary must reset BOTH counters.
Pre-fix, the batch pass reset the shared window on the requests pair and
the tokens pair then saw the fresh window and incremented its stale
value (3 -> 4 > limit 3), so the boundary request got a false 429 that
persisted until the next boundary.
"""
monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60")
monkeypatch.setenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", "false")
_api_key = hash_token("sk-24677-reservation-disabled")
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, rpm_limit=1000, tpm_limit=3)
local_cache = DualCache()
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache),
time_provider=time_controller.now,
)
assert parallel_request_handler.tpm_reservation_enabled is False
# Pin the fake clock so the boundary crossing lands in the same sliver
# production hits: the window-start value has rolled over (integer
# seconds), while the previous window's counters are still cached.
time_controller._current = datetime(2026, 1, 1, 0, 0, 0, 500000)
tokens_key = parallel_request_handler.create_rate_limit_keys("api_key", _api_key, "tokens")
requests_key = parallel_request_handler.create_rate_limit_keys("api_key", _api_key, "requests")
# Steady traffic: one request every 10s fills the TPM window to exactly
# its limit (3 requests, +1 token each); last write at t=20.
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type=""
)
for _ in range(2):
time_controller.advance(10)
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type=""
)
tokens_value = await local_cache.async_get_cache(key=tokens_key)
assert int(tokens_value) == 3, "TPM counter should sit exactly at its limit before the boundary"
# Cross the window boundary (t=59.75, before the counters' TTL lapses):
# the new window must start empty.
time_controller.advance(39.75)
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type=""
)
tokens_value = await local_cache.async_get_cache(key=tokens_key)
requests_value = await local_cache.async_get_cache(key=requests_key)
assert int(tokens_value) == 1, "TPM counter must reset when its window rolls; stale value must not carry over"
assert int(requests_value) == 1, "RPM counter must reset when its window rolls"
# Recovery is immediate, not deferred to the next boundary.
time_controller.advance(0.1)
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type=""
)
tokens_value = await local_cache.async_get_cache(key=tokens_key)
assert int(tokens_value) == 2
@pytest.mark.asyncio
async def test_tpm_reservation_not_poisoned_by_rpm_window_roll(monkeypatch, time_controller):
"""
Reservation enabled (default) + BOTH rpm and tpm limits: the atomic TPM
reservation must not be charged the stale tokens counter left behind by
the RPM pass's window roll.
Pre-fix, the RPM pass rolled the shared window; reserve_tpm_tokens then
read the fresh window together with the previous window's token count
(186) and rejected the boundary request even though the new window was
empty. With per-type windows the tokens window rolls in the reservation
pass itself, so the boundary request is accounted to the new window.
"""
monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60")
_api_key = hash_token("sk-24677-reservation-enabled")
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, rpm_limit=1000, tpm_limit=186)
local_cache = DualCache()
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache),
time_provider=time_controller.now,
)
assert parallel_request_handler.tpm_reservation_enabled is True
# Pin the fake clock so the boundary crossing lands in the same sliver
# production hits: the tokens window-start value has rolled over, while
# the previous window's tokens counter is still cached.
time_controller._current = datetime(2026, 1, 1, 0, 0, 0, 500000)
descriptors = parallel_request_handler._create_rate_limit_descriptors(
user_api_key_dict=user_api_key_dict,
data={},
rpm_limit_type=None,
tpm_limit_type=None,
model_has_failures=False,
)
tokens_key = parallel_request_handler.create_rate_limit_keys("api_key", _api_key, "tokens")
# Steady traffic: six 31-token reservations exactly fill the 186-token
# window; last write at t=50.
for _ in range(6):
if _ > 0:
time_controller.advance(10)
response = await parallel_request_handler.should_rate_limit(
descriptors=descriptors, skip_tpm_check=True
)
assert response["overall_code"] == "OK"
tpm_response = await parallel_request_handler.reserve_tpm_tokens(
descriptors=descriptors, estimated_tokens=31
)
assert tpm_response["overall_code"] == "OK"
tokens_value = await local_cache.async_get_cache(key=tokens_key)
assert int(tokens_value) == 186, "TPM counter should sit exactly at its limit before the boundary"
# First request after the window boundary (t=59.75, before the tokens
# counter's TTL lapses) must be admitted into the new window.
time_controller.advance(9.75)
response = await parallel_request_handler.should_rate_limit(descriptors=descriptors, skip_tpm_check=True)
assert response["overall_code"] == "OK"
tpm_response = await parallel_request_handler.reserve_tpm_tokens(descriptors=descriptors, estimated_tokens=31)
assert tpm_response["overall_code"] == "OK", (
"First reservation after the window boundary must not be limited by the previous window's token count"
)
tokens_value = await local_cache.async_get_cache(key=tokens_key)
assert int(tokens_value) == 31, "Boundary request must be accounted to the new window"
@pytest.mark.asyncio
async def test_over_limit_mirror_does_not_reject_after_window_roll(monkeypatch, time_controller):
"""
The local in-memory over-limit pre-check must ignore a counter snapshot
whose window has already rolled. A genuine over-limit rejection froze an
over-limit value in the local mirror, and the window-blind pre-check then
kept returning 429s for up to a full extra window_size after the boundary.
"""
monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "60")
_api_key = hash_token("sk-24677-stale-mirror")
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key, rpm_limit=2)
local_cache = DualCache()
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache),
time_provider=time_controller.now,
)
requests_key = parallel_request_handler.create_rate_limit_keys("api_key", _api_key, "requests")
# Steady traffic: two requests admitted (t=0, t=10), third at t=20
# genuinely exceeds RPM=2 -> 429, and the local mirror holds the
# over-limit value (3) written at t=20.
for _ in range(2):
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type=""
)
time_controller.advance(10)
with pytest.raises(HTTPException) as exc_info:
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type=""
)
assert exc_info.value.status_code == 429
# After the window boundary (t=59.75, while the stale mirror entry is
# still cached) the key must recover immediately instead of being
# rejected from the stale local mirror.
time_controller.advance(39.75)
await parallel_request_handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict, cache=local_cache, data={}, call_type=""
)
requests_value = await local_cache.async_get_cache(key=requests_key)
assert int(requests_value) == 1, "Counter must reset to 1 for the new window's first request"
def test_is_cache_list_over_limit_ignores_counter_from_expired_window(time_controller):
"""
Unit test for the window-aware pre-check: a counter snapshot paired with
a window_start older than window_size is treated as 0, never as
over-limit.
"""
local_cache = DualCache()
parallel_request_handler = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(local_cache),
time_provider=time_controller.now,
)
window_key = "{api_key:sk-24677-unit}:window:requests"
counter_key = "{api_key:sk-24677-unit}:requests"
keys_to_fetch = [window_key, counter_key]
key_metadata = {
window_key: {
"requests_limit": 2,
"tokens_limit": None,
"window_size": 60,
"descriptor_key": "api_key",
}
}
now_int = int(time_controller.now().timestamp())
# Stale window (started 60s ago) with an over-limit counter value.
response = parallel_request_handler.is_cache_list_over_limit(keys_to_fetch, ["0", 3], key_metadata, now_int=now_int)
assert response["overall_code"] == "OK"
assert response["statuses"][0]["limit_remaining"] == 2
# Same over-limit counter inside a live window is still over-limit.
response = parallel_request_handler.is_cache_list_over_limit(
keys_to_fetch, [str(now_int), 3], key_metadata, now_int=now_int
)
assert response["overall_code"] == "OVER_LIMIT"
# Without now_int the legacy (window-blind) behavior is preserved.
response = parallel_request_handler.is_cache_list_over_limit(keys_to_fetch, ["0", 3], key_metadata)
assert response["overall_code"] == "OVER_LIMIT"