diff --git a/litellm/constants.py b/litellm/constants.py index b9b9c0ba604..81ae98b5883 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1,6 +1,6 @@ import os import sys -from typing import List, Literal, Optional +from typing import Final, List, Literal, Optional from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_or_none @@ -1321,6 +1321,8 @@ ) # by default all litellm proxy keys have a soft budget of 50.0 # makes it clear this is a rate limit error for a litellm virtual key RATE_LIMIT_ERROR_MESSAGE_FOR_VIRTUAL_KEY = "LiteLLM Virtual Key user_api_key_hash" +# rate-limit increment marking a counter as enforced but not advanced +RATE_LIMIT_CHECK_ONLY: Final[Literal["check_only"]] = "check_only" # Python garbage collection threshold configuration # Format: "gen0,gen1,gen2" e.g., "1000,50,50" diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 6e4a6fe1a51..520770fdc61 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -4,7 +4,7 @@ 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 @@ -12,6 +12,7 @@ from litellm import ModelResponse, Router from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache +from litellm.constants import RATE_LIMIT_CHECK_ONLY from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ( @@ -21,6 +22,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( RateLimitDescriptor, RateLimitDescriptorRateLimitObject, + RateLimitIncrementAmounts, _PROXY_MaxParallelRequestsHandler_v3, ) from litellm.proxy.hooks.rate_limiter_utils import ( @@ -442,9 +444,9 @@ async def _check_rate_limits( 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": RATE_LIMIT_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..80c026942b6 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -16,7 +16,9 @@ Dict, List, Literal, + Mapping, Optional, + Sequence, Set, Tuple, TypedDict, @@ -26,7 +28,10 @@ from litellm import DualCache from litellm._logging import verbose_proxy_logger -from litellm.constants import DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE +from litellm.constants import ( + DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE, + RATE_LIMIT_CHECK_ONLY, +) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, @@ -65,6 +70,25 @@ Span = Any InternalUsageCache = Any +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 (`RATE_LIMIT_CHECK_ONLY`). + """ + if isinstance(raw, str): + return 0 if raw == RATE_LIMIT_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 @@ -- 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) -- @@ -140,13 +164,18 @@ ((now - tonumber(window_start)) >= window_size) 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 + 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 +194,9 @@ 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 +1245,7 @@ async def _release_parallel_request_slots( 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 +1267,11 @@ async def atomic_check_and_increment_by_n( 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 | RATE_LIMIT_CHECK_ONLY, "tokens": int | RATE_LIMIT_CHECK_ONLY}`. + A missing or non-positive int means "do not track this + dimension at all"; `RATE_LIMIT_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 +1316,7 @@ async def atomic_check_and_increment_by_n( 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 +1338,11 @@ def _build_descriptor_atomic_payload( 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 +1435,8 @@ async def _refund_applied_descriptor_groups( 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 +1527,7 @@ async def _atomic_check_and_increment_in_memory( 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 +1537,7 @@ async def _atomic_check_and_increment_in_memory( 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 +1555,17 @@ async def _atomic_check_and_increment_in_memory( # 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 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..b711e7a58a9 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 @@ -17,6 +17,7 @@ from litellm import Router from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth +from litellm.constants import RATE_LIMIT_CHECK_ONLY from litellm.proxy.hooks.parallel_request_limiter_v3 import ( MAX_PARALLEL_SLOT_ACQUIRED_KEY, PARALLEL_REQUEST_SLOT_TTL_SECONDS, @@ -4652,3 +4653,116 @@ def _rl_only(headers: Dict[str, Any]) -> Dict[str, Any]: 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 + RATE_LIMIT_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": RATE_LIMIT_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(): + """ + RATE_LIMIT_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": RATE_LIMIT_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": RATE_LIMIT_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"]