diff --git a/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs b/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs index bf6ddfbeb52..974da7c5642 100644 --- a/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs +++ b/litellm-rust/crates/python-bridge/src/redis/store/rate_limits.rs @@ -206,16 +206,27 @@ pub(crate) struct PyAtomicOverLimit { index: usize, current: i64, limit: i64, + window_start: Option, + window_observed_at: Option, } #[pymethods] impl PyAtomicOverLimit { #[new] - fn new(index: usize, current: i64, limit: i64) -> Self { + #[pyo3(signature = (index, current, limit, window_start = None, window_observed_at = None))] + fn new( + index: usize, + current: i64, + limit: i64, + window_start: Option, + window_observed_at: Option, + ) -> Self { Self { index, current, limit, + window_start, + window_observed_at, } } } @@ -313,10 +324,14 @@ fn atomic(reply: AtomicReply) -> PyAtomicReply { index, current, limit, + window_start, + window_observed_at, } => PyAtomicReply::OverLimit(PyAtomicOverLimit { index, current, limit, + window_start: Some(window_start), + window_observed_at: Some(window_observed_at), }), } } diff --git a/litellm-rust/crates/redis-store/src/rate_limits/lua.rs b/litellm-rust/crates/redis-store/src/rate_limits/lua.rs index 96451212d35..a33b7032f0b 100644 --- a/litellm-rust/crates/redis-store/src/rate_limits/lua.rs +++ b/litellm-rust/crates/redis-store/src/rate_limits/lua.rs @@ -60,7 +60,7 @@ pub(super) const CHECK_AND_INCREMENT_BY_N: &str = r#" -- -- 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, now } local time_reply = redis.call('TIME') local now = tonumber(time_reply[1]) local descriptor_count = #KEYS / 2 @@ -94,7 +94,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 } diff --git a/litellm-rust/crates/redis-store/src/rate_limits/scripts.rs b/litellm-rust/crates/redis-store/src/rate_limits/scripts.rs index 04bdad38268..66e17516dd6 100644 --- a/litellm-rust/crates/redis-store/src/rate_limits/scripts.rs +++ b/litellm-rust/crates/redis-store/src/rate_limits/scripts.rs @@ -264,10 +264,21 @@ fn atomic_reply(reply: Value) -> Result { }) .collect::>() .map(AtomicReply::Applied), - Some((Value::Int(1), [index_value, current, limit])) => Ok(AtomicReply::OverLimit { + Some(( + Value::Int(1), + [ + index_value, + current, + limit, + window_start, + window_observed_at, + ], + )) => Ok(AtomicReply::OverLimit { index: index(index_value)?, current: integer(current)?, limit: integer(limit)?, + window_start: integer(window_start)?, + window_observed_at: integer(window_observed_at)?, }), _ => Err(unexpected_reply("rate limiter script reply has no status")), } @@ -301,3 +312,31 @@ fn acquire_reply(reply: Value) -> Result { _ => Err(unexpected_reply("rate limiter script reply has no status")), } } + +#[cfg(test)] +mod tests { + use super::*; + use rstest::rstest; + + #[rstest] + #[case::two_seconds_remaining(9_998)] + #[case::one_second_remaining(9_999)] + fn atomic_rejection_keeps_the_server_window_clock(#[case] observed_at: i64) { + let reply = Value::Array( + [1, 1, 10, 10, 9_910, observed_at] + .into_iter() + .map(Value::Int) + .collect(), + ); + assert_eq!( + atomic_reply(reply).unwrap(), + AtomicReply::OverLimit { + index: 0, + current: 10, + limit: 10, + window_start: 9_910, + window_observed_at: observed_at, + } + ); + } +} diff --git a/litellm-rust/crates/redis-store/src/rate_limits/types.rs b/litellm-rust/crates/redis-store/src/rate_limits/types.rs index 6e81fac7a41..93038e963a6 100644 --- a/litellm-rust/crates/redis-store/src/rate_limits/types.rs +++ b/litellm-rust/crates/redis-store/src/rate_limits/types.rs @@ -30,6 +30,8 @@ pub enum AtomicReply { index: usize, current: i64, limit: i64, + window_start: i64, + window_observed_at: i64, }, } diff --git a/litellm-rust/crates/redis-store/tests/rate_limits.rs b/litellm-rust/crates/redis-store/tests/rate_limits.rs index 52f9169b6da..c1c5bb12081 100644 --- a/litellm-rust/crates/redis-store/tests/rate_limits.rs +++ b/litellm-rust/crates/redis-store/tests/rate_limits.rs @@ -397,18 +397,25 @@ async fn check_and_increment_is_all_or_nothing( assert_eq!(applied.iter().map(|a| a.count).collect::>(), [1, 60]); assert!(applied.iter().all(|a| (a.window_start - now).abs() <= 2)); - assert_eq!( - limits - .check_and_increment(RedisSend::Now, &both(1, 41)) - .unwrap() - .await - .unwrap(), - AtomicReply::OverLimit { - index: 1, - current: 60, - limit: 100 - } - ); + let denied = limits + .check_and_increment(RedisSend::Now, &both(1, 41)) + .unwrap() + .await + .unwrap(); + let AtomicReply::OverLimit { + index, + current, + limit, + window_start, + window_observed_at, + } = denied + else { + panic!("{denied:?}"); + }; + assert_eq!((index, current, limit), (1, 60, 100)); + assert_eq!(window_start, applied[1].window_start); + assert!(window_observed_at >= now); + assert!(window_observed_at <= redis_now(&raw).await); assert_eq!( raw_text(&raw, "{model_per_key:k1:gpt}:requests").await, Some("1".to_owned()) @@ -426,18 +433,21 @@ async fn check_and_increment_is_all_or_nothing( assert!( matches!(at_limit, AtomicReply::Applied(ref applied) if applied[0].count == 2 && applied[1].count == 100) ); - assert_eq!( - limits - .check_and_increment(RedisSend::Now, &both(0, 0)) - .unwrap() - .await - .unwrap(), + let read_only_denied = limits + .check_and_increment(RedisSend::Now, &both(0, 0)) + .unwrap() + .await + .unwrap(); + assert!(matches!( + read_only_denied, AtomicReply::OverLimit { index: 0, current: 2, - limit: 2 - } - ); + limit: 2, + window_start: start, + window_observed_at: observed, + } if start == applied[0].window_start && observed >= window_observed_at && observed <= redis_now(&raw).await + )); } #[rstest] diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 6b36a6ac8bc..fe642a67b9e 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -15,6 +15,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, @@ -412,6 +413,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): @@ -1200,6 +1204,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): cache_values: Sequence[CounterState], key_metadata: Mapping[RateLimitScope, WindowKeyMetadata], read_only: bool = False, + observed_at: float | None = None, ) -> RateLimitResponse: """ Check if the cache values are over the limit. @@ -1208,7 +1213,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): statuses: Final[list[RateLimitStatus]] = [] overall_code = "OK" - for counter, (_window_start, counter_value) in zip(counters, cache_values): + for counter, (window_start, counter_value) in zip(counters, cache_values): item_code = "OK" metadata = key_metadata[counter.scope] current_limit = metadata["requests_limit"] if counter.kind == "requests" else metadata["tokens_limit"] @@ -1221,6 +1226,23 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): 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=counter.kind, + descriptor_key=metadata["descriptor_key"], + descriptor_value=metadata["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, @@ -1412,7 +1434,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): windowed_response = RateLimitResponse(overall_code="OK", statuses=[]) if counters: local_response: Final = self.is_cache_list_over_limit( - counters, self._local_counter_states(counters), key_metadata, read_only=read_only + counters, self._local_counter_states(counters), key_metadata, read_only=read_only, observed_at=now ) if local_response["overall_code"] == "OVER_LIMIT": return local_response @@ -1433,7 +1455,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): window_size=self.window_size, ) - windowed_response = self.is_cache_list_over_limit(counters, cache_values, key_metadata, read_only=read_only) + windowed_response = self.is_cache_list_over_limit( + counters, cache_values, key_metadata, read_only=read_only, observed_at=now + ) if windowed_response["overall_code"] == "OVER_LIMIT": return windowed_response @@ -2074,20 +2098,32 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): level when modifying this code. """ if isinstance(reply, AtomicOverLimit): - meta = per_counter_meta[reply.index] - return RateLimitResponse( - overall_code="OVER_LIMIT", - statuses=[ - RateLimitStatus( - code="OVER_LIMIT", - current_limit=reply.limit, - limit_remaining=max(0, reply.limit - reply.current), - rate_limit_type=meta["rate_limit_type"], - descriptor_key=meta["descriptor_key"], - descriptor_value=meta["descriptor_value"], - ) - ], + over_limit_meta: Final = per_counter_meta[reply.index] + window_start: Final = reply.window_start + window_observed_at: Final = reply.window_observed_at + status: Final = ( + RateLimitStatus( + code="OVER_LIMIT", + current_limit=reply.limit, + limit_remaining=max(0, reply.limit - reply.current), + 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=float(window_start), + window_size=over_limit_meta["window_size"], + window_observed_at=float(window_observed_at), + ) + if window_start is not None and window_observed_at is not None + else RateLimitStatus( + code="OVER_LIMIT", + current_limit=reply.limit, + limit_remaining=max(0, reply.limit - reply.current), + 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]] = [ RateLimitStatus( code="OK", @@ -2121,7 +2157,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]] = [] @@ -2147,6 +2184,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, ) ], ) @@ -3299,9 +3339,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" ) @@ -3316,14 +3367,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/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 773985ad427..2ba30574e02 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -204,13 +204,24 @@ class AtomicApplied: @final class AtomicOverLimit: - def __new__(cls, index: int, current: int, limit: int) -> AtomicOverLimit: ... + def __new__( + cls, + index: int, + current: int, + limit: int, + window_start: int | None = None, + window_observed_at: int | None = None, + ) -> AtomicOverLimit: ... @property def index(self) -> int: ... @property def current(self) -> int: ... @property def limit(self) -> int: ... + @property + def window_start(self) -> int | None: ... + @property + def window_observed_at(self) -> int | None: ... @final class GuardedReply: 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 64e7cde4034..a26c79f17c6 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py @@ -56,6 +56,7 @@ from litellm.proxy.spend_tracking.input_tokens import models_counted_for from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token from litellm.redis.stores.rate_limits import ( + AtomicOverLimit, CounterIncrement, CounterKind, LocalValue, @@ -6854,6 +6855,123 @@ 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}, + ) + counter: Final = handler.rate_limit_counter("api_key", "atomic-tpm-reset", "tokens") + handler.rate_limits.set_local(counter.window, "2000", 60) + handler.rate_limits.set_local(counter, 10, 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}, + ) + _counters, meta = handler._build_descriptor_atomic_payload( + descriptor=descriptor, increment_amounts={"tokens": 1} + ) + response: Final = handler._build_atomic_response( + reply=AtomicOverLimit(0, 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