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:
Kerry Lu 2026-11-12 17:08:02 -08:00 • committed by GitHub
parent df8b3b498e
commit 0324b9d915
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 290 additions and 36 deletions

View file

@ -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}

View file

@ -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({})
)

View file

@ -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