scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(rate-limiting): check apply_to_models against each individual comma-separated model name

Browse source 
apply_to_models was checked against the raw model field as a single string,
so a request going through Router.abatch_completion (comma-separated model,
e.g. "opus-chain,other") never matched a chain-scoped entry's apply_to_models
list, since the joined string is never a member of a list of individual
names. Confirmed live: a chain-scoped concurrency cap already at its limit
for a single-model "opus-chain" request admitted a second "opus-chain,other"
request without ever raising.

_individual_model_names splits the caller's model field the same way
route_llm_request.py does before it actually dispatches through
Router.abatch_completion (call_type == acompletion, comma present).
_classify's admission-time check and admitted_models' success-time
accounting both now test candidate models individually via
_entry_applies_any_candidate_model, a generalization of the existing
admitted-models-only helper. _admitted_models_after now folds in every
individual name from a batch admission, not the joined string, so the
success-time accounting bug closes the same way the admission-time one does.
This commit is contained in:
Deepanshu 2026-09-01 10:47:31 -04:00
parent 06f71891b6
commit 79b8d433da
2 changed files with 163 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -144,18 +144,35 @@ else:
Span: TypeAlias = object
def _entry_applies_any_admitted_model(
entry: TagRateLimitEntry, tags: Sequence[str], key_alias: str | None, admitted_models: frozenset[str]
def _entry_applies_any_candidate_model(
entry: TagRateLimitEntry, tags: Sequence[str], key_alias: str | None, candidate_models: frozenset[str]
) -> bool:
"""Same as `_entry_applies`, except an `apply_to_models`-scoped entry
counts as applying if ANY model an admission attempt for this call_id
saw was in scope -- not just whichever model the call ultimately served.
A `_pre_call_with_fallbacks` retry re-admits with a different model for
the same call_id, and an entry that matched an earlier attempt must
still get its success-time accounting."""
if not admitted_models:
counts as applying if ANY of `candidate_models` is in scope -- not just a
single caller-visible name. Two call sites need this: at admission, a
comma-separated `model` (a non-racing or racing batch dispatch) names
several models at once, any of which should trip a chain-wide cap; at
success-event accounting, a `_pre_call_with_fallbacks` retry re-admits
with a different model for the same call_id, and an entry that matched
an earlier attempt must still get its accounting."""
if not candidate_models:
return _entry_applies(entry, tags, key_alias, None)
return any(_entry_applies(entry, tags, key_alias, model) for model in admitted_models)
return any(_entry_applies(entry, tags, key_alias, model) for model in candidate_models)
def _individual_model_names(model: str | None, call_type: str) -> tuple[str, ...]:
"""`route_llm_request.py` splits a comma-separated `model` on this exact
condition before fanning out through `Router.abatch_completion`/
`abatch_completion_fastest_response` -- an `apply_to_models`-scoped entry
must check each of those individual names, not the raw joined string
(which is never a member of any caller-configured `apply_to_models`
list), or a caller trivially bypasses a chain-wide cap by adding a
second model to the comma list."""
if model is None:
return ()
if call_type != "acompletion" or "," not in model:
return (model,)
return tuple(m.strip() for m in model.split(","))
def _non_racing_batch_width(data: Mapping[str, object], call_type: str) -> int:
@ -447,15 +464,17 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
@staticmethod
def _admitted_models_after(
admitted_models: frozenset[str], model: str | None, renewal_allowed: bool
admitted_models: frozenset[str], candidate_models: frozenset[str], renewal_allowed: bool
) -> frozenset[str]:
"""Only called once this admission attempt has cleared every check
without raising -- a rejected attempt's model must never join
without raising -- a rejected attempt's models must never join
admitted_models, or a later successful attempt's accounting could
wrongly credit an apply_to_models entry that never actually admitted
this request under that model."""
if renewal_allowed and model is not None:
return admitted_models | frozenset((model,))
this request under that model. `candidate_models` is every
individual name a comma-separated `model` names (see
`_individual_model_names`), not the raw joined string."""
if renewal_allowed and candidate_models:
return admitted_models | candidate_models
return admitted_models
@staticmethod
@ -472,7 +491,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
key_alias: str | None,
key_hash: str | None,
now: float,
model: str | None,
candidate_models: frozenset[str],
) -> tuple[_ClassifiedGlobalCheck, ...]:
classified: Final = [] # mutable-ok: sequential accumulator, immediately frozen into a tuple below
for unit in _LIMIT_UNITS:
@ -483,7 +502,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
tag_value = _extract_identity(tags, entry.tag_id)
if tag_value is None:
continue
if not _entry_applies(entry, tags, key_alias, model):
if not _entry_applies_any_candidate_model(entry, tags, key_alias, candidate_models):
continue
effective_key_hash = key_hash if entry.scope_by_key_hash else None
if unit == "concurrency":
@ -609,6 +628,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
key_alias: Final = user_api_key_dict.key_alias
key_hash: Final = user_api_key_dict.api_key
model: Final = data.get("model") if isinstance(data.get("model"), str) else None
candidate_models: Final = frozenset(_individual_model_names(model, call_type))
# First admission for this stash claims ownership; only a later one
# with the same key_hash may renew its charges (see owner_key_hash).
@ -618,9 +638,11 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
now: Final = self._time_provider().timestamp()
stash.admission_time = now
classified: Final = self._classify(config, tags, key_alias, key_hash, now, model)
classified: Final = self._classify(config, tags, key_alias, key_hash, now, candidate_models)
if not classified:
stash.admitted_models = self._admitted_models_after(stash.admitted_models, model, renewal_allowed)
stash.admitted_models = self._admitted_models_after(
stash.admitted_models, candidate_models, renewal_allowed
)
return data
read_only_checks: Final = tuple(c for c in classified if not c.is_atomic)
@ -700,7 +722,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
if request_keys:
stash.charged_request_keys.extend(request_keys) # mutable-ok: see field's own docstring
stash.admitted_models = self._admitted_models_after(stash.admitted_models, model, renewal_allowed)
stash.admitted_models = self._admitted_models_after(stash.admitted_models, candidate_models, renewal_allowed)
return data
async def _release_pending_for_call_id(self, request_kwargs: Mapping[str, object]) -> None:
@ -839,7 +861,7 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o
tag_value = _extract_identity(tags, entry.tag_id)
if tag_value is None:
continue
if not _entry_applies_any_admitted_model(entry, tags, key_alias, admitted_models):
if not _entry_applies_any_candidate_model(entry, tags, key_alias, admitted_models):
continue
increment_value = increment_by_unit[unit]
if increment_value == 0:

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

@ -610,6 +610,127 @@ async def test_apply_to_models_accounts_when_a_fallback_retry_re_admits_with_a_d
)
@pytest.mark.asyncio
async def test_apply_to_models_enforces_for_a_comma_separated_batch_model(time_controller, monkeypatch):
"""Bugbot finding: a caller can dodge a chain-wide apply_to_models cap by
routing through Router.abatch_completion's comma-separated `model` --
admission checks the individual names, not the raw joined string, which
would never match any configured apply_to_models entry."""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{
"concurrency_limits": {
"limits": [
{
"name": "chain_cap",
"tag_id": "end_user_id",
"limit": 1,
"period_seconds": 60,
"apply_to_models": ["opus-chain"],
}
]
}
},
)
hook = _make_hook(time_controller)
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data={**_data(["end_user_id:u1"], call_id="call-1"), "model": "opus-chain"},
call_type="acompletion",
)
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"), "model": "opus-chain,other-model"},
call_type="acompletion",
)
@pytest.mark.asyncio
async def test_apply_to_models_accounts_dollar_spend_from_a_batch_admission(time_controller, monkeypatch):
"""Same gap as above, at success-time accounting: the batch admission's
own admitted_models must record each individual name, not the joined
string, or the spend never lands in the chain-scoped bucket at all."""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{
"dollar_limits": {
"limits": [
{
"name": "chain_spend",
"tag_id": "end_user_id",
"limit": 10.0,
"period_seconds": 86400,
"apply_to_models": ["opus-chain"],
}
]
}
},
)
hook = _make_hook(time_controller)
await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data={**_data(["end_user_id:u1"], call_id="call-1"), "model": "opus-chain,other-model"},
call_type="acompletion",
)
kwargs = {
"litellm_call_id": "call-1",
"metadata": {"tags": ["end_user_id:u1"]},
"standard_logging_object": {"total_tokens": 0, "response_cost": 12.0},
}
await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0)
await asyncio.sleep(0)
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"), "model": "opus-chain"},
call_type="acompletion",
)
@pytest.mark.asyncio
async def test_apply_to_models_batch_split_requires_the_acompletion_call_type(time_controller, monkeypatch):
"""route_llm_request.py only splits a comma-separated model for
call_type acompletion -- any other call type never reaches
Router.abatch_completion, so a comma there is just a literal (if odd)
model name and must not be split."""
monkeypatch.setattr(
litellm,
"global_tag_rate_limits",
{
"concurrency_limits": {
"limits": [
{
"name": "chain_cap",
"tag_id": "end_user_id",
"limit": 1,
"period_seconds": 60,
"apply_to_models": ["opus-chain"],
}
]
}
},
)
hook = _make_hook(time_controller)
result = await hook.async_pre_call_hook(
user_api_key_dict=_key(),
cache=DualCache(),
data={**_data(["end_user_id:u1"], call_id="call-1"), "model": "opus-chain,other-model"},
call_type="embedding",
)
assert result is not None
@pytest.mark.asyncio
async def test_a_rejected_admission_attempts_model_does_not_drive_later_accounting(time_controller, monkeypatch):
"""