scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

fix(dynamic_rate_limiter_v3): a borrower gets its full share back once the model saturates (#51415)

Browse source 
* fix(dynamic_rate_limiter_v3): a borrower gets its full share back once the model saturates

The per-team share counters only count requests admitted while the model is
saturated, so what a team used while the model was idle is not charged
against its share once contention starts. A team below its share is still
admitted ahead of a borrower, and the model's own limits bound the total.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(dynamic_rate_limiter_v3): idle fallback targets and WebSocket turns follow the same share accounting

A router fallback target attempted while idle no longer holds the team's
share, and a Responses WebSocket connection refreshes its share descriptor
on every turn, so idle turns are forgiven and saturated turns count.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(dynamic_rate_limiter_v3): an idle fallback target still records its hold

Settlement finds the hold, gives the admitting model's reservation back and
charges nothing for the idle target, instead of charging the target's
tokens to the admitting model's share. The integration fixture builds its
configuration without mutating the loaded one.

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 15:30:08 -08:00 • committed by GitHub
parent 7130a5b4fd
commit de3f0690e7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 203 additions and 52 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
litellm/proxy/hooks/dynamic_rate_limiter_v3.py
View file

@ -1078,6 +1078,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
saturation: Final = await self._check_model_saturation(requested_model, model_group_info)
saturated: Final = saturation >= _get_priority_settings().saturation_threshold
enforced_team: Final = None if borrowing is None else borrowing.enforced(saturated)
if connection_stash is not None and borrowing is not None:
connection_stash.borrowing_team_model_limit = enforced_team
descriptors: Final = (share_descriptors if saturated else ()) + (
() if enforced_team is None else (enforced_team,)
)
@ -1221,10 +1223,10 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
# All-or-nothing: if any enforced descriptor would exceed its limit,
# no counter is modified and the response carries "OVER_LIMIT".
team_share: Final = self._borrowing_team_share(model=model, user_api_key_dict=user_api_key_dict)
get_or_create_request_stash().borrowing_team_model_limit = None if team_share is None else team_share.share
if team_share is not None:
descriptors_to_check.append(team_share.share)
enforced_team_share: Final = None if team_share is None else team_share.enforced(should_enforce_priority)
get_or_create_request_stash().borrowing_team_model_limit = enforced_team_share
enforced_descriptors: Final[list[RateLimitDescriptor]] = [model_wide_descriptor]
if priority_descriptors and should_enforce_priority:
@ -1355,26 +1357,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
_reserved_window(atomic_response, enforced_team_share) if team_reservation else None
)
self._cap_admitted_output(enforced_team_share, team_reservation, data, call_type)
if team_share is not None and enforced_team_share is None:
await self.v3_limiter.atomic_check_and_increment_by_n(
descriptors=[team_share.share],
increments=[per_request_increment],
parent_otel_span=user_api_key_dict.parent_otel_span,
)
if priority_descriptors and share is not None and not should_enforce_priority:
tracking_only: Final = tuple(
RateLimitDescriptor(
key=descriptor["key"],
value=descriptor["value"],
rate_limit=model_wide_descriptor["rate_limit"],
)
for descriptor in priority_descriptors
)
await self.v3_limiter.atomic_check_and_increment_by_n(
descriptors=list(tracking_only),
increments=[per_request_increment for _ in tracking_only],
parent_otel_span=user_api_key_dict.parent_otel_span,
)
get_or_create_request_stash().rate_limit_response = atomic_response
elif priority_descriptors and not should_enforce_priority:
priority_tracking_response: Final = await self.v3_limiter.should_rate_limit(
@ -1439,8 +1422,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
) -> bool:
"""Hold a router fallback target to the team's share on it once it saturates, and to its borrow ceiling before.
With ``commit`` the attempt is counted toward the target's saturation and the team's counter, holding the
prompt's estimated tokens when the share is enforced, otherwise only checked.
With ``commit`` the attempt is counted toward the target's saturation, and toward the team's counter only
while the share is enforced, holding the prompt's estimated tokens there. An idle target is only checked.
"""
team_share: Final = self._borrowing_team_share(model=model, user_api_key_dict=user_api_key_dict)
model_group_info: Final = self.llm_router.registry().group(model)
@ -1474,13 +1457,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
)
if response["overall_code"] == "OVER_LIMIT":
return False
if enforced is None:
await self.v3_limiter.atomic_check_and_increment_by_n(
descriptors=[team_share.share],
increments=[{"requests": 1, "tokens": reservation}],
parent_otel_span=user_api_key_dict.parent_otel_span,
)
self._hold_fallback_share(model, team_share.share, reservation, _reserved_window(response, enforced))
self._hold_fallback_share(model, team_share.share, enforced, reservation, _reserved_window(response, enforced))
return True
async def _share_has_room(
@ -1496,13 +1473,23 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
)
def _hold_fallback_share(
self, model: str, share: RateLimitDescriptor, reserved_tokens: int, window: ReservationWindowIdentity | None
self,
model: str,
share: RateLimitDescriptor,
charged: RateLimitDescriptor | None,
reserved_tokens: int,
window: ReservationWindowIdentity | None,
) -> None:
"""Remember the hold under the requested name and the resolved group, so settlement finds it either way."""
"""Remember the hold under the requested name and the resolved group, so settlement finds it either way.
``charged`` is the share the target's usage settles to, or None for an idle target whose usage is forgiven.
"""
stash: Final = get_or_create_request_stash()
earlier: Final = stash.borrowing_team_model_fallback_limits.get(model)
carried: Final = earlier.reserved_tokens if earlier is not None and earlier.window == window else 0
hold: Final = BorrowingTeamModelHold(descriptor=share, reserved_tokens=reserved_tokens + carried, window=window)
hold: Final = BorrowingTeamModelHold(
descriptor=charged, reserved_tokens=reserved_tokens + carried, window=window
)
stash.borrowing_team_model_fallback_limits = MappingProxyType(
{**stash.borrowing_team_model_fallback_limits, model: hold, share["value"].rpartition(":")[2]: hold}
)

8
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -572,9 +572,13 @@ class CallTypeRateLimiter(Protocol):
@dataclass(frozen=True, slots=True)
class BorrowingTeamModelHold:
"""A router fallback target's team share with the tokens held on it, settled to actual usage or given back."""
"""A router fallback target's team share with the tokens held on it, settled to actual usage or given back.
descriptor: RateLimitDescriptor
An idle target holds no descriptor: its usage is not charged to the share, and settlement still finds the hold so
the admitting model's reservation is given back instead of taking the target's tokens.
"""
descriptor: RateLimitDescriptor | None
reserved_tokens: int = 0
window: ReservationWindowIdentity | None = None

99
tests/integration/routing/test_dynamic_rate_limiter_borrower_regains_share.py Normal file
View file

@ -0,0 +1,99 @@
import json
import time
import uuid
from collections.abc import Iterator
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import Gateway, eventually, gateway_from_environment
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
OPENAI_MODEL: Final = "gpt-4.1-mini"
MODEL_TPM: Final = 1000
TEAM_TPM: Final = 100
SATURATION_CACHE_TTL: Final = 1
TOKENS_PER_CALL: Final = 30
CALLS_TO_SATURATE: Final = MODEL_TPM // 2 // TOKENS_PER_CALL + 1
UPSTREAM_REPLY: Final = json.dumps(
{
"id": "chatcmpl_borrower_regains_share",
"object": "chat.completion",
"created": 1700000000,
"model": OPENAI_MODEL,
"choices": [{"index": 0, "message": {"role": "assistant", "content": "recovered"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": TOKENS_PER_CALL},
}
).encode()
def _respond(request: Request) -> Reply:
assert request.method == "POST" and request.target.endswith("/v1/chat/completions"), request.target
assert request.headers["authorization"] == "Bearer synthetic-openai-key"
return Reply(body=UPSTREAM_REPLY)
@pytest.fixture(scope="module")
def saturated_model(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, str]]:
directory: Final = tmp_path_factory.mktemp("recovery-window")
base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
configuration: Final = {
**base,
"litellm_settings": {
**base["litellm_settings"],
"callbacks": ["dynamic_rate_limiter_v3"],
"priority_reservation": {"interactive": 0.6, "batch": 0.4},
"priority_reservation_settings": {
"default_priority": 1.0,
"saturation_threshold": 0.5,
"team_model_limit_borrows_idle_capacity_from_other_teams": True,
"saturation_check_cache_ttl": SATURATION_CACHE_TTL,
},
},
}
path: Final = directory / "borrower_regains_share.yaml"
path.write_text(yaml.safe_dump(configuration))
with (
gateway_from_environment() as rig,
wire_server(_respond) as wire,
owned_proxy(rig, directory, {}, config=path) as owned,
):
yield owned, wire.url
def _call(owned: Gateway, model: str, key: str) -> httpx.Response:
return owned.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": f"borrower regains share {uuid.uuid4().hex}"}]},
key=key,
)
def test_a_borrower_gets_its_share_back_once_the_model_saturates_and_a_newcomer_is_served_beside_it(
saturated_model: tuple[Gateway, str],
) -> None:
owned, upstream = saturated_model
with owned.scenario() as scenario:
model: Final = scenario.model(
model=f"openai/{OPENAI_MODEL}", api_base=f"{upstream}/v1", api_key="synthetic-openai-key", tpm=MODEL_TPM
)
borrower: Final = scenario.key(team_id=scenario.team(metadata={"model_tpm_limit": {model: TEAM_TPM}}))
newcomer: Final = scenario.key(team_id=scenario.team(metadata={"model_tpm_limit": {model: TEAM_TPM}}))
borrowed: Final = tuple(_call(owned, model, borrower) for _ in range(CALLS_TO_SATURATE))
assert all(r.status_code == 200 for r in borrowed), [r.status_code for r in borrowed]
time.sleep(SATURATION_CACHE_TTL + 0.5)
forgiven: Final = _call(owned, model, borrower)
assert forgiven.status_code == 200, forgiven.text
served: Final = _call(owned, model, newcomer)
assert served.status_code == 200, served.text
held: Final = eventually(lambda: _call(owned, model, borrower), lambda r: r.status_code == 429, seconds=10)
assert held.headers.get("rate_limited_by") == "model_per_team_saturated", dict(held.headers)
still_served: Final = _call(owned, model, newcomer)
assert still_served.status_code == 200, still_served.text

97
tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py
View file

@ -2163,7 +2163,7 @@ def _borrowing_handler(
async def _admit_then_settle(handler: _PROXY_DynamicRateLimitHandlerV3, auth: UserAPIKeyAuth, tokens: int) -> None:
await _check(handler, auth, saturation=0.0)
await _check(handler, auth, saturation=0.5)
await _settle_tokens(handler, tokens)
@ -2275,13 +2275,39 @@ async def test_team_past_its_model_tpm_limit_is_held_to_it_once_the_model_satura
assert "Team model_tpm_limit: 100, Remaining: 0, Model saturation: 50.0%" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_a_borrowers_usage_while_the_model_was_idle_is_forgiven_once_it_saturates(
monkeypatch: pytest.MonkeyPatch,
):
handler: Final = _borrowing_handler(monkeypatch, borrows=True)
await _check(handler, _team_key(), saturation=0.0)
await _settle_tokens(handler, 150)
assert await _team_tokens_remaining(handler) == 100
await _admit_then_settle(handler, _team_key(), 30)
assert await _team_tokens_remaining(handler) == 70
await _check(handler, _team_key("team-b"), saturation=0.5)
assert await _team_tokens_remaining(handler, _team_key("team-b")) == 100
@pytest.mark.asyncio
async def test_a_borrower_that_spends_its_share_under_saturation_is_held_while_a_newcomer_is_admitted(
monkeypatch: pytest.MonkeyPatch,
):
handler: Final = _borrowing_handler(monkeypatch, borrows=True)
await _admit_then_settle(handler, _team_key(), 150)
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, _team_key(), saturation=0.5)
_held_to_team_limit(exc.value, "tokens", "model_tpm_limit")
await _check(handler, _team_key("team-b"), saturation=0.5)
@pytest.mark.asyncio
async def test_team_past_its_model_rpm_limit_is_held_to_it_once_the_model_saturates(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=True)
auth: Final = _team_key(team_limits={"model_rpm_limit": {"shared-model": 2}})
await _check(handler, auth, saturation=0.0)
await _check(handler, auth, saturation=0.0)
await _check(handler, auth, saturation=0.49)
await _check(handler, auth, saturation=0.5)
await _check(handler, auth, saturation=0.5)
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, auth, saturation=0.5)
_held_to_team_limit(exc.value, "requests", "model_rpm_limit")
@ -2299,7 +2325,7 @@ async def test_requests_admitted_below_saturation_count_only_their_settled_token
@pytest.mark.asyncio
async def test_an_embedding_responses_tokens_count_against_the_team_limit(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=True)
await _check(handler, _team_key(), saturation=0.0)
await _check(handler, _team_key(), saturation=0.5)
await _settle_tokens(
handler,
150,
@ -2315,7 +2341,7 @@ async def test_an_embedding_responses_tokens_count_against_the_team_limit(monkey
@pytest.mark.asyncio
async def test_a_responses_api_calls_tokens_count_against_the_team_limit(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=True)
await _check(handler, _team_key(), saturation=0.0)
await _check(handler, _team_key(), saturation=0.5)
await _settle_tokens(
handler,
150,
@ -2339,7 +2365,7 @@ async def test_a_responses_api_calls_tokens_count_against_the_team_limit(monkeyp
@pytest.mark.asyncio
async def test_a_null_auth_metadata_still_settles_the_teams_tokens(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=True)
await _check(handler, _team_key(), saturation=0.0)
await _check(handler, _team_key(), saturation=0.5)
await handler.async_log_success_event(
kwargs={
"litellm_params": {"metadata": {"model_group": "shared-model"}},
@ -2403,7 +2429,7 @@ async def test_an_exempt_keys_tokens_are_not_charged_to_its_teammates(monkeypatc
@pytest.mark.asyncio
async def test_an_alias_request_counts_against_the_same_team_share_as_its_target(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=True, model_group_alias={"alias-model": "shared-model"})
await _check(handler, _team_key(), saturation=0.0, model="alias-model")
await _check(handler, _team_key(), saturation=0.5, model="alias-model")
await _settle_tokens(handler, 150)
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, _team_key(), saturation=0.5)
@ -2432,7 +2458,7 @@ async def test_a_nested_calls_tokens_are_not_charged_to_the_requests_team_share(
handler = _borrowing_handler(monkeypatch, borrows=True)
claim_request_stash_for_data({"litellm_call_id": "outer-call"})
await _check(handler, _team_key(), saturation=0.0)
await _check(handler, _team_key(), saturation=0.5)
await _settle_tokens(handler, 60, call_id="llm-judge-call")
assert await _team_tokens_remaining(handler) == 100
await _settle_tokens(handler, 60, call_id="outer-call")
@ -2442,7 +2468,7 @@ async def test_a_nested_calls_tokens_are_not_charged_to_the_requests_team_share(
@pytest.mark.asyncio
async def test_a_websocket_turns_tokens_count_against_the_team_limit(monkeypatch: pytest.MonkeyPatch):
handler = _borrowing_handler(monkeypatch, borrows=True)
await _check(handler, _team_key(), saturation=0.0)
await _check(handler, _team_key(), saturation=0.5)
await _settle_tokens(
handler,
150,
@ -2454,6 +2480,29 @@ async def test_a_websocket_turns_tokens_count_against_the_team_limit(monkeypatch
_held_to_team_limit(exc.value, "tokens", "model_tpm_limit")
@pytest.mark.asyncio
async def test_a_websocket_admitted_while_idle_has_its_turns_counted_once_the_model_saturates(
monkeypatch: pytest.MonkeyPatch,
):
handler: Final = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.1)
await _check(handler, _team_key(), saturation=0.0)
await _settle_tokens(
handler,
150,
response=[],
combined_usage_object=litellm.Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150),
)
assert await _team_tokens_remaining(handler) == 100
await handler.enforce_token_limits_for_frame(user_api_key_dict=_team_key(), requested_model="shared-model")
await _settle_tokens(
handler,
90,
response=[],
combined_usage_object=litellm.Usage(prompt_tokens=60, completion_tokens=30, total_tokens=90),
)
assert await _team_tokens_remaining(handler) == 10
@pytest.mark.asyncio
async def test_a_websocket_turn_gets_a_503_when_the_team_counter_cannot_be_read(monkeypatch: pytest.MonkeyPatch):
from litellm.exceptions import ServiceUnavailableError
@ -2513,15 +2562,28 @@ async def test_a_fallback_target_is_refused_once_it_saturates_and_the_teams_shar
@pytest.mark.asyncio
async def test_a_fallback_target_still_borrows_while_idle_and_its_tokens_count_against_that_teams_share(
async def test_a_fallback_target_still_borrows_while_idle_and_its_tokens_are_forgiven_once_it_saturates(
monkeypatch: pytest.MonkeyPatch,
):
handler = _borrowing_handler(monkeypatch, borrows=True)
assert await handler.async_fallback_target_allowed(_team_key(), "shared-model", {}) is True
await _settle_tokens(handler, 150)
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, _team_key(), saturation=0.5)
_held_to_team_limit(exc.value, "tokens", "model_tpm_limit")
await _check(handler, _team_key(), saturation=0.5)
assert await _team_tokens_remaining(handler) == 100
@pytest.mark.asyncio
async def test_an_idle_fallback_target_is_not_charged_and_the_admitting_models_hold_is_given_back(
monkeypatch: pytest.MonkeyPatch,
):
handler: Final = _borrowing_handler(monkeypatch, borrows=True)
key: Final = _team_key(team_limits={"model_tpm_limit": {"shared-model": 100, "other-model": 100}})
await _check(handler, key, saturation=0.5, data=_prompt(20))
assert await _team_tokens_remaining(handler, key) < 100
assert await handler.async_fallback_target_allowed(key, "other-model", _prompt(20)) is True
await _settle_tokens(handler, 150, model_group="other-model")
assert await _team_tokens_remaining(handler, key) == 100
assert await _team_tokens_remaining(handler, key, model="other-model") == 100
@pytest.mark.asyncio
@ -2806,8 +2868,9 @@ async def _burn(
priority: str | None,
tokens: int,
model: str = "shared-model",
saturation: float = 1.0,
):
await _check(handler, auth, saturation=0.0, priority=priority, model=model)
await _check(handler, auth, saturation=saturation, priority=priority, model=model)
await handler.async_log_success_event(
kwargs={
"litellm_params": {"metadata": {"model_group": model}},
@ -3248,7 +3311,7 @@ async def test_a_team_id_spelled_like_a_pool_marker_still_counts_as_a_team(monke
@pytest.mark.asyncio
async def test_requests_sent_while_the_share_was_not_enforced_all_count_once_the_share_grows(
async def test_requests_sent_while_the_share_was_not_enforced_are_forgiven_once_the_model_saturates(
monkeypatch: pytest.MonkeyPatch,
):
handler = _borrowing_handler(monkeypatch, borrows=False, splits=True)
@ -3256,9 +3319,7 @@ async def test_requests_sent_while_the_share_was_not_enforced_all_count_once_the
for _ in range(70):
await _check(handler, _class_team("team-a"), saturation=0.0, priority="interactive")
await _check(handler, _class_team("team-b", "batch"), saturation=0.0, priority="batch")
with pytest.raises(ProxyRateLimitError) as exc:
await _check(handler, _class_team("team-a"), saturation=0.6, priority="interactive")
assert "Active teams in class: 1, Active classes: 2, " in str(exc.value.detail)
await _check(handler, _class_team("team-a"), saturation=0.6, priority="interactive")
@pytest.mark.asyncio