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