mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix(proxy): don't release a live sibling's slot when this hop's own context resolves to nothing configured, and lock in multi-policy release
Cursor Bugbot: _release_own_concurrency_keys still fell back to an unconditional release whenever _resolve_hop_context returned None, which included the common case of a hop whose own model has no tag_rate_limits at all. That hop's own admission never queues a reservation either (it hits the identical "nothing configured" check first), so on abatch_completion the unconditional release could still sweep up a genuinely live sibling's own slot. _resolve_hop_context now folds an empty configured/tags into a resolved-but-empty context instead of returning None for those two cases specifically, since _own_concurrency_keys_for_hop already computes an empty release set from them on its own. None is reserved for the genuinely ambiguous causes (no router, no standard_logging_object, no model_group) where this hop's own admission could have reserved something real but there's no data left to recompute it -- e.g. a different registered CustomLogger rejecting the request before this hop's own call ever ran. Those still need the pre-existing unconditional fallback, which three existing tests already depended on and caught immediately when the first version of this fix collapsed both cases together. Also adds direct empirical coverage (prompted by the sibling PR's global hook finding) confirming a hop matching two different concurrency-scoped entries releases both reservations at once, not just one.
This commit is contained in:
parent
180b5fd69b
commit
f825c03b67
2 changed files with 169 additions and 21 deletions
|
|
@ -1969,17 +1969,25 @@ 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,
|
||||
)
|
||||
if not tags:
|
||||
return None
|
||||
|
||||
# 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.
|
||||
return _HopContext(
|
||||
standard_logging_object=standard_logging_object,
|
||||
configured=configured,
|
||||
|
|
@ -1999,22 +2007,28 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass]
|
|||
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.
|
||||
|
||||
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.
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -3389,6 +3389,140 @@ async def test_batch_sibling_reserving_no_concurrency_slot_does_not_release_a_st
|
|||
)
|
||||
|
||||
|
||||
@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
|
||||
async def test_a_hop_matching_two_concurrency_scoped_entries_releases_both_reservations(time_controller):
|
||||
"""
|
||||
#38347's global hook had a real bug on completely ordinary, non-batch
|
||||
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.
|
||||
"""
|
||||
limiter = _make_limiter(time_controller)
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
_deployment(
|
||||
"grp",
|
||||
"dep-1",
|
||||
{
|
||||
"concurrency_limits": {
|
||||
"limits": [
|
||||
{"name": "per-user", "tag_id": "end_user_id", "limit": 1, "period_seconds": 300},
|
||||
{"name": "per-team", "tag_id": "team_id", "limit": 1, "period_seconds": 300},
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
]
|
||||
)
|
||||
limiter.update_variables(llm_router=router)
|
||||
request_kwargs, kwargs = _call_context(["end_user_id:u1", "team_id:t1"])
|
||||
healthy = router.model_list
|
||||
|
||||
admitted = await limiter.async_filter_deployments(
|
||||
model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs
|
||||
)
|
||||
assert admitted == healthy
|
||||
pending = kwargs.get(_PENDING_CONCURRENCY_KEYS_FIELD)
|
||||
assert pending is not None and len(pending) == 2
|
||||
|
||||
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 not kwargs.get(_PENDING_CONCURRENCY_KEYS_FIELD)
|
||||
|
||||
# Isolated per-policy checks: a follow-up request carrying only one tag
|
||||
# exercises only that one policy, so each must independently show its
|
||||
# own limit=1 slot free again, not just "no exception" from whichever
|
||||
# one entry a partial release happened to free.
|
||||
await limiter.async_filter_deployments(
|
||||
model="grp", healthy_deployments=healthy, messages=None, request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}
|
||||
)
|
||||
await limiter.async_filter_deployments(
|
||||
model="grp", healthy_deployments=healthy, messages=None, request_kwargs={"metadata": {"tags": ["team_id:t1"]}}
|
||||
)
|
||||
|
||||
|
||||
def _request_limit_router(limit: int) -> "litellm.Router":
|
||||
return litellm.Router(
|
||||
model_list=[
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue