From fe457bcb2f2acc52c4aedbe39c76145e1a34c763 Mon Sep 17 00:00:00 2001 From: Shivi Jain Date: Wed, 22 Jul 2026 22:34:31 +0530 Subject: [PATCH] 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. --- .../hooks/parallel_request_limiter_v3.py | 101 ++++++++++----- .../proxy/hooks/test_tpm_concurrent.py | 122 ++++++++++++++++++ 2 files changed, 190 insertions(+), 33 deletions(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index b2f3191db55..7d0f6d2c57e 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -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, diff --git a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py index cbc19c9447a..3f6cbe239bc 100644 --- a/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py +++ b/tests/test_litellm/proxy/hooks/test_tpm_concurrent.py @@ -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"])