diff --git a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py index 1a7cb3f12f7..8ca8b39f0be 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -662,6 +662,21 @@ def _decode_reservations(raw: object) -> tuple[tuple[str, "_PartitionKey"], ...] # request, so its own most recent admission timestamp is the right one. _ADMISSION_TIME_FIELD: Final[str] = "_model_based_tag_rate_limits_admission_time" +# The routing-group membership (candidate_model_names) admission actually +# resolved against, stashed the same way _ADMISSION_TIME_FIELD is. resolve_any +# dedupes divergent per-deployment entries by picking the alphabetically first +# member model_name sharing a signature (resolved_group) -- a pure function of +# this exact candidate set. Router's live routing-group membership can change +# between admission and success (a deployment added or removed mid-request via +# /model/new or a config hot-reload), and success independently re-deriving +# candidate_model_names from *live* membership at that later point can pick a +# different resolved_group than admission did, hashing to a different Redis +# key -- so success accounting silently misses the bucket admission actually +# checked, letting real usage escape the enforced cap. Reusing admission's own +# snapshot keeps resolve_any's output identical at both points regardless of +# what changed in between. +_ROUTING_GROUP_CANDIDATES_FIELD: Final[str] = "_model_based_tag_rate_limits_routing_group_candidates" + class _TagRateLimitIndex: """Rebuilds the limits index when `llm_router.model_list` changes, or at @@ -878,6 +893,25 @@ def _admission_time_or(kwargs: Mapping[str, object], fallback: float) -> float: return recorded if isinstance(recorded, float) else fallback +def _record_routing_group_candidates( + request_kwargs: Mapping[str, object], candidate_model_names: tuple[str, ...] +) -> None: + """Stash the routing-group membership admission resolved against -- see + `_ROUTING_GROUP_CANDIDATES_FIELD`'s docstring for why. Silently a no-op + without a real logging object (defensive only; every real request has + one): success accounting falls back to its own live reconstruction, same + as before this fix existed.""" + logging_obj: Final = request_kwargs.get("litellm_logging_obj") + model_call_details: Final = getattr(logging_obj, "model_call_details", None) + if isinstance(model_call_details, dict): + model_call_details[_ROUTING_GROUP_CANDIDATES_FIELD] = candidate_model_names + + +def _routing_group_candidates_or(kwargs: Mapping[str, object], fallback: tuple[str, ...]) -> tuple[str, ...]: + recorded: Final = kwargs.get(_ROUTING_GROUP_CANDIDATES_FIELD) + return recorded if isinstance(recorded, tuple) else fallback + + @dataclass(frozen=True, slots=True) class _CachePartition: internal_usage_cache: InternalUsageCache @@ -1106,9 +1140,11 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # the purpose of deciding resolved_group, and success accounting has no # way to know which members were healthy at admission time -- it can # only reconstruct the full, static membership (see its own comment - # below). Deriving both sides from the same full-membership source is - # the only way they're guaranteed to dedup to the identical bucket - # regardless of cooldown state at either point in time. + # below). Deriving both sides from the same full-membership source + # handles cooldown-state drift between the two points in time; actual + # membership drift (a deployment added or removed mid-request) still + # needs admission's own snapshot stashed and reused -- see + # _ROUTING_GROUP_CANDIDATES_FIELD's docstring. routing_group_deployments: Final = self.llm_router._get_routing_group_deployments( # pyright: ignore[reportPrivateUsage] # reused across module boundaries, matching resolve_any's own reliance on this method model=model, team_id=team_id ) @@ -1117,6 +1153,7 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] if routing_group_deployments is not None else tuple(name for d in healthy_deployments if isinstance(name := d.get("model_name"), str)) ) + _record_routing_group_candidates(resolved_request_kwargs, candidate_model_names) configured: Final = self._index.get(self.llm_router).resolve_any(model, team_id, candidate_model_names) if not configured: return healthy_deployments @@ -1621,6 +1658,13 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # only when `model_group` isn't a routing group at all (a plain # single-model_name chain, where resolve() already matches directly # and this candidate set is never actually consulted). + # + # Reconstructing live membership here is itself only a fallback: it + # can still disagree with admission's own candidate set if the + # routing group's actual membership changed between the two points + # in time (not just cooldown/health state) -- _routing_group_candidates_or + # below prefers admission's own stashed snapshot whenever one exists. + # See _ROUTING_GROUP_CANDIDATES_FIELD's docstring. 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 @@ -1628,11 +1672,12 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] routing_group_deployments: Final = self.llm_router._get_routing_group_deployments( # pyright: ignore[reportPrivateUsage] # reused across module boundaries, matching resolve_any's own reliance on this method model=model_group, team_id=team_id ) - candidate_model_names: Final = ( + live_candidate_model_names: Final = ( tuple(dep["model_name"] for dep in routing_group_deployments) if routing_group_deployments is not None else ((serving_deployment.model_name,) if serving_deployment is not None else ()) ) + candidate_model_names: Final = _routing_group_candidates_or(kwargs, fallback=live_candidate_model_names) configured: Final = self._index.get(self.llm_router).resolve_any(model_group, team_id, candidate_model_names) if not configured: return diff --git a/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py b/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py index ac24039758e..a3c9e80e502 100644 --- a/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py +++ b/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py @@ -33,7 +33,7 @@ from litellm.proxy.hooks.tag_rate_limits_shared import ( BACKGROUND_TASKS as _BACKGROUND_TASKS, CONCURRENCY_MIN_SAFETY_TTL_SECONDS as _CONCURRENCY_MIN_SAFETY_TTL_SECONDS, ) -from litellm.types.router import RoutingGroup, TagRateLimitEntry, TagRateLimitScope +from litellm.types.router import Deployment, RoutingGroup, TagRateLimitEntry, TagRateLimitScope class TimeController: @@ -1596,6 +1596,77 @@ async def test_log_success_event_accounts_against_the_same_bucket_admission_chec ) +@pytest.mark.asyncio +async def test_log_success_event_uses_admissions_own_candidate_set_when_group_membership_drifts(time_controller): + """ + Bugbot finding: resolve_any's dedup picks the alphabetically first member + model_name sharing a signature as resolved_group, a pure function of + candidate_model_names. Both admission and success independently rebuild + that set from the router's *live* routing-group membership, so a + deployment added mid-request (a hot-reload) whose name sorts earlier can + make success pick a different resolved_group than admission did, + accounting real usage into a bucket admission never checked. + """ + token_limits = { + "token_limits": { + "limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400}] + } + } + router = litellm.Router( + model_list=[ + _deployment("backend-a", "dep-a", token_limits), + _deployment("backend-b", "dep-b", token_limits), + ], + routing_groups=[ + RoutingGroup(group_name="my-group", models=["backend-a", "backend-b"], routing_strategy="simple-shuffle") + ], + ) + limiter = _make_limiter(time_controller) + limiter.update_variables(llm_router=router) + request_kwargs, model_call_details = _call_context(["end_user_id:u1"]) + + healthy = router._get_routing_group_deployments(model="my-group", team_id=None) + admitted = await limiter.async_filter_deployments( + model="my-group", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + assert admitted == healthy + admission_bucket_group = limiter._index.get(router).resolve_any( + "my-group", team_id=None, candidate_model_names=("backend-a", "backend-b") + )[0].resolved_group + + router.get_routing_group("my-group").models.append("backend-0") + router.add_deployment( + Deployment( + model_name="backend-0", + litellm_params={"model": "gpt-4o", "mock_response": "ok"}, # type: ignore + model_info={"id": "dep-0", "tag_rate_limits": token_limits}, + ) + ) + serving_deployment_id = "dep-b" if admission_bucket_group == "backend-a" else "dep-a" + + model_call_details["standard_logging_object"] = { + "model_group": "my-group", + "model_id": serving_deployment_id, + "total_tokens": 42, + "response_cost": 0.01, + } + await limiter.async_log_success_event(kwargs=model_call_details, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + + now = time_controller.now().timestamp() + admission_key = _expected_bucket_key( + "my-group", "tokens", "daily", "end_user_id", "u1", 86400, now, resolved_group=admission_bucket_group, limit=500000 + ) + drifted_key = _expected_bucket_key( + "my-group", "tokens", "daily", "end_user_id", "u1", 86400, now, resolved_group="backend-0", limit=500000 + ) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=admission_key, litellm_parent_otel_span=None)) + == 42.0 + ) + assert await limiter.internal_usage_cache.async_get_cache(key=drifted_key, litellm_parent_otel_span=None) is None + + @pytest.mark.asyncio async def test_admission_dedups_against_the_full_group_not_just_currently_healthy_members(time_controller): """