scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

Merge cd5b3064a3 into 6fa0477d84

Browse source
This commit is contained in:
pnookala-godaddy 2026-11-06 04:58:41 +00:00 • committed by GitHub
commit 47d5c4eafc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 205 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

122
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -16,6 +16,7 @@ from contextlib import asynccontextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field
from datetime import datetime, timezone
from math import ceil
from types import MappingProxyType
from typing import (
TYPE_CHECKING,
@ -317,7 +318,7 @@ CHECK_AND_INCREMENT_BY_N_SCRIPT: Final = """
--
-- Return on success:
-- { 0, new_counter_1, window_start_1, new_counter_2, window_start_2, ... }
-- Return on over-limit: { 1, descriptor_index, current_counter, limit }
-- Return on over-limit: { 1, descriptor_index, current_counter, limit, window_start, redis_now }
local time_reply = redis.call('TIME')
local now = tonumber(time_reply[1])
local descriptor_count = #KEYS / 2
@ -351,7 +352,8 @@ for i = 1, descriptor_count do
blocked = current_counter >= limit
end
if blocked then
return { 1, i, current_counter, limit }
local active_window_start = window_expired and now or tonumber(window_start)
return { 1, i, current_counter, limit, active_window_start, now }
end
descriptor_state[i] = { window_expired, current_counter, window_start }
@ -648,6 +650,9 @@ class RateLimitStatus(TypedDict):
# models' project ITPM/OTPM in one call, or a request carrying multiple
# rate-limited tags, produces statuses sharing the same descriptor_key.
descriptor_value: NotRequired[ReadOnly[str]]
window_start: NotRequired[ReadOnly[float]]
window_size: NotRequired[ReadOnly[int]]
window_observed_at: NotRequired[ReadOnly[float]]
class RateLimitResponse(TypedDict):
@ -1569,6 +1574,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
cache_values: CacheCounterValues,
key_metadata: dict[str, WindowKeyMetadata],
read_only: bool = False,
observed_at: float | None = None,
) -> RateLimitResponse:
"""
Check if the cache values are over the limit.
@ -1581,6 +1587,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
item_code = "OK"
window_key = keys_to_fetch[i]
counter_key = keys_to_fetch[i + 1]
window_start: Final = cache_values[i]
counter_value = cache_values[i + 1]
requests_limit = key_metadata[window_key]["requests_limit"]
tokens_limit = key_metadata[window_key]["tokens_limit"]
@ -1605,15 +1612,32 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
# Only compute limit_remaining if current_limit is not None
limit_remaining = current_limit - int(counter_value) if counter_value is not None else current_limit
if item_code == "OVER_LIMIT" and window_start is not None:
statuses.append(
RateLimitStatus(
code=item_code,
current_limit=current_limit,
limit_remaining=limit_remaining,
rate_limit_type=rate_limit_type,
descriptor_key=key_metadata[window_key]["descriptor_key"],
descriptor_value=key_metadata[window_key]["descriptor_value"],
window_start=float(window_start),
window_size=self.window_size,
window_observed_at=(
observed_at if observed_at is not None else self._get_current_time().timestamp()
),
)
)
continue
statuses.append(
{
"code": item_code,
"current_limit": current_limit,
"limit_remaining": limit_remaining,
"rate_limit_type": rate_limit_type,
"descriptor_key": key_metadata[window_key]["descriptor_key"],
"descriptor_value": key_metadata[window_key]["descriptor_value"],
}
RateLimitStatus(
code=item_code,
current_limit=current_limit,
limit_remaining=limit_remaining,
rate_limit_type=rate_limit_type,
descriptor_key=key_metadata[window_key]["descriptor_key"],
descriptor_value=key_metadata[window_key]["descriptor_value"],
)
)
return RateLimitResponse(overall_code=overall_code, statuses=statuses)
@ -1884,7 +1908,7 @@ 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, read_only=read_only
keys_to_fetch, cache_values, key_metadata, read_only=read_only, observed_at=now
)
if rate_limit_response["overall_code"] == "OVER_LIMIT":
return rate_limit_response
@ -1939,7 +1963,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
windowed_response = self.is_cache_list_over_limit(
keys_to_fetch, cache_values, key_metadata, read_only=read_only
keys_to_fetch, cache_values, key_metadata, read_only=read_only, observed_at=now
)
if windowed_response["overall_code"] == "OVER_LIMIT":
return windowed_response
@ -2679,24 +2703,36 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
status_code: Final = int(raw[0])
if status_code == 1:
# Over limit: { 1, counter_index (1-based), current_counter, limit }
# Over limit: { 1, counter_index (1-based), current_counter, limit, window_start, redis_now }
descriptor_index: Final = int(raw[1]) - 1
current_counter: Final = int(raw[2])
limit: Final = int(raw[3])
meta = per_counter_meta[descriptor_index]
return RateLimitResponse(
overall_code="OVER_LIMIT",
statuses=[
RateLimitStatus(
code="OVER_LIMIT",
current_limit=limit,
limit_remaining=max(0, limit - current_counter),
rate_limit_type=meta["rate_limit_type"],
descriptor_key=meta["descriptor_key"],
descriptor_value=meta["descriptor_value"],
)
],
over_limit_meta: Final = per_counter_meta[descriptor_index]
window_start: Final = float(raw[4]) if len(raw) > 5 else None
window_observed_at: Final = float(raw[5]) if len(raw) > 5 else None
status: Final = (
RateLimitStatus(
code="OVER_LIMIT",
current_limit=limit,
limit_remaining=max(0, limit - current_counter),
rate_limit_type=over_limit_meta["rate_limit_type"],
descriptor_key=over_limit_meta["descriptor_key"],
descriptor_value=over_limit_meta["descriptor_value"],
window_start=window_start,
window_size=over_limit_meta["window_size"],
window_observed_at=window_observed_at,
)
if window_start is not None and window_observed_at is not None
else RateLimitStatus(
code="OVER_LIMIT",
current_limit=limit,
limit_remaining=max(0, limit - current_counter),
rate_limit_type=over_limit_meta["rate_limit_type"],
descriptor_key=over_limit_meta["descriptor_key"],
descriptor_value=over_limit_meta["descriptor_value"],
)
)
return RateLimitResponse(overall_code="OVER_LIMIT", statuses=[status])
statuses: Final[list[RateLimitStatus]] = []
for index, meta in enumerate(per_counter_meta):
@ -2738,7 +2774,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
"""
# Use a single 'now' for the duration of this critical section so all
# descriptors evaluate window expiry consistently.
now_int: Final = int(self._get_current_time().timestamp())
now: Final = self._get_current_time().timestamp()
now_int: Final = int(now)
# Pass 1: read state, validate.
descriptor_state: Final[list[AtomicCounterState]] = []
@ -2776,6 +2813,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
rate_limit_type=meta["rate_limit_type"],
descriptor_key=meta["descriptor_key"],
descriptor_value=meta["descriptor_value"],
window_start=float(now_int if window_expired else int(window_start)),
window_size=meta["window_size"],
window_observed_at=now,
)
],
)
@ -3912,9 +3952,20 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
descriptor_value = matching_descriptor["value"] if matching_descriptor is not None else "unknown"
now = self._get_current_time().timestamp()
reset_time = now + self.window_size
reset_time_formatted = datetime.fromtimestamp(reset_time, tz=timezone.utc).strftime(
window_start: Final = status.get("window_start")
window_size: Final = status.get("window_size")
window_observed_at: Final = status.get("window_observed_at")
reset_time: Final = (
window_start + window_size
if window_start is not None and window_size is not None and window_observed_at is not None
else self._get_current_time().timestamp() + self.window_size
)
retry_after: Final = (
max(1, ceil(reset_time - window_observed_at))
if window_observed_at is not None
else self.window_size
)
reset_time_formatted: Final = datetime.fromtimestamp(reset_time, tz=timezone.utc).strftime(
"%Y-%m-%d %H:%M:%S UTC"
)
@ -3929,14 +3980,15 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
f"Limit resets at: {reset_time_formatted}"
)
headers: Final = {
"retry-after": str(retry_after),
"rate_limit_type": str(status["rate_limit_type"]),
"reset_at": reset_time_formatted,
}
resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(requested_model)
raise ProxyRateLimitError(
detail=detail,
headers={
"retry-after": str(self.window_size),
"rate_limit_type": str(status["rate_limit_type"]),
"reset_at": reset_time_formatted,
},
headers=headers,
rate_limit_type=map_v3_rate_limit_type(status["rate_limit_type"]),
model=resolved_model,
llm_provider=llm_provider,

118
tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -7007,6 +7007,124 @@ async def test_reserve_tpm_tokens_never_evaluates_the_requests_dimension():
assert [s["rate_limit_type"] for s in response["statuses"]] == ["tokens"]
@pytest.mark.asyncio
async def test_rpm_rate_limit_headers_follow_fixed_window_reset(time_controller):
time_controller._current = datetime.fromtimestamp(1_000, tz=timezone.utc)
cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(cache),
time_provider=time_controller.now,
)
auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-reset-rpm"), rpm_limit=1)
await handler.async_pre_call_hook(auth, cache, {}, "")
time_controller.advance(1)
with pytest.raises(HTTPException) as first_denial:
await handler.async_pre_call_hook(auth, cache, {}, "")
time_controller.advance(58)
with pytest.raises(HTTPException) as second_denial:
await handler.async_pre_call_hook(auth, cache, {}, "")
reset_at: Final = datetime.fromtimestamp(1_060, tz=timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
assert first_denial.value.headers["reset_at"] == reset_at
assert second_denial.value.headers["reset_at"] == reset_at
assert first_denial.value.headers["retry-after"] == "59"
assert second_denial.value.headers["retry-after"] == "1"
@pytest.mark.asyncio
async def test_rpm_headers_use_the_enforced_handler_window(time_controller):
time_controller._current = datetime.fromtimestamp(3_000, tz=timezone.utc)
handler: Final = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache()),
time_provider=time_controller.now,
)
descriptor: Final = RateLimitDescriptor(
key="api_key",
value="rpm-handler-window",
rate_limit={"requests_per_unit": 1, "window_size": 90},
)
await handler.should_rate_limit(descriptors=[descriptor])
time_controller.advance(59)
denied: Final = await handler.should_rate_limit(descriptors=[descriptor])
with pytest.raises(HTTPException) as exc_info:
handler._handle_rate_limit_error(denied, [descriptor])
assert exc_info.value.headers["retry-after"] == "1"
assert exc_info.value.headers["reset_at"] == datetime.fromtimestamp(3_060, tz=timezone.utc).strftime(
"%Y-%m-%d %H:%M:%S UTC"
)
time_controller.advance(1.01)
assert (await handler.should_rate_limit(descriptors=[descriptor]))["overall_code"] == "OK"
@pytest.mark.asyncio
async def test_atomic_tpm_in_memory_rejection_has_window_timing(time_controller):
time_controller._current = datetime.fromtimestamp(2_059, tz=timezone.utc)
cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(cache),
time_provider=time_controller.now,
)
descriptor: Final = RateLimitDescriptor(
key="api_key",
value="atomic-tpm-reset",
rate_limit={"tokens_per_unit": 10, "window_size": 60},
)
window_key: Final = "{api_key:atomic-tpm-reset}:window"
counter_key: Final = handler.create_rate_limit_keys("api_key", "atomic-tpm-reset", "tokens")
await cache.async_set_cache(key=window_key, value="2000", ttl=60)
await cache.async_set_cache(key=counter_key, value=10, ttl=60)
response: Final = await handler.atomic_check_and_increment_by_n(
descriptors=[descriptor],
increments=[{"tokens": 1}],
)
with pytest.raises(HTTPException) as exc_info:
handler._handle_rate_limit_error(response, [descriptor])
assert exc_info.value.headers["retry-after"] == "1"
assert exc_info.value.headers["reset_at"] == datetime.fromtimestamp(2_060, tz=timezone.utc).strftime(
"%Y-%m-%d %H:%M:%S UTC"
)
time_controller.advance(1)
assert (await handler.atomic_check_and_increment_by_n(descriptors=[descriptor], increments=[{"tokens": 1}]))[
"overall_code"
] == "OK"
@pytest.mark.parametrize("redis_now, expected_retry_after", [(9_998, "2"), (9_999, "1")])
def test_atomic_redis_rejection_uses_server_window_clock(redis_now: int, expected_retry_after: str) -> None:
handler: Final = _PROXY_MaxParallelRequestsHandler(
internal_usage_cache=InternalUsageCache(DualCache()),
time_provider=lambda: datetime.fromtimestamp(100_000, tz=timezone.utc),
)
descriptor: Final = RateLimitDescriptor(
key="api_key",
value="redis-tpm-reset",
rate_limit={"tokens_per_unit": 10, "window_size": 90},
)
_keys, _args, meta = handler._build_descriptor_atomic_payload(
descriptor=descriptor, increment_amounts={"tokens": 1}
)
response: Final = handler._build_atomic_response(
raw=[1, 1, 10, 10, 9_910, redis_now],
per_counter_meta=meta,
)
with pytest.raises(HTTPException) as exc_info:
handler._handle_rate_limit_error(response, [descriptor])
assert exc_info.value.headers["retry-after"] == expected_retry_after
assert exc_info.value.headers["reset_at"] == datetime.fromtimestamp(10_000, tz=timezone.utc).strftime(
"%Y-%m-%d %H:%M:%S UTC"
)
STATIC_OUTPUT_FLOOR = 1024
ONE_TOKEN_PROMPT = [{"role": "user", "content": "hello"}]
ONE_TOKEN_PROMPT_INPUT_ESTIMATE = 1