mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
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>
This commit is contained in:
parent
df8b3b498e
commit
0324b9d915
3 changed files with 290 additions and 36 deletions
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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({})
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue