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 68238a8459b..7bbf6ee7c14 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -573,18 +573,24 @@ _INDEX_TTL_SECONDS: Final = 5.0 # 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`), but that token is -# only ever trusted at the *next admission's own* stale-hop cleanup (see -# `_release_stale_hop_reservations`), never at the terminal release hooks -# below: a `ContextVar` has the identical non-descendant-task blind spot -# documented above 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, -# leaking every streaming request's reservation instead of releasing it. -# The terminal hooks release unconditionally, accepting the narrower risk -# that a still-live `abatch_completion` sibling's own reservation gets -# freed early, in exchange for actually releasing the far more common -# single-branch (including streaming) case. +# 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. _PENDING_CONCURRENCY_KEYS_FIELD: Final[str] = "_model_based_tag_rate_limits_pending_concurrency_keys" # Identifies which admission call queued a given reservation, scoped by @@ -917,6 +923,69 @@ 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 + overwritten by *this* hop's own attempt just before its callback fires + (see `_resolve_hop_context`'s own docstring for why that's reliable even + though the surrounding `model_call_details` dict is shared across an + `abatch_completion` dispatch's sibling branches).""" + + standard_logging_object: StandardLoggingPayload + configured: tuple[_ConfiguredLimit, ...] + model_group: str + tags: tuple[str, ...] + deployment_id: str | None + key_hash: str | None + key_alias: str | None + + def _resolve_max_in_memory_cache_size() -> int | None: """ `litellm_settings` values reach `litellm.model_based_tag_rate_limits_max_in_memory_cache_size` @@ -1547,10 +1616,9 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] concluded and failed: Router awaits one hop's entire attempt (call plus its own failure handling) before starting the next, and a hop that instead succeeded ends the request there via - async_log_success_event, which already pops everything pending on - this same model_call_details -- so admission is never re-entered, in - that same lineage, while an earlier hop's reservation is still - legitimately in flight. + async_log_success_event, which already released its own reservation + -- so admission is never re-entered, in that same lineage, while an + earlier hop's reservation is still legitimately in flight. The lineage check matters because `model_call_details` is not always scoped to one such chain: `Router.abatch_completion`'s comma-separated @@ -1611,22 +1679,42 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] return frozenset(key for key, _partition_key, token in pending_request_increments if token is admission_token) async def _pop_pending_concurrency_keys( - self, kwargs: Mapping[str, object], *, only_own_lineage: bool = False + self, + kwargs: Mapping[str, object], + *, + only_own_lineage: bool = False, + only_keys: frozenset[str] = frozenset(), ) -> tuple[tuple[str, _PartitionKey], ...]: # Snapshot then remove only those exact entries, never a blanket # clear: a sibling branch sharing this same request's # model_call_details can still be live and appending concurrently # (see the field's own docstring), so wiping the whole list here # would silently strand that branch's reservation instead of - # releasing it later. `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. + # 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 + # `_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. 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 + if (idx := next((i for i, entry in enumerate(pending) if entry[0] == key), None)) is not None + ) snapshot: Final = ( - tuple(entry for entry in pending if entry[2] is _current_admission_token()) + tuple(pending[idx] for idx in matched_indices) + if only_keys + else tuple(entry for entry in pending if entry[2] is _current_admission_token()) if only_own_lineage else tuple(pending) ) @@ -1770,66 +1858,40 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] ) await self._release_keys(release_keys) - async def async_log_failure_event( - self, - kwargs, # noqa: ANN001 # matches CustomLogger.async_log_failure_event; kwargs is dict[str, Any] codebase-wide - response_obj: object, - start_time: object, - end_time: object, - ) -> None: - # No special-case skip for this hook's own tag_rate_limit_exceeded - # rejection: a hop whose own admission rejects never reaches the - # point where a concurrency reservation is queued (see - # async_filter_deployments), so _pop_pending_concurrency_keys already - # returns nothing to release in that case. Skipping release based on - # the exception's error marker alone would be wrong here, since - # global_tag_rate_limits_hook raises 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. - # - # Not `only_own_lineage=True`: a streaming response is consumed (and - # this event fired) from a task the proxy forks independently of - # admission's own, so `_ADMISSION_CONTEXT` never reaches it either -- - # requiring a match here would leak every streaming request's - # reservation until `_CONCURRENCY_MIN_SAFETY_TTL_SECONDS` instead of - # releasing it. `async_release_disconnect_state_hook` stays - # unfiltered for the identical reason; releasing everything here - # keeps a still-live `abatch_completion` sibling's own reservation - # exposed to the same risk that hook already accepts. - release_keys: Final = await self._pop_pending_concurrency_keys(kwargs) - if release_keys: - await self._release_keys(release_keys) - - async def async_log_success_event( - self, - kwargs, # noqa: ANN001 # matches CustomLogger.async_log_success_event; kwargs is dict[str, Any] codebase-wide - response_obj: object, - start_time: object, - end_time: object, - ) -> None: - # Not `only_own_lineage=True`: see async_log_failure_event's own comment above. - 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 - release_task.add_done_callback(_BACKGROUND_TASKS.discard) - + 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 + 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. + """ if self.llm_router is None: - return - + return None standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") if standard_logging_object is None: - return - + return None model_group: Final = standard_logging_object.get("model_group") if not model_group: - return + return None # kwargs here is Logging.model_call_details, not the router's flat # request kwargs admission sees: metadata/litellm_metadata are never # top-level here, only nested under kwargs["litellm_params"] (see # Logging.update_environment_variables). - litellm_params_for_metadata: Final = kwargs.get("litellm_params") or kwargs + litellm_params_raw: Final = kwargs.get("litellm_params") + litellm_params_for_metadata: Final = litellm_params_raw if isinstance(litellm_params_raw, Mapping) else kwargs metadata_variable_name: Final = _resolve_authoritative_metadata_variable_name(litellm_params_for_metadata) team_id: Final = _extract_team_id(litellm_params_for_metadata, metadata_variable_name) key_hash: Final = _extract_key_hash(litellm_params_for_metadata, metadata_variable_name) @@ -1873,7 +1935,7 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] 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 + return None tags: Final = _order_tags_for_identity_resolution( _get_tags_from_request_kwargs(kwargs, metadata_variable_name=metadata_variable_name), @@ -1881,22 +1943,108 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass] metadata_variable_name, ) if not tags: + return None + + return _HopContext( + standard_logging_object=standard_logging_object, + configured=configured, + model_group=model_group, + tags=tags, + deployment_id=deployment_id if isinstance(deployment_id, str) else None, + key_hash=key_hash, + 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). + Falls back to the pre-existing unconditional release only when + `context` itself couldn't be resolved (nothing configured for this + tag/model, or an identity-extraction edge case) -- in which case there + 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. + """ + 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 own_concurrency_keys: + return await self._pop_pending_concurrency_keys(kwargs, only_keys=own_concurrency_keys) + return await self._pop_pending_concurrency_keys(kwargs) + + async def async_log_failure_event( + self, + kwargs, # noqa: ANN001 # matches CustomLogger.async_log_failure_event; kwargs is dict[str, Any] codebase-wide + response_obj: object, + start_time: object, + end_time: object, + ) -> None: + # No special-case skip for this hook's own tag_rate_limit_exceeded + # rejection: a hop whose own admission rejects never reaches the + # point where a concurrency reservation is queued (see + # async_filter_deployments), so nothing is ever found to release in + # that case. Skipping release based on the exception's error marker + # alone would be wrong here, since global_tag_rate_limits_hook raises + # 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)) + if release_keys: + await self._release_keys(release_keys) + + async def async_log_success_event( + self, + kwargs, # noqa: ANN001 # matches CustomLogger.async_log_success_event; kwargs is dict[str, Any] codebase-wide + response_obj: object, + start_time: object, + end_time: object, + ) -> None: + context: Final = self._resolve_hop_context(kwargs) + release_keys: Final = await self._release_own_concurrency_keys(kwargs, context) + 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 + release_task.add_done_callback(_BACKGROUND_TASKS.discard) + + if context is None: return now: Final = _admission_time_or(kwargs, fallback=self._time_provider().timestamp()) increment_by_unit: Final[Mapping[_LimitUnit, float]] = MappingProxyType( { - "tokens": float(standard_logging_object.get("total_tokens") or 0), - "dollars": float(standard_logging_object.get("response_cost") or 0), + "tokens": float(context.standard_logging_object.get("total_tokens") or 0), + "dollars": float(context.standard_logging_object.get("response_cost") or 0), } ) operation_by_limit: Final = tuple( (configured_limit, operation) - for configured_limit in configured + for configured_limit in context.configured if ( operation := _increment_operation_for_limit( - configured_limit, model_group, tags, deployment_id, key_hash, key_alias, increment_by_unit, now + configured_limit, + context.model_group, + context.tags, + context.deployment_id, + context.key_hash, + context.key_alias, + increment_by_unit, + now, ) ) is not None 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 3e34715a606..67d00685410 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 @@ -27,11 +27,15 @@ from litellm.proxy.hooks.model_based_tag_rate_limits_hook import ( _current_admission_token, _extract_team_id, _inflight_key, + _own_concurrency_key_for_limit, + _own_concurrency_keys_for_hop, _pending_reservations_cache_key, _PROXY_ModelBasedTagRateLimitsHook, ) from litellm.proxy.hooks.tag_rate_limits_shared import ( BACKGROUND_TASKS as _BACKGROUND_TASKS, +) +from litellm.proxy.hooks.tag_rate_limits_shared import ( CONCURRENCY_MIN_SAFETY_TTL_SECONDS as _CONCURRENCY_MIN_SAFETY_TTL_SECONDS, ) from litellm.types.router import ( @@ -3076,6 +3080,112 @@ async def test_concurrent_batch_siblings_do_not_bypass_a_concurrency_limit(time_ assert len(rejections) == 1 +@pytest.mark.asyncio +async def test_fast_failing_batch_sibling_does_not_release_a_still_executing_siblings_slot(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. + """ + 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 + ) + + # 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()) + + # 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", + "total_tokens": 0, + "response_cost": 0, + } + 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. + 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"]}}, + ) + + def _request_limit_router(limit: int) -> "litellm.Router": return litellm.Router( model_list=[ @@ -4255,6 +4365,76 @@ 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), @@ -4336,6 +4516,50 @@ 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) # --------------------------------------------------------------------------- @@ -4779,12 +5003,16 @@ async def test_concurrency_scope_by_key_hash_gives_independent_reservations_per_ ) async def _release(key: str, kwargs: dict): + # Mirrors real Router dispatch: the deployment-specific litellm_params + # this hop actually attempted with carries the same authenticated + # user_api_key admission itself read, which is what + # _resolve_hop_context recomputes this hop's own concurrency key from. + kwargs["metadata"]["user_api_key"] = key kwargs["standard_logging_object"] = { "model_group": "grp", "model_id": "dep-1", "total_tokens": 0, "response_cost": 0, - "metadata": {"user_api_key_hash": key}, } await limiter.async_log_success_event( kwargs=kwargs,