fix(rate_limiter): enforce TPM pre-call in dynamic_rate_limiter_v3 priority pools

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"

litellm/proxy/hooks/dynamic_rate_limiter_v3.py

@@ -4,14 +4,15 @@
 
 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
 
 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 (
@@ -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,

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(

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