From d3b8df8b284ba5bc58c3572f36d36d52e5f98eb6 Mon Sep 17 00:00:00 2001 From: michelligabriele Date: Mon, 25 May 2026 22:24:43 +0200 Subject: [PATCH] fix(v3 limiter): cap no-max_tokens TPM floor at smallest configured limit --- .../hooks/parallel_request_limiter_v3.py | 32 +++++-- .../proxy/hooks/test_tpm_concurrent.py | 93 +++++++++++++++++++ 2 files changed, 117 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 283a3d8d10b..405db4cff7d 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -344,6 +344,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self, data: dict, model: Optional[str] = None, + min_configured_tpm_limit: Optional[int] = None, ) -> int: """ Estimate total tokens this request will consume so we can reserve them @@ -351,6 +352,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): estimated = input_tokens + max_tokens. Supports chat (messages), completions (prompt), and embeddings (input). + + ``min_configured_tpm_limit`` is the smallest ``tokens_per_unit`` among + the TPM-bearing descriptors this request will be charged against. When + provided, the no-``max_tokens`` output-budget floor is capped at a + fraction of that limit so small TPM caps remain usable. Omit to + preserve the unconstrained floor. """ messages = data.get("messages") prompt = data.get("prompt") @@ -394,11 +401,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): case _: # No max_tokens specified — reserve at least the input size with a # conservative floor so a stream of small concurrent requests can't - # collectively bypass the limit. - max_tokens_estimate = max( - estimated_input_tokens, - DEFAULT_MAX_TOKENS_ESTIMATE // 4, - ) + # collectively bypass the limit. Cap the floor by a fraction of + # the smallest TPM limit this request will be charged against, + # so a small per-tenant TPM cap can't be tripped by the floor + # alone. + output_floor = DEFAULT_MAX_TOKENS_ESTIMATE // 4 + if min_configured_tpm_limit is not None: + output_floor = min( + output_floor, + max(1, min_configured_tpm_limit // 4), + ) + max_tokens_estimate = max(estimated_input_tokens, output_floor) total_estimated = estimated_input_tokens + max_tokens_estimate @@ -2009,10 +2022,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # in-memory check otherwise — single-worker protection still holds # even without Redis. # ---------------------------------------------------------------- - has_tpm_limits = any( - (d.get("rate_limit") or {}).get("tokens_per_unit") is not None + configured_tpm_limits = [ + (d.get("rate_limit") or {}).get("tokens_per_unit") for d in descriptors - ) + if (d.get("rate_limit") or {}).get("tokens_per_unit") is not None + ] + has_tpm_limits = bool(configured_tpm_limits) if has_tpm_limits: # Floor at 1 token so contentless requests (/responses, @@ -2026,6 +2041,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self._estimate_tokens_for_request( data=data, model=requested_model, + min_configured_tpm_limit=min(configured_tpm_limits), ), 1, ) diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index e294d1471db..2c74d5d877d 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -995,5 +995,98 @@ async def test_token_rate_limit_headers_present_in_stored_response(rate_limiter) assert api_key_tokens["limit_remaining"] >= 0 +@pytest.mark.asyncio +async def test_estimate_tokens_floor_caps_at_smallest_configured_tpm(rate_limiter): + """ + Regression: with a small configured TPM cap and no max_tokens, the + output-budget floor must be capped at a fraction of that limit so the + reservation alone can't trip the limit. + """ + handler, _cache = rate_limiter + + estimate = handler._estimate_tokens_for_request( + data={"messages": [{"role": "user", "content": "hello"}]}, + min_configured_tpm_limit=1000, + ) + # input ~= 5//4 = 1 token; output floor capped at 1000//4 = 250; + # total ~= 251 (well under 1000). + assert ( + estimate <= 1000 // 2 + ), f"With TPM=1000, reservation must stay well under the limit; got {estimate}" + assert estimate >= 1, "Estimate must be at least the call-site floor of 1" + + +@pytest.mark.asyncio +async def test_estimate_tokens_floor_unchanged_for_large_tpm(rate_limiter): + """ + Large TPM budgets must keep the 1024-token floor so a stream of small + concurrent requests can't collectively bypass the limit. + """ + handler, _cache = rate_limiter + + estimate = handler._estimate_tokens_for_request( + data={"messages": [{"role": "user", "content": "hello"}]}, + min_configured_tpm_limit=100_000, + ) + # input ~= 1; output floor = min(1024, 100_000//4=25_000) = 1024; + # total ~= 1025. + assert estimate == 1 + 1024 + + +@pytest.mark.asyncio +async def test_estimate_tokens_floor_unchanged_when_kwarg_omitted(rate_limiter): + """ + Callers that don't pass min_configured_tpm_limit (legacy path, tests that + stub the estimator) must observe the pre-fix floor. + """ + handler, _cache = rate_limiter + + estimate = handler._estimate_tokens_for_request( + data={"messages": [{"role": "user", "content": "hello"}]}, + ) + assert estimate == 1 + 1024 + + +@pytest.mark.asyncio +async def test_small_tpm_cap_admits_no_max_tokens_request(rate_limiter): + """ + Regression (end-to-end at the hook level): a project-level model_tpm_limit + of 1000 with a tiny no-max_tokens request must not 429 on the first call. + Pre-fix the 1024-token floor tripped OVER_LIMIT against the 1000-token cap + on every request. + """ + handler, cache = rate_limiter + + api_key = hash_token("sk-small-tpm") + user_api_key_dict = UserAPIKeyAuth( + api_key=api_key, + project_id="proj-small-tpm", + project_metadata={ + "model_tpm_limit": {"gpt-3.5-turbo": 1000}, + "model_rpm_limit": {"gpt-3.5-turbo": 60}, + }, + ) + + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "hello"}], + } + + # Must not raise — pre-fix this was a 429. + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="", + ) + + reserved = (data.get("metadata") or {}).get(TPM_RESERVED_TOKENS_KEY) + assert reserved is not None, "Reservation should have been stashed" + assert reserved <= 1000 // 2, ( + f"Capped floor must keep the reservation well under the 1000 TPM " + f"cap; got {reserved}" + ) + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"])