mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
fix(proxy): reserve exactly one concurrency unit per abatch_completion dispatch, not one per branch
Live verification (following up on the sibling global-hook PR) confirmed a
foundational error in this hook's whole batch-completion concurrency model:
Router.abatch_completion's comma-separated branches share one litellm_logging_obj,
and therefore one Logging instance, so litellm's own has_logged_{event_type}
dedup means only ONE terminal success/failure event EVER fires for the whole
dispatch, never one per branch. The "reserve one unit per matching branch,
release exactly one per branch's own terminal event" model built up over this
session's earlier rounds was solving a race that cannot happen; the real bug
was that N units got reserved but only 1 ever got released, leaking N-1 units
until their safety TTL on every ordinary multi-model batch call.
Admission (async_filter_deployments) now dedupes concurrency reservations by
key for the whole dispatch: a key already pending (staked by an earlier
branch of the same dispatch) is skipped entirely, riding along on that one
reservation for free rather than re-checking and re-incrementing. The claim
is staked synchronously, before the atomic check-and-increment it's part of,
closing the race a genuinely concurrent sibling branch could otherwise hit
(asyncio only switches tasks at an await); rolled back if that same atomic
batch is then rejected by a different check, since nothing was actually
incremented for it.
Release goes back to fully unconditional (removing _release_own_concurrency_keys,
_resolve_hop_context's per-hop key recomputation for concurrency, and
_pop_pending_concurrency_keys's only_keys filter): with at most one
reservation per key ever pending, whichever single terminal event fires
correctly releases exactly that one entry, matching the model already
accepted for Router.abatch_completion_fastest_response.
Token/dollar accounting has no such fix available within this hook: with
only one terminal event ever firing, only one branch's own usage is ever
accounted, and every other branch's real usage silently escapes accounting
entirely. This is a real, pre-existing gap this session's fixes did not
introduce and cannot close from inside this hook -- documented in
_resolve_hop_context's own docstring and covered by a test describing the
gap concretely; a real fix needs Router-level changes (a distinct Logging
instance per branch, or aggregating usage before the one terminal event).
Adds a real-pipeline regression test (common_processing_pre_call_logic then
route_request, matching what route_llm_request.py does in production) that
drives the actual hook through litellm's own logging worker and confirms a
3-model batch dispatch against a concurrency limit of 1 admits all three
and fully releases the single reservation.
This commit is contained in:
Deepanshu
parent
f825c03b67
commit
c817e0150c
2 changed files with 423 additions and 540 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue