scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix(rate-limiting): release concurrency slots on a cross-hook rejection

Browse source 
Both tag-rate-limit hooks raise the identical ProxyRateLimitError shape
(detail.error == tag_rate_limit_exceeded). async_log_failure_event fires on
every registered CustomLogger regardless of which one raised, so each hook's
own skip-on-that-marker special case incorrectly suppressed releasing a
reservation the *other* hook's rejection left pending, leaking the slot
until the safety TTL. Removes the special case from both hooks: a hook's
own rejection never reaches the point where a reservation is queued, so the
existing empty-pending-keys check already no-ops correctly for that case.

Found by Veria AI on PR #36541.
This commit is contained in:
Deepanshu 2026-08-24 16:17:28 -04:00
parent 095321afe4
commit cae03ad043
4 changed files with 106 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/proxy/hooks/global_tag_rate_limits_hook.py
View file

@ -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

16
litellm/proxy/hooks/model_based_tag_rate_limits_hook.py
View file

@ -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)

44
tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py
View file

@ -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(

44
tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py
View file

@ -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):
"""