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): """