scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

fix(proxy): fix double refund and multimodal undercounting in project io token reservation

Browse source
This commit is contained in:
Shivi Jain 2026-07-22 23:09:49 +05:30
parent fe457bcb2f
commit 4e20fd1382
2 changed files with 168 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

55
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -2542,6 +2542,41 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
llm_provider=llm_provider,
)
def _estimate_precise_input_tokens(self, data: dict, model: str | None) -> int:
"""
Model-aware input token estimate for the project ITPM reservation,
using ``litellm.token_counter`` -- the same approach the
deployment-level itpm/otpm check uses in
``io_token_rate_limit_check.py``. Unlike the cheap char-count
estimate the combined-TPM path uses, this accounts for image/tool
content so a burst of multimodal or tool-heavy requests can't each
reserve only the one-token floor and blow past ITPM before post-call
reconciliation catches up.
Falls back to the cheap char-count estimate if ``token_counter``
can't resolve a tokenizer for this model (e.g. an unrecognized
custom model name) or otherwise raises.
"""
from litellm import token_counter
try:
return max(
0,
int(
token_counter(
model=model or "",
messages=data.get("messages"),
text=data.get("prompt") or data.get("input"),
tools=data.get("tools"),
tool_choice=data.get("tool_choice"),
use_default_image_token_count=True,
)
),
)
except Exception: # noqa: BLE001 - any tokenizer/model-resolution failure degrades to the cheap estimate, never a 500
estimated_input_tokens, _ = self._estimate_input_and_output_tokens(data=data)
return estimated_input_tokens
async def _reserve_project_io_tokens_or_raise(
self,
descriptors: list[RateLimitDescriptor],
@ -2575,10 +2610,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
]
min_configured_otpm_limit = min(configured_otpm_limits) if configured_otpm_limits else None
estimated_input_tokens, estimated_output_tokens = self._estimate_input_and_output_tokens(
_, estimated_output_tokens = self._estimate_input_and_output_tokens(
data=data,
min_configured_tpm_limit=min_configured_otpm_limit,
)
estimated_input_tokens = self._estimate_precise_input_tokens(data=data, model=requested_model)
estimated_input_tokens = max(estimated_input_tokens, 1)
estimated_output_tokens = max(estimated_output_tokens, 1)
@ -2605,12 +2641,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
if io_response["overall_code"] == "OVER_LIMIT":
# A combined-TPM reservation may have already succeeded above for
# this same request; refund it too, or its counter stays inflated
# until the window's TTL expires.
await self._refund_reserved_tokens(
scopes=tpm_reservation_scopes,
amount=tpm_reservation_amount,
parent_otel_span=user_api_key_dict.parent_otel_span,
)
# until the window's TTL expires. Mark it released so the
# ProxyRateLimitError we're about to raise doesn't get refunded
# a second time when async_post_call_failure_hook sees the same
# (still-stashed) reservation and refunds it again.
if tpm_reservation_amount > 0:
await self._refund_reserved_tokens(
scopes=tpm_reservation_scopes,
amount=tpm_reservation_amount,
parent_otel_span=user_api_key_dict.parent_otel_span,
)
self._mark_reservation_released(data)
acquisition = self._get_parallel_slot_acquisition(kwargs=data)
if acquisition is not None:
await self._release_parallel_request_slots(

120
tests/test_litellm/proxy/hooks/test_tpm_concurrent.py
View file

@ -1523,5 +1523,125 @@ async def test_proxy_rejection_refunds_itpm_otpm_only_reservation_with_no_combin
assert otpm_after == 0, f"OTPM-only reservation leaked on proxy rejection: {otpm_after}"
@pytest.mark.asyncio
async def test_otpm_rejection_does_not_double_refund_combined_tpm(rate_limiter):
"""
Regression for a High-severity review finding: when the project ITPM
reservation succeeds but OTPM is then over limit,
_reserve_project_io_tokens_or_raise rolls back the combined-TPM
reservation that already succeeded earlier in the same pre-call, then
raises. If it doesn't also mark that reservation released,
async_post_call_failure_hook -- which fires next in the real request
lifecycle, since raising from async_pre_call_hook triggers it -- sees
the same still-stashed reservation and refunds it a second time,
driving the combined TPM counter negative and letting a caller push
past the project's real TPM budget by repeatedly triggering OTPM
rejections.
"""
handler, cache = rate_limiter
api_key = hash_token("sk-double-refund")
user_api_key_dict = UserAPIKeyAuth(
api_key=api_key,
project_id="proj-double-refund",
project_metadata={
"model_tpm_limit": {"bedrock_mantle/claude-opus": 100000},
"model_itpm_limit": {"bedrock_mantle/claude-opus": 100000},
"model_otpm_limit": {"bedrock_mantle/claude-opus": 5},
},
)
data = {
"model": "bedrock_mantle/claude-opus",
"messages": [{"role": "user", "content": "hello there, this is a test message"}],
"max_tokens": 60, # blows past the 5-token OTPM limit
}
tpm_counter_key = handler.create_rate_limit_keys(
key="model_per_project", value="proj-double-refund:bedrock_mantle/claude-opus", rate_limit_type="tokens"
)
with pytest.raises(Exception) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=cache,
data=data,
call_type="",
)
assert getattr(exc_info.value, "status_code", None) == 429
tpm_after_pre_call = int(await cache.async_get_cache(key=tpm_counter_key, local_only=True) or 0)
assert tpm_after_pre_call == 0, f"combined TPM reservation not rolled back: {tpm_after_pre_call}"
# In the real request lifecycle, async_post_call_failure_hook fires next
# for a pre-call rejection. It must not refund the same reservation again.
await handler.async_post_call_failure_hook(
request_data=data,
original_exception=exc_info.value,
user_api_key_dict=user_api_key_dict,
)
tpm_after_failure_hook = int(await cache.async_get_cache(key=tpm_counter_key, local_only=True) or 0)
assert tpm_after_failure_hook == 0, (
f"combined TPM counter went negative from a double refund: {tpm_after_failure_hook}"
)
@pytest.mark.asyncio
async def test_itpm_reservation_accounts_for_image_content_not_just_text(rate_limiter):
"""
Regression for a Medium-severity review finding: the ITPM reservation
used to estimate input tokens from message text length alone, so a
request with a tiny text prompt but real image content would reserve
almost nothing, letting a burst of such requests blow past the
configured ITPM limit before post-call usage reconciliation catches up.
_estimate_precise_input_tokens uses litellm.token_counter (with
use_default_image_token_count=True, so this test makes no network call)
to account for image content instead of only text length.
"""
handler, cache = rate_limiter
api_key = hash_token("sk-multimodal-itpm")
user_api_key_dict = UserAPIKeyAuth(
api_key=api_key,
project_id="proj-multimodal",
project_metadata={
# Tighter than the ~250-token default image estimate, but far
# bigger than the handful of tokens the bare text "hi" would
# cost under the old char-count-only estimate.
"model_itpm_limit": {"bedrock_mantle/claude-opus": 50},
},
)
data = {
"model": "bedrock_mantle/claude-opus",
"messages": [
{
"role": "user",
"content": [
{"type": "text", "text": "hi"},
{
"type": "image_url",
"image_url": {"url": "https://example.com/some-image.png"},
},
],
}
],
}
with pytest.raises(Exception) as exc_info:
await handler.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=cache,
data=data,
call_type="",
)
assert getattr(exc_info.value, "status_code", None) == 429, (
"Expected the image content to push the ITPM reservation over the "
"50-token limit; if this doesn't raise, the estimate is undercounting "
"non-text content again."
)
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])