mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(proxy): stop a batch sibling's terminal event from releasing another branch's live concurrency slot
Bugbot finding: async_log_success_event/async_log_failure_event still called _pop_pending_concurrency_keys with no filter, so the first finishing branch of an abatch_completion dispatch released every reservation on the shared model_call_details, including a still-live sibling branch's own slot. Reservations are now tagged with an admission-scoped token from a ContextVar rather than a raw asyncio.Task: create_task snapshots the current Context, so a task explicitly forked from within one hop's own admission (litellm's own logging dispatch explicitly propagates context) still reads back the same token, while abatch_completion's sibling branches, forked before any of them call admission, each mint their own distinct token on first use. Both the stale-hop cleanup and the terminal release hooks (success/failure) now require this token to match before releasing; the disconnect hook is left unfiltered, since its own scenario (mid-stream disconnect) cannot co-occur with abatch_completion's combined, non-streaming response.
This commit is contained in:
Deepanshu
parent
744733ac2d
commit
c2d881894e
2 changed files with 256 additions and 87 deletions
|
|
@ -12,6 +12,7 @@ per-deployment dedup.
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
import contextvars
|
||||
import json
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from dataclasses import dataclass, replace
|
||||
|
|
@ -558,15 +559,42 @@ _INDEX_TTL_SECONDS: Final = 5.0
|
|||
# can't be forged or guessed.
|
||||
#
|
||||
# "shared across that request's own fallback hops" is not the same as
|
||||
# "scoped to one asyncio Task": `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
|
||||
# each entry also carries the Task that queued it (see
|
||||
# `_queue_pending_reservations`), letting `_release_stale_hop_reservations`
|
||||
# tell a genuinely stale same-task hop apart from a still-live sibling
|
||||
# branch's own reservation.
|
||||
# "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`) so a release can
|
||||
# tell a genuinely stale same-lineage hop apart from a still-live sibling
|
||||
# branch's own reservation, without reintroducing the non-descendant-task
|
||||
# blind spot a bare `ContextVar` has for the data itself (see above): the
|
||||
# token is only ever compared for *identity*, never relied on to carry the
|
||||
# reservation across a task boundary the way `model_call_details` does.
|
||||
_PENDING_CONCURRENCY_KEYS_FIELD: Final[str] = "_model_based_tag_rate_limits_pending_concurrency_keys"
|
||||
|
||||
# Identifies which admission call queued a given reservation, scoped by
|
||||
# asyncio Context rather than by Task identity: `asyncio.create_task`
|
||||
# snapshots the *current* Context into the new Task, so a task explicitly
|
||||
# forked from within one hop's own admission (e.g. litellm's own
|
||||
# `LoggingWorker.enqueue` dispatch, which explicitly propagates the calling
|
||||
# context) still reads back the same token, while `abatch_completion`'s
|
||||
# sibling branches -- forked *before* any of them ever called admission --
|
||||
# each start from an unset ContextVar and mint their own distinct token on
|
||||
# first use, never matching each other's.
|
||||
_ADMISSION_CONTEXT: Final[contextvars.ContextVar[object | None]] = contextvars.ContextVar(
|
||||
"_model_based_tag_rate_limits_admission_context", default=None
|
||||
)
|
||||
|
||||
|
||||
def _current_admission_token() -> object:
|
||||
token: Final = _ADMISSION_CONTEXT.get()
|
||||
if token is not None:
|
||||
return token
|
||||
fresh_token: Final = object()
|
||||
_ADMISSION_CONTEXT.set(fresh_token)
|
||||
return fresh_token
|
||||
|
||||
|
||||
# Same `model_call_details`-stashing rationale as the field above, for a
|
||||
# different unit: "requests" is atomic and admitted once per hop (see
|
||||
# _ATOMIC_UNITS), same as concurrency, but a "requests" limit is meant to cap
|
||||
|
|
@ -898,20 +926,20 @@ def _queue_pending_reservations(
|
|||
) -> None:
|
||||
"""Stash reservations on the request's own `model_call_details`, under
|
||||
`field` -- see `_PENDING_CONCURRENCY_KEYS_FIELD`'s docstring for why this,
|
||||
not a ContextVar or `litellm_call_id`. Silently a no-op without a real
|
||||
logging object (defensive only; every real request has one): a queued
|
||||
concurrency reservation still self-heals via
|
||||
`_CONCURRENCY_MIN_SAFETY_TTL_SECONDS`, just later.
|
||||
not a `litellm_call_id`. Silently a no-op without a real logging object
|
||||
(defensive only; every real request has one): a queued concurrency
|
||||
reservation still self-heals via `_CONCURRENCY_MIN_SAFETY_TTL_SECONDS`,
|
||||
just later.
|
||||
|
||||
Each entry is stamped with the queueing coroutine's own `asyncio.Task`:
|
||||
`Router.abatch_completion`'s comma-separated multi-model dispatch runs
|
||||
several `acompletion` calls concurrently as *separate* tasks that all
|
||||
share one `model_call_details` (the proxy attaches one `litellm_logging_obj`
|
||||
to the request before the comma-split, and every branch inherits that
|
||||
same reference), so this field is no longer scoped to one logical
|
||||
request's own serial fallback chain the way its docstring assumes.
|
||||
`_release_stale_hop_reservations` uses the stamp to tell "an earlier hop
|
||||
of *this* chain, safe to reclaim" apart from "a concurrent sibling
|
||||
Each entry is stamped with `_current_admission_token()`: `Router.abatch_completion`'s
|
||||
comma-separated multi-model dispatch runs several `acompletion` calls
|
||||
concurrently as *separate* tasks that all share one `model_call_details`
|
||||
(the proxy attaches one `litellm_logging_obj` to the request before the
|
||||
comma-split, and every branch inherits that same reference), so this
|
||||
field is no longer scoped to one logical request's own serial fallback
|
||||
chain the way its docstring assumes. `_release_stale_hop_reservations`
|
||||
and the terminal release hooks use the stamp to tell "this same
|
||||
admission lineage, safe to reclaim" apart from "a concurrent sibling
|
||||
branch's own still-live reservation," which must never be touched here.
|
||||
"""
|
||||
logging_obj: Final = request_kwargs.get("litellm_logging_obj")
|
||||
|
|
@ -922,9 +950,9 @@ def _queue_pending_reservations(
|
|||
if pending is None:
|
||||
pending = [] # mutable-ok: shared, request-scoped accumulator; see field's own docstring # rebind-ok: lazily initialized only when absent
|
||||
model_call_details[field] = pending
|
||||
current_task: Final = asyncio.current_task()
|
||||
admission_token: Final = _current_admission_token()
|
||||
pending.extend(
|
||||
(key, partition_key, current_task) for key, partition_key in reservations
|
||||
(key, partition_key, admission_token) for key, partition_key in reservations
|
||||
) # mutable-ok: see comment above
|
||||
|
||||
|
||||
|
|
@ -1483,24 +1511,24 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass]
|
|||
async def _release_stale_hop_reservations(self, request_kwargs: Mapping[str, object]) -> frozenset[str]:
|
||||
"""
|
||||
A concurrency reservation still queued when a *new* hop's admission
|
||||
runs, *within the same asyncio Task*, can only belong to an earlier
|
||||
hop of this same request's own fallback chain that already concluded
|
||||
and failed: Router awaits one hop's entire attempt (call plus its own
|
||||
failure handling) before starting the next, and a hop that instead
|
||||
succeeded ends the request there via async_log_success_event, which
|
||||
already pops everything -- so admission is never re-entered, in that
|
||||
same task, while an earlier hop's reservation is still legitimately
|
||||
in flight.
|
||||
runs, *within the same admission lineage*, can only belong to an
|
||||
earlier hop of this same request's own fallback chain that already
|
||||
concluded and failed: Router awaits one hop's entire attempt (call
|
||||
plus its own failure handling) before starting the next, and a hop
|
||||
that instead succeeded ends the request there via
|
||||
async_log_success_event, which already pops its own lineage's
|
||||
entries -- so admission is never re-entered, in that same lineage,
|
||||
while an earlier hop's reservation is still legitimately in flight.
|
||||
|
||||
The task check matters because `model_call_details` is not always
|
||||
The lineage check matters because `model_call_details` is not always
|
||||
scoped to one such chain: `Router.abatch_completion`'s comma-separated
|
||||
multi-model dispatch runs several branches concurrently, each its own
|
||||
Task, but every branch shares the identical `litellm_logging_obj` (see
|
||||
`_queue_pending_reservations`'s own docstring) -- so a reservation
|
||||
queued by a still-running sibling branch can be sitting here too, and
|
||||
releasing it out from under that branch would let more calls through
|
||||
a concurrency limit than it allows. Only entries this exact Task
|
||||
queued are safe to treat as stale; anything else is left for its own
|
||||
a concurrency limit than it allows. Only entries `_current_admission_token()`
|
||||
stamped are safe to treat as stale; anything else is left for its own
|
||||
branch to release.
|
||||
|
||||
LiteLLM only invokes a request's CustomLogger.async_log_failure_event
|
||||
|
|
@ -1541,17 +1569,17 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass]
|
|||
model_call_details: Final = getattr(logging_obj, "model_call_details", None)
|
||||
if not isinstance(model_call_details, dict):
|
||||
return frozenset()
|
||||
release_keys: Final = await self._pop_pending_concurrency_keys(model_call_details, only_current_task=True)
|
||||
release_keys: Final = await self._pop_pending_concurrency_keys(model_call_details, only_own_lineage=True)
|
||||
if release_keys:
|
||||
await self._release_keys(release_keys)
|
||||
pending_request_increments: Final = model_call_details.get(_PENDING_REQUEST_INCREMENTS_FIELD)
|
||||
if not isinstance(pending_request_increments, list):
|
||||
return frozenset()
|
||||
current_task: Final = asyncio.current_task()
|
||||
return frozenset(key for key, _partition_key, task in pending_request_increments if task is current_task)
|
||||
admission_token: Final = _current_admission_token()
|
||||
return frozenset(key for key, _partition_key, token in pending_request_increments if token is admission_token)
|
||||
|
||||
async def _pop_pending_concurrency_keys(
|
||||
self, kwargs: Mapping[str, object], *, only_current_task: bool = False
|
||||
self, kwargs: Mapping[str, object], *, only_own_lineage: bool = False
|
||||
) -> tuple[tuple[str, _PartitionKey], ...]:
|
||||
# Every caller of this method is itself a normal release path, so
|
||||
# also clear the async_post_call_failure_hook cache mirror for the
|
||||
|
|
@ -1585,15 +1613,20 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass]
|
|||
# model_call_details can still be live and appending concurrently
|
||||
# (see the field's own docstring), so wiping the whole list here
|
||||
# would silently strand that branch's reservation instead of
|
||||
# releasing it later. `only_current_task` additionally excludes any
|
||||
# entry a *different*, still-running Task queued -- see
|
||||
# releasing it later. `only_own_lineage` additionally excludes any
|
||||
# entry a *different* admission lineage queued -- see
|
||||
# `_release_stale_hop_reservations`'s own docstring for why that
|
||||
# distinction, not just presence, decides what's actually stale.
|
||||
# distinction, not just presence, decides what's actually stale;
|
||||
# the same distinction applies to a terminal release (success/failure)
|
||||
# racing a still-live sibling branch's own reservation.
|
||||
pending: Final = kwargs.get(_PENDING_CONCURRENCY_KEYS_FIELD)
|
||||
if not isinstance(pending, list) or not pending:
|
||||
return ()
|
||||
current_task: Final = asyncio.current_task()
|
||||
snapshot: Final = tuple(entry for entry in pending if not only_current_task or entry[2] is current_task)
|
||||
snapshot: Final = (
|
||||
tuple(entry for entry in pending if entry[2] is _current_admission_token())
|
||||
if only_own_lineage
|
||||
else tuple(pending)
|
||||
)
|
||||
for entry in snapshot:
|
||||
try:
|
||||
pending.remove(entry)
|
||||
|
|
@ -1610,6 +1643,12 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass]
|
|||
Without this, the reservation would sit held until _CONCURRENCY_MIN_SAFETY_TTL_SECONDS
|
||||
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.
|
||||
"""
|
||||
logging_obj: Final = request_data.get("litellm_logging_obj")
|
||||
model_call_details: Final = getattr(logging_obj, "model_call_details", None)
|
||||
|
|
@ -1685,13 +1724,18 @@ class _PROXY_ModelBasedTagRateLimitsHook( # pyright: ignore[reportUnusedClass]
|
|||
# the exception's error marker alone would be wrong here, since
|
||||
# global_tag_rate_limits_hook raises the identical marker -- that
|
||||
# rejection can land after this hook already reserved a slot for the
|
||||
# same request, and that slot must still be released.
|
||||
release_keys: Final = await self._pop_pending_concurrency_keys(kwargs)
|
||||
# same request, and that slot must still be released. only_own_lineage
|
||||
# keeps this from releasing a still-live sibling branch's own
|
||||
# reservation when model_call_details is shared across an
|
||||
# abatch_completion dispatch -- see _PENDING_CONCURRENCY_KEYS_FIELD's
|
||||
# docstring.
|
||||
release_keys: Final = await self._pop_pending_concurrency_keys(kwargs, only_own_lineage=True)
|
||||
if release_keys:
|
||||
await self._release_keys(release_keys)
|
||||
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None:
|
||||
release_keys: Final = await self._pop_pending_concurrency_keys(kwargs)
|
||||
# only_own_lineage: see async_log_failure_event's own comment above.
|
||||
release_keys: Final = await self._pop_pending_concurrency_keys(kwargs, only_own_lineage=True)
|
||||
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
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ from litellm.proxy.hooks.model_based_tag_rate_limits_hook import (
|
|||
_build_group_limits,
|
||||
_build_limits_index,
|
||||
_ConfiguredLimit,
|
||||
_current_admission_token,
|
||||
_extract_team_id,
|
||||
_inflight_key,
|
||||
_pending_reservations_cache_key,
|
||||
|
|
@ -33,7 +34,14 @@ from litellm.proxy.hooks.tag_rate_limits_shared import (
|
|||
BACKGROUND_TASKS as _BACKGROUND_TASKS,
|
||||
CONCURRENCY_MIN_SAFETY_TTL_SECONDS as _CONCURRENCY_MIN_SAFETY_TTL_SECONDS,
|
||||
)
|
||||
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, RoutingGroup, TagRateLimitEntry, TagRateLimitScope
|
||||
from litellm.types.router import (
|
||||
Deployment,
|
||||
LiteLLM_Params,
|
||||
ModelInfo,
|
||||
RoutingGroup,
|
||||
TagRateLimitEntry,
|
||||
TagRateLimitScope,
|
||||
)
|
||||
|
||||
|
||||
class TimeController:
|
||||
|
|
@ -230,7 +238,11 @@ async def test_filter_deployments_ignores_a_forged_empty_litellm_metadata_key(ti
|
|||
deployment = _deployment(
|
||||
"grp",
|
||||
"dep-1",
|
||||
{"request_limits": {"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}]}},
|
||||
{
|
||||
"request_limits": {
|
||||
"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}]
|
||||
}
|
||||
},
|
||||
)
|
||||
router = litellm.Router(model_list=[deployment])
|
||||
limiter.update_variables(llm_router=router)
|
||||
|
|
@ -258,7 +270,11 @@ async def test_filter_deployments_reads_metadata_when_litellm_metadata_is_presen
|
|||
deployment = _deployment(
|
||||
"grp",
|
||||
"dep-1",
|
||||
{"request_limits": {"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}]}},
|
||||
{
|
||||
"request_limits": {
|
||||
"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}]
|
||||
}
|
||||
},
|
||||
)
|
||||
router = litellm.Router(model_list=[deployment])
|
||||
limiter.update_variables(llm_router=router)
|
||||
|
|
@ -288,7 +304,11 @@ async def test_filter_deployments_ignores_a_forged_populated_litellm_metadata_ke
|
|||
deployment = _deployment(
|
||||
"grp",
|
||||
"dep-1",
|
||||
{"request_limits": {"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}]}},
|
||||
{
|
||||
"request_limits": {
|
||||
"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 1, "period_seconds": 86400}]
|
||||
}
|
||||
},
|
||||
)
|
||||
router = litellm.Router(model_list=[deployment])
|
||||
limiter.update_variables(llm_router=router)
|
||||
|
|
@ -1730,7 +1750,15 @@ async def test_log_success_event_accounts_against_the_same_bucket_admission_chec
|
|||
|
||||
now = time_controller.now().timestamp()
|
||||
token_key = _expected_bucket_key(
|
||||
"my-group", "tokens", "daily", "end_user_id", "u1", 86400, now, resolved_group=admission_bucket_group, limit=500000
|
||||
"my-group",
|
||||
"tokens",
|
||||
"daily",
|
||||
"end_user_id",
|
||||
"u1",
|
||||
86400,
|
||||
now,
|
||||
resolved_group=admission_bucket_group,
|
||||
limit=500000,
|
||||
)
|
||||
assert (
|
||||
float(await limiter.internal_usage_cache.async_get_cache(key=token_key, litellm_parent_otel_span=None)) == 42.0
|
||||
|
|
@ -1772,9 +1800,11 @@ async def test_log_success_event_uses_admissions_own_candidate_set_when_group_me
|
|||
model="my-group", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs
|
||||
)
|
||||
assert admitted == healthy
|
||||
admission_bucket_group = limiter._index.get(router).resolve_any(
|
||||
"my-group", team_id=None, candidate_model_names=("backend-a", "backend-b")
|
||||
)[0].resolved_group
|
||||
admission_bucket_group = (
|
||||
limiter._index.get(router)
|
||||
.resolve_any("my-group", team_id=None, candidate_model_names=("backend-a", "backend-b"))[0]
|
||||
.resolved_group
|
||||
)
|
||||
|
||||
routing_group = router.get_routing_group("my-group")
|
||||
assert routing_group is not None
|
||||
|
|
@ -1799,7 +1829,15 @@ async def test_log_success_event_uses_admissions_own_candidate_set_when_group_me
|
|||
|
||||
now = time_controller.now().timestamp()
|
||||
admission_key = _expected_bucket_key(
|
||||
"my-group", "tokens", "daily", "end_user_id", "u1", 86400, now, resolved_group=admission_bucket_group, limit=500000
|
||||
"my-group",
|
||||
"tokens",
|
||||
"daily",
|
||||
"end_user_id",
|
||||
"u1",
|
||||
86400,
|
||||
now,
|
||||
resolved_group=admission_bucket_group,
|
||||
limit=500000,
|
||||
)
|
||||
drifted_key = _expected_bucket_key(
|
||||
"my-group", "tokens", "daily", "end_user_id", "u1", 86400, now, resolved_group="backend-0", limit=500000
|
||||
|
|
@ -1881,7 +1919,13 @@ async def test_log_success_event_accounts_against_the_key_hash_admission_checked
|
|||
token_limits = {
|
||||
"token_limits": {
|
||||
"limits": [
|
||||
{"name": "daily", "tag_id": "end_user_id", "limit": 500000, "period_seconds": 86400, "scope_by_key_hash": True}
|
||||
{
|
||||
"name": "daily",
|
||||
"tag_id": "end_user_id",
|
||||
"limit": 500000,
|
||||
"period_seconds": 86400,
|
||||
"scope_by_key_hash": True,
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
|
|
@ -1933,7 +1977,9 @@ async def test_log_success_event_charges_the_window_admission_checked_not_a_late
|
|||
current when the response finishes.
|
||||
"""
|
||||
token_limits = {
|
||||
"token_limits": {"limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 500, "period_seconds": 60}]}
|
||||
"token_limits": {
|
||||
"limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 500, "period_seconds": 60}]
|
||||
}
|
||||
}
|
||||
router = litellm.Router(model_list=[_deployment("grp", "dep-1", token_limits)])
|
||||
limiter = _make_limiter(time_controller)
|
||||
|
|
@ -1967,7 +2013,9 @@ async def test_log_success_event_charges_the_window_admission_checked_not_a_late
|
|||
)
|
||||
assert (
|
||||
float(
|
||||
await limiter.internal_usage_cache.async_get_cache(key=admitted_window_bucket, litellm_parent_otel_span=None)
|
||||
await limiter.internal_usage_cache.async_get_cache(
|
||||
key=admitted_window_bucket, litellm_parent_otel_span=None
|
||||
)
|
||||
)
|
||||
== 42.0
|
||||
)
|
||||
|
|
@ -1993,7 +2041,11 @@ async def test_log_success_event_accounts_against_the_team_id_admission_checked(
|
|||
deployment = _deployment(
|
||||
"real-model-name",
|
||||
"dep-1",
|
||||
{"token_limits": {"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500, "period_seconds": 86400}]}},
|
||||
{
|
||||
"token_limits": {
|
||||
"limits": [{"name": "daily", "tag_id": "end_user_id", "limit": 500, "period_seconds": 86400}]
|
||||
}
|
||||
},
|
||||
)
|
||||
deployment["model_info"]["team_id"] = "team-1"
|
||||
deployment["model_info"]["team_public_model_name"] = "team-alias-name"
|
||||
|
|
@ -2958,6 +3010,63 @@ async def test_concurrent_batch_siblings_do_not_bypass_a_concurrency_limit(time_
|
|||
assert len(rejections) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_batch_siblings_terminal_event_does_not_release_a_live_siblings_slot(time_controller):
|
||||
"""
|
||||
Bugbot finding: async_log_success_event/async_log_failure_event still
|
||||
released every pending reservation unconditionally, so the first
|
||||
finishing abatch_completion branch freed a still-live sibling branch's
|
||||
own concurrency slot too, letting a third, unrelated caller admit past
|
||||
a limit that branch was still genuinely occupying.
|
||||
"""
|
||||
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"])
|
||||
branch_two_admitted = asyncio.Event()
|
||||
|
||||
async def _branch_two() -> None:
|
||||
await limiter.async_filter_deployments(
|
||||
model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs
|
||||
)
|
||||
branch_two_admitted.set()
|
||||
# Still "in flight" while branch one below finishes and releases.
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
task_two = asyncio.create_task(_branch_two())
|
||||
|
||||
await limiter.async_filter_deployments(
|
||||
model="grp", healthy_deployments=healthy, messages=None, request_kwargs=request_kwargs
|
||||
)
|
||||
await branch_two_admitted.wait()
|
||||
kwargs["standard_logging_object"] = {
|
||||
"model_group": "grp",
|
||||
"model_id": "dep-1",
|
||||
"total_tokens": 0,
|
||||
"response_cost": 0,
|
||||
}
|
||||
await limiter.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0)
|
||||
await asyncio.sleep(0)
|
||||
|
||||
# Only branch one's own slot was freed; branch two's is still live, so
|
||||
# exactly one more caller fits before the limit of 2 is hit again.
|
||||
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"]}},
|
||||
)
|
||||
await task_two
|
||||
|
||||
|
||||
def _request_limit_router(limit: int) -> "litellm.Router":
|
||||
return litellm.Router(
|
||||
model_list=[
|
||||
|
|
@ -2966,7 +3075,9 @@ def _request_limit_router(limit: int) -> "litellm.Router":
|
|||
"dep-1",
|
||||
{
|
||||
"request_limits": {
|
||||
"limits": [{"name": "per_period", "tag_id": "end_user_id", "limit": limit, "period_seconds": 300}]
|
||||
"limits": [
|
||||
{"name": "per_period", "tag_id": "end_user_id", "limit": limit, "period_seconds": 300}
|
||||
]
|
||||
}
|
||||
},
|
||||
)
|
||||
|
|
@ -3419,7 +3530,9 @@ def _redis_limiter(time_controller: TimeController):
|
|||
pytest.skip("Redis environment variables (REDIS_HOST, REDIS_PORT) not set")
|
||||
redis_cache = RedisCache(host=redis_host, port=int(redis_port), password=os.getenv("REDIS_PASSWORD"))
|
||||
dual_cache = DualCache(redis_cache=redis_cache)
|
||||
return _PROXY_ModelBasedTagRateLimitsHook(internal_usage_cache=dual_cache, time_provider=time_controller.now), redis_cache
|
||||
return _PROXY_ModelBasedTagRateLimitsHook(
|
||||
internal_usage_cache=dual_cache, time_provider=time_controller.now
|
||||
), redis_cache
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -3545,7 +3658,11 @@ async def test_redis_backed_token_admission_sees_increments_the_in_memory_cache_
|
|||
_deployment(
|
||||
"grp",
|
||||
"dep-1",
|
||||
{"token_limits": {"limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 100, "period_seconds": 60}]}},
|
||||
{
|
||||
"token_limits": {
|
||||
"limits": [{"name": "per_minute", "tag_id": "end_user_id", "limit": 100, "period_seconds": 60}]
|
||||
}
|
||||
},
|
||||
)
|
||||
]
|
||||
)
|
||||
|
|
@ -3626,7 +3743,9 @@ async def test_redis_backed_concurrency_ttl_refreshes_on_every_admission(time_co
|
|||
# would make a real-time before/after comparison too slow to assert
|
||||
# on deterministically) with refresh_ttl=True, matching how a
|
||||
# concurrency check is actually admitted.
|
||||
admitted, _ = await limiter._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True)
|
||||
admitted, _ = await limiter._check_and_increment_one(
|
||||
cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True
|
||||
)
|
||||
assert admitted
|
||||
ttl_after_first_admission = await redis_cache.redis_async_client.ttl(key)
|
||||
assert ttl_after_first_admission > 0
|
||||
|
|
@ -3636,7 +3755,9 @@ async def test_redis_backed_concurrency_ttl_refreshes_on_every_admission(time_co
|
|||
# A second admission on the same still-live key, most of the way
|
||||
# through the first admission's ttl, must push the ttl back out to
|
||||
# the full window again, not leave it counting down toward zero.
|
||||
admitted, _ = await limiter._check_and_increment_one(cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True)
|
||||
admitted, _ = await limiter._check_and_increment_one(
|
||||
cache, key, limit=100, increment=1.0, ttl=3, refresh_ttl=True
|
||||
)
|
||||
assert admitted
|
||||
ttl_after_second_admission = await redis_cache.redis_async_client.ttl(key)
|
||||
assert ttl_after_second_admission >= 2
|
||||
|
|
@ -4112,9 +4233,9 @@ def test_concurrency_ttl_floor_does_not_shorten_a_longer_period_seconds():
|
|||
@pytest.mark.asyncio
|
||||
async def test_release_in_a_forked_task_is_visible_to_the_parent_context(time_controller):
|
||||
limiter = _make_limiter(time_controller)
|
||||
# Entries are (key, partition_key, queueing_task) triples in production
|
||||
# (see _queue_pending_reservations); the task is irrelevant to this
|
||||
# specific release path (only_current_task defaults False here).
|
||||
# Entries are (key, partition_key, admission_token) triples in production
|
||||
# (see _queue_pending_reservations); the token is irrelevant to this
|
||||
# specific release path (only_own_lineage defaults False here).
|
||||
model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: [("key1", None, None)]}
|
||||
|
||||
async def detached_release():
|
||||
|
|
@ -4156,28 +4277,26 @@ async def test_release_is_not_repeated_for_the_same_snapshot(time_controller):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_only_current_task_leaves_a_concurrent_siblings_reservation_alone(time_controller):
|
||||
async def test_release_only_own_lineage_leaves_a_concurrent_siblings_reservation_alone(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 shares one litellm_logging_obj (the proxy attaches it to the
|
||||
request before the comma-split), so a new hop's admission could see a
|
||||
still-live sibling branch's own reservation sitting in the same
|
||||
model_call_details and wrongly sweep it up as "stale". only_current_task
|
||||
must leave a differently-tasked entry untouched.
|
||||
Bugbot/Veria AI finding: Router.abatch_completion's comma-separated
|
||||
multi-model dispatch runs each model concurrently, each its own asyncio
|
||||
Task, but every branch shares one litellm_logging_obj (the proxy attaches
|
||||
it to the request before the comma-split), so a still-live sibling
|
||||
branch's own reservation can sit in the same model_call_details.
|
||||
only_own_lineage must release this context's own entry while leaving a
|
||||
differently-lineaged (sibling branch's) entry untouched.
|
||||
"""
|
||||
limiter = _make_limiter(time_controller)
|
||||
own_token = _current_admission_token()
|
||||
sibling_token = object()
|
||||
model_call_details: dict = {
|
||||
_PENDING_CONCURRENCY_KEYS_FIELD: [("own-key", None, own_token), ("sibling-key", None, sibling_token)]
|
||||
}
|
||||
|
||||
async def _reserve_as_a_separate_task() -> None:
|
||||
pass # the task object itself is the fixture; body is irrelevant
|
||||
|
||||
sibling_task = asyncio.create_task(_reserve_as_a_separate_task())
|
||||
await sibling_task
|
||||
model_call_details: dict = {_PENDING_CONCURRENCY_KEYS_FIELD: [("sibling-key", None, sibling_task)]}
|
||||
|
||||
released = await limiter._pop_pending_concurrency_keys(model_call_details, only_current_task=True)
|
||||
assert released == ()
|
||||
assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == [("sibling-key", None, sibling_task)]
|
||||
released = await limiter._pop_pending_concurrency_keys(model_call_details, only_own_lineage=True)
|
||||
assert released == (("own-key", None),)
|
||||
assert model_call_details[_PENDING_CONCURRENCY_KEYS_FIELD] == [("sibling-key", None, sibling_token)]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -4230,7 +4349,9 @@ async def test_cross_unit_refund_leaves_no_phantom_increment_in_memory(time_cont
|
|||
)
|
||||
|
||||
now = time_controller.now().timestamp()
|
||||
request_key = _expected_bucket_key("grp", "requests", "per_minute", "end_user_id", "refund-check", 60, now, limit=10)
|
||||
request_key = _expected_bucket_key(
|
||||
"grp", "requests", "per_minute", "end_user_id", "refund-check", 60, now, limit=10
|
||||
)
|
||||
value = await limiter.internal_usage_cache.async_get_cache(key=request_key, litellm_parent_otel_span=None)
|
||||
assert (float(value) if value is not None else 0.0) == 1.0
|
||||
|
||||
|
|
@ -4304,7 +4425,9 @@ async def test_exception_mid_batch_refunds_every_earlier_admission_before_propag
|
|||
raising_key = "{tag_rl:test:exception-refund:b}:requests"
|
||||
|
||||
class _FlakyLimiter(_PROXY_ModelBasedTagRateLimitsHook):
|
||||
async def _check_and_increment_one(self, cache, key: str, limit: float, increment: float, ttl: int, refresh_ttl: bool):
|
||||
async def _check_and_increment_one(
|
||||
self, cache, key: str, limit: float, increment: float, ttl: int, refresh_ttl: bool
|
||||
):
|
||||
if key == raising_key:
|
||||
raise RuntimeError("simulated transient redis failure")
|
||||
return await super()._check_and_increment_one(cache, key, limit, increment, ttl, refresh_ttl)
|
||||
|
|
@ -4342,7 +4465,9 @@ async def test_a_raising_keys_own_ambiguous_outcome_is_never_refunded(time_contr
|
|||
raising_key = "{tag_rl:test:ambiguous-no-refund:b}:requests"
|
||||
|
||||
class _FlakyLimiter(_PROXY_ModelBasedTagRateLimitsHook):
|
||||
async def _check_and_increment_one(self, cache, key: str, limit: float, increment: float, ttl: int, refresh_ttl: bool):
|
||||
async def _check_and_increment_one(
|
||||
self, cache, key: str, limit: float, increment: float, ttl: int, refresh_ttl: bool
|
||||
):
|
||||
if key == raising_key:
|
||||
# Simulate Redis committing the increment before the
|
||||
# response is lost: the write actually happens...
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue