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

fix(rate-limit): settle a fallback's tokens on the model group that served it (#51082)

Browse source 
* fix(rate-limit): settle a fallback's tokens on the model group that served it

The v3 limiter pinned post-call TPM settlement to the model reserved at
admission, so when the router fell back to another model group the
requested model's counter took the usage and the serving model's counter
stayed empty. When the served model group differs from the admitted one,
charge the served group and refund the abandoned group's per-model
reservation. Same-model retries and alias reloads keep the admission target.

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

* fix(rate-limit): refund an abandoned fallback reservation only on counters the served model does not settle

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

* test(rate-limit): wait for the primary refund and the fallback charge before reading counters

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

* fix(rate-limit): only treat a limiter-admitted fallback target as the served model at settlement

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

* test(rate-limit): drop unrelated formatting from the v3 limiter unit test file

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

* test(rate-limit): remove helper copies duplicated while appending the fallback settlement tests

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-10 18:36:55 -08:00 • committed by GitHub
parent 373fe799e7
commit 917fd11966
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 277 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -500,6 +500,7 @@ class RequestRateLimiterStash:
reserved_tokens: int = 0
reserved_model: RateLimitedModel | None = None
reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset)
fallback_models: frozenset[str] = field(default_factory=frozenset)
ptu_ceiling: PTUTeamCeiling | None = None
ptu_reserved_tokens: int = 0
ptu_reserved_window_identities: frozenset[ReservationWindowIdentity] = field(default_factory=frozenset)
@ -3786,6 +3787,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
) -> bool:
"""Reserve the caller's per-model limits on a router fallback target. False when any is exhausted."""
data: Final = MappingProxyType({**request_data, "model": model})
stash: Final = _request_stash.get()
hidden_stash: Final = _request_stash.set(None)
try:
built: Final = await self._build_request_rate_limit_descriptors(user_api_key_dict, data, None)
@ -3794,10 +3796,16 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
descriptors: Final = [
descriptor for descriptor in built if descriptor["key"] in PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS
]
if not descriptors:
return True
response: Final = await self.should_rate_limit(descriptors, parent_otel_span=user_api_key_dict.parent_otel_span)
return response["overall_code"] == "OK"
allowed: Final = (
not descriptors
or (await self.should_rate_limit(descriptors, parent_otel_span=user_api_key_dict.parent_otel_span))[
"overall_code"
]
== "OK"
)
if allowed and stash is not None:
stash.fallback_models = stash.fallback_models | {model}
return allowed
@asynccontextmanager
async def request_capacity(
@ -4720,12 +4728,21 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
reserved_tokens: Final = open_reservation.reserved_tokens if open_reservation is not None else 0
reserved_model: Final = stash.reserved_model if stash is not None else None
reserved_scopes: Final[frozenset[tuple[str, str]]] = stash.reserved_scopes if stash is not None else frozenset()
# Reconciliation must target the same model-scoped counter that the
# pre-call reservation incremented. If a reservation was made,
served_model: Final = self._rate_limited_model(model_group)
# ``reserved_model`` (resolved at admission, so an alias map reload
# mid-flight cannot move the charge) is authoritative; otherwise fall
# back to the router's ``model_group`` (the no-reservation charge path).
reconcile_model: Final = reserved_model if reserved_model is not None else self._rate_limited_model(model_group)
# mid-flight cannot move the charge) settles the request unless the
# router fell back to a group this limiter admitted as a fallback
# target: then the served group is charged and the abandoned group's
# per-model reservation is refunded. Client-supplied ``model_group``
# metadata alone never moves the charge.
fell_back: Final = (
stash is not None
and reserved_model is not None
and served_model is not None
and served_model.requested != reserved_model.requested
and served_model.requested in stash.fallback_models
)
reconcile_model: Final = served_model if fell_back or reserved_model is None else reserved_model
reconcile_models: Final = self._rate_limited_models(reconcile_model)
pipeline_operations: Final[list[ReservationOperation]] = []
@ -4764,6 +4781,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
)
)
charged_targets: Final = [target for target in targets if target not in key_owned_team_scopes]
abandoned_scopes: Final = (
tuple(
sorted(
scope
for scope in reserved_scopes
if scope[0] in PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS and scope not in charged_targets
)
)
if fell_back
else ()
)
if reserved_tokens > 0 and total_tokens < reserved_tokens:
verbose_proxy_logger.debug(
"Releasing unused TPM budget on success: reserved=%s, actual=%s, release=%s",
@ -4779,6 +4807,14 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
reserved_tokens=reserved_tokens,
)
)
pipeline_operations.extend(
self._build_reservation_aware_tpm_ops(
targets=abandoned_scopes,
reserved_scopes=reserved_scopes,
actual_tokens=0,
reserved_tokens=reserved_tokens,
)
)
call_type: Final[object] = kwargs.get("call_type")
ptu_skips_socket_usage: Final = call_type == CallTypes.aresponses_websocket.value
pipeline_operations.extend(

91
tests/integration/routing/test_fallback_settles_serving_model.py Normal file
View file

@ -0,0 +1,91 @@
from __future__ import annotations
import os
import time
import uuid
from collections.abc import Iterator
from contextlib import ExitStack
from dataclasses import dataclass
from typing import Final
import pytest
from integration._support.client import Gateway, gateway_from_environment
from integration._support.process import OwnedState, graceful_stop_seconds, owned_proxy_process
from integration._support.wire import wire_server
from integration.routing.test_fallback_rate_limit_check import provider_error_reply
from integration.routing.test_per_model_rate_limit_fallbacks import (
assert_served,
chat,
healthy_reply,
model_entry,
proxy_config,
)
from redis import Redis
pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120)
ADMISSION_PROBE_TOKENS: Final = 1
SETTLEMENT_DEADLINE_SECONDS: Final = 15.0
@dataclass(frozen=True, slots=True)
class Rig:
gateway: Gateway
primary: str
fallback: str
def key_model_tokens(group: str) -> int:
with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), socket_timeout=2) as cache:
counters: Final = tuple(cache.scan_iter(match=f"*model_per_key:*:{group}}}:tokens"))
return sum(int(cache.get(counter) or 0) for counter in counters)
def settled_key_model_tokens(primary: str, fallback: str, total_tokens: int) -> tuple[int, int]:
deadline: Final = time.monotonic() + SETTLEMENT_DEADLINE_SECONDS
while True:
charged: Final = (key_model_tokens(primary), key_model_tokens(fallback))
if (charged[0] == 0 and charged[1] >= total_tokens) or time.monotonic() >= deadline:
return charged
time.sleep(0.5)
@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]:
suffix: Final = uuid.uuid4().hex[:10]
primary: Final = f"settle-primary-{suffix}"
fallback: Final = f"settle-fallback-{suffix}"
directory: Final = tmp_path_factory.mktemp("fallback-settles-serving-model")
with ExitStack() as stack:
rig_gateway: Final = stack.enter_context(gateway_from_environment())
primary_wire: Final = stack.enter_context(wire_server(provider_error_reply))
fallback_wire: Final = stack.enter_context(wire_server(healthy_reply))
config: Final = proxy_config(
directory,
model_list=(model_entry(primary, primary_wire), model_entry(fallback, fallback_wire)),
fallbacks=({primary: [fallback]},),
general_settings={},
)
owned: Final = stack.enter_context(
owned_proxy_process(
rig_gateway, directory, {}, config=config, workers=2, state=OwnedState.SHARED_WITH_WORKER
)
)
yield Rig(owned.gateway, primary, fallback)
def test_provider_error_fallback_charges_its_tokens_to_the_model_that_served(rig: Rig) -> None:
with rig.gateway.scenario() as scenario:
key: Final = scenario.key(model_tpm_limit={rig.primary: 100_000, rig.fallback: 1_000})
response: Final = chat(rig.gateway, key, rig.primary, f"pmrl-{uuid.uuid4().hex}")
assert_served(response, rig.fallback, "served pmrl-")
total_tokens: Final = int(response.json()["usage"]["total_tokens"])
assert total_tokens > 0, response.text
primary_tokens, fallback_tokens = settled_key_model_tokens(rig.primary, rig.fallback, total_tokens)
assert primary_tokens == 0, (primary_tokens, fallback_tokens)
assert total_tokens <= fallback_tokens <= total_tokens + ADMISSION_PROBE_TOKENS, (
primary_tokens,
fallback_tokens,
)

141
tests/unit/proxy/hooks/test_parallel_request_limiter_v3.py
View file

@ -9941,3 +9941,144 @@ async def test_end_user_model_rpm_limit_is_shared_across_keys_and_isolated_per_c
await _rpm_request(handler, cache, globex, "target")
await _rpm_request(handler, cache, acme_via_key_one, "other-model")
@pytest.mark.asyncio
async def test_success_tpm_accounting_moves_the_charge_to_the_model_group_that_served_a_fallback() -> None:
cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
key_metadata: Final = {"model_tpm_limit": {"primary": 1000, "fallback": 1000}}
auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-fallback-settle"), metadata=key_metadata)
await handler.async_pre_call_hook(
user_api_key_dict=auth,
cache=cache,
data={"model": "primary", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 10},
call_type="acompletion",
)
stash: Final = get_request_stash()
assert stash is not None
assert stash.reserved_model == RateLimitedModel(requested="primary", group="primary")
assert stash.reserved_tokens > 0
assert await handler.async_fallback_target_allowed(
auth, "fallback", {"messages": [{"role": "user", "content": "hello"}]}
)
response: Final = ModelResponse(
id="fallback-settle",
object="chat.completion",
created=int(datetime.now().timestamp()),
model="fallback",
usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150),
choices=[],
)
kwargs: Final = {
"standard_logging_object": {"metadata": {"user_api_key_hash": auth.api_key}},
"litellm_params": {
"metadata": {
"model_group": "fallback",
"original_model_group": "primary",
"user_api_key_metadata": key_metadata,
}
},
"model": "fallback",
}
ops: Final = handler._build_success_event_pipeline_operations(
kwargs=kwargs, response_obj=response, rate_limit_type="total"
)
charged: Final = {_increment_of(op).counter: _increment_of(op).amount for op in ops}
assert charged[_counter("model_per_key", f"{auth.api_key}:fallback")] == 150
assert charged[_counter("model_per_key", f"{auth.api_key}:primary")] == -stash.reserved_tokens
@pytest.mark.asyncio
async def test_success_tpm_accounting_settles_an_alias_to_target_fallback_once_per_counter(
collision_limiter: tuple[DualCache, PROXY_MaxParallelRequestsHandler_v3, str, str],
) -> None:
cache, handler, alias, target = collision_limiter
key_metadata: Final = {"model_tpm_limit": {alias: 1000, target: 1000}}
auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-alias-fallback-settle"), metadata=key_metadata)
await handler.async_pre_call_hook(
user_api_key_dict=auth,
cache=cache,
data={"model": alias, "messages": [{"role": "user", "content": "hello"}], "max_tokens": 10},
call_type="acompletion",
)
stash: Final = get_request_stash()
assert stash is not None
assert stash.reserved_model == RateLimitedModel(requested=alias, group=target)
assert stash.reserved_tokens > 0
assert await handler.async_fallback_target_allowed(
auth, target, {"messages": [{"role": "user", "content": "hello"}]}
)
response: Final = ModelResponse(
id="alias-fallback-settle",
object="chat.completion",
created=int(datetime.now().timestamp()),
model=target,
usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150),
choices=[],
)
kwargs: Final = {
"standard_logging_object": {"metadata": {"user_api_key_hash": auth.api_key}},
"litellm_params": {
"metadata": {"model_group": target, "original_model_group": alias, "user_api_key_metadata": key_metadata}
},
"model": target,
}
ops: Final = handler._build_success_event_pipeline_operations(
kwargs=kwargs, response_obj=response, rate_limit_type="total"
)
increments: Final = tuple(_increment_of(op) for op in ops)
target_counter: Final = _counter("model_per_key", f"{auth.api_key}:{target}")
alias_counter: Final = _counter("model_per_key", f"{auth.api_key}:{alias}")
assert [inc.amount for inc in increments if inc.counter == target_counter] == [150 - stash.reserved_tokens]
assert [inc.amount for inc in increments if inc.counter == alias_counter] == []
@pytest.mark.asyncio
async def test_success_tpm_accounting_ignores_a_client_supplied_model_group_without_a_router_fallback() -> None:
cache: Final = DualCache()
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(cache))
key_metadata: Final = {"model_tpm_limit": {"primary": 1000}}
auth: Final = UserAPIKeyAuth(api_key=hash_token("sk-client-model-group"), metadata=key_metadata)
await handler.async_pre_call_hook(
user_api_key_dict=auth,
cache=cache,
data={"model": "primary", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 10},
call_type="acompletion",
)
stash: Final = get_request_stash()
assert stash is not None
assert stash.reserved_tokens > 0
response: Final = ModelResponse(
id="client-model-group",
object="chat.completion",
created=int(datetime.now().timestamp()),
model="primary",
usage=Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150),
choices=[],
)
kwargs: Final = {
"standard_logging_object": {"metadata": {"user_api_key_hash": auth.api_key}},
"litellm_params": {"metadata": {"model_group": "uncapped", "user_api_key_metadata": key_metadata}},
"model": "primary",
}
ops: Final = handler._build_success_event_pipeline_operations(
kwargs=kwargs, response_obj=response, rate_limit_type="total"
)
charged: Final = {_increment_of(op).counter: _increment_of(op).amount for op in ops}
assert charged[_counter("model_per_key", f"{auth.api_key}:primary")] == 150 - stash.reserved_tokens
assert not any(":uncapped" in counter.scope.descriptor_value for counter in charged)