From de3f0690e7e85d4c1d6bbcbd784ce2b93ab3ea9c Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Thu, 12 Nov 2026 15:30:08 -0800 Subject: [PATCH] fix(dynamic_rate_limiter_v3): a borrower gets its full share back once the model saturates (#51415) * fix(dynamic_rate_limiter_v3): a borrower gets its full share back once the model saturates The per-team share counters only count requests admitted while the model is saturated, so what a team used while the model was idle is not charged against its share once contention starts. A team below its share is still admitted ahead of a borrower, and the model's own limits bound the total. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(dynamic_rate_limiter_v3): idle fallback targets and WebSocket turns follow the same share accounting A router fallback target attempted while idle no longer holds the team's share, and a Responses WebSocket connection refreshes its share descriptor on every turn, so idle turns are forgiven and saturated turns count. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(dynamic_rate_limiter_v3): an idle fallback target still records its hold Settlement finds the hold, gives the admitting model's reservation back and charges nothing for the idle target, instead of charging the target's tokens to the admitting model's share. The integration fixture builds its configuration without mutating the loaded one. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/hooks/dynamic_rate_limiter_v3.py | 51 ++++------ .../hooks/parallel_request_limiter_v3.py | 8 +- ...mic_rate_limiter_borrower_regains_share.py | 99 +++++++++++++++++++ .../hooks/test_dynamic_rate_limiter_v3.py | 97 ++++++++++++++---- 4 files changed, 203 insertions(+), 52 deletions(-) create mode 100644 tests/integration/routing/test_dynamic_rate_limiter_borrower_regains_share.py diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 2873121cf56..a58c6feb99a 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -1078,6 +1078,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): saturation: Final = await self._check_model_saturation(requested_model, model_group_info) saturated: Final = saturation >= _get_priority_settings().saturation_threshold enforced_team: Final = None if borrowing is None else borrowing.enforced(saturated) + if connection_stash is not None and borrowing is not None: + connection_stash.borrowing_team_model_limit = enforced_team descriptors: Final = (share_descriptors if saturated else ()) + ( () if enforced_team is None else (enforced_team,) ) @@ -1221,10 +1223,10 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): # All-or-nothing: if any enforced descriptor would exceed its limit, # no counter is modified and the response carries "OVER_LIMIT". team_share: Final = self._borrowing_team_share(model=model, user_api_key_dict=user_api_key_dict) - get_or_create_request_stash().borrowing_team_model_limit = None if team_share is None else team_share.share if team_share is not None: descriptors_to_check.append(team_share.share) enforced_team_share: Final = None if team_share is None else team_share.enforced(should_enforce_priority) + get_or_create_request_stash().borrowing_team_model_limit = enforced_team_share enforced_descriptors: Final[list[RateLimitDescriptor]] = [model_wide_descriptor] if priority_descriptors and should_enforce_priority: @@ -1355,26 +1357,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): _reserved_window(atomic_response, enforced_team_share) if team_reservation else None ) self._cap_admitted_output(enforced_team_share, team_reservation, data, call_type) - if team_share is not None and enforced_team_share is None: - await self.v3_limiter.atomic_check_and_increment_by_n( - descriptors=[team_share.share], - increments=[per_request_increment], - parent_otel_span=user_api_key_dict.parent_otel_span, - ) if priority_descriptors and share is not None and not should_enforce_priority: - tracking_only: Final = tuple( - RateLimitDescriptor( - key=descriptor["key"], - value=descriptor["value"], - rate_limit=model_wide_descriptor["rate_limit"], - ) - for descriptor in priority_descriptors - ) - await self.v3_limiter.atomic_check_and_increment_by_n( - descriptors=list(tracking_only), - increments=[per_request_increment for _ in tracking_only], - parent_otel_span=user_api_key_dict.parent_otel_span, - ) get_or_create_request_stash().rate_limit_response = atomic_response elif priority_descriptors and not should_enforce_priority: priority_tracking_response: Final = await self.v3_limiter.should_rate_limit( @@ -1439,8 +1422,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): ) -> bool: """Hold a router fallback target to the team's share on it once it saturates, and to its borrow ceiling before. - With ``commit`` the attempt is counted toward the target's saturation and the team's counter, holding the - prompt's estimated tokens when the share is enforced, otherwise only checked. + With ``commit`` the attempt is counted toward the target's saturation, and toward the team's counter only + while the share is enforced, holding the prompt's estimated tokens there. An idle target is only checked. """ team_share: Final = self._borrowing_team_share(model=model, user_api_key_dict=user_api_key_dict) model_group_info: Final = self.llm_router.registry().group(model) @@ -1474,13 +1457,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): ) if response["overall_code"] == "OVER_LIMIT": return False - if enforced is None: - await self.v3_limiter.atomic_check_and_increment_by_n( - descriptors=[team_share.share], - increments=[{"requests": 1, "tokens": reservation}], - parent_otel_span=user_api_key_dict.parent_otel_span, - ) - self._hold_fallback_share(model, team_share.share, reservation, _reserved_window(response, enforced)) + self._hold_fallback_share(model, team_share.share, enforced, reservation, _reserved_window(response, enforced)) return True async def _share_has_room( @@ -1496,13 +1473,23 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): ) def _hold_fallback_share( - self, model: str, share: RateLimitDescriptor, reserved_tokens: int, window: ReservationWindowIdentity | None + self, + model: str, + share: RateLimitDescriptor, + charged: RateLimitDescriptor | None, + reserved_tokens: int, + window: ReservationWindowIdentity | None, ) -> None: - """Remember the hold under the requested name and the resolved group, so settlement finds it either way.""" + """Remember the hold under the requested name and the resolved group, so settlement finds it either way. + + ``charged`` is the share the target's usage settles to, or None for an idle target whose usage is forgiven. + """ stash: Final = get_or_create_request_stash() earlier: Final = stash.borrowing_team_model_fallback_limits.get(model) carried: Final = earlier.reserved_tokens if earlier is not None and earlier.window == window else 0 - hold: Final = BorrowingTeamModelHold(descriptor=share, reserved_tokens=reserved_tokens + carried, window=window) + hold: Final = BorrowingTeamModelHold( + descriptor=charged, reserved_tokens=reserved_tokens + carried, window=window + ) stash.borrowing_team_model_fallback_limits = MappingProxyType( {**stash.borrowing_team_model_fallback_limits, model: hold, share["value"].rpartition(":")[2]: hold} ) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 7d5daff71cd..e520a8d7d92 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -572,9 +572,13 @@ class CallTypeRateLimiter(Protocol): @dataclass(frozen=True, slots=True) class BorrowingTeamModelHold: - """A router fallback target's team share with the tokens held on it, settled to actual usage or given back.""" + """A router fallback target's team share with the tokens held on it, settled to actual usage or given back. - descriptor: RateLimitDescriptor + An idle target holds no descriptor: its usage is not charged to the share, and settlement still finds the hold so + the admitting model's reservation is given back instead of taking the target's tokens. + """ + + descriptor: RateLimitDescriptor | None reserved_tokens: int = 0 window: ReservationWindowIdentity | None = None diff --git a/tests/integration/routing/test_dynamic_rate_limiter_borrower_regains_share.py b/tests/integration/routing/test_dynamic_rate_limiter_borrower_regains_share.py new file mode 100644 index 00000000000..0e87ed0bf13 --- /dev/null +++ b/tests/integration/routing/test_dynamic_rate_limiter_borrower_regains_share.py @@ -0,0 +1,99 @@ +import json +import time +import uuid +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import httpx +import pytest +import yaml +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.process import owned_proxy +from integration._support.wire import Reply, Request, wire_server + +OPENAI_MODEL: Final = "gpt-4.1-mini" +MODEL_TPM: Final = 1000 +TEAM_TPM: Final = 100 +SATURATION_CACHE_TTL: Final = 1 +TOKENS_PER_CALL: Final = 30 +CALLS_TO_SATURATE: Final = MODEL_TPM // 2 // TOKENS_PER_CALL + 1 +UPSTREAM_REPLY: Final = json.dumps( + { + "id": "chatcmpl_borrower_regains_share", + "object": "chat.completion", + "created": 1700000000, + "model": OPENAI_MODEL, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "recovered"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": TOKENS_PER_CALL}, + } +).encode() + + +def _respond(request: Request) -> Reply: + assert request.method == "POST" and request.target.endswith("/v1/chat/completions"), request.target + assert request.headers["authorization"] == "Bearer synthetic-openai-key" + return Reply(body=UPSTREAM_REPLY) + + +@pytest.fixture(scope="module") +def saturated_model(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, str]]: + directory: Final = tmp_path_factory.mktemp("recovery-window") + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration: Final = { + **base, + "litellm_settings": { + **base["litellm_settings"], + "callbacks": ["dynamic_rate_limiter_v3"], + "priority_reservation": {"interactive": 0.6, "batch": 0.4}, + "priority_reservation_settings": { + "default_priority": 1.0, + "saturation_threshold": 0.5, + "team_model_limit_borrows_idle_capacity_from_other_teams": True, + "saturation_check_cache_ttl": SATURATION_CACHE_TTL, + }, + }, + } + path: Final = directory / "borrower_regains_share.yaml" + path.write_text(yaml.safe_dump(configuration)) + with ( + gateway_from_environment() as rig, + wire_server(_respond) as wire, + owned_proxy(rig, directory, {}, config=path) as owned, + ): + yield owned, wire.url + + +def _call(owned: Gateway, model: str, key: str) -> httpx.Response: + return owned.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"borrower regains share {uuid.uuid4().hex}"}]}, + key=key, + ) + + +def test_a_borrower_gets_its_share_back_once_the_model_saturates_and_a_newcomer_is_served_beside_it( + saturated_model: tuple[Gateway, str], +) -> None: + owned, upstream = saturated_model + with owned.scenario() as scenario: + model: Final = scenario.model( + model=f"openai/{OPENAI_MODEL}", api_base=f"{upstream}/v1", api_key="synthetic-openai-key", tpm=MODEL_TPM + ) + borrower: Final = scenario.key(team_id=scenario.team(metadata={"model_tpm_limit": {model: TEAM_TPM}})) + newcomer: Final = scenario.key(team_id=scenario.team(metadata={"model_tpm_limit": {model: TEAM_TPM}})) + + borrowed: Final = tuple(_call(owned, model, borrower) for _ in range(CALLS_TO_SATURATE)) + assert all(r.status_code == 200 for r in borrowed), [r.status_code for r in borrowed] + time.sleep(SATURATION_CACHE_TTL + 0.5) + + forgiven: Final = _call(owned, model, borrower) + assert forgiven.status_code == 200, forgiven.text + served: Final = _call(owned, model, newcomer) + assert served.status_code == 200, served.text + + held: Final = eventually(lambda: _call(owned, model, borrower), lambda r: r.status_code == 429, seconds=10) + assert held.headers.get("rate_limited_by") == "model_per_team_saturated", dict(held.headers) + still_served: Final = _call(owned, model, newcomer) + assert still_served.status_code == 200, still_served.text diff --git a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py index 1e4d89b011a..554a8bb3215 100644 --- a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -2163,7 +2163,7 @@ def _borrowing_handler( async def _admit_then_settle(handler: _PROXY_DynamicRateLimitHandlerV3, auth: UserAPIKeyAuth, tokens: int) -> None: - await _check(handler, auth, saturation=0.0) + await _check(handler, auth, saturation=0.5) await _settle_tokens(handler, tokens) @@ -2275,13 +2275,39 @@ async def test_team_past_its_model_tpm_limit_is_held_to_it_once_the_model_satura assert "Team model_tpm_limit: 100, Remaining: 0, Model saturation: 50.0%" in str(exc.value.detail) +@pytest.mark.asyncio +async def test_a_borrowers_usage_while_the_model_was_idle_is_forgiven_once_it_saturates( + monkeypatch: pytest.MonkeyPatch, +): + handler: Final = _borrowing_handler(monkeypatch, borrows=True) + await _check(handler, _team_key(), saturation=0.0) + await _settle_tokens(handler, 150) + assert await _team_tokens_remaining(handler) == 100 + await _admit_then_settle(handler, _team_key(), 30) + assert await _team_tokens_remaining(handler) == 70 + await _check(handler, _team_key("team-b"), saturation=0.5) + assert await _team_tokens_remaining(handler, _team_key("team-b")) == 100 + + +@pytest.mark.asyncio +async def test_a_borrower_that_spends_its_share_under_saturation_is_held_while_a_newcomer_is_admitted( + monkeypatch: pytest.MonkeyPatch, +): + handler: Final = _borrowing_handler(monkeypatch, borrows=True) + await _admit_then_settle(handler, _team_key(), 150) + with pytest.raises(ProxyRateLimitError) as exc: + await _check(handler, _team_key(), saturation=0.5) + _held_to_team_limit(exc.value, "tokens", "model_tpm_limit") + await _check(handler, _team_key("team-b"), saturation=0.5) + + @pytest.mark.asyncio async def test_team_past_its_model_rpm_limit_is_held_to_it_once_the_model_saturates(monkeypatch: pytest.MonkeyPatch): handler = _borrowing_handler(monkeypatch, borrows=True) auth: Final = _team_key(team_limits={"model_rpm_limit": {"shared-model": 2}}) - await _check(handler, auth, saturation=0.0) - await _check(handler, auth, saturation=0.0) await _check(handler, auth, saturation=0.49) + await _check(handler, auth, saturation=0.5) + await _check(handler, auth, saturation=0.5) with pytest.raises(ProxyRateLimitError) as exc: await _check(handler, auth, saturation=0.5) _held_to_team_limit(exc.value, "requests", "model_rpm_limit") @@ -2299,7 +2325,7 @@ async def test_requests_admitted_below_saturation_count_only_their_settled_token @pytest.mark.asyncio async def test_an_embedding_responses_tokens_count_against_the_team_limit(monkeypatch: pytest.MonkeyPatch): handler = _borrowing_handler(monkeypatch, borrows=True) - await _check(handler, _team_key(), saturation=0.0) + await _check(handler, _team_key(), saturation=0.5) await _settle_tokens( handler, 150, @@ -2315,7 +2341,7 @@ async def test_an_embedding_responses_tokens_count_against_the_team_limit(monkey @pytest.mark.asyncio async def test_a_responses_api_calls_tokens_count_against_the_team_limit(monkeypatch: pytest.MonkeyPatch): handler = _borrowing_handler(monkeypatch, borrows=True) - await _check(handler, _team_key(), saturation=0.0) + await _check(handler, _team_key(), saturation=0.5) await _settle_tokens( handler, 150, @@ -2339,7 +2365,7 @@ async def test_a_responses_api_calls_tokens_count_against_the_team_limit(monkeyp @pytest.mark.asyncio async def test_a_null_auth_metadata_still_settles_the_teams_tokens(monkeypatch: pytest.MonkeyPatch): handler = _borrowing_handler(monkeypatch, borrows=True) - await _check(handler, _team_key(), saturation=0.0) + await _check(handler, _team_key(), saturation=0.5) await handler.async_log_success_event( kwargs={ "litellm_params": {"metadata": {"model_group": "shared-model"}}, @@ -2403,7 +2429,7 @@ async def test_an_exempt_keys_tokens_are_not_charged_to_its_teammates(monkeypatc @pytest.mark.asyncio async def test_an_alias_request_counts_against_the_same_team_share_as_its_target(monkeypatch: pytest.MonkeyPatch): handler = _borrowing_handler(monkeypatch, borrows=True, model_group_alias={"alias-model": "shared-model"}) - await _check(handler, _team_key(), saturation=0.0, model="alias-model") + await _check(handler, _team_key(), saturation=0.5, model="alias-model") await _settle_tokens(handler, 150) with pytest.raises(ProxyRateLimitError) as exc: await _check(handler, _team_key(), saturation=0.5) @@ -2432,7 +2458,7 @@ async def test_a_nested_calls_tokens_are_not_charged_to_the_requests_team_share( handler = _borrowing_handler(monkeypatch, borrows=True) claim_request_stash_for_data({"litellm_call_id": "outer-call"}) - await _check(handler, _team_key(), saturation=0.0) + await _check(handler, _team_key(), saturation=0.5) await _settle_tokens(handler, 60, call_id="llm-judge-call") assert await _team_tokens_remaining(handler) == 100 await _settle_tokens(handler, 60, call_id="outer-call") @@ -2442,7 +2468,7 @@ async def test_a_nested_calls_tokens_are_not_charged_to_the_requests_team_share( @pytest.mark.asyncio async def test_a_websocket_turns_tokens_count_against_the_team_limit(monkeypatch: pytest.MonkeyPatch): handler = _borrowing_handler(monkeypatch, borrows=True) - await _check(handler, _team_key(), saturation=0.0) + await _check(handler, _team_key(), saturation=0.5) await _settle_tokens( handler, 150, @@ -2454,6 +2480,29 @@ async def test_a_websocket_turns_tokens_count_against_the_team_limit(monkeypatch _held_to_team_limit(exc.value, "tokens", "model_tpm_limit") +@pytest.mark.asyncio +async def test_a_websocket_admitted_while_idle_has_its_turns_counted_once_the_model_saturates( + monkeypatch: pytest.MonkeyPatch, +): + handler: Final = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.1) + await _check(handler, _team_key(), saturation=0.0) + await _settle_tokens( + handler, + 150, + response=[], + combined_usage_object=litellm.Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150), + ) + assert await _team_tokens_remaining(handler) == 100 + await handler.enforce_token_limits_for_frame(user_api_key_dict=_team_key(), requested_model="shared-model") + await _settle_tokens( + handler, + 90, + response=[], + combined_usage_object=litellm.Usage(prompt_tokens=60, completion_tokens=30, total_tokens=90), + ) + assert await _team_tokens_remaining(handler) == 10 + + @pytest.mark.asyncio async def test_a_websocket_turn_gets_a_503_when_the_team_counter_cannot_be_read(monkeypatch: pytest.MonkeyPatch): from litellm.exceptions import ServiceUnavailableError @@ -2513,15 +2562,28 @@ async def test_a_fallback_target_is_refused_once_it_saturates_and_the_teams_shar @pytest.mark.asyncio -async def test_a_fallback_target_still_borrows_while_idle_and_its_tokens_count_against_that_teams_share( +async def test_a_fallback_target_still_borrows_while_idle_and_its_tokens_are_forgiven_once_it_saturates( monkeypatch: pytest.MonkeyPatch, ): handler = _borrowing_handler(monkeypatch, borrows=True) assert await handler.async_fallback_target_allowed(_team_key(), "shared-model", {}) is True await _settle_tokens(handler, 150) - with pytest.raises(ProxyRateLimitError) as exc: - await _check(handler, _team_key(), saturation=0.5) - _held_to_team_limit(exc.value, "tokens", "model_tpm_limit") + await _check(handler, _team_key(), saturation=0.5) + assert await _team_tokens_remaining(handler) == 100 + + +@pytest.mark.asyncio +async def test_an_idle_fallback_target_is_not_charged_and_the_admitting_models_hold_is_given_back( + monkeypatch: pytest.MonkeyPatch, +): + handler: Final = _borrowing_handler(monkeypatch, borrows=True) + key: Final = _team_key(team_limits={"model_tpm_limit": {"shared-model": 100, "other-model": 100}}) + await _check(handler, key, saturation=0.5, data=_prompt(20)) + assert await _team_tokens_remaining(handler, key) < 100 + assert await handler.async_fallback_target_allowed(key, "other-model", _prompt(20)) is True + await _settle_tokens(handler, 150, model_group="other-model") + assert await _team_tokens_remaining(handler, key) == 100 + assert await _team_tokens_remaining(handler, key, model="other-model") == 100 @pytest.mark.asyncio @@ -2806,8 +2868,9 @@ async def _burn( priority: str | None, tokens: int, model: str = "shared-model", + saturation: float = 1.0, ): - await _check(handler, auth, saturation=0.0, priority=priority, model=model) + await _check(handler, auth, saturation=saturation, priority=priority, model=model) await handler.async_log_success_event( kwargs={ "litellm_params": {"metadata": {"model_group": model}}, @@ -3248,7 +3311,7 @@ async def test_a_team_id_spelled_like_a_pool_marker_still_counts_as_a_team(monke @pytest.mark.asyncio -async def test_requests_sent_while_the_share_was_not_enforced_all_count_once_the_share_grows( +async def test_requests_sent_while_the_share_was_not_enforced_are_forgiven_once_the_model_saturates( monkeypatch: pytest.MonkeyPatch, ): handler = _borrowing_handler(monkeypatch, borrows=False, splits=True) @@ -3256,9 +3319,7 @@ async def test_requests_sent_while_the_share_was_not_enforced_all_count_once_the for _ in range(70): await _check(handler, _class_team("team-a"), saturation=0.0, priority="interactive") await _check(handler, _class_team("team-b", "batch"), saturation=0.0, priority="batch") - with pytest.raises(ProxyRateLimitError) as exc: - await _check(handler, _class_team("team-a"), saturation=0.6, priority="interactive") - assert "Active teams in class: 1, Active classes: 2, " in str(exc.value.detail) + await _check(handler, _class_team("team-a"), saturation=0.6, priority="interactive") @pytest.mark.asyncio