From 25fc49bdc3cb3b6aabd13a28cd8a391b7378795e Mon Sep 17 00:00:00 2001 From: Deepanshu Date: Thu, 3 Sep 2026 10:18:37 -0400 Subject: [PATCH] fix(rate-limiting): reserve one unit for the whole non-racing batch dispatch, not one per model Router.abatch_completion shares one litellm_logging_obj across every branch (attached in common_request_processing.py before Router ever sees the request, then threaded unchanged through every branch's own kwargs). Logging.should_run_logging dedups per event type against that shared object, so at most one success and one failure callback ever fire for the whole dispatch, never once per branch. Confirmed live with a real Router.abatch_completion and a real, shared litellm_logging_obj attached the way production actually does it: a 3-model batch reserving one unit per model left 2 of 3 units permanently leaked after every branch genuinely succeeded, since only one release event ever arrived to pop them. Reverts the per-branch reservation this hook added when porting #38347's batch-width fix, back to reserving exactly one unit for the whole dispatch and releasing everything pending on whichever single event fires first -- the same accepted tradeoff already made for the racing abatch_completion_fastest_response variant, whose cancelled losers never fire a terminal event at all. A second event firing (the mixed success-and-failure case) finds nothing left and is a safe no-op. Verified this doesn't apply to model_based_tag_rate_limits_hook's own per-Router-hop admission the same way: each hop's own routing/admission call is independent of the shared logging object (Router splits the batch and routes each branch separately before any of this is relevant), so its existing own-hop-key recomputation on release already never over-releases; the branches it can't get a release signal for are bounded by the same safety-TTL tradeoff, not an unbounded leak. Added a real end-to-end regression test for both hooks using the actual shared litellm_logging_obj construction, not an idealized one-event-per-branch assumption. --- .../hooks/global_tag_rate_limits_hook.py | 131 +++------- .../hooks/test_global_tag_rate_limits_hook.py | 234 +++++++----------- .../test_model_based_tag_rate_limits_hook.py | 91 +++++++ 3 files changed, 221 insertions(+), 235 deletions(-) diff --git a/litellm/proxy/hooks/global_tag_rate_limits_hook.py b/litellm/proxy/hooks/global_tag_rate_limits_hook.py index 3c912c6e2a7..347077d5f45 100644 --- a/litellm/proxy/hooks/global_tag_rate_limits_hook.py +++ b/litellm/proxy/hooks/global_tag_rate_limits_hook.py @@ -69,18 +69,27 @@ reservations early. `Router.abatch_completion`'s comma-separated `model` fans this one admission out into several real, independent LLM calls below this hook entirely (it never re-runs `async_pre_call_hook` per branch the way -`model_based_tag_rate_limits_hook`'s per-Router-hop admission does), each of -which reliably fires its own terminal success/failure event -- confirmed -live against the real dispatch. A single concurrency reservation for the -whole batch would be released by whichever branch finishes first, letting -the still-running siblings push real concurrent calls past the configured -cap, so admission instead reserves one unit per comma-separated model (see -`_non_racing_batch_width`) and each branch's own event releases just its own -share (see `_release_one_pending_for_call_id`). The racing -`abatch_completion_fastest_response` variant is excluded from that: it -cancels every losing branch without ever firing a terminal event for it -(confirmed live), so reserving more than the single unit it already does -would leak every losing branch's share until the safety TTL. +`model_based_tag_rate_limits_hook`'s per-Router-hop admission does). Every +branch shares the one `litellm_logging_obj` this same request already +carries (attached in `common_request_processing.py` before Router ever +sees the request, then threaded unchanged through every branch's own +`**kwargs`), so `Logging.should_run_logging`'s per-event-type dedup fires +this hook's own success/failure callback at most once for a success and at +most once for a failure across the *whole* dispatch, never once per +branch -- confirmed live against the real dispatch, including with a real, +shared `litellm_logging_obj` attached the way production actually does it. +Reserving one unit per model and expecting each to release on its own +branch's event (an earlier version of this hook did exactly that) is +therefore not just occasionally racy: it guarantees every batch of width +more than 2 leaks the surplus for the safety TTL, every single call. This +hook instead reserves exactly one unit for the whole non-racing dispatch, +same as the racing `abatch_completion_fastest_response` variant below, +and releases it on whichever single event fires first -- a second event +firing (the mixed success-and-failure case) finds nothing left and is a +safe no-op. That undercounts the dispatch's real, brief concurrent +provider load, the same accepted tradeoff already made for +`abatch_completion_fastest_response`, whose cancelled losers never fire a +terminal event at all. """ import asyncio @@ -169,34 +178,6 @@ def _individual_model_names(model: str | None, call_type: str) -> tuple[str, ... return tuple(m.strip() for m in model.split(",")) -def _non_racing_batch_width(data: Mapping[str, object], call_type: str) -> int: - """`Router.abatch_completion` (not `abatch_completion_fastest_response`) - fans this one admission out into one independent real LLM call per - comma-separated model in `model`, and every one of those branches - reliably fires its own terminal success/failure event -- confirmed - empirically, no cancellation involved, unlike the racing - `fastest_response` variant, which cancels every losing branch without - ever firing a terminal event for it and must keep reserving a single - unit released by whichever branch finishes first. - - A single concurrency reservation for the whole non-racing dispatch would - get released by whichever branch finishes first, letting the - still-running siblings push real concurrent calls past the configured - cap. Reserving one unit per branch instead, released one at a time as - each branch's own event fires, keeps the count accurate. - - Mirrors the exact condition `route_llm_request.py` uses to route into - `abatch_completion` in the first place, so this only ever fires for a - request that will actually take that path. - """ - if call_type != "acompletion" or data.get("fastest_response"): - return 1 - model_field: Final = data.get("model") - if not isinstance(model_field, str) or "," not in model_field: - return 1 - return len(model_field.split(",")) - - def _hash_tag(entry: TagRateLimitEntry, unit: _LimitUnit, tag_value: str, key_hash: str | None) -> str: """ Global-hook equivalent of `model_based_tag_rate_limits_hook._hash_tag`, @@ -268,11 +249,6 @@ class _GlobalTagRateLimitStash: # an earlier attempt must still get its accounting at success time even # though the request ultimately serves from a later attempt's model. admitted_models: frozenset[str] = field(default_factory=frozenset) - # One entry per reserved unit, not one entry per distinct key: a - # non-racing batch dispatch (see _non_racing_batch_width) reserves and - # appends batch_width entries for the same key, and each branch's own - # terminal event pops exactly one of them -- see - # _release_one_pending_for_call_id. pending_concurrency_keys: list[tuple[str, _PartitionKey]] = field(default_factory=list) # mutable-ok: queue # "requests" keys already charged for this call_id -- veria-ai finding: # ProxyBaseLLMRequestProcessing._pre_call_with_fallbacks reruns the whole @@ -690,7 +666,6 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o already_reserved_concurrency_keys: Final = frozenset( key for key, _partition_key in stash.pending_concurrency_keys ) - non_racing_batch_width: Final = _non_racing_batch_width(data, call_type) failing_index, values = await self._atomic_check_and_increment( tuple( ( @@ -706,17 +681,14 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o # rejecting) refunds that zero-cost renewal as a # genuine no-op, same reasoning as # model_based_tag_rate_limits_hook's identical fix - # for its own per-hop retries. Otherwise a - # concurrency check reserves one unit per non-racing - # batch branch (see _non_racing_batch_width), not - # just one for the whole dispatch. + # for its own per-hop retries. 0.0 if renewal_allowed and ( (check.unit == "requests" and check.key in stash.charged_request_keys) or (check.unit == "concurrency" and check.key in already_reserved_concurrency_keys) ) - else (float(non_racing_batch_width) if check.unit == "concurrency" else 1.0), + else 1.0, self._ttl_for(check.unit, check.entry), check.unit == "concurrency", ) @@ -734,14 +706,10 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o # renewed at zero net cost above, so re-adding it here would # make release (which decrements once per queued entry) decrement # twice for a counter that was only ever incremented once. - # One entry per reserved unit (non_racing_batch_width of them, - # ordinarily 1) -- see pending_concurrency_keys's own docstring - # for why a non-racing batch dispatch needs one entry per branch. concurrency_reservations: Final = tuple( (check.key, _partition_key(check.entry)) for check in atomic_checks if check.unit == "concurrency" and check.key not in already_reserved_concurrency_keys - for _ in range(non_racing_batch_width) ) if concurrency_reservations: stash.pending_concurrency_keys.extend(concurrency_reservations) # mutable-ok: see field's own docstring @@ -769,14 +737,15 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o async def _release_pending_for_call_id(self, request_kwargs: Mapping[str, object]) -> None: """Releases every reservation still pending for this call_id at once. - Correct for a disconnect or a chain-exhausted failure: either aborts - every not-yet-completed branch of a non-racing batch together (see - _non_racing_batch_width), so none of them will ever fire its own - terminal event to release its own share individually -- whatever's - still pending here is exactly what those abandoned branches - reserved, no more (any branch that already completed already - popped its own entry via _release_one_pending_for_call_id) and no - less.""" + Every reservation is fungible under one call_id, and at most one + success and one failure callback ever fires for the whole call_id + (see the module docstring's explanation of `Logging.should_run_logging`'s + per-event-type dedup across a batch dispatch's shared logging + object) -- releasing everything still pending here is therefore + always correct, never a still-genuinely-running branch's share, + whether this fires from a disconnect, a chain-exhausted failure, or + a terminal callback that already ran once for this call_id. + """ stash: Final = _stash_for_call(_call_id_from_kwargs(request_kwargs)) if stash is None or not stash.pending_concurrency_keys: return @@ -784,21 +753,6 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o stash.pending_concurrency_keys.clear() await self._release_keys(release_keys) - async def _release_one_pending_for_call_id( - self, request_kwargs: Mapping[str, object] - ) -> tuple[str, _PartitionKey] | None: - """Releases exactly one reservation, not every reservation - currently pending: a non-racing batch dispatch reserves one unit - per comma-separated model (see _non_racing_batch_width), and each - branch's own terminal event must only release its own share, not a - still-running sibling's. Entries under one call_id are otherwise - fungible (same key repeated), so which single entry gets popped - doesn't matter.""" - stash: Final = _stash_for_call(_call_id_from_kwargs(request_kwargs)) - if stash is None or not stash.pending_concurrency_keys: - return None - return stash.pending_concurrency_keys.pop() # mutable-ok: see field's own docstring - async def async_release_disconnect_state_hook(self, request_data: Mapping[str, object]) -> None: await self._release_pending_for_call_id(request_data) @@ -826,26 +780,21 @@ class _PROXY_GlobalTagRateLimitsHook( # pyright: ignore[reportUnusedClass] # o self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object ) -> None: # Always release regardless of which hook raised: this hook's own - # rejection never reserves a slot, so there's nothing pending to - # release in that case; a rejection from model_based_tag_rate_limits_hook - # (same error marker) can still land after this hook already - # reserved its own slot. Only this one branch's own share, not - # every reservation still pending for a non-racing batch's other, - # still-running branches -- see _release_one_pending_for_call_id. - released_entry: Final = await self._release_one_pending_for_call_id(kwargs) - if released_entry is not None: - await self._release_keys((released_entry,)) + # rejection never reserves a slot, so pending_concurrency_keys is + # already empty in that case and the check below no-ops; a rejection + # from model_based_tag_rate_limits_hook (same error marker) can still + # land after this hook already reserved its own slot. + await self._release_pending_for_call_id(kwargs) async def async_log_success_event( self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object ) -> None: - released_entry: Final = await self._release_one_pending_for_call_id(kwargs) - if released_entry is not None: - release_task: Final = asyncio.create_task(self._release_keys((released_entry,))) + stash: Final = _stash_for_call(_call_id_from_kwargs(kwargs)) + if stash is not None and stash.pending_concurrency_keys: + release_task: Final = asyncio.create_task(self._release_pending_for_call_id(kwargs)) _BACKGROUND_TASKS.add(release_task) # mutable-ok: see _BACKGROUND_TASKS's own docstring release_task.add_done_callback(_BACKGROUND_TASKS.discard) - stash: Final = _stash_for_call(_call_id_from_kwargs(kwargs)) config: Final = self._refresh_config() if config is None: return diff --git a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py index 3a0e6863e29..83bd800fc35 100644 --- a/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py +++ b/tests/test_litellm/proxy/hooks/test_global_tag_rate_limits_hook.py @@ -1303,12 +1303,18 @@ async def test_concurrent_requests_do_not_share_each_others_reservation_state(ti # --------------------------------------------------------------------------- # Non-racing batch dispatch (Router.abatch_completion via a comma-separated -# `model`): one admission fans out into several real, independent LLM calls, -# each firing its own terminal event -- confirmed live against the real -# Router.abatch_completion (not just this hook in isolation) that every -# branch always emits exactly one success or failure event, unlike the -# racing abatch_completion_fastest_response variant, whose cancelled losers -# never emit any terminal event at all. +# `model`): every branch shares the one litellm_logging_obj this request +# already carries (attached before Router ever sees it, then threaded +# unchanged through every branch's own kwargs), so Logging.should_run_logging's +# per-event-type dedup means at most one success and one failure callback +# ever fire for the whole dispatch -- confirmed live against the real +# Router.abatch_completion with a real, shared litellm_logging_obj attached +# the way production actually does it, not an idealized one-event-per-branch +# assumption. Reserving more than one unit per dispatch would leak the rest +# until the safety TTL every single time a batch has more than 1-2 branches, +# not just occasionally -- so this hook reserves exactly one unit for the +# whole dispatch, same as the racing abatch_completion_fastest_response +# variant below. # --------------------------------------------------------------------------- @@ -1320,16 +1326,17 @@ def _batch_data(tags: list[str], model: str, call_id: str = "call-1", fastest_re @pytest.mark.asyncio -async def test_non_racing_batch_reserves_one_unit_per_model(time_controller, monkeypatch): - """A single admission for a 3-model comma-separated batch must reserve 3 - units, not 1 -- otherwise the cap is measured against the wrong number of - real concurrent LLM calls the batch actually makes.""" +async def test_non_racing_batch_reserves_only_one_unit_for_the_whole_dispatch(time_controller, monkeypatch): + """A single admission for a comma-separated batch must reserve exactly + one unit, not one per model: only one terminal callback (at most) will + ever fire for this call_id regardless of how many models are in the + batch, so reserving more would never get released until the safety TTL.""" monkeypatch.setattr( litellm, "global_tag_rate_limits", { "concurrency_limits": { - "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 3, "period_seconds": 60}] + "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] } }, ) @@ -1341,85 +1348,69 @@ async def test_non_racing_batch_reserves_one_unit_per_model(time_controller, mon data=_batch_data(["end_user_id:u1"], model="model-a,model-b,model-c", call_id="call-1"), call_type="acompletion", ) - # The batch alone already occupies every unit of the cap of 3. - with pytest.raises(ProxyRateLimitError): - await hook.async_pre_call_hook( - user_api_key_dict=_key(), - cache=DualCache(), - data=_data(["end_user_id:u1"], call_id="call-2"), - call_type="completion", - ) - - -@pytest.mark.asyncio -async def test_non_racing_batch_releases_one_unit_per_branch(time_controller, monkeypatch): - """Each branch's own success/failure event must release only its own - share of a non-racing batch's reservation -- releasing the whole - reservation on the first branch to finish would let a new, unrelated - request in while the batch's own remaining branches are still genuinely - in flight, exceeding the configured cap.""" - monkeypatch.setattr( - litellm, - "global_tag_rate_limits", - { - "concurrency_limits": { - "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 3, "period_seconds": 60}] - } - }, - ) - hook = _make_hook(time_controller) - - await hook.async_pre_call_hook( - user_api_key_dict=_key(), - cache=DualCache(), - data=_batch_data(["end_user_id:u1"], model="model-a,model-b,model-c", call_id="call-1"), - call_type="acompletion", - ) - batch_kwargs = {"litellm_call_id": "call-1", "metadata": {"tags": ["end_user_id:u1"]}} await hook.async_log_success_event(kwargs=batch_kwargs, response_obj=None, start_time=0, end_time=0) await asyncio.sleep(0) - # Only 1 of 3 branches finished -- the other 2 still hold their own - # units, so a brand-new request needing all 3 remaining units of the cap - # must still be rejected. - with pytest.raises(ProxyRateLimitError): - await hook.async_pre_call_hook( - user_api_key_dict=_key(), - cache=DualCache(), - data=_batch_data(["end_user_id:u1"], model="model-x,model-y,model-z", call_id="call-2"), - call_type="acompletion", - ) - - # The other 2 branches finish (one success, one failure). - await hook.async_log_success_event(kwargs=batch_kwargs, response_obj=None, start_time=0, end_time=0) - await hook.async_log_failure_event(kwargs=batch_kwargs, response_obj=None, start_time=0, end_time=0) - await asyncio.sleep(0) - - # Now every unit from the first batch is back -- no leak, no - # over-release along the way. + # The single reserved unit is fully released by that one event -- a + # fresh request needing that same single unit must be admitted, not + # rejected by 2 phantom leftover units. result = await hook.async_pre_call_hook( user_api_key_dict=_key(), cache=DualCache(), - data=_data(["end_user_id:u1"], call_id="call-3"), + data=_data(["end_user_id:u1"], call_id="call-2"), call_type="completion", ) assert result is not None @pytest.mark.asyncio -async def test_non_racing_batch_disconnect_releases_only_what_remains(time_controller, monkeypatch): - """A client disconnect mid-batch cancels every branch that hasn't - already finished together, so none of those cancelled branches ever - fires its own terminal event. The disconnect hook must release exactly - what's left -- not double-release a branch that already released its - own share via a real success/failure event first.""" +async def test_non_racing_batch_second_terminal_event_is_a_safe_no_op(time_controller, monkeypatch): + """should_run_logging dedups per event type, not once overall, so a + batch with a mix of successful and failed branches fires both a success + and a failure callback for the same call_id. Whichever fires second must + find nothing left pending and no-op instead of double-releasing.""" monkeypatch.setattr( litellm, "global_tag_rate_limits", { "concurrency_limits": { - "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 3, "period_seconds": 60}] + "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] + } + }, + ) + hook = _make_hook(time_controller) + + await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_batch_data(["end_user_id:u1"], model="model-a,model-b", call_id="call-1"), + call_type="acompletion", + ) + batch_kwargs = {"litellm_call_id": "call-1", "metadata": {"tags": ["end_user_id:u1"]}} + await hook.async_log_success_event(kwargs=batch_kwargs, response_obj=None, start_time=0, end_time=0) + await asyncio.sleep(0) + await hook.async_log_failure_event(kwargs=batch_kwargs, response_obj=None, start_time=0, end_time=0) + + result = await hook.async_pre_call_hook( + user_api_key_dict=_key(), + cache=DualCache(), + data=_data(["end_user_id:u1"], call_id="call-2"), + call_type="completion", + ) + assert result is not None + + +@pytest.mark.asyncio +async def test_non_racing_batch_disconnect_releases_the_single_reservation(time_controller, monkeypatch): + """A client disconnect mid-batch must release the one reservation the + whole dispatch holds, the same as any other call_id.""" + monkeypatch.setattr( + litellm, + "global_tag_rate_limits", + { + "concurrency_limits": { + "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] } }, ) @@ -1431,10 +1422,6 @@ async def test_non_racing_batch_disconnect_releases_only_what_remains(time_contr data=_batch_data(["end_user_id:u1"], model="model-a,model-b,model-c", call_id="call-1"), call_type="acompletion", ) - batch_kwargs = {"litellm_call_id": "call-1", "metadata": {"tags": ["end_user_id:u1"]}} - await hook.async_log_success_event(kwargs=batch_kwargs, response_obj=None, start_time=0, end_time=0) - await asyncio.sleep(0) - await hook.async_release_disconnect_state_hook({"litellm_call_id": "call-1"}) result = await hook.async_pre_call_hook( @@ -1470,8 +1457,6 @@ async def test_fastest_response_batch_still_reserves_only_one_unit(time_controll data=_batch_data(["end_user_id:u1"], model="model-a,model-b", call_id="call-1", fastest_response=True), call_type="acompletion", ) - # If this admission had reserved 2 units instead of 1, this second, - # unrelated request (cap of 1) would already be rejected here too. kwargs = {"litellm_call_id": "call-1", "metadata": {"tags": ["end_user_id:u1"]}} await hook.async_log_success_event(kwargs=kwargs, response_obj=None, start_time=0, end_time=0) await asyncio.sleep(0) @@ -1486,18 +1471,21 @@ async def test_fastest_response_batch_still_reserves_only_one_unit(time_controll @pytest.mark.asyncio -async def test_real_abatch_completion_admits_once_and_releases_per_branch(time_controller, monkeypatch): - """End-to-end against the real Router.abatch_completion (mocked LLM - responses, no network): confirms admission fires exactly once for the - whole dispatch -- not once per branch the way model_based_tag_rate_limits_hook's - per-Router-hop admission does -- and that a cap sized to the batch width - is never exceeded even while some branches are still genuinely running.""" +async def test_real_abatch_completion_with_a_shared_logging_obj_does_not_leak(time_controller, monkeypatch): + """End-to-end against the real Router.abatch_completion with a real, + shared litellm_logging_obj attached the same way common_request_processing.py + does it before Router ever sees the request -- not an idealized + one-event-per-branch assumption that doesn't hold in production. Every + branch here genuinely succeeds, but Logging.should_run_logging still + suppresses every callback after the first, so this proves the single + reservation this hook takes for the whole dispatch is still fully + released, not just 1 of however many it might otherwise have reserved.""" monkeypatch.setattr( litellm, "global_tag_rate_limits", { "concurrency_limits": { - "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 3, "period_seconds": 60}] + "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] } }, ) @@ -1512,11 +1500,11 @@ async def test_real_abatch_completion_admits_once_and_releases_per_branch(time_c }, { "model_name": "model-slow-1", - "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "slow-1-done", "mock_delay": 0.3}, + "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "slow-1-done", "mock_delay": 0.05}, }, { "model_name": "model-slow-2", - "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "slow-2-done", "mock_delay": 0.3}, + "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "slow-2-done", "mock_delay": 0.05}, }, ] ) @@ -1524,80 +1512,38 @@ async def test_real_abatch_completion_admits_once_and_releases_per_branch(time_c admitted_data = await hook.async_pre_call_hook( user_api_key_dict=_key(), cache=DualCache(), - data=_batch_data(["end_user_id:u1"], model="model-fast,model-slow-1,model-slow-2", call_id="call-1"), + data={ + **_batch_data(["end_user_id:u1"], model="model-fast,model-slow-1,model-slow-2", call_id="call-1"), + "messages": [{"role": "user", "content": "hello"}], + }, call_type="acompletion", ) - batch_task = asyncio.create_task( - router.abatch_completion( - models=["model-fast", "model-slow-1", "model-slow-2"], - messages=[{"role": "user", "content": "hello"}], - metadata=admitted_data["metadata"], - litellm_call_id=admitted_data["litellm_call_id"], - ) + logging_obj, admitted_data = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=time_controller.now(), + **admitted_data, ) + admitted_data["litellm_logging_obj"] = logging_obj - # Give only the fast branch (0.01s mock_delay) time to finish; the two - # slow branches (0.3s) are still genuinely in flight. + models = [m.strip() for m in admitted_data.pop("model").split(",")] + responses = await router.abatch_completion(models=models, **admitted_data) await asyncio.sleep(0.1) - # 1 of 3 units released, 2 still held by the still-running slow - # branches -- a brand-new request needing all 3 units of the cap must - # still be rejected. - with pytest.raises(ProxyRateLimitError): - await hook.async_pre_call_hook( - user_api_key_dict=_key(), - cache=DualCache(), - data=_batch_data(["end_user_id:u1"], model="model-x,model-y,model-z", call_id="call-2"), - call_type="acompletion", - ) - - responses = await batch_task assert [r.choices[0].message.content for r in responses] == ["fast-done", "slow-1-done", "slow-2-done"] - await asyncio.sleep(0) - # Every branch has now finished and released its own share -- no leak. + # No leak: the single reserved unit is back, so a fresh request needing + # it is admitted, not wrongly rejected by leftover phantom reservations. result = await hook.async_pre_call_hook( user_api_key_dict=_key(), cache=DualCache(), - data=_data(["end_user_id:u1"], call_id="call-3"), + data=_data(["end_user_id:u1"], call_id="call-2"), call_type="completion", ) assert result is not None -@pytest.mark.asyncio -async def test_comma_separated_model_outside_acompletion_call_type_is_not_a_batch(time_controller, monkeypatch): - """A comma in `model` only means a non-racing batch dispatch for the - exact call_type route_llm_request.py itself gates on -- any other - call_type must reserve just 1 unit, matching whatever that call_type - actually does under the hood.""" - monkeypatch.setattr( - litellm, - "global_tag_rate_limits", - { - "concurrency_limits": { - "limits": [{"name": "conc", "tag_id": "end_user_id", "limit": 1, "period_seconds": 60}] - } - }, - ) - hook = _make_hook(time_controller) - - await hook.async_pre_call_hook( - user_api_key_dict=_key(), - cache=DualCache(), - data=_batch_data(["end_user_id:u1"], model="model-a,model-b,model-c", call_id="call-1"), - call_type="completion", - ) - with pytest.raises(ProxyRateLimitError): - await hook.async_pre_call_hook( - user_api_key_dict=_key(), - cache=DualCache(), - data=_data(["end_user_id:u1"], call_id="call-2"), - call_type="completion", - ) - - # --------------------------------------------------------------------------- # Accounting: tokens/dollars via async_log_success_event # --------------------------------------------------------------------------- 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 7efc4b5232a..c4d1259a50e 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 @@ -4908,6 +4908,97 @@ async def test_fast_failing_batch_sibling_does_not_release_a_still_executing_sib ) +@pytest.mark.asyncio +async def test_real_abatch_completion_with_a_shared_logging_obj_never_over_releases(time_controller, monkeypatch): + """ + End-to-end against the real Router.abatch_completion with a real, shared + litellm_logging_obj attached the way common_request_processing.py does + it before Router ever sees the request -- not the two tests above, which + manually fire one callback and assert on that in isolation. Every branch + here reserves its own deployment-scoped slot at its own admission hop + (unaffected by logging-object sharing, since routing happens per branch + regardless), but Logging.should_run_logging's per-event-type dedup means + at most one success callback fires for the whole dispatch, no matter how + many branches actually completed. This hook's own admission token/hop + recomputation must never release more than that one branch's own slot -- + the other branches' slots are expected to sit until the safety TTL, the + same accepted tradeoff already made for abatch_completion_fastest_response's + cancelled losers, not an unbounded or over-eager release. + """ + limiter = _make_limiter(time_controller) + router = litellm.Router( + model_list=[ + _deployment( + "model-fast", + "dep-fast", + { + "concurrency_limits": { + "limits": [{"name": "c", "tag_id": "end_user_id", "limit": 1, "period_seconds": 300}] + } + }, + ), + _deployment( + "model-slow-1", + "dep-slow-1", + { + "concurrency_limits": { + "limits": [{"name": "c", "tag_id": "end_user_id", "limit": 1, "period_seconds": 300}] + } + }, + ), + _deployment( + "model-slow-2", + "dep-slow-2", + { + "concurrency_limits": { + "limits": [{"name": "c", "tag_id": "end_user_id", "limit": 1, "period_seconds": 300}] + } + }, + ), + ] + ) + limiter.update_variables(llm_router=router) + monkeypatch.setattr(litellm, "callbacks", [limiter]) + + data = { + "model": "model-fast,model-slow-1,model-slow-2", + "messages": [{"role": "user", "content": "hello"}], + "litellm_call_id": "call-1", + "metadata": {"tags": ["end_user_id:u1"]}, + } + logging_obj, data = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + **data, + ) + data["litellm_logging_obj"] = logging_obj + + models = [m.strip() for m in data.pop("model").split(",")] + responses = await router.abatch_completion(models=models, **data) + await asyncio.sleep(0.1) + + assert [r.choices[0].message.content for r in responses] == ["ok", "ok", "ok"] + + # At most one of the 3 branches' own slot was released via its own + # callback; the rest must still be held, not silently freed by a + # sibling's event -- a fresh request for each of those models must + # still be rejected against the cap of 1. + released_count = 0 + for model_name in ("model-fast", "model-slow-1", "model-slow-2"): + try: + await limiter.async_filter_deployments( + model=model_name, + healthy_deployments=router.model_list, + messages=None, + request_kwargs={"metadata": {"tags": ["end_user_id:u1"]}}, + ) + released_count += 1 + except ProxyRateLimitError: + pass + assert released_count <= 1 + + @pytest.mark.asyncio async def test_fast_succeeding_batch_sibling_does_not_release_a_still_executing_siblings_slot(time_controller): """