From 6735375de7ef0c58fe783ff7fc1a42aa11f0107c Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 25 Jul 2026 02:02:54 +0000 Subject: [PATCH] refactor(rate_limiter): move check-only sentinel into constants Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 4 +++- .../proxy/hooks/dynamic_rate_limiter_v3.py | 4 ++-- .../hooks/parallel_request_limiter_v3.py | 21 +++++++------------ .../hooks/test_parallel_request_limiter_v3.py | 12 +++++------ 4 files changed, 19 insertions(+), 22 deletions(-) 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 @@ DEFAULT_SOFT_BUDGET = float( ) # 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 735979d3b38..520770fdc61 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -12,6 +12,7 @@ import litellm 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 ( @@ -19,7 +20,6 @@ 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, @@ -446,7 +446,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): per_request_increment: RateLimitIncrementAmounts = { "requests": 1, - "tokens": CHECK_ONLY, + "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 fbdf68f5d81..80c026942b6 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -14,7 +14,6 @@ from typing import ( Any, Callable, Dict, - Final, List, Literal, Mapping, @@ -29,7 +28,10 @@ from typing import ( 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, @@ -68,8 +70,6 @@ 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] @@ -80,10 +80,10 @@ def _resolve_increment(raw: RateLimitIncrement | None) -> int | None: 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`). + advancing it (`RATE_LIMIT_CHECK_ONLY`). """ if isinstance(raw, str): - return 0 if raw == CHECK_ONLY else None + return 0 if raw == RATE_LIMIT_CHECK_ONLY else None if raw is None or raw <= 0: return None return raw @@ -163,9 +163,6 @@ 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 and increment > 0 then current_counter = 0 @@ -173,8 +170,6 @@ for i = 1, descriptor_count do current_counter = tonumber(redis.call('GET', counter_key) or 0) end - -- 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 @@ -1273,9 +1268,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): descriptors: rate-limit descriptors to check increments: per-descriptor increment amounts, indexed parallel to `descriptors`. Each entry is - `{"requests": int | CHECK_ONLY, "tokens": int | CHECK_ONLY}`. + `{"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"; `CHECK_ONLY` means "enforce this dimension + dimension at all"; `RATE_LIMIT_CHECK_ONLY` means "enforce this dimension against usage already recorded, without advancing it". Returns: 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 f834cfa9961..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,8 +17,8 @@ import litellm 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 ( - CHECK_ONLY, MAX_PARALLEL_SLOT_ACQUIRED_KEY, PARALLEL_REQUEST_SLOT_TTL_SECONDS, ) @@ -4666,7 +4666,7 @@ def _tpm_only_descriptor(limit: int) -> Dict[str, Any]: 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 + 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( @@ -4682,7 +4682,7 @@ def test_check_only_increment_still_emits_a_token_key(): keys, args, meta = handler._build_descriptor_atomic_payload( descriptor=descriptor, - increment_amounts={"requests": 1, "tokens": CHECK_ONLY}, + increment_amounts={"requests": 1, "tokens": RATE_LIMIT_CHECK_ONLY}, ) assert keys == [ "{priority_model:my-model:team_a}:window", @@ -4695,7 +4695,7 @@ def test_check_only_increment_still_emits_a_token_key(): @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 + 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()) @@ -4723,7 +4723,7 @@ async def test_check_only_tokens_reject_at_limit_without_advancing_counter(): await set_tokens(179) under_limit = await handler.atomic_check_and_increment_by_n( descriptors=[descriptor], - increments=[{"tokens": CHECK_ONLY}], + increments=[{"tokens": RATE_LIMIT_CHECK_ONLY}], ) assert under_limit["overall_code"] == "OK" assert await read_tokens() == 179 @@ -4731,7 +4731,7 @@ async def test_check_only_tokens_reject_at_limit_without_advancing_counter(): await set_tokens(180) at_limit = await handler.atomic_check_and_increment_by_n( descriptors=[descriptor], - increments=[{"tokens": CHECK_ONLY}], + increments=[{"tokens": RATE_LIMIT_CHECK_ONLY}], ) assert at_limit["overall_code"] == "OVER_LIMIT" assert at_limit["statuses"][0]["rate_limit_type"] == "tokens"