mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(proxy): reserve the larger declared output budget for TPM limits (#37001)
The TPM pre-call reservation read `max_tokens or max_completion_tokens`, so a request declaring both was charged for whichever field came first. A caller sending `max_tokens=1` with `max_completion_tokens=10000` reserved 2 tokens and was then free to consume ten thousand, since the provider honours the modern field and litellm's own param mapping drops the legacy one for the gpt-5 and o-series families. Reserve against the larger of the declared budgets instead. Over-reserving is the safe direction for a limiter: post-call reconciliation refunds the difference between the reservation and actual usage, while under-reserving lets the window be exceeded before anything notices.
This commit is contained in:
parent
592564db23
commit
33e9f54dc8
2 changed files with 108 additions and 1 deletions
|
|
@ -462,6 +462,23 @@ def _call_id_from_callback_kwargs(kwargs: object) -> str | None:
|
|||
return call_id if isinstance(call_id, str) else None
|
||||
|
||||
|
||||
def _declared_output_budget(value: object) -> int | None:
|
||||
"""Coerce a declared output budget to tokens, or None when it names no budget.
|
||||
|
||||
Accepts every shape the pre-existing ``int(...)`` coercion did, floats and numeric
|
||||
strings included, because a budget this cannot read is a budget this cannot reserve
|
||||
against, which is the bypass the caller-declared limits are checked for.
|
||||
"""
|
||||
if isinstance(value, (int, float)):
|
||||
return int(value)
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return int(float(value))
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
||||
def __init__(
|
||||
self,
|
||||
|
|
@ -604,7 +621,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
|
||||
estimated_input_tokens: Final = max(1, total_chars // DEFAULT_CHARS_PER_TOKEN) if total_chars > 0 else 0
|
||||
|
||||
explicit_max_tokens: Final = data.get("max_tokens") or data.get("max_completion_tokens")
|
||||
# Both spellings can arrive together, e.g. a deployment-level max_tokens default under a
|
||||
# client-supplied max_completion_tokens. Reserving against the larger keeps the estimate an
|
||||
# upper bound on what the provider can emit, whichever one it ends up honouring.
|
||||
declared_output_budgets: Final = tuple(
|
||||
budget
|
||||
for budget in (
|
||||
_declared_output_budget(data.get("max_tokens")),
|
||||
_declared_output_budget(data.get("max_completion_tokens")),
|
||||
)
|
||||
if budget is not None
|
||||
)
|
||||
explicit_max_tokens: Final = max(declared_output_budgets) if declared_output_budgets else None
|
||||
|
||||
match (explicit_max_tokens, input_text):
|
||||
case (mt, _) if mt is not None:
|
||||
|
|
|
|||
|
|
@ -5593,3 +5593,82 @@ def test_internal_call_origin_success_ops_are_skipped():
|
|||
|
||||
assert charged
|
||||
assert skipped == []
|
||||
|
||||
|
||||
def _conflicting_budget_bodies() -> Dict[str, Dict[str, object]]:
|
||||
"""The same request, three ways of declaring the output budget."""
|
||||
base = {"model": "gpt-5-chat", "messages": [{"role": "user", "content": "hi"}]}
|
||||
return {
|
||||
"both": {**base, "max_tokens": 1, "max_completion_tokens": 10000},
|
||||
"only_large": {**base, "max_completion_tokens": 10000},
|
||||
"only_small": {**base, "max_tokens": 1},
|
||||
}
|
||||
|
||||
|
||||
def test_conflicting_token_limits_reserve_the_larger_declared_budget():
|
||||
"""Both spellings together must reserve the larger budget, not whichever is read first.
|
||||
|
||||
A request declaring max_tokens=1 alongside max_completion_tokens=10000 previously
|
||||
reserved one output token while the provider stayed free to emit ten thousand.
|
||||
"""
|
||||
handler = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(DualCache())
|
||||
)
|
||||
bodies = _conflicting_budget_bodies()
|
||||
|
||||
reserved = {
|
||||
label: handler._estimate_tokens_for_request(data=body)
|
||||
for label, body in bodies.items()
|
||||
}
|
||||
|
||||
assert reserved["both"] == reserved["only_large"]
|
||||
assert reserved["both"] > reserved["only_small"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("declared", [10000, 10000.0, "10000"])
|
||||
def test_non_integer_output_budgets_still_reserve_their_declared_size(declared):
|
||||
"""A budget litellm cannot read is a budget it cannot reserve against.
|
||||
|
||||
A float or numeric-string max_tokens is explicit enough to suppress the capped
|
||||
output floor, so dropping it from the estimate under-reserves and reopens the
|
||||
same TPM bypass that reading both spellings was meant to close.
|
||||
"""
|
||||
handler = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(DualCache())
|
||||
)
|
||||
base = {"model": "gpt-5-chat", "messages": [{"role": "user", "content": "hi"}]}
|
||||
|
||||
reserved = handler._estimate_tokens_for_request(data={**base, "max_tokens": declared})
|
||||
reserved_int = handler._estimate_tokens_for_request(data={**base, "max_tokens": 10000})
|
||||
|
||||
assert reserved == reserved_int
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conflicting_token_limits_cannot_bypass_tpm_reservation():
|
||||
"""The pre-call hook must refuse a request whose larger declared budget exceeds the TPM limit."""
|
||||
local_cache = DualCache()
|
||||
handler = _PROXY_MaxParallelRequestsHandler(
|
||||
internal_usage_cache=InternalUsageCache(local_cache)
|
||||
)
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key=hash_token("sk-conflicting-budgets"), tpm_limit=100, models=[]
|
||||
)
|
||||
bodies = _conflicting_budget_bodies()
|
||||
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=local_cache,
|
||||
data=dict(bodies["only_small"]),
|
||||
call_type="",
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=local_cache,
|
||||
data=dict(bodies["both"]),
|
||||
call_type="",
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 429
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue