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