From 0324b9d915bdeee8fb4910b3b90259ba8636c1e5 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Thu, 12 Nov 2026 17:08:02 -0800 Subject: [PATCH] fix(dynamic_rate_limiter_v3): a capped borrower keeps its cap on its own counter and gets a fresh share at saturation (#51494) * fix(dynamic_rate_limiter_v3): give a capped borrower its own cap counter so it also gets a fresh share at saturation A team with team_model_limit_max_borrow_multiple enforced its cap on the same counter as its share, so its idle-time usage kept counting and it got no fresh share once the model saturated. The cap now runs on its own counter (model_per_team_max_borrow_multiple) that counts all of the team's usage this window, while the share counter only counts usage once the model is saturated. Also repairs the six share-metric unit tests that main broke when #51415 and #51417 landed together, by saturating the model from a separate priority class before the measured calls. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(dynamic_rate_limiter_v3): size a capped team's token reservation by its share, idle or saturated The reservation for a prompt without max_tokens was bounded by whichever descriptor was being held, the cap while idle and the share once saturated, so the same request was refused at one moment and admitted the next with the same counters. Reserve against the share, the tighter of the two, in both states. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(dynamic_rate_limiter_v3): track the cap's reservation apart from the share's The cap and the share shared one reserved-token figure and the fallback carry keyed on the share's window, so an idle admission's hold on the cap was refunded from a share that never carried it once a WebSocket turn saturated, stacked idle fallback holds could carry across cap windows, and an idle fallback served for a capped team published no share gauges. Each counter now carries its own reservation and window, and a fallback hold keeps the target's share so the gauges read the cap counter against it. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(dynamic_rate_limiter_v3): keep a WebSocket's share hold until it settles and read capped idle borrowing against the team share A frame checked after the model went idle again replaced the connection's share descriptor with None while the admission's hold on that share was still outstanding, so settlement never gave the hold back. The share is now kept while its reservation is open. With the class share split on, the share gauges picked the wider class share over a capped idle team's own share, understating its borrowing, so the tighter of the two is published. 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 | 171 ++++++++++++++---- .../hooks/parallel_request_limiter_v3.py | 8 + .../hooks/test_dynamic_rate_limiter_v3.py | 147 ++++++++++++++- 3 files changed, 290 insertions(+), 36 deletions(-) diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index a58c6feb99a..cadc07a35df 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -390,6 +390,20 @@ def _tightest_team_share(descriptors: Iterable[RateLimitDescriptor | None]) -> t return min(shares, key=lambda share: share[1], default=None) +def _tightest_of(*shares: tuple[RateLimitDescriptor, int] | None) -> tuple[RateLimitDescriptor, int] | None: + return min((share for share in shares if share is not None), key=lambda share: share[1], default=None) + + +def _capped_idle_share( + share: RateLimitDescriptor | None, cap: RateLimitDescriptor | None +) -> tuple[RateLimitDescriptor, int] | None: + """A capped team's usage while the model is idle sits on its cap counter, read against its share.""" + tightest: Final = _tightest_team_share((share,)) + if tightest is None or cap is None: + return None + return cap, tightest[1] + + def _usage_of(kwargs: Mapping[str, object], response_obj: object) -> Usage | None: """Usage of a modeled response, a Responses WebSocket turn's combined usage, or its raw event list.""" usage: Final = getattr(response_obj, "usage", None) @@ -428,7 +442,7 @@ class _BorrowingTeamShare: ceiling: RateLimitDescriptor | None def enforced(self, saturated: bool) -> RateLimitDescriptor | None: - return self.share if saturated else self.ceiling + return self.share if saturated else None def _team_limit_name(rate_limit_type: str) -> str: @@ -480,15 +494,25 @@ def _team_model_settlement( return () served: Final = stash.borrowing_team_model_fallback_limits.get(model_group) if model_group is not None else None admitted: Final = stash.borrowing_team_model_limit + cap: Final = stash.borrowing_team_model_cap reserved_tokens: Final = stash.borrowing_team_model_reserved_tokens + cap_reserved_tokens: Final = stash.borrowing_team_model_cap_reserved_tokens window: Final = stash.borrowing_team_model_reserved_window + cap_window: Final = stash.borrowing_team_model_cap_reserved_window if served is None: - return _window_guarded(_team_model_token_increment(admitted, total_tokens - reserved_tokens, ttl), window) + return ( + *_window_guarded(_team_model_token_increment(admitted, total_tokens - reserved_tokens, ttl), window), + *_window_guarded(_team_model_token_increment(cap, total_tokens - cap_reserved_tokens, ttl), cap_window), + ) return ( *_window_guarded(_team_model_token_increment(admitted, -reserved_tokens, ttl), window), + *_window_guarded(_team_model_token_increment(cap, -cap_reserved_tokens, ttl), cap_window), *_window_guarded( _team_model_token_increment(served.descriptor, total_tokens - served.reserved_tokens, ttl), served.window ), + *_window_guarded( + _team_model_token_increment(served.cap, total_tokens - served.cap_reserved_tokens, ttl), served.cap_window + ), ) @@ -858,12 +882,16 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): return _BorrowingTeamShare( limits=limits, share=self._team_share_descriptor(value, limits), - ceiling=None if ceiling is None else self._team_share_descriptor(value, ceiling), + ceiling=None + if ceiling is None + else self._team_share_descriptor(value, ceiling, TEAM_MODEL_LIMIT_MAX_BORROW_MULTIPLE), ) - def _team_share_descriptor(self, value: str, limits: BorrowingTeamModelLimits) -> RateLimitDescriptor: + def _team_share_descriptor( + self, value: str, limits: BorrowingTeamModelLimits, key: str = TEAM_MODEL_LIMIT_AT_SATURATION + ) -> RateLimitDescriptor: return RateLimitDescriptor( - key=TEAM_MODEL_LIMIT_AT_SATURATION, + key=key, value=value, rate_limit={ "tokens_per_unit": limits.tpm, @@ -881,7 +909,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): user_api_key_dict: UserAPIKeyAuth, team_share: _BorrowingTeamShare, ) -> ProxyRateLimitError: - if saturation >= _get_priority_settings().saturation_threshold: + if status["descriptor_key"] != TEAM_MODEL_LIMIT_MAX_BORROW_MULTIPLE: return self._team_limit_error(model, status, saturation, priority, user_api_key_dict) return self._team_borrow_ceiling_error( model, status, saturation, priority, user_api_key_dict, team_share.limits @@ -922,8 +950,13 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): loggers: Final = _prometheus_loggers() if not loggers or stash is None or model_group is None or settlement.user_api_key_team_id is None: return - share: Final = _tightest_team_share( - (stash.borrowing_team_model_limit, stash.priority_class_share) if hold is None else (hold.descriptor,) + share: Final = ( + _tightest_of(_tightest_team_share((hold.descriptor,)), _capped_idle_share(hold.share, hold.cap)) + if hold is not None + else _tightest_of( + _tightest_team_share((stash.borrowing_team_model_limit, stash.priority_class_share)), + _capped_idle_share(stash.borrowing_team_model_share, stash.borrowing_team_model_cap), + ) ) if share is None: return @@ -1078,10 +1111,19 @@ 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) + team_cap: Final = None if borrowing is None else borrowing.ceiling 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,) + connection_stash.borrowing_team_model_limit = ( + connection_stash.borrowing_team_model_limit + if enforced_team is None and connection_stash.borrowing_team_model_reserved_tokens + else enforced_team + ) + connection_stash.borrowing_team_model_share = borrowing.share + connection_stash.borrowing_team_model_cap = team_cap + descriptors: Final = ( + (share_descriptors if saturated else ()) + + (() if enforced_team is None else (enforced_team,)) + + (() if team_cap is None else (team_cap,)) ) if not descriptors: return @@ -1226,28 +1268,34 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): 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) + team_cap: Final = None if team_share is None else team_share.ceiling + held_team_share: Final = enforced_team_share if enforced_team_share is not None else team_cap get_or_create_request_stash().borrowing_team_model_limit = enforced_team_share + get_or_create_request_stash().borrowing_team_model_share = None if team_share is None else team_share.share + get_or_create_request_stash().borrowing_team_model_cap = team_cap enforced_descriptors: Final[list[RateLimitDescriptor]] = [model_wide_descriptor] if priority_descriptors and should_enforce_priority: enforced_descriptors.extend(priority_descriptors) if enforced_team_share is not None: enforced_descriptors.append(enforced_team_share) + if team_cap is not None: + enforced_descriptors.append(team_cap) per_request_increment: Final[dict[Literal["requests", "tokens"], int]] = { "requests": 1, "tokens": 0, } team_reservation: Final = ( - await self._team_token_reservation(enforced_team_share, data, model, call_type, user_api_key_dict) - if enforced_team_share is not None + await self._team_token_reservation(team_share.share, data, model, call_type, user_api_key_dict) + if team_share is not None and held_team_share is not None else 0 ) atomic_response: Final = await self.v3_limiter.atomic_check_and_increment_by_n( descriptors=enforced_descriptors, increments=[ {"requests": 1, "tokens": team_reservation} - if descriptor is enforced_team_share + if descriptor is enforced_team_share or descriptor is team_cap else per_request_increment for descriptor in enforced_descriptors ], @@ -1307,7 +1355,10 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): llm_provider=llm_provider, ) - if descriptor_key == TEAM_MODEL_LIMIT_AT_SATURATION and team_share is not None: + if ( + descriptor_key in (TEAM_MODEL_LIMIT_AT_SATURATION, TEAM_MODEL_LIMIT_MAX_BORROW_MULTIPLE) + and team_share is not None + ): raise self._team_share_error(model, status, saturation, priority, user_api_key_dict, team_share) if descriptor_key == PRIORITY_CLASS_SHARE_AT_SATURATION and share is not None: raise self._class_share_error(model, status, saturation, priority, user_api_key_dict, share) @@ -1352,11 +1403,15 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): # enforced. The increment-then-check semantics of should_rate_limit # are fine because we ignore the OVER_LIMIT response. stash: Final = get_or_create_request_stash() - stash.borrowing_team_model_reserved_tokens = team_reservation + stash.borrowing_team_model_reserved_tokens = team_reservation if enforced_team_share is not None else 0 + stash.borrowing_team_model_cap_reserved_tokens = team_reservation if team_cap is not None else 0 stash.borrowing_team_model_reserved_window = ( _reserved_window(atomic_response, enforced_team_share) if team_reservation else None ) - self._cap_admitted_output(enforced_team_share, team_reservation, data, call_type) + stash.borrowing_team_model_cap_reserved_window = ( + _reserved_window(atomic_response, team_cap) if team_reservation else None + ) + self._cap_admitted_output(None if team_share is None else team_share.share, team_reservation, data, call_type) if priority_descriptors and share is not None and not should_enforce_priority: get_or_create_request_stash().rate_limit_response = atomic_response elif priority_descriptors and not should_enforce_priority: @@ -1431,33 +1486,48 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): return True saturation: Final = await self._check_model_saturation(model, model_group_info) enforced: Final = team_share.enforced(saturation >= _get_priority_settings().saturation_threshold) + cap: Final = team_share.ceiling + held: Final = enforced if enforced is not None else cap reservation: Final = ( await self._team_token_reservation( - enforced, + team_share.share, {**request_data, "model": model}, model, get_or_create_request_stash().call_type, user_api_key_dict, ) - if enforced is not None + if held is not None else 0 ) + team_descriptors: Final = tuple(d for d in (enforced, cap) if d is not None) if not commit: - return enforced is None or await self._share_has_room(enforced, reservation, user_api_key_dict) + for descriptor in team_descriptors: + if not await self._share_has_room(descriptor, reservation, user_api_key_dict): + return False + return True model_wide_descriptor: Final = self._create_model_tracking_descriptor( model=model, model_group_info=model_group_info, high_limit_multiplier=1 ) response: Final = await self.v3_limiter.atomic_check_and_increment_by_n( - descriptors=[model_wide_descriptor, *(() if enforced is None else (enforced,))], + descriptors=[model_wide_descriptor, *team_descriptors], increments=[ {"requests": 1, "tokens": 0}, - *(() if enforced is None else ({"requests": 1, "tokens": reservation},)), + *({"requests": 1, "tokens": reservation} for _ in team_descriptors), ], parent_otel_span=user_api_key_dict.parent_otel_span, ) if response["overall_code"] == "OVER_LIMIT": return False - self._hold_fallback_share(model, team_share.share, enforced, reservation, _reserved_window(response, enforced)) + self._hold_fallback_share( + model, + team_share.share, + enforced, + reservation if enforced is not None else 0, + _reserved_window(response, enforced), + cap, + reservation if cap is not None else 0, + _reserved_window(response, cap), + ) return True async def _share_has_room( @@ -1479,16 +1549,34 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): charged: RateLimitDescriptor | None, reserved_tokens: int, window: ReservationWindowIdentity | None, + cap: RateLimitDescriptor | None = None, + cap_reserved_tokens: int = 0, + cap_window: ReservationWindowIdentity | None = None, ) -> None: """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. + ``cap`` is the team's borrow ceiling on the target, charged whether or not the target is saturated. Each + counter carries an earlier hold on the same target only while that hold sits in the counter's current window. """ 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 + carried: Final = ( + earlier.reserved_tokens if earlier is not None and window is not None and earlier.window == window else 0 + ) + carried_cap: Final = ( + earlier.cap_reserved_tokens + if earlier is not None and cap_window is not None and earlier.cap_window == cap_window + else 0 + ) hold: Final = BorrowingTeamModelHold( - descriptor=charged, reserved_tokens=reserved_tokens + carried, window=window + descriptor=charged, + reserved_tokens=reserved_tokens + carried, + window=window, + cap=cap, + cap_reserved_tokens=cap_reserved_tokens + carried_cap, + cap_window=cap_window, + share=share, ) stash.borrowing_team_model_fallback_limits = MappingProxyType( {**stash.borrowing_team_model_fallback_limits, model: hold, share["value"].rpartition(":")[2]: hold} @@ -1811,7 +1899,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): ) if stash is not None: stash.borrowing_team_model_reserved_tokens = 0 + stash.borrowing_team_model_cap_reserved_tokens = 0 stash.borrowing_team_model_reserved_window = None + stash.borrowing_team_model_cap_reserved_window = None stash.borrowing_team_model_fallback_limits = MappingProxyType( {k: h for k, h in stash.borrowing_team_model_fallback_limits.items() if h is not served} ) @@ -1915,28 +2005,41 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): _usage_of(callback_kwargs, None if isinstance(response_obj, BaseException) else response_obj), self.v3_limiter.get_rate_limit_type(), ) + own_usage: Final = used if failed_hold is None else 0 + own_delta: Final = own_usage - stash.borrowing_team_model_reserved_tokens + own_cap_delta: Final = own_usage - stash.borrowing_team_model_cap_reserved_tokens refund: Final = ( *_window_guarded( - _team_model_token_increment( - stash.borrowing_team_model_limit, - (used if failed_hold is None else 0) - stash.borrowing_team_model_reserved_tokens, - self.v3_limiter.window_size, - ), + _team_model_token_increment(stash.borrowing_team_model_limit, own_delta, self.v3_limiter.window_size), stash.borrowing_team_model_reserved_window, ), + *_window_guarded( + _team_model_token_increment(stash.borrowing_team_model_cap, own_cap_delta, self.v3_limiter.window_size), + stash.borrowing_team_model_cap_reserved_window, + ), *( () if failed_hold is None - else _window_guarded( - _team_model_token_increment( - failed_hold.descriptor, used - failed_hold.reserved_tokens, self.v3_limiter.window_size + else ( + *_window_guarded( + _team_model_token_increment( + failed_hold.descriptor, used - failed_hold.reserved_tokens, self.v3_limiter.window_size + ), + failed_hold.window, + ), + *_window_guarded( + _team_model_token_increment( + failed_hold.cap, used - failed_hold.cap_reserved_tokens, self.v3_limiter.window_size + ), + failed_hold.cap_window, ), - failed_hold.window, ) ), ) stash.borrowing_team_model_reserved_tokens = 0 + stash.borrowing_team_model_cap_reserved_tokens = 0 stash.borrowing_team_model_reserved_window = None + stash.borrowing_team_model_cap_reserved_window = None if failed_hold is not None: stash.borrowing_team_model_fallback_limits = MappingProxyType( {k: h for k, h in stash.borrowing_team_model_fallback_limits.items() if h is not failed_hold} diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index e520a8d7d92..bf6703443c2 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -581,6 +581,10 @@ class BorrowingTeamModelHold: descriptor: RateLimitDescriptor | None reserved_tokens: int = 0 window: ReservationWindowIdentity | None = None + cap: RateLimitDescriptor | None = None + cap_reserved_tokens: int = 0 + cap_window: ReservationWindowIdentity | None = None + share: RateLimitDescriptor | None = None @dataclass(slots=True) @@ -631,6 +635,10 @@ class RequestRateLimiterStash: borrowing_team_model_limit: RateLimitDescriptor | None = None borrowing_team_model_reserved_tokens: int = 0 borrowing_team_model_reserved_window: ReservationWindowIdentity | None = None + borrowing_team_model_share: RateLimitDescriptor | None = None + borrowing_team_model_cap: RateLimitDescriptor | None = None + borrowing_team_model_cap_reserved_tokens: int = 0 + borrowing_team_model_cap_reserved_window: ReservationWindowIdentity | None = None borrowing_team_model_fallback_limits: Mapping[str, BorrowingTeamModelHold] = field( default_factory=lambda: MappingProxyType({}) ) 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 a54912c44ba..44af7ddc08f 100644 --- a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -3140,6 +3140,62 @@ def _held_to_borrow_cap(exc: ProxyRateLimitError, rate_limit_type: str, limit_se ), exc.detail +async def _cap_tokens_remaining(handler: _PROXY_DynamicRateLimitHandlerV3, auth: UserAPIKeyAuth) -> int: + share: Final = handler._borrowing_team_share("shared-model", auth) + assert share is not None and share.ceiling is not None + response: Final = await handler.v3_limiter.should_rate_limit( + [share.ceiling], read_only=True, skip_parallel_check=True + ) + return next(s["limit_remaining"] for s in response["statuses"] if s["rate_limit_type"] == "tokens") + + +@pytest.mark.asyncio +async def test_a_capped_borrowers_idle_usage_counts_toward_its_cap_but_not_toward_its_share( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _check(handler, _capped_team_key(2), saturation=0.0) + await _settle_tokens(handler, 150) + assert await _cap_tokens_remaining(handler, _capped_team_key(2)) == 50 + assert await _team_tokens_remaining(handler, _capped_team_key(2)) == 100 + await _check(handler, _capped_team_key(2), saturation=0.5) + await _settle_tokens(handler, 30) + assert await _cap_tokens_remaining(handler, _capped_team_key(2)) == 20 + assert await _team_tokens_remaining(handler, _capped_team_key(2)) == 70 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("saturation", [0.0, 0.5], ids=["idle", "saturated"]) +async def test_a_capped_borrowers_reservation_is_sized_by_its_share_not_its_cap( + monkeypatch: pytest.MonkeyPatch, saturation: float +): + """A prompt without max_tokens reserves a floor that scales with the limit it is sized against. Sizing it + against the cap would refuse a request its own share would let through, and would admit different + requests depending on whether the model happened to be saturated.""" + handler: Final = _borrowing_handler(monkeypatch, borrows=True) + await _admit_at_the_hook(handler, _capped_team_key(2), tokens=160) + short_prompt_without_max_tokens: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "hi"}]} + await _check(handler, _capped_team_key(2), saturation=saturation, data=short_prompt_without_max_tokens) + assert short_prompt_without_max_tokens["max_tokens"] == handler.v3_limiter.no_max_tokens_output_floor(100) + + +@pytest.mark.asyncio +async def test_a_capped_borrower_is_held_to_its_cap_while_idle_after_a_saturated_stretch( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _check(handler, _capped_team_key(2), saturation=0.0) + await _settle_tokens(handler, 150) + await _check(handler, _capped_team_key(2), saturation=0.5) + await _settle_tokens(handler, 60) + with pytest.raises(ProxyRateLimitError) as exc: + await _check(handler, _capped_team_key(2), saturation=0.49) + _held_to_borrow_cap(exc.value, "tokens", "model_tpm_limit") + with pytest.raises(ProxyRateLimitError) as saturated: + await _check(handler, _capped_team_key(2), saturation=0.5) + _held_to_borrow_cap(saturated.value, "tokens", "model_tpm_limit") + + @pytest.mark.asyncio async def test_team_borrows_past_its_share_while_it_stays_under_its_max_borrow_multiple( monkeypatch: pytest.MonkeyPatch, @@ -3248,6 +3304,44 @@ async def test_a_websocket_turn_is_refused_at_the_teams_max_borrow_multiple_whil _held_to_borrow_cap(exc.value, "tokens", "model_tpm_limit") +@pytest.mark.asyncio +async def test_a_websocket_admitted_while_saturated_gives_its_share_hold_back_when_the_model_goes_idle( + monkeypatch: pytest.MonkeyPatch, +): + handler: Final = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.5) + await _check(handler, _capped_team_key(2), saturation=0.5, data=_prompt(20)) + assert await _team_tokens_remaining(handler, _capped_team_key(2)) < 100 + await handler.enforce_token_limits_for_frame(user_api_key_dict=_capped_team_key(2), 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, _capped_team_key(2)) == 10 + assert await _cap_tokens_remaining(handler, _capped_team_key(2)) == 110 + + +@pytest.mark.asyncio +async def test_a_capped_websocket_admitted_while_idle_settles_its_first_saturated_turn_on_a_fresh_share( + monkeypatch: pytest.MonkeyPatch, +): + """The idle admission held tokens on the cap only, so the share the turn settles to once the model saturates + starts from zero rather than being refunded a hold it never carried.""" + handler: Final = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.1) + await _check(handler, _capped_team_key(2), saturation=0.0, data=_prompt(20)) + await asyncio.create_task(_saturate_model_from_another_request(handler, 200)) + await handler.enforce_token_limits_for_frame(user_api_key_dict=_capped_team_key(2), 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, _capped_team_key(2)) == 10 + assert await _cap_tokens_remaining(handler, _capped_team_key(2)) == 110 + + @pytest.mark.asyncio async def test_a_websocket_turn_under_the_teams_max_borrow_multiple_still_borrows(monkeypatch: pytest.MonkeyPatch): handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.5) @@ -3444,6 +3538,22 @@ async def _admit_at_the_hook(handler: _PROXY_DynamicRateLimitHandlerV3, auth: Us await _settle_tokens(handler, tokens) +async def _saturate_model_from_another_request(handler: _PROXY_DynamicRateLimitHandlerV3, tokens: int) -> None: + """``_saturate_model`` on a request stash of its own, as a second request in the proxy would have, so the + stash of the request under test keeps its holds.""" + from litellm.proxy.hooks import parallel_request_limiter_v3 + + parallel_request_limiter_v3._request_stash.set(None) # pyright: ignore[reportPrivateUsage] # fresh request context + await _saturate_model(handler, tokens) + + +async def _saturate_model(handler: _PROXY_DynamicRateLimitHandlerV3, tokens: int) -> None: + """Idle usage that pushes the model past its saturation threshold without touching any team's share, from a + priority class of its own so the default pool the measured team sits in stays untouched.""" + warm_up: Final = UserAPIKeyAuth(api_key="sk-warm-up", team_id="team-warm-up", team_metadata={"priority": "batch"}) + await _admit_at_the_hook(handler, warm_up, tokens=tokens) + + async def _refused_at_the_hook( handler: _PROXY_DynamicRateLimitHandlerV3, auth: UserAPIKeyAuth, model: str = "shared-model" ) -> ProxyRateLimitError: @@ -3507,7 +3617,7 @@ async def test_a_failed_request_settles_the_team_share_gauges_too(monkeypatch: p "auth, burns, saturation_threshold, reason, enforced_share_tokens", [ (_team_key(), (60, 70), 0.0, "model_per_team_saturated", 100), - (_capped_team_key(2), (150, 60), 0.5, "model_per_team_max_borrow_multiple", 200), + (_capped_team_key(2), (150, 60), 0.5, "model_per_team_max_borrow_multiple", 100), ], ids=["saturated", "borrow_cap"], ) @@ -3535,7 +3645,8 @@ async def test_a_refused_request_is_counted_under_the_reason_its_429_header_carr ) assert _sample(registry, "litellm_dynamic_limiter_share_tokens", **TEAM_LABELS) == enforced_share_tokens assert ( - _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **TEAM_LABELS) == sum(burns) - enforced_share_tokens + _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **TEAM_LABELS) + == sum(burns) - enforced_share_tokens ) @@ -3607,6 +3718,7 @@ async def test_share_gauges_and_throttles_carry_the_teams_priority_so_the_series team_id="team-a", team_metadata={"model_tpm_limit": {"shared-model": 100}, "priority": "interactive"}, ) + await _saturate_model(handler, 200) await _admit_at_the_hook(handler, team, tokens=60) await _admit_at_the_hook(handler, team, tokens=70) labels: Final = {"model_group": "shared-model", "team": "team-a", "priority": "interactive"} @@ -3626,6 +3738,7 @@ async def test_a_fallback_served_request_publishes_the_teams_share_on_the_fallba handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.0) registry = _prometheus_sink(monkeypatch) team: Final = _team_key(team_limits={"model_tpm_limit": {"other-model": 40, "shared-model": 100}}) + await _saturate_model(handler, 500) assert await handler.async_pre_call_hook(team, DualCache(), {"model": "other-model"}, "completion") is None assert await handler.async_fallback_target_allowed(team, "shared-model", {}) is True await _settle_tokens(handler, 150, model_group="shared-model") @@ -3635,6 +3748,35 @@ async def test_a_fallback_served_request_publishes_the_teams_share_on_the_fallba assert _sample(registry, "litellm_dynamic_limiter_share_tokens", model_group="other-model") is None +@pytest.mark.asyncio +async def test_a_capped_idle_teams_borrowing_is_read_against_its_share_not_the_wider_class_share( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True, splits=True) + registry = _prometheus_sink(monkeypatch) + await _admit_at_the_hook(handler, _capped_team_key(2), tokens=150) + assert _sample(registry, "litellm_dynamic_limiter_share_tokens", **TEAM_LABELS) == 100 + assert _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **TEAM_LABELS) == 50 + + +@pytest.mark.asyncio +async def test_an_idle_fallback_served_for_a_capped_team_still_publishes_its_borrowing(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + registry = _prometheus_sink(monkeypatch) + team: Final = _team_key( + team_limits={ + "model_tpm_limit": {"other-model": 40, "shared-model": 100}, + "team_model_limit_max_borrow_multiple": 2, + } + ) + assert await handler.async_pre_call_hook(team, DualCache(), {"model": "other-model"}, "completion") is None + assert await handler.async_fallback_target_allowed(team, "shared-model", {}) is True + await _settle_tokens(handler, 150, model_group="shared-model") + shared: Final = {"model_group": "shared-model", "team": "team-a", "priority": "default"} + assert _sample(registry, "litellm_dynamic_limiter_share_tokens", **shared) == 100 + assert _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **shared) == 50 + + @pytest.mark.asyncio async def test_a_failing_metrics_sink_never_changes_admission(monkeypatch: pytest.MonkeyPatch): from litellm.integrations.prometheus import PrometheusLogger @@ -3648,6 +3790,7 @@ async def test_a_failing_metrics_sink_never_changes_admission(monkeypatch: pytes handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.0) _prometheus_sink(monkeypatch, BrokenSink) + await _saturate_model(handler, 200) await _admit_at_the_hook(handler, _team_key(), tokens=60) await _admit_at_the_hook(handler, _team_key(), tokens=70) assert await _team_tokens_remaining(handler) == -30