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