From 4296e2f221297228d949caedb908fc1b27408645 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 25 Jul 2026 01:35:52 +0000 Subject: [PATCH] fix(rate_limiter): enforce TPM pre-call in dynamic_rate_limiter_v3 priority pools Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/dynamic_rate_limiter_v3.py | 8 +- .../hooks/parallel_request_limiter_v3.py | 78 +++++++++-- .../hooks/test_dynamic_rate_limiter_v3.py | 121 ++++++++++++++++++ .../hooks/test_parallel_request_limiter_v3.py | 114 +++++++++++++++++ 4 files changed, 305 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 6e4a6fe1a51..735979d3b38 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -4,7 +4,7 @@ Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting import os from datetime import datetime -from typing import TYPE_CHECKING, Callable, Dict, List, Literal, Optional, Union +from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Union from fastapi import HTTPException @@ -19,8 +19,10 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ( map_v3_rate_limit_type, ) from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + CHECK_ONLY, RateLimitDescriptor, RateLimitDescriptorRateLimitObject, + RateLimitIncrementAmounts, _PROXY_MaxParallelRequestsHandler_v3, ) from litellm.proxy.hooks.rate_limiter_utils import ( @@ -442,9 +444,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): if priority_descriptors and should_enforce_priority: enforced_descriptors.extend(priority_descriptors) - per_request_increment: Dict[Literal["requests", "tokens"], int] = { + per_request_increment: RateLimitIncrementAmounts = { "requests": 1, - "tokens": 0, + "tokens": CHECK_ONLY, } atomic_response = await self.v3_limiter.atomic_check_and_increment_by_n( descriptors=enforced_descriptors, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index b2216488db2..fbdf68f5d81 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -14,9 +14,12 @@ from typing import ( Any, Callable, Dict, + Final, List, Literal, + Mapping, Optional, + Sequence, Set, Tuple, TypedDict, @@ -65,6 +68,27 @@ else: Span = Any InternalUsageCache = Any +CHECK_ONLY: Final[Literal["check_only"]] = "check_only" + +RateLimitIncrement = Union[int, Literal["check_only"]] +RateLimitIncrementAmounts = Mapping[Literal["requests", "tokens"], RateLimitIncrement] + + +def _resolve_increment(raw: RateLimitIncrement | None) -> int | None: + """ + Map a requested increment onto the Lua/in-memory ARGV increment. + + Returns None when the counter should not be tracked or enforced at all, + and 0 when it should be enforced against usage recorded so far without + advancing it (`CHECK_ONLY`). + """ + if isinstance(raw, str): + return 0 if raw == CHECK_ONLY else None + if raw is None or raw <= 0: + return None + return raw + + BATCH_RATE_LIMITER_SCRIPT = """ local results = {} local now = tonumber(ARGV[1]) @@ -115,7 +139,7 @@ CHECK_AND_INCREMENT_BY_N_SCRIPT = """ -- KEYS layout: pairs of (window_key, counter_key), one pair per descriptor. -- ARGV layout: per-descriptor 4-tuple, starting at ARGV[1]: -- ARGV[(i-1)*4 + 1] = limit --- ARGV[(i-1)*4 + 2] = increment +-- ARGV[(i-1)*4 + 2] = increment (0 = check the counter without advancing it) -- ARGV[(i-1)*4 + 3] = ttl_seconds (counter TTL when window resets) -- ARGV[(i-1)*4 + 4] = window_size_seconds (sliding-window length) -- @@ -139,14 +163,24 @@ for i = 1, descriptor_count do local window_expired = (not window_start) or ((now - tonumber(window_start)) >= window_size) + -- Token counters are written post-call by the success logger, which + -- never touches the window key, so a check-only counter reads the raw + -- value and relies on the counter's TTL to bound staleness. local current_counter - if window_expired then + if window_expired and increment > 0 then current_counter = 0 else current_counter = tonumber(redis.call('GET', counter_key) or 0) end - if current_counter + increment > limit then + -- A check-only counter (increment 0) rejects at `current >= limit`, the + -- same point an increment of 1 rejects at. + local check_increment = increment + if check_increment < 1 then + check_increment = 1 + end + + if current_counter + check_increment > limit then return { 1, i, current_counter, limit } end @@ -165,7 +199,9 @@ for i = 1, descriptor_count do local window_expired = descriptor_state[i][1] - if window_expired then + if increment == 0 then + table.insert(results, descriptor_state[i][2]) + elseif window_expired then redis.call('SET', window_key, tostring(now)) redis.call('SET', counter_key, increment) redis.call('EXPIRE', window_key, window_size) @@ -1214,7 +1250,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): async def atomic_check_and_increment_by_n( self, descriptors: List[RateLimitDescriptor], - increments: List[Dict[Literal["requests", "tokens"], int]], + increments: Sequence[RateLimitIncrementAmounts], parent_otel_span: Optional[Span] = None, ) -> RateLimitResponse: """ @@ -1236,8 +1272,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Args: descriptors: rate-limit descriptors to check increments: per-descriptor increment amounts, indexed parallel to - `descriptors`. Each entry is `{"requests": int, "tokens": int}` - — values default to 0 when a descriptor has no matching limit. + `descriptors`. Each entry is + `{"requests": int | CHECK_ONLY, "tokens": int | CHECK_ONLY}`. + A missing or non-positive int means "do not track this + dimension at all"; `CHECK_ONLY` means "enforce this dimension + against usage already recorded, without advancing it". Returns: RateLimitResponse with one status per (descriptor, rate_limit_type) @@ -1282,7 +1321,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): def _build_descriptor_atomic_payload( self, descriptor: RateLimitDescriptor, - increment_amounts: Dict[Literal["requests", "tokens"], int], + increment_amounts: RateLimitIncrementAmounts, ) -> Tuple[List[str], List[Any], List[Dict[str, Any]]]: """ Build (KEYS, ARGV, per-counter meta) for a single descriptor's Lua @@ -1304,11 +1343,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): rlt: Literal["requests", "tokens"] = cast(Literal["requests", "tokens"], rate_limit_type) if rlt == "requests": limit_value = rate_limit.get("requests_per_unit") - inc_amount = int(increment_amounts.get("requests", 0) or 0) + inc_amount = _resolve_increment(increment_amounts.get("requests", 0)) else: limit_value = rate_limit.get("tokens_per_unit") - inc_amount = int(increment_amounts.get("tokens", 0) or 0) - if limit_value is None or inc_amount <= 0: + inc_amount = _resolve_increment(increment_amounts.get("tokens", 0)) + if limit_value is None or inc_amount is None: continue counter_key = self.create_rate_limit_keys(descriptor_key, descriptor_value, rlt) # Counter-key TTL and window_size are conceptually distinct @@ -1401,6 +1440,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return for group_meta in applied: for entry in group_meta: + if entry["increment"] == 0: + continue try: await redis_cache.async_increment( key=entry["counter_key"], @@ -1491,7 +1532,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): window_expired = window_start is None or (now_int - int(window_start)) >= window_size current_counter = ( 0 - if window_expired + if window_expired and meta["increment"] > 0 else int( await self.internal_usage_cache.async_get_cache( key=meta["counter_key"], @@ -1501,7 +1542,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): or 0 ) ) - if current_counter + meta["increment"] > meta["current_limit"]: + if current_counter + max(meta["increment"], 1) > meta["current_limit"]: return RateLimitResponse( overall_code="OVER_LIMIT", statuses=[ @@ -1519,6 +1560,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # Pass 2: apply increments. statuses: List[RateLimitStatus] = [] for meta, state in zip(per_counter_meta, descriptor_state): + if meta["increment"] == 0: + statuses.append( + RateLimitStatus( + code="OK", + current_limit=meta["current_limit"], + limit_remaining=max(0, meta["current_limit"] - state["current"]), + rate_limit_type=meta["rate_limit_type"], + descriptor_key=meta["descriptor_key"], + ) + ) + continue new_counter = meta["increment"] if state["window_expired"] else state["current"] + meta["increment"] if state["window_expired"]: await self.internal_usage_cache.async_set_cache( diff --git a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py index 00ed7e8cd6c..bfe3f09b25b 100644 --- a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -12,6 +12,7 @@ from datetime import datetime, timedelta from unittest.mock import AsyncMock, patch import pytest +from fastapi import HTTPException sys.path.insert(0, os.path.abspath("../../../..")) @@ -1772,3 +1773,123 @@ async def test_priority_429_includes_model_name_and_configured_limits(): assert "Priority: prod" in error_msg, error_msg assert "Rate limit type: tokens" in error_msg, error_msg assert "Model saturation:" in error_msg, error_msg + + +async def _seed_token_counter(handler, descriptor_key: str, descriptor_value: str, tokens: int) -> None: + await handler.internal_usage_cache.async_set_cache( + key=handler.v3_limiter.create_rate_limit_keys( + key=descriptor_key, + value=descriptor_value, + rate_limit_type="tokens", + ), + value=tokens, + litellm_parent_otel_span=None, + local_only=True, + ) + + +@pytest.mark.asyncio +async def test_priority_pool_over_tpm_reservation_is_rejected(): + """ + A priority pool that has already burned more tokens than its reservation + must be rejected pre-call, even though the request's own token count is + unknown at that point (LIT-4800: the TPM branch never ran). + """ + os.environ["LITELLM_LICENSE"] = "test-license-key" + litellm.priority_reservation = {"team_a": 0.6, "team_b": 0.4} + litellm.priority_reservation_settings.saturation_threshold = 0.5 + + model = "tpm-only-priority-model" + total_tpm = 300 + + handler = DynamicRateLimitHandler(internal_usage_cache=DualCache()) + handler.update_variables( + llm_router=Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "tpm": total_tpm, + }, + } + ] + ) + ) + + # team_a reserved TPM is 180; the pool is over it while the model as a + # whole (200/300) still has headroom, so the priority pool is the only + # descriptor that can reject. + await _seed_token_counter(handler, "model_saturation_check", model, 200) + await _seed_token_counter(handler, "priority_model", f"{model}:team_a", 200) + + user_api_key_dict = UserAPIKeyAuth(api_key="sk-team-a", metadata={"priority": "team_a"}) + + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={"model": model}, + call_type="completion", + ) + + assert exc_info.value.status_code == 429 + detail = exc_info.value.detail + assert isinstance(detail, dict) + assert "Priority-based rate limit exceeded" in detail["error"], detail + assert "Rate limit type: tokens" in detail["error"], detail + + +@pytest.mark.asyncio +async def test_priority_pool_under_tpm_reservation_is_allowed_without_advancing_tokens(): + """ + The pre-call TPM check is read-only: a pool under its reservation is + admitted and its token counter is left untouched, since real usage is + only known post-call. + """ + os.environ["LITELLM_LICENSE"] = "test-license-key" + litellm.priority_reservation = {"team_a": 0.6, "team_b": 0.4} + litellm.priority_reservation_settings.saturation_threshold = 0.5 + + model = "tpm-only-priority-model-under-limit" + total_tpm = 300 + + handler = DynamicRateLimitHandler(internal_usage_cache=DualCache()) + handler.update_variables( + llm_router=Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "tpm": total_tpm, + }, + } + ] + ) + ) + + await _seed_token_counter(handler, "model_saturation_check", model, 200) + await _seed_token_counter(handler, "priority_model", f"{model}:team_a", 100) + + user_api_key_dict = UserAPIKeyAuth(api_key="sk-team-a", metadata={"priority": "team_a"}) + + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=DualCache(), + data={"model": model}, + call_type="completion", + ) + + pool_tokens = await handler.internal_usage_cache.async_get_cache( + key=handler.v3_limiter.create_rate_limit_keys( + key="priority_model", + value=f"{model}:team_a", + rate_limit_type="tokens", + ), + litellm_parent_otel_span=None, + local_only=True, + ) + assert int(pool_tokens) == 100 diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index c76e1a60afd..f834cfa9961 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -18,6 +18,7 @@ from litellm import Router from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + CHECK_ONLY, MAX_PARALLEL_SLOT_ACQUIRED_KEY, PARALLEL_REQUEST_SLOT_TTL_SECONDS, ) @@ -4652,3 +4653,116 @@ async def test_streaming_mirror_matches_non_streaming_header_shape(monkeypatch): f" non_streaming={_rl_only(non_stream_headers)}" ) assert "x-ratelimit-model_per_key-remaining-requests" in stream_slp_headers + + +def _tpm_only_descriptor(limit: int) -> Dict[str, Any]: + return { + "key": "priority_model", + "value": "my-model:team_a", + "rate_limit": {"tokens_per_unit": limit, "window_size": 60}, + } + + +def test_check_only_increment_still_emits_a_token_key(): + """ + A non-positive int increment means "don't track this dimension", but + CHECK_ONLY means "enforce it without advancing it" and must still put the + token counter on the Lua KEYS list (LIT-4800). + """ + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + descriptor = _tpm_only_descriptor(180) + + untracked_keys, _, _ = handler._build_descriptor_atomic_payload( + descriptor=descriptor, + increment_amounts={"requests": 1, "tokens": 0}, + ) + assert untracked_keys == [] + + keys, args, meta = handler._build_descriptor_atomic_payload( + descriptor=descriptor, + increment_amounts={"requests": 1, "tokens": CHECK_ONLY}, + ) + assert keys == [ + "{priority_model:my-model:team_a}:window", + "{priority_model:my-model:team_a}:tokens", + ] + assert args == [180, 0, 60, 60] + assert meta[0]["increment"] == 0 + + +@pytest.mark.asyncio +async def test_check_only_tokens_reject_at_limit_without_advancing_counter(): + """ + CHECK_ONLY rejects at `current >= limit`, the same point an increment of 1 + rejects at, and leaves the counter untouched when it admits. + """ + internal_usage_cache = InternalUsageCache(DualCache()) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=internal_usage_cache) + descriptor = _tpm_only_descriptor(180) + counter_key = "{priority_model:my-model:team_a}:tokens" + + async def set_tokens(value: int) -> None: + await internal_usage_cache.async_set_cache( + key=counter_key, + value=value, + litellm_parent_otel_span=None, + local_only=True, + ) + + async def read_tokens() -> int: + return int( + await internal_usage_cache.async_get_cache( + key=counter_key, + litellm_parent_otel_span=None, + local_only=True, + ) + ) + + await set_tokens(179) + under_limit = await handler.atomic_check_and_increment_by_n( + descriptors=[descriptor], + increments=[{"tokens": CHECK_ONLY}], + ) + assert under_limit["overall_code"] == "OK" + assert await read_tokens() == 179 + + await set_tokens(180) + at_limit = await handler.atomic_check_and_increment_by_n( + descriptors=[descriptor], + increments=[{"tokens": CHECK_ONLY}], + ) + assert at_limit["overall_code"] == "OVER_LIMIT" + assert at_limit["statuses"][0]["rate_limit_type"] == "tokens" + assert await read_tokens() == 180 + + +@pytest.mark.asyncio +async def test_token_only_increment_does_not_enforce_rpm(): + """ + check_and_increment_tokens passes tokens only; RPM is enforced separately + by should_rate_limit, so a zero requests increment must stay untracked + rather than double-enforcing. + """ + internal_usage_cache = InternalUsageCache(DualCache()) + handler = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=internal_usage_cache) + descriptor: Dict[str, Any] = { + "key": "key", + "value": "sk-1", + "rate_limit": {"requests_per_unit": 1, "tokens_per_unit": 1000, "window_size": 60}, + } + await internal_usage_cache.async_set_cache( + key="{key:sk-1}:requests", + value=1, + litellm_parent_otel_span=None, + local_only=True, + ) + + response = await handler.atomic_check_and_increment_by_n( + descriptors=[descriptor], + increments=[{"tokens": 10}], + ) + + assert response["overall_code"] == "OK" + assert [status["rate_limit_type"] for status in response["statuses"]] == ["tokens"]