scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

revert(rate-limiting): drop per-branch concurrency reservation width, abatch_completion never fires per-branch terminal events

Browse source 
Reconciliation finding from #36541: verified against the real proxy pipeline
(common_processing_pre_call_logic -> route_request -> abatch_completion, not
a hand-built abatch_completion call with its own fresh kwargs) that every
branch of a comma-separated abatch_completion dispatch shares the identical
litellm_logging_obj/model_call_details the proxy attaches before the
dispatch. Logging.should_run_logging gates each event type on that shared
object's own has_logged_{event_type} flag, so at most one success event and
at most one failure event ever fire for the whole dispatch, never one per
branch -- confirmed live, including the mixed-outcome case (one branch
succeeds, another fails) where both can still fire once each since the two
flags gate independently.

This invalidates the batch-width reservation design from earlier rounds
(_non_racing_batch_width, its nested-messages accounting, and the
one-entry-per-branch release split): there is no per-branch multiplicity of
terminal events to drain, so reserving more than one unit for the whole
dispatch only makes a real batch request get rejected against a cap sized
for a single admission. Reverts admission to reserve exactly one
concurrency unit for the whole dispatch and release it on whichever
terminal event fires, the same model already used for
abatch_completion_fastest_response.

_own_concurrency_keys/_pop_matching_keys_for_call_id (releasing every
concurrency policy a request matched, not just one) and the
apply_to_models individual-model-name matching are unaffected: both are
about identity/policy matching, not branch counting, and remain necessary
regardless of this reconciliation.
This commit is contained in:
Deepanshu 2026-09-03 11:03:19 -04:00
parent 71a27935fc
commit 6312aebc17
2 changed files with 74 additions and 357 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -36,18 +36,20 @@ correctly releases its concurrency reservation through this mechanism.
`Router.abatch_completion`'s comma-separated `model` fans this one admission
out into several real, independent LLM calls below this hook entirely (it
never re-runs `async_pre_call_hook` per branch the way
`model_based_tag_rate_limits_hook`'s per-Router-hop admission does), each of
which reliably fires its own terminal success/failure event -- confirmed
live against the real dispatch. A single concurrency reservation for the
whole batch would be released by whichever branch finishes first, letting
the still-running siblings push real concurrent calls past the configured
cap, so admission instead reserves one unit per comma-separated model (see
`_non_racing_batch_width`) and each branch's own event releases just its own
share (see `_release_own_share`). The racing
`abatch_completion_fastest_response` variant is excluded from that: it
cancels every losing branch without ever firing a terminal event for it
(confirmed live), so reserving more than the single unit it already does
would leak every losing branch's share until the safety TTL.
`model_based_tag_rate_limits_hook`'s per-Router-hop admission does). Every
branch shares the identical `litellm_logging_obj`/`model_call_details` the
proxy attached before the dispatch, and `Logging.should_run_logging` gates
each event type (`async_success`, `async_failure`) on that same shared
dict's own `has_logged_{event_type}` flag -- confirmed live against the real
pipeline (`common_processing_pre_call_logic` -> `route_request` ->
`abatch_completion`, not a hand-built `abatch_completion` call with its own
fresh kwargs, which never shares that flag and was mistaken for per-branch
firing in an earlier round): at most one success event and at most one
failure event ever fire for the whole dispatch, never one per branch.
Reserving a single unit for the whole dispatch and releasing it on whichever
of those fires is therefore correct and sufficient -- the same model already
used for the racing `abatch_completion_fastest_response` variant, whose
cancelled losing branches never fire a terminal event at all.
"""
import asyncio
@ -199,44 +201,6 @@ def _individual_model_names(model: str | None, call_type: str) -> tuple[str, ...
return tuple(m.strip() for m in model.split(","))
def _non_racing_batch_width(data: Mapping[str, object], call_type: str) -> int:
"""`Router.abatch_completion` (not `abatch_completion_fastest_response`)
fans this one admission out into one independent real LLM call per
comma-separated model in `model`, and every one of those branches
reliably fires its own terminal success/failure event -- confirmed
empirically, no cancellation involved, unlike the racing
`fastest_response` variant, which cancels every losing branch without
ever firing a terminal event for it and must keep reserving a single
unit released by whichever branch finishes first.
A single concurrency reservation for the whole non-racing dispatch would
get released by whichever branch finishes first, letting the
still-running siblings push real concurrent calls past the configured
cap. Reserving one unit per branch instead, released one at a time as
each branch's own event fires, keeps the count accurate.
Mirrors the exact condition `route_llm_request.py` uses to route into
`abatch_completion` in the first place, so this only ever fires for a
request that will actually take that path. `abatch_completion` itself
also accepts a nested `messages: list[list[...]]` ("N requests to M
models") and dispatches one branch per (message, model) pair, not one
per model -- veria-ai's finding on this file.
"""
if call_type != "acompletion" or data.get("fastest_response"):
return 1
model_field: Final = data.get("model")
if not isinstance(model_field, str) or "," not in model_field:
return 1
model_count: Final = len(model_field.split(","))
messages_field: Final = data.get("messages")
message_list_count: Final = (
len(messages_field)
if isinstance(messages_field, list) and messages_field and all(isinstance(m, list) for m in messages_field)
else 1
)
return model_count * message_list_count
def _hash_tag(entry: TagRateLimitEntry, unit: _LimitUnit, tag_value: str, key_hash: str | None) -> str:
"""Namespaced under `tag_rl:global:` so it never collides with
`model_based_tag_rate_limits_hook`'s own `tag_rl:{model_group}:...` keys."""
@ -290,11 +254,11 @@ class _GlobalTagRateLimitStash:
# an earlier attempt must still get its accounting at success time even
# though the request ultimately serves from a later attempt's model.
admitted_models: frozenset[str] = field(default_factory=frozenset)
# One entry per reserved unit, not one entry per distinct key: a
# non-racing batch dispatch (see _non_racing_batch_width) reserves and
# appends `batch_width` entries per matching policy, and each branch's
# own terminal event pops exactly one entry per its own matching
# policies -- see _release_own_share.
# One entry per matching concurrency-scoped policy, not one entry total:
# a request can match more than one concurrency entry (a global cap and
# a named per-team cap on the same tag, for example), and the terminal
# event that releases this admission must release every one of them --
# see _release_own_share.
pending_concurrency_keys: list[tuple[str, _PartitionKey]] = field(default_factory=list) # mutable-ok: queue
# Keys already charged for this call_id, so a fallback retry (same
# litellm_call_id, different model) renews instead of double-charging.
@ -695,7 +659,6 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
already_reserved_concurrency_keys: Final = frozenset(
key for key, _partition_key in stash.pending_concurrency_keys
)
non_racing_batch_width: Final = _non_racing_batch_width(data, call_type)
failing_index, values = await self._atomic_check_and_increment(
tuple(
(
@ -705,16 +668,13 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
# A key already charged/reserved for this call_id (an
# earlier fallback attempt for the same request) renews
# at zero net cost instead of charging a second unit.
# Otherwise a concurrency check reserves one unit per
# non-racing batch branch (see _non_racing_batch_width),
# not just one for the whole dispatch.
0.0
if renewal_allowed
and (
(check.unit == "requests" and check.key in stash.charged_request_keys)
or (check.unit == "concurrency" and check.key in already_reserved_concurrency_keys)
)
else (float(non_racing_batch_width) if check.unit == "concurrency" else 1.0),
else 1.0,
self._ttl_for(check.unit, check.entry),
check.unit == "concurrency",
)
@ -729,14 +689,11 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
# Exclude already_reserved_concurrency_keys: that key renewed at
# zero cost above, so re-adding it would make release decrement
# twice for a counter only ever incremented once. One entry per
# reserved unit (non_racing_batch_width of them, ordinarily 1) --
# see pending_concurrency_keys's own docstring for why.
# twice for a counter only ever incremented once.
concurrency_reservations: Final = tuple(
(check.key, _partition_key(check.entry))
for check in atomic_checks
if check.unit == "concurrency" and check.key not in already_reserved_concurrency_keys
for _ in range(non_racing_batch_width)
)
if concurrency_reservations:
stash.pending_concurrency_keys.extend(concurrency_reservations) # mutable-ok: see field's own docstring
@ -761,13 +718,13 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
async def _release_pending_for_call_id(self, request_kwargs: Mapping[str, object]) -> None:
"""Releases every reservation still pending for this call_id at once.
Correct for a disconnect or a chain-exhausted failure: either aborts
every not-yet-completed branch of a non-racing batch together (see
_non_racing_batch_width), so none of them will ever fire its own
terminal event to release its own share individually -- whatever's
still pending here is exactly what those abandoned branches reserved,
no more (any branch that already completed already popped its own
entry via _release_one_pending_for_call_id) and no less."""
A single admission per call_id reserves at most one concurrency
entry per matching policy (see pending_concurrency_keys's own
docstring), never more than once per policy, so releasing everything
pending here is always exactly this one admission's own share --
correct for a disconnect, a chain-exhausted failure, or as the
fallback in _release_own_share when this hop's own identity can't be
resolved at release time."""
stash: Final = _stash_for_call(_call_id_from_kwargs(request_kwargs))
if stash is None or not stash.pending_concurrency_keys:
return
@ -775,21 +732,6 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
stash.pending_concurrency_keys.clear()
await self._release_keys(release_keys)
async def _release_one_pending_for_call_id(
self, request_kwargs: Mapping[str, object]
) -> tuple[str, _PartitionKey] | None:
"""Releases exactly one reservation, not every reservation currently
pending. Fallback only, for when this hop's own identity can't be
resolved at release time (see _own_concurrency_keys_for_release):
entries under one call_id are otherwise fungible (same key
repeated), so releasing one arbitrary entry is safe even without
knowing which policy it belongs to, but releasing every entry would
risk sweeping up a still-running sibling branch's own share."""
stash: Final = _stash_for_call(_call_id_from_kwargs(request_kwargs))
if stash is None or not stash.pending_concurrency_keys:
return None
return stash.pending_concurrency_keys.pop() # mutable-ok: see field's own docstring
async def _pop_matching_keys_for_call_id(
self, request_kwargs: Mapping[str, object], only_keys: frozenset[str]
) -> tuple[tuple[str, _PartitionKey], ...]:
@ -797,10 +739,8 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
reservation currently pending: a request can match more than one
concurrency-scoped entry (see _own_concurrency_keys), and each
terminal event must release every one of its own matches, not just
one -- while still leaving a still-running sibling branch's own
share (a different key, or another fungible entry under a shared
key) untouched. An empty only_keys means this hop's own identity
resolved to zero matching policies, so it pops nothing."""
one. An empty only_keys means this hop's own identity resolved to
zero matching policies, so it pops nothing."""
stash: Final = _stash_for_call(_call_id_from_kwargs(request_kwargs))
if stash is None or not stash.pending_concurrency_keys:
return ()
@ -856,8 +796,8 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
own_concurrency_keys: Final = self._own_concurrency_keys_for_release(kwargs)
if own_concurrency_keys is not None:
return await self._pop_matching_keys_for_call_id(kwargs, own_concurrency_keys)
one_entry: Final = await self._release_one_pending_for_call_id(kwargs)
return (one_entry,) if one_entry is not None else ()
await self._release_pending_for_call_id(kwargs)
return ()
async def async_release_disconnect_state_hook(self, request_data: Mapping[str, object]) -> None:
await self._release_pending_for_call_id(request_data)
@ -893,9 +833,8 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
# rejection never reserves a slot, so there is nothing to pop in
# that case; a rejection from model_based_tag_rate_limits_hook (same
# error marker) can still land after this hook already reserved its
# own slot. Only this one branch's own share, not every reservation
# still pending for a non-racing batch's other, still-running
# branches -- see _release_own_share.
# own slot. Every concurrency policy this admission matched, not
# just one -- see _release_own_share.
released_entries: Final = await self._release_own_share(kwargs)
if released_entries:
await self._release_keys(released_entries)

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

@ -18,7 +18,6 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitErro
from litellm.proxy.hooks.global_tag_rate_limits_hook import (
_PROXY_GlobalTagRateLimitsHook,
)
from litellm.router import Router
class TimeController:
@ -1233,157 +1232,55 @@ async def test_concurrent_requests_do_not_share_each_others_reservation_state(ti
# ---------------------------------------------------------------------------
# Non-racing batch dispatch (Router.abatch_completion via a comma-separated
# `model`): one admission fans out into several real, independent LLM calls,
# each firing its own terminal event -- confirmed live against the real
# Router.abatch_completion (not just this hook in isolation) that every
# branch always emits exactly one success or failure event, unlike the
# racing abatch_completion_fastest_response variant, whose cancelled losers
# never emit any terminal event at all.
# Batch dispatch (Router.abatch_completion via a comma-separated `model`):
# every branch shares the identical litellm_logging_obj/model_call_details
# the proxy attaches before the dispatch. Verified against the real pipeline
# (common_processing_pre_call_logic -> route_request -> abatch_completion,
# not a hand-built abatch_completion call with its own fresh kwargs, which
# never shares that object and looks like independent per-branch firing):
# Logging.should_run_logging gates each event type on that shared dict's own
# has_logged_{event_type} flag, so at most one success event and at most one
# failure event ever fire for the whole dispatch, never one per branch.
# Reserving a single unit for the whole dispatch and releasing it on
# whichever of those fires is correct -- the same model already used for the
# racing abatch_completion_fastest_response variant, whose cancelled losing
# branches never fire a terminal event at all.
# ---------------------------------------------------------------------------
def _batch_data(tags: list[str], model: str, call_id: str = "call-1", fastest_response: bool = False) -> dict:
data = {**_data(tags, call_id=call_id), "model": model}
if fastest_response:
data["fastest_response"] = True
return data
@pytest.mark.asyncio
async def test_non_racing_batch_reserves_one_unit_per_model(time_controller, monkeypatch):
"""A single admission for a 3-model comma-separated batch must reserve 3
units, not 1 -- otherwise the cap is measured against the wrong number of
real concurrent LLM calls the batch actually makes."""
async def test_batch_dispatch_reserves_a_single_concurrency_unit(time_controller, monkeypatch):
"""Reserving one unit per comma-separated model would reject this
3-model batch outright against a cap of 1, since the real dispatch only
ever produces at most one real terminal event total (see module
docstring), not one per model."""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{
"concurrency_limits": {
"limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 3, "period_seconds": 60}]
"limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}]
}
},
)
hook = _make_hook(time_controller)
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_batch_data(["end_user_id:u1"], model="model-a,model-b,model-c", call_id="call-1"),
call_type="acompletion",
)
# The batch alone already occupies every unit of the cap of 3.
with pytest.raises(ProxyRateLimitError):
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",
)
@pytest.mark.asyncio
async def test_non_racing_batch_releases_one_unit_per_branch(time_controller, monkeypatch):
"""Each branch's own success/failure event must release only its own
share of a non-racing batch's reservation -- releasing the whole
reservation on the first branch to finish would let a new, unrelated
request in while the batch's own remaining branches are still genuinely
in flight, exceeding the configured cap."""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{
"concurrency_limits": {
"limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 3, "period_seconds": 60}]
}
},
)
hook = _make_hook(time_controller)
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_batch_data(["end_user_id:u1"], model="model-a,model-b,model-c", call_id="call-1"),
call_type="acompletion",
)
batch_kwargs = {"litellm_call_id": "call-1", "metadata": {"tags": ["end_user_id:u1"]}}
await hook.async_log_success_event(kwargs=batch_kwargs, response_obj=None, start_time=0, end_time=0)
await asyncio.sleep(0)
# Only 1 of 3 branches finished -- the other 2 still hold their own
# units, so a brand-new request needing all 3 remaining units of the cap
# must still be rejected.
with pytest.raises(ProxyRateLimitError):
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_batch_data(["end_user_id:u1"], model="model-x,model-y,model-z", call_id="call-2"),
call_type="acompletion",
)
# The other 2 branches finish (one success, one failure).
await hook.async_log_success_event(kwargs=batch_kwargs, response_obj=None, start_time=0, end_time=0)
await hook.async_log_failure_event(kwargs=batch_kwargs, response_obj=None, start_time=0, end_time=0)
await asyncio.sleep(0)
# Now every unit from the first batch is back -- no leak, no
# over-release along the way.
result = await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_data(["end_user_id:u1"], call_id="call-3"),
call_type="completion",
data={**_data(["end_user_id:u1"], call_id="call-1"), "model": "model-a,model-b,model-c"},
call_type="acompletion",
)
assert result is not None
@pytest.mark.asyncio
async def test_non_racing_batch_disconnect_releases_only_what_remains(time_controller, monkeypatch):
"""A client disconnect mid-batch cancels every branch that hasn't
already finished together, so none of those cancelled branches ever
fires its own terminal event. The disconnect hook must release exactly
what's left -- not double-release a branch that already released its
own share via a real success/failure event first."""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{
"concurrency_limits": {
"limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 3, "period_seconds": 60}]
}
},
)
hook = _make_hook(time_controller)
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_batch_data(["end_user_id:u1"], model="model-a,model-b,model-c", call_id="call-1"),
call_type="acompletion",
)
batch_kwargs = {"litellm_call_id": "call-1", "metadata": {"tags": ["end_user_id:u1"]}}
await hook.async_log_success_event(kwargs=batch_kwargs, response_obj=None, start_time=0, end_time=0)
await asyncio.sleep(0)
await hook.async_release_disconnect_state_hook({"litellm_call_id": "call-1"})
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_fastest_response_batch_still_reserves_only_one_unit(time_controller, monkeypatch):
"""abatch_completion_fastest_response cancels every losing branch
without ever firing a terminal event for it, so it must keep reserving a
single unit for the whole dispatch -- reserving one per model here would
leak every losing branch's share until the safety TTL on every single
call."""
async def test_batch_dispatch_release_is_safe_across_both_real_terminal_events(time_controller, monkeypatch):
"""has_logged_async_success and has_logged_async_failure gate
independently, so a mixed-outcome dispatch (one branch succeeds,
another fails) can fire both a real success event and a real failure
event for the same shared reservation. Releasing on the first must
leave the second nothing to release, not double-release or error."""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
@ -1398,11 +1295,10 @@ async def test_fastest_response_batch_still_reserves_only_one_unit(time_controll
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_batch_data(["end_user_id:u1"], model="model-a,model-b", call_id="call-1", fastest_response=True),
data={**_data(["end_user_id:u1"], call_id="call-1"), "model": "model-a,model-b"},
call_type="acompletion",
)
# If this admission had reserved 2 units instead of 1, this second,
# unrelated request (cap of 1) would already be rejected here too.
kwargs = {"litellm_call_id": "call-1", "metadata": {"tags": ["end_user_id:u1"]}}
await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0)
await asyncio.sleep(0)
@ -1415,94 +1311,13 @@ async def test_fastest_response_batch_still_reserves_only_one_unit(time_controll
)
assert result is not None
@pytest.mark.asyncio
async def test_real_abatch_completion_admits_once_and_releases_per_branch(time_controller, monkeypatch):
"""End-to-end against the real Router.abatch_completion (mocked LLM
responses, no network): confirms admission fires exactly once for the
whole dispatch -- not once per branch the way model_based_tag_rate_limits_hook's
per-Router-hop admission does -- and that a cap sized to the batch width
is never exceeded even while some branches are still genuinely running."""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{
"concurrency_limits": {
"limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 3, "period_seconds": 60}]
}
},
)
hook = _make_hook(time_controller)
monkeypatch.setattr(litellm, "callbacks", [hook])
router = Router(
model_list=[
{
"model_name": "model-fast",
"litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "fast-done", "mock_delay": 0.01},
},
{
"model_name": "model-slow-1",
"litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "slow-1-done", "mock_delay": 0.3},
},
{
"model_name": "model-slow-2",
"litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "slow-2-done", "mock_delay": 0.3},
},
]
)
admitted_data = await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_batch_data(["end_user_id:u1"], model="model-fast,model-slow-1,model-slow-2", call_id="call-1"),
call_type="acompletion",
)
batch_task = asyncio.create_task(
router.abatch_completion(
models=["model-fast", "model-slow-1", "model-slow-2"],
messages=[{"role": "user", "content": "hello"}],
metadata=admitted_data["metadata"],
litellm_call_id=admitted_data["litellm_call_id"],
)
)
# Give only the fast branch (0.01s mock_delay) time to finish; the two
# slow branches (0.3s) are still genuinely in flight.
await asyncio.sleep(0.1)
# 1 of 3 units released, 2 still held by the still-running slow
# branches -- a brand-new request needing all 3 units of the cap must
# still be rejected.
with pytest.raises(ProxyRateLimitError):
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_batch_data(["end_user_id:u1"], model="model-x,model-y,model-z", call_id="call-2"),
call_type="acompletion",
)
responses = await batch_task
assert [r.choices[0].message.content for r in responses] == ["fast-done", "slow-1-done", "slow-2-done"]
await asyncio.sleep(0)
# Every branch has now finished and released its own share -- no leak.
result = await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_data(["end_user_id:u1"], call_id="call-3"),
call_type="completion",
)
assert result is not None
await hook.async_log_failure_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0)
@pytest.mark.asyncio
async def test_comma_separated_model_outside_acompletion_call_type_is_not_a_batch(time_controller, monkeypatch):
"""A comma in `model` only means a non-racing batch dispatch for the
exact call_type route_llm_request.py itself gates on -- any other
call_type must reserve just 1 unit, matching whatever that call_type
actually does under the hood."""
async def test_fastest_response_batch_also_reserves_a_single_unit(time_controller, monkeypatch):
"""abatch_completion_fastest_response's own comma-separated `model` must
reserve exactly one unit too, same as the non-racing dispatch above."""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
@ -1514,54 +1329,17 @@ async def test_comma_separated_model_outside_acompletion_call_type_is_not_a_batc
)
hook = _make_hook(time_controller)
await hook.async_pre_call_hook(
result = await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data=_batch_data(["end_user_id:u1"], model="model-a,model-b,model-c", call_id="call-1"),
call_type="completion",
)
with pytest.raises(ProxyRateLimitError):
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",
)
@pytest.mark.asyncio
async def test_non_racing_batch_width_accounts_for_nested_message_lists(time_controller, monkeypatch):
"""veria-ai finding: Router.abatch_completion's "N requests to M models"
mode (a nested `messages: list[list[...]]`) dispatches one real branch
per (message, model) pair, not one per model -- admission must reserve
model_count * message_list_count units, not just model_count."""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{
"concurrency_limits": {
"limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 4, "period_seconds": 60}]
}
data={
**_data(["end_user_id:u1"], call_id="call-1"),
"model": "model-a,model-b",
"fastest_response": True,
},
call_type="acompletion",
)
hook = _make_hook(time_controller)
data = {
**_data(["end_user_id:u1"], call_id="call-1"),
"model": "model-a,model-b",
"messages": [[{"role": "user", "content": "q1"}], [{"role": "user", "content": "q2"}]],
}
await hook.async_pre_call_hook(user_api_key_dict=_key(), cache=DualCache(), data=data, call_type="acompletion")
# 2 models x 2 message-lists = 4 real branches, exactly at the cap of 4 --
# a fifth unit for anything else must be rejected.
with pytest.raises(ProxyRateLimitError):
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