From c62d79f7dfd1eb8121dfa11f9b66e6ebcae6efc7 Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Fri, 14 Aug 2026 13:05:49 -0400 Subject: [PATCH] fix(rate-limiting): resolve limits by deployment model_name for routing-group calls 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. --- litellm/proxy/hooks/tag_rate_limiter.py | 63 +++++++- .../proxy/hooks/test_tag_rate_limiter.py | 153 ++++++++++++++++++ 2 files changed, 211 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/hooks/tag_rate_limiter.py b/litellm/proxy/hooks/tag_rate_limiter.py index bf3f231788f..875fc8338a7 100644 --- a/litellm/proxy/hooks/tag_rate_limiter.py +++ b/litellm/proxy/hooks/tag_rate_limiter.py @@ -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( { diff --git a/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py b/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py index 856883c62cc..e1fcf76ab5e 100644 --- a/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_tag_rate_limiter.py @@ -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 # ---------------------------------------------------------------------------