mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(proxy): refund itpm/otpm reservations correctly on proxy-side rejection
async_post_call_failure_hook had its own refund path, separate from async_log_failure_event, for requests rejected before the LLM call ever ran (e.g. a downstream guardrail block). That path predates itpm/otpm and only knew about the combined TPM reservation: it refunded every token descriptor by the flat combined amount, which is wrong once the project itpm/otpm descriptors are in that same list, and it returned early whenever the combined reservation was zero, which is exactly the case when a project only sets model_itpm_limit/model_otpm_limit with no model_tpm_limit at all.
This commit is contained in:
Shivi Jain
parent
aa130f4e4c
commit
fe457bcb2f
2 changed files with 190 additions and 33 deletions
|
|
@ -3990,17 +3990,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
traceback_str: Optional[str] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Release the parallel-request slot and any TPM reservation when the
|
||||
request is rejected after the pre-call hook acquired them but before
|
||||
the LLM call ran (e.g. a downstream guardrail/auth hook raised).
|
||||
Without this, those resources are stranded — async_log_failure_event
|
||||
is a litellm completion-level callback and never fires for proxy-side
|
||||
rejections, so a leaked slot would occupy the gauge for the full
|
||||
PARALLEL_REQUEST_SLOT_TTL_SECONDS.
|
||||
Release the parallel-request slot and any TPM/ITPM/OTPM reservation
|
||||
when the request is rejected after the pre-call hook acquired them
|
||||
but before the LLM call ran (e.g. a downstream guardrail/auth hook
|
||||
raised). Without this, those resources are stranded —
|
||||
async_log_failure_event is a litellm completion-level callback and
|
||||
never fires for proxy-side rejections, so a leaked slot would occupy
|
||||
the gauge for the full PARALLEL_REQUEST_SLOT_TTL_SECONDS.
|
||||
|
||||
Idempotent: the slot release clears the acquisition marker (and slot
|
||||
removal is a no-op ZREM on a second run), and the TPM refund is
|
||||
guarded by TPM_RESERVATION_RELEASED_KEY — if both this hook and
|
||||
removal is a no-op ZREM on a second run), and the reservation refund
|
||||
is guarded by TPM_RESERVATION_RELEASED_KEY (shared across the
|
||||
TPM/ITPM/OTPM buckets) — if both this hook and
|
||||
async_log_failure_event end up running in the same flow, only the
|
||||
first release/refund applies.
|
||||
"""
|
||||
|
|
@ -4016,36 +4017,70 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
if self._is_reservation_released(kwargs=request_data):
|
||||
return
|
||||
reserved_tokens = self._get_reserved_tokens_from_kwargs(kwargs=request_data)
|
||||
if reserved_tokens <= 0:
|
||||
itpm_reserved = self._get_reserved_itpm_tokens_from_kwargs(kwargs=request_data)
|
||||
otpm_reserved = self._get_reserved_otpm_tokens_from_kwargs(kwargs=request_data)
|
||||
if reserved_tokens <= 0 and itpm_reserved <= 0 and otpm_reserved <= 0:
|
||||
return
|
||||
|
||||
# Refund directly against the descriptors we reserved against —
|
||||
# the pre-call hook stashes them in the request-data metadata
|
||||
# channels before success/failure callbacks run.
|
||||
stashed = self._lookup_stashed_value(
|
||||
kwargs=request_data,
|
||||
standard_logging_metadata=None,
|
||||
key=RATE_LIMIT_DESCRIPTORS_KEY,
|
||||
)
|
||||
descriptors: List[RateLimitDescriptor] = stashed if isinstance(stashed, list) else []
|
||||
ops: List[RedisPipelineIncrementOperation] = []
|
||||
for descriptor in descriptors:
|
||||
rate_limit = descriptor.get("rate_limit") or {}
|
||||
if rate_limit.get("tokens_per_unit") is None:
|
||||
continue
|
||||
ops.append(
|
||||
RedisPipelineIncrementOperation(
|
||||
key=self.create_rate_limit_keys(
|
||||
descriptor["key"],
|
||||
descriptor["value"],
|
||||
"tokens",
|
||||
),
|
||||
increment_value=-reserved_tokens,
|
||||
ttl=self.window_size,
|
||||
|
||||
if reserved_tokens > 0:
|
||||
# Refund directly against the descriptors we reserved
|
||||
# against — the pre-call hook stashes them in the
|
||||
# request-data metadata channels before success/failure
|
||||
# callbacks run. Excludes the project ITPM/OTPM descriptors,
|
||||
# which are reserved from different amounts and refunded
|
||||
# separately below.
|
||||
stashed = self._lookup_stashed_value(
|
||||
kwargs=request_data,
|
||||
standard_logging_metadata=None,
|
||||
key=RATE_LIMIT_DESCRIPTORS_KEY,
|
||||
)
|
||||
descriptors: List[RateLimitDescriptor] = stashed if isinstance(stashed, list) else []
|
||||
for descriptor in descriptors:
|
||||
if descriptor["key"] in (PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY):
|
||||
continue
|
||||
rate_limit = descriptor.get("rate_limit") or {}
|
||||
if rate_limit.get("tokens_per_unit") is None:
|
||||
continue
|
||||
ops.append(
|
||||
RedisPipelineIncrementOperation(
|
||||
key=self.create_rate_limit_keys(
|
||||
descriptor["key"],
|
||||
descriptor["value"],
|
||||
"tokens",
|
||||
),
|
||||
increment_value=-reserved_tokens,
|
||||
ttl=self.window_size,
|
||||
)
|
||||
)
|
||||
|
||||
if itpm_reserved > 0:
|
||||
itpm_scopes = self._get_reserved_itpm_scopes_from_kwargs(kwargs=request_data)
|
||||
ops.extend(
|
||||
self._build_reservation_aware_tpm_ops(
|
||||
targets=list(itpm_scopes),
|
||||
reserved_scopes=itpm_scopes,
|
||||
actual_tokens=0,
|
||||
reserved_tokens=itpm_reserved,
|
||||
)
|
||||
)
|
||||
if otpm_reserved > 0:
|
||||
otpm_scopes = self._get_reserved_otpm_scopes_from_kwargs(kwargs=request_data)
|
||||
ops.extend(
|
||||
self._build_reservation_aware_tpm_ops(
|
||||
targets=list(otpm_scopes),
|
||||
reserved_scopes=otpm_scopes,
|
||||
actual_tokens=0,
|
||||
reserved_tokens=otpm_reserved,
|
||||
)
|
||||
)
|
||||
|
||||
if ops:
|
||||
verbose_proxy_logger.debug(f"Releasing reserved TPM tokens on proxy-level rejection: {reserved_tokens}")
|
||||
verbose_proxy_logger.debug(
|
||||
f"Releasing reserved tokens on proxy-level rejection: "
|
||||
f"tpm={reserved_tokens}, itpm={itpm_reserved}, otpm={otpm_reserved}"
|
||||
)
|
||||
await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline(
|
||||
increment_list=ops,
|
||||
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
|
|
|
|||
|
|
@ -1401,5 +1401,127 @@ async def test_project_itpm_otpm_released_on_failure(rate_limiter):
|
|||
assert any(i["increment"] == -60 for i in otpm_releases), otpm_releases
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_rejection_refunds_itpm_otpm_by_their_own_amount_not_combined(rate_limiter):
|
||||
"""
|
||||
Regression for a Greptile-flagged bug: when a project configures both a
|
||||
combined model_tpm_limit and split model_itpm_limit/model_otpm_limit for
|
||||
the same model, async_post_call_failure_hook's proxy-side refund path
|
||||
used to decrement every token descriptor -- including the ITPM/OTPM
|
||||
ones -- by the flat combined reservation amount, instead of each
|
||||
bucket's own reserved amount. That drives the split counters negative
|
||||
(or under-refunds them) instead of returning them to exactly zero.
|
||||
"""
|
||||
handler, cache = rate_limiter
|
||||
|
||||
api_key = hash_token("sk-mixed-tpm-io")
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key=api_key,
|
||||
project_id="proj-mixed",
|
||||
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": 100000},
|
||||
},
|
||||
)
|
||||
|
||||
data = {
|
||||
"model": "bedrock_mantle/claude-opus",
|
||||
"messages": [{"role": "user", "content": "hello there, this is a test message"}],
|
||||
"max_tokens": 60,
|
||||
}
|
||||
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=cache,
|
||||
data=data,
|
||||
call_type="",
|
||||
)
|
||||
|
||||
tpm_counter_key = handler.create_rate_limit_keys(
|
||||
key="model_per_project", value="proj-mixed:bedrock_mantle/claude-opus", rate_limit_type="tokens"
|
||||
)
|
||||
itpm_counter_key = handler.create_rate_limit_keys(
|
||||
key="model_per_project_itpm", value="proj-mixed:bedrock_mantle/claude-opus", rate_limit_type="tokens"
|
||||
)
|
||||
otpm_counter_key = handler.create_rate_limit_keys(
|
||||
key="model_per_project_otpm", value="proj-mixed:bedrock_mantle/claude-opus", rate_limit_type="tokens"
|
||||
)
|
||||
|
||||
tpm_reserved = int(await cache.async_get_cache(key=tpm_counter_key, local_only=True) or 0)
|
||||
itpm_reserved = int(await cache.async_get_cache(key=itpm_counter_key, local_only=True) or 0)
|
||||
otpm_reserved = int(await cache.async_get_cache(key=otpm_counter_key, local_only=True) or 0)
|
||||
assert tpm_reserved > 0 and itpm_reserved > 0 and otpm_reserved > 0
|
||||
|
||||
await handler.async_post_call_failure_hook(
|
||||
request_data=data,
|
||||
original_exception=Exception("guardrail rejected"),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
tpm_after = int(await cache.async_get_cache(key=tpm_counter_key, local_only=True) or 0)
|
||||
itpm_after = int(await cache.async_get_cache(key=itpm_counter_key, local_only=True) or 0)
|
||||
otpm_after = int(await cache.async_get_cache(key=otpm_counter_key, local_only=True) or 0)
|
||||
|
||||
assert tpm_after == 0, f"combined TPM counter leaked: {tpm_after}"
|
||||
assert itpm_after == 0, f"ITPM counter corrupted by combined-amount refund: {itpm_after}"
|
||||
assert otpm_after == 0, f"OTPM counter corrupted by combined-amount refund: {otpm_after}"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_rejection_refunds_itpm_otpm_only_reservation_with_no_combined_tpm(rate_limiter):
|
||||
"""
|
||||
Regression for the second half of the same bug: with only
|
||||
model_itpm_limit/model_otpm_limit configured (no model_tpm_limit), the
|
||||
combined reserved_tokens is 0, and the proxy-side refund path used to
|
||||
return immediately on that -- leaking the ITPM/OTPM reservations until
|
||||
the rate-limit window's TTL expired.
|
||||
"""
|
||||
handler, cache = rate_limiter
|
||||
|
||||
api_key = hash_token("sk-io-only")
|
||||
user_api_key_dict = UserAPIKeyAuth(
|
||||
api_key=api_key,
|
||||
project_id="proj-io-only",
|
||||
project_metadata={
|
||||
"model_itpm_limit": {"bedrock_mantle/claude-opus": 100000},
|
||||
"model_otpm_limit": {"bedrock_mantle/claude-opus": 100000},
|
||||
},
|
||||
)
|
||||
|
||||
data = {
|
||||
"model": "bedrock_mantle/claude-opus",
|
||||
"messages": [{"role": "user", "content": "hello there, this is a test message"}],
|
||||
"max_tokens": 60,
|
||||
}
|
||||
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=cache,
|
||||
data=data,
|
||||
call_type="",
|
||||
)
|
||||
|
||||
itpm_counter_key = handler.create_rate_limit_keys(
|
||||
key="model_per_project_itpm", value="proj-io-only:bedrock_mantle/claude-opus", rate_limit_type="tokens"
|
||||
)
|
||||
otpm_counter_key = handler.create_rate_limit_keys(
|
||||
key="model_per_project_otpm", value="proj-io-only:bedrock_mantle/claude-opus", rate_limit_type="tokens"
|
||||
)
|
||||
assert int(await cache.async_get_cache(key=itpm_counter_key, local_only=True) or 0) > 0
|
||||
assert int(await cache.async_get_cache(key=otpm_counter_key, local_only=True) or 0) > 0
|
||||
|
||||
await handler.async_post_call_failure_hook(
|
||||
request_data=data,
|
||||
original_exception=Exception("guardrail rejected"),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
itpm_after = int(await cache.async_get_cache(key=itpm_counter_key, local_only=True) or 0)
|
||||
otpm_after = int(await cache.async_get_cache(key=otpm_counter_key, local_only=True) or 0)
|
||||
assert itpm_after == 0, f"ITPM-only reservation leaked on proxy rejection: {itpm_after}"
|
||||
assert otpm_after == 0, f"OTPM-only reservation leaked on proxy rejection: {otpm_after}"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "-s"])
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue