diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index ba595ad47ee..42e6d24c7b8 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -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, diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py index f9c234be55a..806875db537 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -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