diff --git a/litellm/proxy/hooks/global_tag_rate_limits_hook.py b/litellm/proxy/hooks/global_tag_rate_limits_hook.py index c6a6ab556dd..29a0939151e 100644 --- a/litellm/proxy/hooks/global_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/global_tag_rate_limits_hook.py @@ -61,7 +61,6 @@ from litellm.proxy.hooks.model_based_tag_rate_limits_hook import ( _ATOMIC_UNITS, # pyright: ignore[reportPrivateUsage] # reused across module boundaries, see module docstring _BACKGROUND_TASKS, # pyright: ignore[reportPrivateUsage] # reused across module boundaries, see module docstring _CONCURRENCY_MIN_SAFETY_TTL_SECONDS, # pyright: ignore[reportPrivateUsage] # reused across module boundaries, see module docstring - _EMPTY_MAPPING, # pyright: ignore[reportPrivateUsage] # reused across module boundaries, see module docstring _LIMIT_UNITS, # pyright: ignore[reportPrivateUsage] # reused across module boundaries, see module docstring _UNIT_TO_GROUP_FIELD, # pyright: ignore[reportPrivateUsage] # reused across module boundaries, see module docstring _UNIT_TO_RATE_LIMIT_TYPE, # pyright: ignore[reportPrivateUsage] # reused across module boundaries, see module docstring @@ -513,13 +512,15 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o await self._release_keys(release_keys) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: - if isinstance(kwargs.get("exception"), ProxyRateLimitError): - detail: Final = ( - kwargs["exception"].detail if isinstance(kwargs["exception"].detail, dict) else _EMPTY_MAPPING - ) - if detail.get("error") == "tag_rate_limit_exceeded": - return - + # No special-case skip for this hook's own tag_rate_limit_exceeded + # rejection: that rejection never reaches the point where a + # concurrency reservation is queued (see async_pre_call_hook), so + # stash.pending_concurrency_keys is already empty in that case and + # the check below naturally no-ops. Skipping release based on the + # exception's error marker alone would be wrong here, since + # model_based_tag_rate_limits_hook raises the identical marker -- + # that rejection can land after this hook already reserved a slot + # for this same request, and that slot must still be released. stash: Final = _stash_for_call(_call_id_from_kwargs(kwargs)) if stash is None or not stash.pending_concurrency_keys: return 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 01e981a8407..f18f4a82d66 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -1440,13 +1440,15 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] await self._release_keys(release_keys) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: - if isinstance(kwargs.get("exception"), ProxyRateLimitError): - detail: Final = ( - kwargs["exception"].detail if isinstance(kwargs["exception"].detail, dict) else _EMPTY_MAPPING - ) - if detail.get("error") == "tag_rate_limit_exceeded": - return - + # No special-case skip for this hook's own tag_rate_limit_exceeded + # rejection: a hop whose own admission rejects never reaches the + # point where a concurrency reservation is queued (see + # async_filter_deployments), so _pop_pending_concurrency_keys already + # returns nothing to release in that case. Skipping release based on + # the exception's error marker alone would be wrong here, since + # global_tag_rate_limits_hook raises the identical marker -- that + # rejection can land after this hook already reserved a slot for the + # same request, and that slot must still be released. release_keys: Final = self._pop_pending_concurrency_keys(kwargs) if release_keys: await self._release_keys(release_keys) diff --git a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py index 934f177a352..de8543159dc 100644 --- a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py +++ b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py @@ -319,6 +319,50 @@ async def test_concurrency_reservation_released_on_success(time_controller, monk assert result is not None +@pytest.mark.asyncio +async def test_concurrency_reservation_released_when_a_different_hook_rejects_the_request(time_controller, monkeypatch): + """ + model_based_tag_rate_limits_hook raises the identical ProxyRateLimitError + shape (detail["error"] == "tag_rate_limit_exceeded") this hook's own + admission does, since both hooks share the same rejection marker. + async_log_failure_event fires on every registered CustomLogger regardless + of which one raised, so this hook must still release its own successfully + reserved concurrency slot when the *other* hook is what rejected the + request -- skipping release just because the marker matches would leak + this hook's own slot until the safety TTL, even though nothing about this + hook's own admission failed. + """ + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + {"concurrency_limits": {"limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}]}}, + ) + hook = _make_hook(time_controller) + + data = _data(["end_user_id:u1"], call_id="call-1") + await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="completion") + + other_hooks_rejection = ProxyRateLimitError( + detail={"error": "tag_rate_limit_exceeded", "type": "requests", "tag_id": "end_user_id"}, + headers={"retry-after": "60"}, + rate_limit_type=None, + model="gpt-4o", + llm_provider="litellm_proxy", + ) + kwargs = {"litellm_call_id": "call-1", "exception": other_hooks_rejection} + await hook.async_log_failure_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + + # The slot was released despite the shared rejection marker, so a fresh + # request must be admitted again. + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + assert result is not None + + @pytest.mark.asyncio async def test_concurrency_reservation_released_on_disconnect(time_controller, monkeypatch): monkeypatch.setattr( 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 662fb366718..6dc7da61c56 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 @@ -2234,6 +2234,50 @@ async def test_concurrency_slot_released_on_failure_frees_capacity(time_controll assert result == healthy +@pytest.mark.asyncio +async def test_concurrency_slot_released_when_a_different_hook_rejects_the_request(time_controller): + """ + global_tag_rate_limits_hook raises the identical ProxyRateLimitError + shape (detail["error"] == "tag_rate_limit_exceeded") this hook's own + admission does. async_log_failure_event fires on every registered + CustomLogger regardless of which one raised, so this hook must still + release its own successfully reserved concurrency slot when the *other* + hook is what rejected the request -- skipping release based on the + shared marker alone would leak this hook's own slot until the safety + TTL, even though nothing about this hook's own admission failed. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + 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, + ) + + other_hooks_rejection = ProxyRateLimitError( + detail={"error": "tag_rate_limit_exceeded", "type": "requests", "tag_id": "end_user_id"}, + headers={"retry-after": "60"}, + rate_limit_type=None, + model="grp", + llm_provider="litellm_proxy", + ) + kwargs["exception"] = other_hooks_rejection + await limiter.async_log_failure_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) + + 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_on_fallback_recovered_hop_failure(time_controller): """