scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

fix(proxy): release exactly the concurrency reservation a completing hop's own admission made, not every pending entry

Browse source 
Live verification found the abatch_completion sibling-slot leak still real
and reproducible: async_log_success_event/async_log_failure_event released
every pending reservation on model_call_details unconditionally, so a
branch that finished fast released a still-executing sibling branch's own
slot too, even after the earlier task/ContextVar-based fix (which only
ever protected the admission-time stale-hop cleanup, never these two
terminal hooks -- they had to stay unconditional or a streaming response's
own release, firing from a task the proxy forks independently of
admission's, would never match and leak until TTL instead).

Root cause traced properly this time: Router.abatch_completion's sibling
branches share one model_call_details, litellm_call_id, and top-level
metadata/litellm_metadata dict by reference (confirmed live), so stashing
a new per-admission id into any of those would hit the identical
last-write-wins race the original bug already exploits. What actually
stays reliably per-hop, even while that surrounding object is shared, is
standard_logging_object and litellm_params: 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 the success/failure callback firing, so a sibling branch
never gets a chance to interleave and overwrite it first. Token/dollar
accounting already depended on exactly this being reliable; concurrency
release now reuses the identical identity resolution (extracted into
_resolve_hop_context) to recompute the exact key(s) the completing hop's
own admission reserved, and releases at most one matching entry per key --
reservations sharing a key are fungible, so this never touches a sibling's
entry under a different key, and never over-releases a shared one either.

Falls back to the pre-existing unconditional release only when that
recomputation itself fails (nothing configured for this tag/model, or an
identity-extraction edge case): there is no way to tell a sibling's
reservation apart from this hop's own in that case either, so it only
matters for the single-branch case release already handled correctly
before abatch_completion existed.

_release_stale_hop_reservations (admission-time stale-hop cleanup) keeps
its task/ContextVar-based filter unchanged -- it only ever runs
synchronously within admission's own hop sequence, never a task fork, so
it was never exposed to the streaming-callback blind spot this fix closes
for the two terminal hooks.
This commit is contained in:
Deepanshu 2026-08-31 12:03:29 -04:00
parent 0b8c9bd489
commit 780def9194
2 changed files with 454 additions and 78 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

302
litellm/proxy/hooks/model_based_tag_rate_limits_hook.py
View file

@ -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

230
tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py
View file

@ -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,