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 6fb6e354b12..d99896af513 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -571,26 +571,27 @@ _INDEX_TTL_SECONDS: Final = 5.0 # "shared across every task that happens to touch model_call_details": # `Router.abatch_completion`'s comma-separated multi-model dispatch runs # several branches concurrently, each its own Task, but hands every branch -# the identical `litellm_logging_obj` -- so a still-live sibling branch's -# own reservation can sit in this same list. Each entry also carries an -# admission-scoped token (see `_current_admission_token`), trusted only at -# the *next admission's own* stale-hop cleanup (see -# `_release_stale_hop_reservations`): a `ContextVar` has the same -# non-descendant-task blind spot the comment above documents for the -# reservation data itself, so a streaming response's own success event -- -# fired from a task the proxy forked independently of admission's -- would -# never see a matching token either. -# -# The terminal release hooks (`async_log_success_event`/ -# `async_log_failure_event`) use a different filter instead, immune to that -# blind spot: `_release_own_concurrency_keys` recomputes exactly the key(s) -# *this* hop's own admission reserved from data reliably specific to this -# hop even while `model_call_details` is shared (see -# `_resolve_hop_context`'s own docstring), and releases at most one entry -# per matching key -- reservations sharing a key are fungible, so this -# never touches a still-live sibling branch's own entry under a different -# key, and releases exactly one unit under a shared key even when a sibling -# holds another. +# the identical `litellm_logging_obj`, and therefore the identical `Logging` +# instance -- confirmed live: only ONE terminal success/failure event ever +# fires for the *whole* dispatch, never one per branch, because every +# branch's own completion checks the identical instance's +# `has_logged_{event_type}` flag (`Logging.should_run_logging`) before +# dispatching to any registered `CustomLogger`, and whichever branch gets +# there first flips it for all the others. Reserving one unit per matching +# branch and relying on that many independent releases was the wrong model: +# with only one release ever happening, every other branch's own +# reservation would leak until `_CONCURRENCY_MIN_SAFETY_TTL_SECONDS` on +# every ordinary multi-model batch call, not just a race. Admission +# (`async_filter_deployments`, via `_pending_concurrency_keys`) instead +# reserves at most one unit per key for the whole dispatch regardless of +# how many branches match it -- every branch after the first rides along on +# that one reservation for free -- so the single release that does happen +# always exactly balances what was reserved. Each entry also carries an +# admission-scoped token (see `_current_admission_token`), used only by +# `_release_stale_hop_reservations`'s *own* admission-time cleanup of a +# prior hop of the identical serial fallback chain, never by the terminal +# release hooks, which release everything still pending unconditionally -- +# safe now that at most one reservation per key ever exists at a time. _PENDING_CONCURRENCY_KEYS_FIELD: Final[str] = "_model_based_tag_rate_limits_pending_concurrency_keys" # Identifies which admission call queued a given reservation, scoped by @@ -936,52 +937,6 @@ def _increment_operation_for_limit( ) -def _own_concurrency_key_for_limit( - configured_limit: _ConfiguredLimit, - model_group: str, - tags: Sequence[str], - deployment_id: str | None, - key_hash: str | None, - key_alias: str | None, -) -> str | None: - """The exact `_inflight_key` this hop's own admission would have reserved - for `configured_limit`, or `None` if `configured_limit` doesn't apply to - this hop at all -- mirrors `_increment_operation_for_limit`'s own - deployment_scope/tag_value/entry_applies checks, restricted to the - "concurrency" unit `_increment_operation_for_limit` itself skips.""" - if configured_limit.unit != "concurrency": - return None - if configured_limit.deployment_scope is not None and deployment_id not in configured_limit.deployment_scope: - return None - tag_value: Final = _extract_identity(tags, configured_limit.entry.tag_id) - if tag_value is None: - return None - if not _entry_applies(configured_limit.entry, tags, key_alias, model_group): - return None - key_hash_for_limit: Final = key_hash if configured_limit.entry.scope_by_key_hash else None - return _inflight_key(model_group, configured_limit, tag_value, key_hash=key_hash_for_limit) - - -def _own_concurrency_keys_for_hop( - configured: Sequence[_ConfiguredLimit], - model_group: str, - tags: Sequence[str], - deployment_id: str | None, - key_hash: str | None, - key_alias: str | None, -) -> frozenset[str]: - return frozenset( - key - for configured_limit in configured - if ( - key := _own_concurrency_key_for_limit( - configured_limit, model_group, tags, deployment_id, key_hash, key_alias - ) - ) - is not None - ) - - class _HopContext(NamedTuple): """The identity of the one hop whose success/failure event this is -- resolved from `standard_logging_object`/`litellm_params`, both freshly @@ -1056,6 +1011,42 @@ def _queue_pending_reservations( ) # mutable-ok: see comment above +def _pending_concurrency_keys(request_kwargs: Mapping[str, object]) -> frozenset[str]: + """Every concurrency key already queued for this call, by any branch of + an `abatch_completion` dispatch -- see `_PENDING_CONCURRENCY_KEYS_FIELD`'s + docstring for why a dispatch reserves at most one unit per key regardless + of how many branches match it.""" + logging_obj: Final = request_kwargs.get("litellm_logging_obj") + model_call_details: Final = getattr(logging_obj, "model_call_details", None) + if not isinstance(model_call_details, dict): + return frozenset() + pending: Final = model_call_details.get(_PENDING_CONCURRENCY_KEYS_FIELD) + if not isinstance(pending, list): + return frozenset() + return frozenset(entry[0] for entry in pending) + + +def _discard_pending_concurrency_keys(request_kwargs: Mapping[str, object], keys: Iterable[str]) -> None: + """Rolls back this hop's own just-staked claim(s) for `keys` -- used + when the atomic batch they were staked ahead of ends up rejected, since + `_atomic_check_and_increment`'s all-or-nothing contract means nothing + was actually incremented for them; leaving the claim in place would + both corrupt release-time bookkeeping (nothing to release against) and + make a genuinely later sibling wrongly skip its own real reservation, + believing this one already covers it.""" + logging_obj: Final = request_kwargs.get("litellm_logging_obj") + model_call_details: Final = getattr(logging_obj, "model_call_details", None) + if not isinstance(model_call_details, dict): + return + pending: Final = model_call_details.get(_PENDING_CONCURRENCY_KEYS_FIELD) + if not isinstance(pending, list): + return + for key in keys: + entry = next((candidate for candidate in pending if candidate[0] == key), None) + if entry is not None: + pending.remove(entry) # mutable-ok: shared, request-scoped accumulator; see field's own docstring + + 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 @@ -1375,7 +1366,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, model, now) - classified: Final = tuple( + raw_classified: Final = tuple( check for configured_limit in configured if ( @@ -1392,6 +1383,34 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] ) is not None ) + # A dispatch reserves at most one concurrency unit per key regardless + # of how many abatch_completion branches match it -- see + # _PENDING_CONCURRENCY_KEYS_FIELD's own docstring for why. Staked + # synchronously below, before this function's next `await`: asyncio + # only switches tasks at an `await`, so a sibling branch admitting + # concurrently can never observe this hop mid-decision, only either + # fully before or fully after it -- closing the race that would + # otherwise let two branches both see "not yet reserved" and each + # stake their own. + already_pending_keys: Final = _pending_concurrency_keys(resolved_request_kwargs) + own_new_concurrency_claims: Final = [] # mutable-ok: staked synchronously below before any further await; see comment above + deduped_classified: Final = [] # mutable-ok: see comment above + for check in raw_classified: + # not Final: rebound each loop iteration + already_claimed = check.configured_limit.unit == "concurrency" and ( + check.key in already_pending_keys or any(key == check.key for key, _ in own_new_concurrency_claims) + ) + if already_claimed: + continue + if check.configured_limit.unit == "concurrency": + own_new_concurrency_claims.append((check.key, _partition_key(check.configured_limit.entry))) + deduped_classified.append(check) + classified: Final = tuple(deduped_classified) + if own_new_concurrency_claims: + _queue_pending_reservations( + resolved_request_kwargs, _PENDING_CONCURRENCY_KEYS_FIELD, own_new_concurrency_claims + ) + read_only_checks: Final = tuple((c.configured_limit, c.tag_value, c.key) for c in classified if not c.is_atomic) atomic_checks: Final = tuple((c.configured_limit, c.tag_value, c.key) for c in classified if c.is_atomic) @@ -1427,22 +1446,23 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] ) ) if failing_index is not None: + if own_new_concurrency_claims: + # Nothing was actually incremented for these (the whole + # batch is all-or-nothing), so the claim staked above + # must be rolled back -- otherwise it would both corrupt + # release-time bookkeeping and make a genuinely later + # sibling wrongly skip its own real reservation. + _discard_pending_concurrency_keys( + resolved_request_kwargs, (key for key, _partition_key in own_new_concurrency_claims) + ) failing_limit, failing_tag_value, _ = atomic_checks[failing_index] self._raise_over_limit(failing_limit, failing_tag_value, model, current=values[0]) - concurrency_reservations: Final = tuple( - (key, _partition_key(configured_limit.entry)) - for configured_limit, _tag_value, key in atomic_checks - if configured_limit.unit == "concurrency" - ) - if concurrency_reservations: - _queue_pending_reservations( - resolved_request_kwargs, _PENDING_CONCURRENCY_KEYS_FIELD, concurrency_reservations - ) + if own_new_concurrency_claims: await self._mirror_pending_reservations( resolved_request_kwargs.get("litellm_call_id"), _extract_key_hash(resolved_request_kwargs, metadata_variable_name), - concurrency_reservations, + tuple(own_new_concurrency_claims), ) # Only genuinely new keys, never one already in stale_request_keys: @@ -1714,7 +1734,6 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] kwargs: Mapping[str, object], *, only_own_lineage: bool = False, - 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 @@ -1723,31 +1742,20 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # would silently strand that branch's reservation instead of # releasing it later. # - # `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`. - # - # `only_own_lineage` additionally excludes any entry a *different* - # admission lineage queued -- see `_release_stale_hop_reservations`'s - # own docstring for why that distinction, not just presence, decides - # what's actually stale. + # `only_own_lineage` excludes any entry a *different* admission + # lineage queued -- see `_release_stale_hop_reservations`'s own + # docstring for why that distinction, not just presence, decides + # what's actually stale. The terminal release hooks + # (`async_log_success_event`/`async_log_failure_event`) never pass + # it: with admission now reserving at most one unit per key for the + # whole dispatch (see `_PENDING_CONCURRENCY_KEYS_FIELD`'s docstring), + # everything still pending when the one terminal event for this + # dispatch fires is safe to release unconditionally. pending: Final = kwargs.get(_PENDING_CONCURRENCY_KEYS_FIELD) if not isinstance(pending, list) or not pending: return () - matched_indices: Final = tuple( - idx - 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 is not None - else tuple(entry for entry in pending if entry[2] is _current_admission_token()) + tuple(entry for entry in pending if entry[2] is _current_admission_token()) if only_own_lineage else tuple(pending) ) @@ -1820,11 +1828,10 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] expires, letting a caller who repeatedly opens and immediately drops streaming requests exhaust their own tag's concurrency limit for free. - Deliberately not `only_own_lineage=True`: `Router.abatch_completion` - returns every branch's response together rather than a single stream, - so its concurrent-sibling-branch race this hook otherwise guards - against (see `_PENDING_CONCURRENCY_KEYS_FIELD`'s docstring) cannot - co-occur with a mid-stream disconnect here. + Deliberately not `only_own_lineage=True`: at most one reservation + per key is ever pending for the whole request (see + `_PENDING_CONCURRENCY_KEYS_FIELD`'s docstring), so releasing + everything still pending here is exactly releasing this hop's own. """ logging_obj: Final = request_data.get("litellm_logging_obj") model_call_details: Final = getattr(logging_obj, "model_call_details", None) @@ -1893,22 +1900,28 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] def _resolve_hop_context(self, kwargs: Mapping[str, object]) -> _HopContext | None: """ - `standard_logging_object` and `litellm_params` are freshly overwritten - by *this* hop's own attempt immediately before its success/failure + Resolves the identity this hop's own token/dollar accounting needs + (`async_log_success_event`'s own tail, below) from + `standard_logging_object`/`litellm_params`, freshly overwritten by + *this* hop's own attempt immediately before its success/failure callback fires -- Router builds each hop's own `litellm_params` from whichever deployment it actually attempted, and litellm's dispatch writes `standard_logging_object` with no `await` between that write - and this callback firing, so a sibling branch of an `abatch_completion` - dispatch has no chance to interleave and overwrite it first even - though the surrounding `model_call_details` dict (`kwargs` here) is - the identical object shared across every branch (confirmed live: it, - `litellm_call_id`, and the top-level `metadata`/`litellm_metadata` - dicts are all one shared object across a comma-separated dispatch's - branches). Token/dollar accounting below already depends on this - being reliably per-hop; reused here to recompute exactly the - concurrency key(s) this hop's own admission reserved, rather than a - value that would have to survive being read back from a different - branch or task. + and this callback firing. + + Only one such callback ever fires per `abatch_completion` dispatch, + not one per branch (see `_PENDING_CONCURRENCY_KEYS_FIELD`'s + docstring): every branch shares the identical `Logging` instance, so + whichever branch's completion reaches `should_run_logging` first + wins for the whole dispatch, and its own `standard_logging_object` + is whatever the shared `model_call_details` happens to hold at that + moment -- reliably this hop's own for a genuine single-hop or + serial-fallback request, but for a multi-model batch, only the + winning branch's own usage is ever accounted; every other branch's + real token/dollar usage is silently never counted at all. Fixing + that needs Router-level changes (a distinct `Logging` instance per + branch, or aggregating usage before the one terminal event) -- + tracked as a known limitation, not attempted here. """ if self.llm_router is None: return None @@ -1969,25 +1982,17 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] 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 + tags: Final = _order_tags_for_identity_resolution( _get_tags_from_request_kwargs(kwargs, metadata_variable_name=metadata_variable_name), kwargs, metadata_variable_name, ) - # Deliberately not an early `return None` when `configured` or - # `tags` is empty (Cursor Bugbot finding: that used to make - # `_release_own_concurrency_keys` fall back to an unconditional - # release): `async_filter_deployments` hits this exact same - # "nothing configured" / "no matching tag" check on *this hop's own - # admission* before it could ever queue a concurrency reservation, - # so an empty `configured` or `tags` here means this hop reserved - # nothing of its own to release either -- `_own_concurrency_keys_for_hop` - # already returns an empty set for either case on its own, without - # needing this function to special-case it. `None` is reserved for - # the genuinely ambiguous causes above (no router, no - # `standard_logging_object`, no `model_group`), where this hop's own - # admission *could* have reserved something we have no way to - # recompute -- those still need the unconditional fallback. + if not tags: + return None + return _HopContext( standard_logging_object=standard_logging_object, configured=configured, @@ -1998,50 +2003,6 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] key_alias=key_alias, ) - async def _release_own_concurrency_keys( - self, kwargs: Mapping[str, object], context: "_HopContext | None" - ) -> tuple[tuple[str, _PartitionKey], ...]: - """ - Releases exactly the concurrency reservation(s) this hop's own - admission made, computed via `context` the same way admission itself - computed them -- never "every reservation currently pending", which - would also release a still-live `abatch_completion` sibling branch's - own reservation (see `_PENDING_CONCURRENCY_KEYS_FIELD`'s docstring). - - A resolved `context` whose own hop matched no concurrency-unit limit - at all (Cursor Bugbot finding: "no `tag_rate_limits` configured for - this model" is the common case) must release nothing, not fall - through to the unconditional path below: `async_filter_deployments` - hits that identical "nothing configured"/"no matching tag" check on - *this exact hop's own admission*, before it could ever queue a - concurrency reservation, so this hop reserved nothing of its own to - release either -- `_resolve_hop_context` folds both into an - empty-but-resolved `context` rather than `None` for exactly this - reason (see its own docstring). - - `context is None` means `_resolve_hop_context` hit one of its - genuinely ambiguous causes instead (no router, no - `standard_logging_object`, no `model_group`) -- cases where this - hop's own admission *could* have reserved something, but there is - no reliable data left to recompute what. Only those fall back to - the pre-existing unconditional release, the same as before - `abatch_completion` existed: e.g. a *different* registered - CustomLogger rejecting the request before this hop's own call ever - ran (so no `standard_logging_object` was ever built) must still - release this hook's own successfully reserved slot. - """ - 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, - ) - return await self._pop_pending_concurrency_keys(kwargs, only_keys=own_concurrency_keys) - async def async_log_failure_event( self, kwargs, # noqa: ANN001 # matches CustomLogger.async_log_failure_event; kwargs is dict[str, Any] codebase-wide @@ -2058,7 +2019,13 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # the identical marker -- that rejection can land after this hook # already reserved a slot for the same request, and that slot must # still be released. - release_keys: Final = await self._release_own_concurrency_keys(kwargs, self._resolve_hop_context(kwargs)) + # + # Unconditional, not filtered to this hop's own lineage: only one + # terminal event ever fires per abatch_completion dispatch (see + # _PENDING_CONCURRENCY_KEYS_FIELD's docstring), and admission now + # reserves at most one unit per key for the whole dispatch, so + # whatever's still pending here is exactly that one reservation. + release_keys: Final = await self._pop_pending_concurrency_keys(kwargs) if release_keys: await self._release_keys(release_keys) @@ -2070,7 +2037,7 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] end_time: object, ) -> None: context: Final = self._resolve_hop_context(kwargs) - release_keys: Final = await self._release_own_concurrency_keys(kwargs, context) + release_keys: Final = await self._pop_pending_concurrency_keys(kwargs) if release_keys: release_task: Final = asyncio.create_task(self._release_keys(release_keys)) _BACKGROUND_TASKS.add(release_task) # mutable-ok: see _BACKGROUND_TASKS's own docstring 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 e1eb107b82a..306e1ed03e6 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 @@ -26,10 +26,10 @@ from litellm.proxy.hooks.model_based_tag_rate_limits_hook import ( _build_limits_index, _ConfiguredLimit, _current_admission_token, + _discard_pending_concurrency_keys, _extract_team_id, _inflight_key, - _own_concurrency_key_for_limit, - _own_concurrency_keys_for_hop, + _pending_concurrency_keys, _pending_reservations_cache_key, _PROXY_ModelBasedTagRateLimitsHook, _record_admission_time, @@ -3168,60 +3168,124 @@ async def test_next_hops_admission_releases_a_prior_hops_leaked_reservation(time @pytest.mark.asyncio -async def test_concurrent_batch_siblings_do_not_bypass_a_concurrency_limit(time_controller): +async def test_concurrent_batch_siblings_dedupe_to_exactly_one_concurrency_reservation(time_controller): """ - Veria AI finding: Router.abatch_completion's comma-separated multi-model - dispatch runs each model concurrently as its own asyncio.Task, but every - branch is handed the identical litellm_logging_obj (the proxy attaches - one to the request before the comma-split), so two genuinely concurrent - branches share one model_call_details. Before the fix, a second branch's - admission-time stale-reservation cleanup couldn't tell that apart from - an earlier, already-failed hop of its own retry chain, so it released - the first branch's still-live reservation and let both branches through - a concurrency limit of 1. + Router.abatch_completion's comma-separated multi-model dispatch runs + each model concurrently as its own asyncio.Task, but every branch is + handed the identical litellm_logging_obj (the proxy attaches one to the + request before the comma-split), so two genuinely concurrent branches + share one model_call_details -- and, confirmed live, that shared + Logging instance means only ONE terminal success/failure event will + ever fire for the whole dispatch (litellm's own has_logged_{event_type} + dedup), never one per branch. Reserving one unit per branch and relying + on that many independent releases was the wrong model: with only one + release ever happening, every other branch's own reservation would leak + until its safety TTL on every ordinary multi-model batch call. Admission + instead reserves at most one unit per key for the whole dispatch: both + branches here admit successfully under a limit of 1, and only one real + reservation exists regardless. """ limiter = _make_limiter(time_controller) router = _concurrency_router(limit=1) limiter.update_variables(llm_router=router) healthy = router.model_list request_kwargs, _kwargs = _call_context(["end_user_id:u1"]) - first_admitted = asyncio.Event() - async def _branch_one() -> None: - await limiter.async_filter_deployments( - model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs - ) - first_admitted.set() - # Held "in flight" while the second branch's admission runs, exactly - # like two concurrently in-flight provider calls. - await asyncio.sleep(0.05) - - async def _branch_two() -> None: - await first_admitted.wait() + async def _admit() -> None: await limiter.async_filter_deployments( model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs ) results = await asyncio.gather( - asyncio.create_task(_branch_one()), asyncio.create_task(_branch_two()), return_exceptions=True + asyncio.create_task(_admit()), asyncio.create_task(_admit()), return_exceptions=True ) rejections = [result for result in results if isinstance(result, ProxyRateLimitError)] - assert len(rejections) == 1 + assert rejections == [] + assert len(_pending_concurrency_keys(request_kwargs)) == 1 @pytest.mark.asyncio -async def test_fast_failing_batch_sibling_does_not_release_a_still_executing_siblings_slot(time_controller): +async def test_three_batch_siblings_still_dedupe_to_exactly_one_concurrency_reservation(time_controller): + """A dispatch's own reservation count does not scale with how many + models are in its comma-separated list.""" + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + request_kwargs, _kwargs = _call_context(["end_user_id:u1"]) + + async def _admit() -> None: + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + results = await asyncio.gather( + *(asyncio.create_task(_admit()) for _ in range(3)), return_exceptions=True + ) + rejections = [result for result in results if isinstance(result, ProxyRateLimitError)] + assert rejections == [] + assert len(_pending_concurrency_keys(request_kwargs)) == 1 + + +@pytest.mark.asyncio +async def test_rejected_atomic_batch_rolls_back_its_own_newly_staked_concurrency_claim(time_controller): """ - Live repro: racing two comma-separated abatch_completion branches sharing - one model_call_details, one bad-keyed for a fast 401, showed the - genuinely-executing sibling's own inflight key drop before its real - completion finished, admitting a third same-tag caller past a - concurrency=1 cap. async_log_failure_event used to release every pending - reservation unconditionally regardless of which hop it belonged to; it - must release only the failing branch's own slot. + A concurrency claim is staked synchronously before the atomic + check-and-increment it's part of even runs, to close the race two + concurrent siblings could otherwise hit (see + _PENDING_CONCURRENCY_KEYS_FIELD's docstring). If that batch then + rejects because of a *different* check in it, nothing was actually + incremented for the concurrency claim either -- the whole batch is + all-or-nothing -- so the staked claim must be rolled back. Left in + place, it would falsely tell a later admission this key is already + covered by a real reservation that never actually happened. """ limiter = _make_limiter(time_controller) - router = _concurrency_router(limit=2) + router = litellm.Router( + model_list=[ + _deployment( + "grp", + "dep-1", + { + "concurrency_limits": { + "limits": [{"name": "inflight", "tag_id": "shared_pool", "limit": 2, "period_seconds": 300}] + }, + "request_limits": { + "limits": [{"name": "per_period", "tag_id": "end_user_id", "limit": 1, "period_seconds": 300}] + }, + }, + ) + ] + ) + limiter.update_variables(llm_router=router) + healthy = router.model_list + + first_request_kwargs, _first_kwargs = _call_context(["end_user_id:u1", "shared_pool:pool-a"]) + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=first_request_kwargs + ) + + second_request_kwargs, _second_kwargs = _call_context(["end_user_id:u1", "shared_pool:pool-a"]) + with pytest.raises(ProxyRateLimitError): + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=second_request_kwargs + ) + + assert _pending_concurrency_keys(second_request_kwargs) == frozenset() + + +@pytest.mark.asyncio +async def test_the_dispatchs_one_surviving_success_event_fully_releases_its_single_reservation(time_controller): + """ + Only one terminal event -- success or failure -- ever fires for the + whole abatch_completion dispatch. Since admission reserves exactly one + unit for the dispatch regardless of how many branches matched it, that + one event's own unconditional release always exactly balances it: + whichever branch's data happens to be reflected when it fires, the + concurrency slot is fully freed for the request as a whole. + """ + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) limiter.update_variables(llm_router=router) healthy = router.model_list request_kwargs, kwargs = _call_context(["end_user_id:u1"]) @@ -3231,17 +3295,46 @@ async def test_fast_failing_batch_sibling_does_not_release_a_still_executing_sib model="grp", 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 -- - # exactly like Router.abatch_completion's real dispatch, and needed for - # _release_stale_hop_reservations' own admission-lineage token to treat - # them as two independent lineages rather than two hops of one chain. await asyncio.create_task(_admit()) await asyncio.create_task(_admit()) + assert len(_pending_concurrency_keys(request_kwargs)) == 1 + + kwargs["standard_logging_object"] = { + "model_group": "grp", + "model_id": "dep-1", + "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) + + assert _pending_concurrency_keys(request_kwargs) == frozenset() + await limiter.async_filter_deployments( + model="grp", + healthy_deployments=healthy, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + + +@pytest.mark.asyncio +async def test_the_dispatchs_one_surviving_failure_event_fully_releases_its_single_reservation(time_controller): + """Same as the success-event version above, but for the failure path.""" + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + healthy = router.model_list + request_kwargs, kwargs = _call_context(["end_user_id:u1"]) + + async def _admit() -> None: + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + await asyncio.create_task(_admit()) + await asyncio.create_task(_admit()) + assert len(_pending_concurrency_keys(request_kwargs)) == 1 - # The second branch fails fast (e.g. a bad key on that specific model) -- - # its own failure event must release only its own slot, not the first - # branch's, which is still genuinely executing. kwargs["standard_logging_object"] = { "model_group": "grp", "model_id": "dep-1", @@ -3250,214 +3343,13 @@ async def test_fast_failing_batch_sibling_does_not_release_a_still_executing_sib } await limiter.async_log_failure_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) - # One slot freed (the failing branch's), one still held (the executing - # branch's): a fresh request fits, a second one does not. + assert _pending_concurrency_keys(request_kwargs) == frozenset() await limiter.async_filter_deployments( model="grp", healthy_deployments=healthy, messages=None, request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, ) - with pytest.raises(ProxyRateLimitError): - await limiter.async_filter_deployments( - model="grp", - healthy_deployments=healthy, - messages=None, - request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, - ) - - -@pytest.mark.asyncio -async def test_fast_succeeding_batch_sibling_does_not_release_a_still_executing_siblings_slot(time_controller): - """ - Same live repro as the failure-path version above, but for the more - common case: one branch of a comma-separated abatch_completion dispatch - finishes (successfully) well before its sibling. async_log_success_event - must release only that one branch's own slot. - """ - limiter = _make_limiter(time_controller) - router = _concurrency_router(limit=2) - limiter.update_variables(llm_router=router) - healthy = router.model_list - request_kwargs, kwargs = _call_context(["end_user_id:u1"]) - - async def _admit() -> None: - await limiter.async_filter_deployments( - model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs - ) - - # Each branch its own asyncio.Task -- see the failure-path test above - # for why that's required to model two independent admission lineages. - await asyncio.create_task(_admit()) - await asyncio.create_task(_admit()) - - kwargs["standard_logging_object"] = { - "model_group": "grp", - "model_id": "dep-1", - "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) - - await limiter.async_filter_deployments( - model="grp", - healthy_deployments=healthy, - messages=None, - request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, - ) - with pytest.raises(ProxyRateLimitError): - await limiter.async_filter_deployments( - model="grp", - healthy_deployments=healthy, - messages=None, - request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, - ) - - -@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"]}}, - ) - - -@pytest.mark.asyncio -async def test_batch_sibling_with_no_tag_rate_limits_at_all_does_not_release_a_still_executing_siblings_slot( - time_controller, -): - """ - Cursor Bugbot finding: `_release_own_concurrency_keys` still fell back to - an unconditional release whenever `_resolve_hop_context` returned `None` - -- including the common case of a hop whose own model has no - `tag_rate_limits` configured at all, so `resolve_any` finds nothing and - `_resolve_hop_context` bails out. That same "nothing configured" check - is also the *first* thing this hop's own admission does, before it could - ever queue a concurrency reservation -- so a `None` context means this - hop reserved nothing of its own, the same as a resolved context with no - matching concurrency-unit limit. On `abatch_completion` the still-live - sibling's own reservation must survive this branch's completion. - """ - limiter = _make_limiter(time_controller) - 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-none", "dep-none", {}), - ] - ) - 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"] - none_healthy = [d for d in router.model_list if d["model_info"]["id"] == "dep-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 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-none", none_healthy)) - - # The unconfigured branch finishes first. Its own success event resolves - # no context at all (nothing configured for "grp-none"), and reserved no - # concurrency slot of its own -- it must not touch the sibling's. - kwargs["standard_logging_object"] = { - "model_group": "grp-none", - "model_id": "dep-none", - "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"]}}, - ) @pytest.mark.asyncio @@ -3467,10 +3359,9 @@ async def test_a_hop_matching_two_concurrency_scoped_entries_releases_both_reser traffic: a request matching multiple concurrency-scoped TagRateLimitEntry policies reserved one key per matching policy at admission but a terminal event released only one entry total, leaking every other match - until its safety TTL. Verifies this hook's own release path -- which - computes `only_keys` as the full set of every matching policy's own key, - not a single arbitrary one -- actually releases all of them, not just - the first. + until its safety TTL. Verifies this hook's own release path (an + unconditional pop of everything still pending for the call) actually + releases all of them, not just the first. """ limiter = _make_limiter(time_controller) router = litellm.Router( @@ -3523,6 +3414,145 @@ async def test_a_hop_matching_two_concurrency_scoped_entries_releases_both_reser ) +@pytest.mark.asyncio +async def test_real_pipeline_abatch_completion_reserves_and_releases_exactly_one_concurrency_unit( + time_controller, monkeypatch +): + """ + Real-pipeline regression test: drives an actual comma-separated + abatch_completion dispatch through the real proxy pre-call pipeline -- + common_processing_pre_call_logic followed by route_request, exactly + what route_llm_request.py does in production -- which is what + pre-attaches one shared litellm_logging_obj to every branch before the + comma-split. Confirmed live: that shared Logging instance means only + ONE terminal success/failure event ever fires for the whole dispatch, + never one per branch, since every branch checks the identical + instance's has_logged_{event_type} flag before dispatching to any + CustomLogger. This drives the real hook (registered in litellm.callbacks, + invoked by Router.async_callback_filter_deployments and litellm's own + logging worker, not called directly) through a real 3-branch dispatch + of the same model against a concurrency limit of 1: without dedup, the + second and third branches would each be rejected outright (confirmed + empirically without the admission-side fix); with it, all three succeed, + and the single reservation is fully released once the one surviving + terminal event fires -- checked after an explicit flush of litellm's + own async logging worker, not a fixed sleep, since that worker (unlike + this hook's own success/failure hooks in the other tests here) runs the + completion callback on its own background task, not inline. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + from litellm.proxy import common_request_processing as cpr_module + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.route_llm_request import route_request + from litellm.proxy.utils import ProxyLogging + + limiter = _make_limiter(time_controller) + router = _concurrency_router(limit=1) + limiter.update_variables(llm_router=router) + monkeypatch.setattr(litellm, "callbacks", [limiter]) + + model_string = "grp,grp,grp" + base_data = { + "model": model_string, + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"tags": ["end_user_id:u1"]}, + } + + async def mock_add_litellm_data_to_request(*_args, **_kwargs): + return dict(base_data) + + monkeypatch.setattr(cpr_module, "add_litellm_data_to_request", mock_add_litellm_data_to_request) + + async def passthrough_pre_call_hook(user_api_key_dict, data, call_type): + return data + + proxy_logging_obj = MagicMock(spec=ProxyLogging) + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=passthrough_pre_call_hook) + + processing_obj = ProxyBaseLLMRequestProcessing(data=dict(base_data)) + mock_request = MagicMock() + mock_request.headers = {} + data, _logging_obj = await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=UserAPIKeyAuth(), + proxy_logging_obj=proxy_logging_obj, + proxy_config=MagicMock(spec=ProxyConfig), + route_type="acompletion", + ) + + dispatch_result = route_request( + data=data, + llm_router=router, + user_model=None, + route_type="acompletion", + user_api_key_dict=UserAPIKeyAuth(), + ) + resolved = await dispatch_result + if asyncio.iscoroutine(resolved): + resolved = await resolved + assert len(resolved) == 3 + assert not any(isinstance(branch_result, Exception) for branch_result in resolved) + + await GLOBAL_LOGGING_WORKER.flush() + await asyncio.sleep(0.2) + + assert _pending_concurrency_keys(data) == frozenset() + + +@pytest.mark.asyncio +async def test_known_limitation_batch_dispatch_only_accounts_one_branchs_token_usage(time_controller): + """ + Known, accepted limitation (see `_resolve_hop_context`'s own docstring): + only one terminal event ever fires for the whole `abatch_completion` + dispatch, so only whichever one branch's own `standard_logging_object` + is reflected in the shared `model_call_details` at that moment gets + accounted -- never the sum of every branch's real usage. Fixing this + needs Router-level changes (a distinct Logging instance per branch, or + aggregating usage before the one terminal event fires), out of scope + for this hook. + + Documents the gap concretely: two branches each genuinely used 100 + tokens (200 total), but only one terminal event ever fires in reality, + so only 100 tokens are ever accounted, not 200. This is not a bug this + hook can fix; it's a pre-existing accounting gap the batch-completion + fixes elsewhere in this file did not introduce and cannot close. + """ + 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", "dep-1", token_limits)]) + limiter.update_variables(llm_router=router) + request_kwargs, kwargs = _call_context(["end_user_id:u1"]) + healthy = router.model_list + + await limiter.async_filter_deployments( + model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs + ) + + # Simulates the one terminal event that actually fires for the whole + # dispatch in reality: whichever branch's data happens to be reflected, + # each branch here genuinely used 100 tokens, so a correct sum would be + # 200 -- only this one branch's 100 is ever visible to account. + kwargs["standard_logging_object"] = { + "model_group": "grp", + "model_id": "dep-1", + "total_tokens": 100, + "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 = _expected_bucket_key("grp", "tokens", "daily", "end_user_id", "u1", 86400, now, limit=500000) + accounted = float(await limiter.internal_usage_cache.async_get_cache(key=token_key, litellm_parent_otel_span=None)) + assert accounted == 100.0 # the sibling branch's own real 100 tokens are never accounted at all + + def _request_limit_router(limit: int) -> "litellm.Router": return litellm.Router( model_list=[ @@ -4702,76 +4732,6 @@ def test_concurrency_ttl_floor_does_not_shorten_a_longer_period_seconds(): assert _PROXY_ModelBasedTagRateLimitsHook._ttl_for(configured_limit) == _CONCURRENCY_MIN_SAFETY_TTL_SECONDS + 100 -# --------------------------------------------------------------------------- -# _own_concurrency_key_for_limit / _own_concurrency_keys_for_hop -- the -# terminal release paths' own recomputation of exactly which reservation(s) -# a completing hop is entitled to release, mirroring -# _increment_operation_for_limit's own deployment_scope/tag_value/ -# entry_applies checks. -# --------------------------------------------------------------------------- - - -def test_own_concurrency_key_for_limit_matches_the_key_admission_would_reserve(): - entry = TagRateLimitEntry(name="inflight", tag_id="end_user_id", limit=1, period_seconds=300) - configured_limit = _ConfiguredLimit(unit="concurrency", entry=entry, deployment_scope=None) - admission_key = _inflight_key("grp", configured_limit, "u1", key_hash=None) - assert _own_concurrency_key_for_limit(configured_limit, "grp", ["end_user_id:u1"], "dep-1", None, None) == ( - admission_key - ) - - -def test_own_concurrency_key_for_limit_is_none_for_a_non_concurrency_unit(): - entry = TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=500, period_seconds=86400) - configured_limit = _ConfiguredLimit(unit="tokens", entry=entry, deployment_scope=None) - assert _own_concurrency_key_for_limit(configured_limit, "grp", ["end_user_id:u1"], "dep-1", None, None) is None - - -def test_own_concurrency_key_for_limit_is_none_outside_its_deployment_scope(): - entry = TagRateLimitEntry(name="inflight", tag_id="end_user_id", limit=1, period_seconds=300) - configured_limit = _ConfiguredLimit(unit="concurrency", entry=entry, deployment_scope=("dep-1",)) - assert _own_concurrency_key_for_limit(configured_limit, "grp", ["end_user_id:u1"], "dep-2", None, None) is None - assert _own_concurrency_key_for_limit(configured_limit, "grp", ["end_user_id:u1"], "dep-1", None, None) is not None - - -def test_own_concurrency_key_for_limit_is_none_without_a_matching_tag(): - entry = TagRateLimitEntry(name="inflight", tag_id="end_user_id", limit=1, period_seconds=300) - configured_limit = _ConfiguredLimit(unit="concurrency", entry=entry, deployment_scope=None) - assert _own_concurrency_key_for_limit(configured_limit, "grp", ["team_id:t1"], "dep-1", None, None) is None - - -def test_own_concurrency_key_for_limit_folds_in_key_hash_only_when_scoped(): - scoped_entry = TagRateLimitEntry( - name="inflight", tag_id="end_user_id", limit=1, period_seconds=300, scope_by_key_hash=True - ) - scoped_limit = _ConfiguredLimit(unit="concurrency", entry=scoped_entry, deployment_scope=None) - key_with_hash = _own_concurrency_key_for_limit(scoped_limit, "grp", ["end_user_id:u1"], "dep-1", "hashA", None) - key_with_different_hash = _own_concurrency_key_for_limit( - scoped_limit, "grp", ["end_user_id:u1"], "dep-1", "hashB", None - ) - assert key_with_hash != key_with_different_hash - - unscoped_entry = TagRateLimitEntry(name="inflight", tag_id="end_user_id", limit=1, period_seconds=300) - unscoped_limit = _ConfiguredLimit(unit="concurrency", entry=unscoped_entry, deployment_scope=None) - key_ignoring_hash_a = _own_concurrency_key_for_limit( - unscoped_limit, "grp", ["end_user_id:u1"], "dep-1", "hashA", None - ) - key_ignoring_hash_b = _own_concurrency_key_for_limit( - unscoped_limit, "grp", ["end_user_id:u1"], "dep-1", "hashB", None - ) - assert key_ignoring_hash_a == key_ignoring_hash_b - - -def test_own_concurrency_keys_for_hop_only_collects_concurrency_unit_entries(): - concurrency_entry = TagRateLimitEntry(name="inflight", tag_id="end_user_id", limit=1, period_seconds=300) - token_entry = TagRateLimitEntry(name="daily", tag_id="end_user_id", limit=500, period_seconds=86400) - configured = ( - _ConfiguredLimit(unit="concurrency", entry=concurrency_entry, deployment_scope=None), - _ConfiguredLimit(unit="tokens", entry=token_entry, deployment_scope=None), - ) - keys = _own_concurrency_keys_for_hop(configured, "grp", ["end_user_id:u1"], "dep-1", None, None) - assert len(keys) == 1 - - # --------------------------------------------------------------------------- # pending-concurrency-key field on model_call_details must survive a detached # asyncio.create_task fork (e.g. litellm's own failure-logging dispatch), @@ -4853,50 +4813,6 @@ async def test_release_only_own_lineage_leaves_a_concurrent_siblings_reservation assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == [("sibling-key", None, sibling_token)] -@pytest.mark.asyncio -async def test_pop_pending_concurrency_keys_only_keys_releases_at_most_one_per_key(time_controller): - """ - only_keys is the terminal (success/failure) release paths' own filter: - reservations sharing a key are fungible, so a matching key releases - exactly one entry, never every entry sharing that key -- two branches - admitted under the identical key must each release their own unit - independently, not have one release both at once. - """ - limiter = _make_limiter(time_controller) - model_call_details: dict = { - _PENDING_CONCURRENCY_KEYS_FIELD: [ - ("shared-key", None, object()), - ("shared-key", None, object()), - ("other-key", None, object()), - ] - } - - first_release = await limiter._pop_pending_concurrency_keys(model_call_details, only_keys=frozenset({"shared-key"})) - assert first_release == (("shared-key", None),) - assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == [ - ("shared-key", None, model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD][0][2]), - ("other-key", None, model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD][1][2]), - ] - - second_release = await limiter._pop_pending_concurrency_keys( - model_call_details, only_keys=frozenset({"shared-key"}) - ) - assert second_release == (("shared-key", None),) - assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == [ - ("other-key", None, model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD][0][2]) - ] - - -@pytest.mark.asyncio -async def test_pop_pending_concurrency_keys_only_keys_ignores_a_non_matching_key(time_controller): - limiter = _make_limiter(time_controller) - model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: [("sibling-key", None, object())]} - - released = await limiter._pop_pending_concurrency_keys(model_call_details, only_keys=frozenset({"my-own-key"})) - assert released == () - assert len(model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD]) == 1 - - # --------------------------------------------------------------------------- # refund-on-rollback across differently-hash-tagged keys (Redis Cluster safety) # ---------------------------------------------------------------------------