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): resolve limits by deployment model_name for routing-group calls

Browse source 
Router deliberately keeps a callable routing-group name (and, per its
own design, a direct deployment-id address) distinct from every member
deployment's own model_name (Router._get_routing_group_deployments).
Since _LimitsIndex only keys by model_name and team alias, a
group-addressed call previously matched neither table and the limiter
silently no-opped for both admission (async_filter_deployments) and
success accounting (async_log_success_event), even though the member
deployments carried real tag_rate_limits under their own model_name.

Add _LimitsIndex.resolve_any(), which falls back to resolving via each
candidate deployment's own model_name when the caller-visible name
matches neither table, stamping each result with the model_name it
came from (_ConfiguredLimit.resolved_group) so hashing stays
namespaced per underlying model_name -- otherwise two different
model_names sharing one routing group with an identically-named,
identically-configured limit would collide onto one Redis counter.

Direct deployment-id addressing has a related but separate, more severe
gap: Router.async_get_healthy_deployments returns a single dict (not a
list) for that path and short-circuits before Router.async_callback_
filter_deployments is ever called, so every CustomLogger.async_filter_
deployments-based hook is skipped, not just this one. That is a
Router-level structural issue affecting many hooks and is out of scope
for this PR.
This commit is contained in:
Deepanshu 2026-08-14 13:05:49 -04:00
parent 8d4528ba37
commit c62d79f7df
2 changed files with 211 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

63
litellm/proxy/hooks/tag_rate_limiter.py
View file

@ -2,7 +2,7 @@
import asyncio
import contextvars
from collections.abc import Callable, Mapping, Sequence
from collections.abc import Callable, Iterable, Mapping, Sequence
from dataclasses import dataclass, replace
from datetime import datetime
from itertools import groupby
@ -132,6 +132,17 @@ class _ConfiguredLimit:
# same Redis counter despite the index itself correctly scoping the
# lookup by (team_id, alias).
team_scope: str | None = None
# The real model_name this limit was found under when `resolve()`'s
# direct lookup by the caller-visible model string missed and
# `resolve_any()` fell back to resolving via a candidate deployment's
# own model_name instead (routing groups, and any other indirection
# where Router deliberately keeps the caller-visible name distinct from
# every deployment's own model_name). None when resolved directly, in
# which case the caller-visible name is already unambiguous and safe to
# hash by. Set, this overrides the caller-visible name in the bucket key
# so limits from two different underlying model_names sharing one
# routing group never collide on one counter.
resolved_group: str | None = None
def _extract_identity(tags: Sequence[str], tag_id: str) -> str | None:
@ -313,6 +324,30 @@ class _LimitsIndex:
return scoped
return self.by_model_name.get(model, ())
def resolve_any(
self, model: str, team_id: str | None, candidate_model_names: Iterable[str]
) -> tuple[_ConfiguredLimit, ...]:
"""
Like `resolve`, but falls back to each candidate deployment's own
`model_name` when `model` itself matches neither table -- Router
deliberately keeps `model` as a callable routing-group name distinct
from every member deployment's own `model_name` (see
`Router._get_routing_group_deployments`), so a group-addressed call
would otherwise never match this index at all despite its member
deployments carrying real `tag_rate_limits`. Each fallback result is
stamped with the `model_name` it actually came from (`resolved_group`)
so hashing stays namespaced per underlying group even when the
candidates span more than one `model_name`.
"""
direct: Final = self.resolve(model, team_id)
if direct:
return direct
return tuple(
replace(limit, resolved_group=name)
for name in frozenset(candidate_model_names)
for limit in self.by_model_name.get(name, ())
)
def _team_alias_key(deployment: Mapping[str, object]) -> tuple[str, str] | None:
model_info: Final = deployment.get("model_info") or _EMPTY_MAPPING
@ -464,6 +499,13 @@ def _scope_suffix(deployment_scope: tuple[str, ...] | None) -> str:
def _hash_tag(model_group: str, configured: _ConfiguredLimit, tag_value: str, key_hash: str | None) -> str:
# resolved_group overrides the caller-visible model_group when this
# limit was found via resolve_any()'s per-deployment fallback (routing
# groups): the caller-visible name is ambiguous there (shared by every
# member model_name), so hashing by it would collide two different
# underlying model_names' identically-named limits onto one counter.
# See _ConfiguredLimit.resolved_group.
effective_model_group: Final = configured.resolved_group if configured.resolved_group is not None else model_group
scope: Final = _scope_suffix(configured.deployment_scope)
# team_scope disambiguates two teams that publish the identical
# team_public_model_name alias with identically-configured limits --
@ -472,7 +514,7 @@ def _hash_tag(model_group: str, configured: _ConfiguredLimit, tag_value: str, ke
team_suffix: Final = f":team:{configured.team_scope}" if configured.team_scope is not None else ""
key_suffix: Final = f":key:{key_hash}" if key_hash is not None else ""
return (
f"tag_rl:{model_group}:{configured.unit}:{configured.entry.name}:{configured.entry.tag_id}:"
f"tag_rl:{effective_model_group}:{configured.unit}:{configured.entry.name}:{configured.entry.tag_id}:"
f"{scope}{team_suffix}:{tag_value}{key_suffix}"
)
@ -690,7 +732,10 @@ class _PROXY_TagRateLimiter( # pyright: ignore[reportUnusedClass] # only refer
resolved_request_kwargs: Final = request_kwargs or _EMPTY_MAPPING
metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(resolved_request_kwargs)
team_id: Final = _extract_team_id(resolved_request_kwargs, metadata_variable_name)
configured: Final = self._index.get(self.llm_router).resolve(model, team_id)
candidate_model_names: Final = tuple(
name for d in healthy_deployments if isinstance(name := d.get("model_name"), str)
)
configured: Final = self._index.get(self.llm_router).resolve_any(model, team_id, candidate_model_names)
if not configured:
return healthy_deployments
@ -882,7 +927,16 @@ class _PROXY_TagRateLimiter( # pyright: ignore[reportUnusedClass] # only refer
standard_logging_metadata: Final = standard_logging_object.get("metadata") or _EMPTY_MAPPING
team_id: Final = standard_logging_metadata.get("user_api_key_team_id")
key_hash: Final = standard_logging_metadata.get("user_api_key_hash")
configured: Final = self._index.get(self.llm_router).resolve(model_group, team_id)
# model_group is the caller-visible name, which Router deliberately
# keeps distinct from the serving deployment's own model_name for a
# routing-group call (see resolve_any's docstring); fall back to the
# one deployment that actually served this hop.
deployment_id: Final = standard_logging_object.get("model_id")
serving_deployment: Final = (
self.llm_router.get_deployment(deployment_id) if isinstance(deployment_id, str) else None
)
candidate_model_names: Final = (serving_deployment.model_name,) if serving_deployment is not None else ()
configured: Final = self._index.get(self.llm_router).resolve_any(model_group, team_id, candidate_model_names)
if not configured:
return
@ -900,7 +954,6 @@ class _PROXY_TagRateLimiter( # pyright: ignore[reportUnusedClass] # only refer
if not tags:
return
deployment_id: Final = standard_logging_object.get("model_id")
now: Final = self._time_provider().timestamp()
increment_by_unit: Final[Mapping[_LimitUnit, float]] = MappingProxyType(
{

153
tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py
View file

@ -257,6 +257,113 @@ async def test_filter_deployments_allows_under_limit_and_rejects_at_limit(time_c
assert exc_info.value.detail["limit_name"] == "per_minute"
@pytest.mark.asyncio
async def test_filter_deployments_falls_back_to_deployment_model_name_for_routing_group_calls(time_controller):
"""
Router keeps a callable routing-group name distinct from every member
deployment's own model_name (see Router._get_routing_group_deployments),
so async_filter_deployments can be called with model="my-group" while
healthy_deployments carries the group's real member deployments. The
limiter must still resolve and enforce each member's own configured
limits rather than silently no-opping because "my-group" itself never
appears in the index.
"""
limiter = _make_limiter(time_controller)
router = litellm.Router(
model_list=[
_deployment(
"backend-a",
"dep-1",
{
"request_limits": {
"limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 2, "period_seconds": 60}]
}
},
)
]
)
limiter.update_variables(llm_router=router)
healthy = router.model_list
for _ in range(2):
result = await limiter.async_filter_deployments(
model="my-group",
healthy_deployments=healthy,
messages=None,
request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}},
)
assert result == healthy
with pytest.raises(ProxyRateLimitError):
await limiter.async_filter_deployments(
model="my-group",
healthy_deployments=healthy,
messages=None,
request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}},
)
@pytest.mark.asyncio
async def test_filter_deployments_routing_group_does_not_collide_across_different_model_names(time_controller):
"""
A routing group can span deployments from different model_names that
happen to declare an identically-named, identically-configured limit.
Each must get its own bucket (keyed by its own model_name via
resolved_group), not share one just because the caller addressed both
through the same group name.
"""
limiter = _make_limiter(time_controller)
router = litellm.Router(
model_list=[
_deployment(
"backend-a",
"dep-a",
{
"request_limits": {
"limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}]
}
},
),
_deployment(
"backend-b",
"dep-b",
{
"request_limits": {
"limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}]
}
},
),
]
)
limiter.update_variables(llm_router=router)
healthy = router.model_list
# Exhaust backend-a's limit (limit=1) via the group-addressed call.
await limiter.async_filter_deployments(
model="my-group",
healthy_deployments=[healthy[0]],
messages=None,
request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}},
)
with pytest.raises(ProxyRateLimitError):
await limiter.async_filter_deployments(
model="my-group",
healthy_deployments=[healthy[0]],
messages=None,
request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}},
)
# backend-b's own bucket must be untouched -- same group, same tag, same
# limit name, but a different underlying model_name.
result = await limiter.async_filter_deployments(
model="my-group",
healthy_deployments=[healthy[1]],
messages=None,
request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}},
)
assert result == [healthy[1]]
@pytest.mark.asyncio
async def test_filter_deployments_per_entry_fail_open_when_tag_absent(time_controller):
"""
@ -522,6 +629,52 @@ async def test_log_success_event_reads_nested_litellm_metadata_when_that_is_auth
)
@pytest.mark.asyncio
async def test_log_success_event_falls_back_to_serving_deployment_model_name_for_routing_group_calls(
time_controller,
):
"""
standard_logging_object["model_group"] is the caller-visible name from
Router._update_kwargs_before_fallbacks -- for a routing-group call this
is the group name too, which never appears in the index. Success
accounting must fall back to the model_name of the deployment that
actually served this hop (standard_logging_object["model_id"]).
"""
limiter = _make_limiter(time_controller)
router = litellm.Router(
model_list=[
_deployment(
"backend-a",
"dep-1",
{
"token_limits": {
"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400}]
}
},
)
]
)
limiter.update_variables(llm_router=router)
kwargs = {
"metadata": {"tags": ["end_user_id:u1"]},
"standard_logging_object": {
"model_group": "my-group",
"model_id": "dep-1",
"total_tokens": 42,
"response_cost": 0.01,
},
}
await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0)
await asyncio.sleep(0)
now = time_controller.now().timestamp()
token_key = f"{{tag_rl:backend-a:tokens:daily:end_user_id:chain:u1}}:{int(now) // 86400}"
assert (
float(await limiter.internal_usage_cache.async_get_cache(key=token_key, litellm_parent_otel_span=None)) == 42.0
)
# ---------------------------------------------------------------------------
# concurrency limits -- reserve at admission, release on success/failure
# ---------------------------------------------------------------------------