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>

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,

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(

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

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"]