diff --git a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py index 4c32f822b42..a1bba1af4d2 100644 --- a/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/model_based_tag_rate_limits_hook.py @@ -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 diff --git a/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py b/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py index 60df49728f1..67a8c34bf29 100644 --- a/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py +++ b/tests/test_litellm/proxy/hooks/test_model_based_tag_rate_limits_hook.py @@ -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...