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 7bbf6ee7c14..ebe468b4aae 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -743,24 +743,37 @@ def _decode_reservations(raw: object) -> tuple[tuple[str, "_PartitionKey"], ...] # admission actually checked, letting a burst of calls admitted against one # (still-under-limit) window get charged entirely into the next window's # fresh, unrelated counter -- silently bypassing the limit right around each -# rollover. Overwritten by each hop's own admission (last-write-wins), which -# is correct: success only ever fires for whichever hop actually served the -# request, so its own most recent admission timestamp is the right one. +# rollover. +# +# Keyed by model_group, not a single scalar: `Router.abatch_completion`'s +# comma-separated branches share this identical model_call_details object +# (see `_PENDING_CONCURRENCY_KEYS_FIELD`'s docstring), each addressing its own +# model_group concurrently, so a single last-write-wins value would have one +# branch's admission timestamp overwritten by whichever sibling's admission +# happened to run second -- letting success accounting classify a bucket +# against a completely different branch's timing. A hop's own model_group is +# reliably per-hop even though the surrounding dict is shared (see +# `_resolve_hop_context`'s own docstring), so keying by it lets concurrent +# siblings addressing different model_groups each own a distinct slot; two +# siblings addressing the identical model_group (a literal `"grp,grp"` dial) +# still race on that one slot, same residual as the single-branch case this +# field originally handled. _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. +# resolved against, stashed the same way _ADMISSION_TIME_FIELD is (also keyed +# by model_group, for the identical `abatch_completion` cross-branch reason). +# 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" @@ -1043,38 +1056,56 @@ def _queue_pending_reservations( ) # mutable-ok: see comment above -def _record_admission_time(request_kwargs: Mapping[str, object], now: float) -> None: - """Stash this hop's admission timestamp -- see `_ADMISSION_TIME_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 fresh timestamp, same as before this fix existed.""" +def _record_admission_time(request_kwargs: Mapping[str, object], model_group: str, now: float) -> None: + """Stash this hop's admission timestamp under its own model_group -- see + `_ADMISSION_TIME_FIELD`'s docstring for why keyed, not scalar. Silently a + no-op without a real logging object (defensive only; every real request + has one): success accounting falls back to its own fresh timestamp, 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[_ADMISSION_TIME_FIELD] = now + if not isinstance(model_call_details, dict): + return + by_model_group = model_call_details.get(_ADMISSION_TIME_FIELD) # rebind-ok: lazily initialized below when absent + if not isinstance(by_model_group, dict): + by_model_group = {} # mutable-ok: shared, request-scoped accumulator; see field's own docstring # rebind-ok: lazily initialized only when absent + model_call_details[_ADMISSION_TIME_FIELD] = by_model_group + by_model_group[model_group] = now # mutable-ok: see comment above -def _admission_time_or(kwargs: Mapping[str, object], fallback: float) -> float: - recorded: Final = kwargs.get(_ADMISSION_TIME_FIELD) +def _admission_time_or(kwargs: Mapping[str, object], model_group: str, fallback: float) -> float: + by_model_group: Final = kwargs.get(_ADMISSION_TIME_FIELD) + recorded: Final = by_model_group.get(model_group) if isinstance(by_model_group, dict) else None return recorded if isinstance(recorded, float) else fallback def _record_routing_group_candidates( - request_kwargs: Mapping[str, object], candidate_model_names: tuple[str, ...] + request_kwargs: Mapping[str, object], model_group: str, 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.""" + """Stash the routing-group membership admission resolved against, under + this hop's own model_group -- see `_ROUTING_GROUP_CANDIDATES_FIELD`'s + docstring for why keyed, not scalar. 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 + if not isinstance(model_call_details, dict): + return + by_model_group = model_call_details.get( # rebind-ok: lazily initialized below when absent + _ROUTING_GROUP_CANDIDATES_FIELD + ) + if not isinstance(by_model_group, dict): + by_model_group = {} # mutable-ok: shared, request-scoped accumulator; see field's own docstring # rebind-ok: lazily initialized only when absent + model_call_details[_ROUTING_GROUP_CANDIDATES_FIELD] = by_model_group + by_model_group[model_group] = candidate_model_names # mutable-ok: see comment above -def _routing_group_candidates_or(kwargs: Mapping[str, object], fallback: tuple[str, ...]) -> tuple[str, ...]: - recorded: Final = kwargs.get(_ROUTING_GROUP_CANDIDATES_FIELD) +def _routing_group_candidates_or( + kwargs: Mapping[str, object], model_group: str, fallback: tuple[str, ...] +) -> tuple[str, ...]: + by_model_group: Final = kwargs.get(_ROUTING_GROUP_CANDIDATES_FIELD) + recorded: Final = by_model_group.get(model_group) if isinstance(by_model_group, dict) else None return recorded if isinstance(recorded, tuple) else fallback @@ -1326,7 +1357,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) + _record_routing_group_candidates(resolved_request_kwargs, model, 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 @@ -1343,7 +1374,7 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] key_alias: Final = _extract_key_alias(resolved_request_kwargs, metadata_variable_name) now: Final = self._time_provider().timestamp() - _record_admission_time(resolved_request_kwargs, now) + _record_admission_time(resolved_request_kwargs, model, now) classified: Final = tuple( check for configured_limit in configured @@ -1683,7 +1714,7 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] kwargs: Mapping[str, object], *, only_own_lineage: bool = False, - only_keys: frozenset[str] = frozenset(), + only_keys: frozenset[str] | None = None, ) -> tuple[tuple[str, _PartitionKey], ...]: # Snapshot then remove only those exact entries, never a blanket # clear: a sibling branch sharing this same request's @@ -1692,9 +1723,11 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # would silently strand that branch's reservation instead of # releasing it later. # - # `only_keys`, when non-empty, takes priority: at most one entry per - # requested key, since reservations sharing a key are fungible (any - # one of them represents the same +1 to the same counter) -- see + # `only_keys`, when passed (even an empty frozenset -- a hop that + # itself reserved no concurrency slot must release nothing, not fall + # through to every pending entry), takes priority: at most one entry + # per requested key, since reservations sharing a key are fungible + # (any one of them represents the same +1 to the same counter) -- see # `_release_own_concurrency_keys`'s own docstring for why this is the # terminal (success/failure) release paths' own filter, not # `only_own_lineage`. @@ -1708,12 +1741,12 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] return () matched_indices: Final = tuple( idx - for key in only_keys + for key in (only_keys or frozenset()) if (idx := next((i for i, entry in enumerate(pending) if entry[0] == key), None)) is not None ) snapshot: Final = ( tuple(pending[idx] for idx in matched_indices) - if only_keys + if only_keys is not None else tuple(entry for entry in pending if entry[2] is _current_admission_token()) if only_own_lineage else tuple(pending) @@ -1932,7 +1965,9 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] 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) + candidate_model_names: Final = _routing_group_candidates_or( + kwargs, model_group, 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 None @@ -1970,22 +2005,28 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] is no way to tell a sibling's reservation apart from this hop's own anyway, so this only ever matters for the single-branch case that release already handled correctly before `abatch_completion` existed. + + A resolved `context` whose own hop reserved *no* concurrency slot at + all (no concurrency-unit limit matches its model/tags/deployment -- + e.g. only a token or dollar limit is configured for this tag) must + release nothing, not fall through to the unconditional path: on + `abatch_completion` that shared pending list can still hold a + genuinely live sibling branch's own reservation (a different + deployment/model_group whose admission did reserve one), and this + hop reserving zero keys is not the same signal as `context` failing + to resolve at all. """ - own_concurrency_keys: Final = ( - _own_concurrency_keys_for_hop( - context.configured, - context.model_group, - context.tags, - context.deployment_id, - context.key_hash, - context.key_alias, - ) - if context is not None - else frozenset() + if context is None: + return await self._pop_pending_concurrency_keys(kwargs) + own_concurrency_keys: Final = _own_concurrency_keys_for_hop( + context.configured, + context.model_group, + context.tags, + context.deployment_id, + context.key_hash, + context.key_alias, ) - if own_concurrency_keys: - return await self._pop_pending_concurrency_keys(kwargs, only_keys=own_concurrency_keys) - return await self._pop_pending_concurrency_keys(kwargs) + return await self._pop_pending_concurrency_keys(kwargs, only_keys=own_concurrency_keys) async def async_log_failure_event( self, @@ -2024,7 +2065,7 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] if context is None: return - now: Final = _admission_time_or(kwargs, fallback=self._time_provider().timestamp()) + now: Final = _admission_time_or(kwargs, context.model_group, fallback=self._time_provider().timestamp()) increment_by_unit: Final[Mapping[_LimitUnit, float]] = MappingProxyType( { "tokens": float(context.standard_logging_object.get("total_tokens") or 0), 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 67d00685410..e216cd19c7a 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 @@ -19,6 +19,7 @@ from litellm.caching.dual_cache import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.model_based_tag_rate_limits_hook import ( + _admission_time_or, _PENDING_CONCURRENCY_KEYS_FIELD, _bucket_key, _build_group_limits, @@ -31,6 +32,9 @@ from litellm.proxy.hooks.model_based_tag_rate_limits_hook import ( _own_concurrency_keys_for_hop, _pending_reservations_cache_key, _PROXY_ModelBasedTagRateLimitsHook, + _record_admission_time, + _record_routing_group_candidates, + _routing_group_candidates_or, ) from litellm.proxy.hooks.tag_rate_limits_shared import ( BACKGROUND_TASKS as _BACKGROUND_TASKS, @@ -1853,6 +1857,131 @@ async def test_log_success_event_uses_admissions_own_candidate_set_when_group_me assert await limiter.internal_usage_cache.async_get_cache(key=drifted_key, litellm_parent_otel_span=None) is None +def test_admission_time_and_routing_group_candidates_are_keyed_per_model_group(): + """ + Cursor Bugbot finding: `_ADMISSION_TIME_FIELD`/`_ROUTING_GROUP_CANDIDATES_FIELD` + were a single last-write-wins value on the shared model_call_details, so + two `abatch_completion` branches admitting concurrently against + *different* model_groups had one branch's own snapshot overwritten by + whichever admitted second, regardless of which model_group either + addressed. Keyed by model_group instead: each branch's own admission + owns a distinct slot, so a sibling addressing a different model_group + can never clobber it. + """ + model_call_details: dict = {} + request_kwargs = {"litellm_logging_obj": SimpleNamespace(model_call_details=model_call_details)} + + _record_admission_time(request_kwargs, "grp-a", 100.0) + _record_routing_group_candidates(request_kwargs, "grp-a", ("backend-a1", "backend-a2")) + _record_admission_time(request_kwargs, "grp-b", 200.0) + _record_routing_group_candidates(request_kwargs, "grp-b", ("backend-b1", "backend-b2")) + + assert _admission_time_or(model_call_details, "grp-a", fallback=-1.0) == 100.0 + assert _admission_time_or(model_call_details, "grp-b", fallback=-1.0) == 200.0 + assert _routing_group_candidates_or(model_call_details, "grp-a", fallback=()) == ("backend-a1", "backend-a2") + assert _routing_group_candidates_or(model_call_details, "grp-b", fallback=()) == ("backend-b1", "backend-b2") + + +@pytest.mark.asyncio +async def test_success_accounting_is_not_contaminated_by_a_concurrent_siblings_routing_group_candidates( + time_controller, +): + """ + End-to-end version of the unit test above: two `abatch_completion` + branches of one comma-separated dispatch (e.g. "group-a,group-b") admit + concurrently under the identical shared model_call_details, each + against its own routing group. A single last-write-wins snapshot field + let whichever branch admitted second overwrite the other's, so the + first branch's own success event reconstructed the *second* branch's + candidate set -- resolve_any's fallback then resolved against a + completely different routing group's own deployments, hashing usage + into a bucket admission for this branch never checked (and that + branch's real bucket was never touched at all). + """ + token_limits_a = { + "token_limits": {"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400}]} + } + token_limits_b = { + "token_limits": {"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 999, "period_seconds": 86400}]} + } + router = litellm.Router( + model_list=[ + _deployment("backend-a1", "dep-a1", token_limits_a), + _deployment("backend-a2", "dep-a2", token_limits_a), + _deployment("backend-b1", "dep-b1", token_limits_b), + _deployment("backend-b2", "dep-b2", token_limits_b), + ], + routing_groups=[ + RoutingGroup(group_name="group-a", models=["backend-a1", "backend-a2"], routing_strategy="simple-shuffle"), + RoutingGroup(group_name="group-b", models=["backend-b1", "backend-b2"], 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_a = router._get_routing_group_deployments(model="group-a", team_id=None) + healthy_b = router._get_routing_group_deployments(model="group-b", team_id=None) + assert healthy_a is not None and healthy_b is not None + + async def _admit(model: str, healthy: list) -> None: + await limiter.async_filter_deployments( + model=model, healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + # Both branches admit concurrently, each its own asyncio.Task, sharing + # one model_call_details -- exactly Router.abatch_completion's real + # dispatch for a comma-separated "group-a,group-b" model string. + await asyncio.create_task(_admit("group-a", healthy_a)) + await asyncio.create_task(_admit("group-b", healthy_b)) + + admission_bucket_group_a = ( + limiter._index.get(router) + .resolve_any("group-a", team_id=None, candidate_model_names=("backend-a1", "backend-a2"))[0] + .resolved_group + ) + + model_call_details["standard_logging_object"] = { + "model_group": "group-a", + "model_id": "dep-a1" if admission_bucket_group_a == "backend-a1" else "dep-a2", + "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() + correct_bucket_key = _expected_bucket_key( + "group-a", + "tokens", + "daily", + "end_user_id", + "u1", + 86400, + now, + resolved_group=admission_bucket_group_a, + limit=500000, + ) + wrong_bucket_key_b1 = _expected_bucket_key( + "group-a", "tokens", "daily", "end_user_id", "u1", 86400, now, resolved_group="backend-b1", limit=999 + ) + wrong_bucket_key_b2 = _expected_bucket_key( + "group-a", "tokens", "daily", "end_user_id", "u1", 86400, now, resolved_group="backend-b2", limit=999 + ) + assert ( + float(await limiter.internal_usage_cache.async_get_cache(key=correct_bucket_key, litellm_parent_otel_span=None)) + == 42.0 + ) + assert ( + await limiter.internal_usage_cache.async_get_cache(key=wrong_bucket_key_b1, litellm_parent_otel_span=None) + is None + ) + assert ( + await limiter.internal_usage_cache.async_get_cache(key=wrong_bucket_key_b2, 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): """ @@ -3186,6 +3315,80 @@ async def test_fast_succeeding_batch_sibling_does_not_release_a_still_executing_ ) +@pytest.mark.asyncio +async def test_batch_sibling_reserving_no_concurrency_slot_does_not_release_a_still_executing_siblings_slot( + time_controller, +): + """ + Cursor Bugbot finding: `_release_own_concurrency_keys` fell back to + releasing every pending reservation whenever this hop's own + `own_concurrency_keys` was empty -- including when this hop's context + resolved successfully but simply reserved no concurrency slot at all + (its own model/tag only has a token limit configured, not a concurrency + one). On `abatch_completion` that shared pending list can still hold a + genuinely live sibling branch's own reservation (a different + model_group whose own admission did reserve a concurrency slot), so the + token-only branch's completion used to release the concurrency-limited + branch's still-executing slot. + """ + limiter = _make_limiter(time_controller) + token_limits = { + "token_limits": {"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400}]} + } + router = litellm.Router( + model_list=[ + _deployment( + "grp-conc", + "dep-conc", + { + "concurrency_limits": { + "limits": [{"name": "inflight", "tag_id": "end_user_id", "limit": 1, "period_seconds": 300}] + } + }, + ), + _deployment("grp-token", "dep-token", token_limits), + ] + ) + limiter.update_variables(llm_router=router) + request_kwargs, kwargs = _call_context(["end_user_id:u1"]) + conc_healthy = [d for d in router.model_list if d["model_info"]["id"] == "dep-conc"] + token_healthy = [d for d in router.model_list if d["model_info"]["id"] == "dep-token"] + + async def _admit(model: str, healthy: list) -> None: + await limiter.async_filter_deployments( + model=model, healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + # Both branches of one abatch_completion dispatch admit under the + # identical shared model_call_details, each its own asyncio.Task -- see + # the concurrency-sibling tests above for why that models two + # independent admission lineages. + await asyncio.create_task(_admit("grp-conc", conc_healthy)) + await asyncio.create_task(_admit("grp-token", token_healthy)) + + # The token-only branch finishes first. Its own success event resolves a + # real context (a token limit matches its tag/model), but reserves no + # concurrency slot of its own -- it must not touch the sibling's. + kwargs["standard_logging_object"] = { + "model_group": "grp-token", + "model_id": "dep-token", + "total_tokens": 10, + "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) + + # The concurrency-limited branch's own slot must still be held: a fresh + # request against the same tag/model is rejected under the limit=1 cap. + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp-conc", + healthy_deployments=conc_healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + + def _request_limit_router(limit: int) -> "litellm.Router": return litellm.Router( model_list=[