diff --git a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py index 2026706f6e3..2149e254e9a 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.core_helpers import ( _get_parent_otel_span_from_kwargs, # pyright: ignore[reportPrivateUsage] # reused across module boundaries, matching dynamic_rate_limiter_v3's identical import get_metadata_variable_name_from_kwargs, ) +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.parallel_request_limiter_v3 import ( _PROXY_MaxParallelRequestsHandler_v3, # pyright: ignore[reportPrivateUsage] # this hook explicitly reuses its Redis/TTL-preserving increment machinery, see module docstring @@ -1453,6 +1454,35 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] if release_keys: await self._release_keys(release_keys) + async def async_post_call_failure_hook( + self, + request_data: dict, # mutable-ok: must match CustomLogger.async_post_call_failure_hook's own base signature exactly + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + """ + litellm's Logging object sets has_logged_async_failure=True after + the first hop of a fallback chain fails, which blocks + async_log_failure_event for every later hop (see + fallback_event_handlers.py's own docstring) -- so a chain's own + final, chain-exhausting failure never reaches that callback at all, + and _release_stale_hop_reservations only cleans up a stale + reservation when a *next* hop's admission runs, which never happens + after the last one. This hook fires exactly once per proxy request, + at the point the proxy gives up and returns an error to the caller, + regardless of how many hops ran or whether the completion-level + callback was suppressed for this one -- the one reliable place left + to release whatever reservation is still pending. + """ + logging_obj: Final = request_data.get("litellm_logging_obj") + model_call_details: Final = getattr(logging_obj, "model_call_details", None) + if not isinstance(model_call_details, dict): + return + release_keys: Final = self._pop_pending_concurrency_keys(model_call_details) + if release_keys: + await self._release_keys(release_keys) + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: # No special-case skip for this hook's own tag_rate_limit_exceeded # rejection: a hop whose own admission rejects never reaches the diff --git a/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py b/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py index 95e751277cf..1ab22959358 100644 --- a/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py +++ b/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py @@ -16,6 +16,7 @@ from pydantic import ValidationError import litellm from litellm.caching.dual_cache import DualCache +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.model_based_tag_rate_limits_hook import ( _BACKGROUND_TASKS, @@ -2347,6 +2348,48 @@ async def test_concurrency_slot_released_on_failure_frees_capacity(time_controll assert result == healthy +@pytest.mark.asyncio +async def test_concurrency_slot_released_by_post_call_failure_hook_on_the_final_fallback_hop(time_controller): + """ + litellm's Logging object sets has_logged_async_failure=True after the + first hop's failure and blocks async_log_failure_event for every later + hop (see fallback_event_handlers.py), so a fallback chain's own final, + chain-exhausting failure never reaches async_log_failure_event at all -- + _release_stale_hop_reservations only cleans up a stale reservation when + a *next* hop's admission runs, and there is no next hop after the last + one. async_post_call_failure_hook fires exactly once, at the point the + proxy gives up and returns an error to the caller, regardless of how + many hops ran or whether the completion-level callback was suppressed -- + it must release whatever reservation is still pending at that point. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + # This hop's admission reserves the slot; its own failure is the chain's + # final one, so async_log_failure_event never fires for it (simulating + # litellm's has_logged_async_failure dedup blocking the callback here). + request_kwargs, _kwargs = _call_context(["end_user_id:u1"]) + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + await limiter.async_post_call_failure_hook( + request_data=request_kwargs, + original_exception=Exception("all deployments failed"), + user_api_key_dict=UserAPIKeyAuth(api_key="hash"), + ) + + result = await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + assert result == healthy + + @pytest.mark.asyncio async def test_concurrency_slot_released_when_a_different_hook_rejects_the_request(time_controller): """