fix(router): enforce deployment rpm on passthrough routes (#45987)

_ageneric_api_call_with_fallbacks_helper (used by /v1/messages and other
passthrough routes) never recorded selection-time usage, so with
router_settings.enable_pre_call_checks: true a deployment's rpm limit was
only enforced on /v1/chat/completions (_acompletion), not on /v1/messages.

Track the picked deployment at selection time the same way _acompletion
does, so _pre_call_checks filters deployments that reached their per-minute
rpm budget on every route.

Signed-off-by: apex-mochen <2756823972@qq.com>
This commit is contained in:
叮当猫 2026-10-12 11:14:24 +08:00
parent 102619d47a
commit 8e71fa5c63
2 changed files with 106 additions and 2 deletions

View file

@ -5492,6 +5492,12 @@ class Router:
self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name=function_name)
# Count the picked deployment against its per-minute rpm budget at
# selection time, the same way _acompletion does, so
# _pre_call_checks enforces rpm limits on passthrough routes
# (/v1/messages and friends) that go through this helper.
self._track_deployment_metrics(deployment=deployment, parent_otel_span=parent_otel_span)
data: Final = deployment["litellm_params"].copy()
model_name: Final = data["model"]
self.total_calls[model_name] += 1
@ -13786,7 +13792,7 @@ class Router:
messages: list[dict[str, str]] | None = None,
input: str | list | None = None,
specific_deployment: bool | None = False,
):
) -> dict[str, Any]:
"""
Async implementation of 'get_available_deployments'.
@ -15199,7 +15205,9 @@ class Router:
return pass_through_deployments
def _track_deployment_metrics(self, deployment, parent_otel_span: Span | None, response=None):
def _track_deployment_metrics(
self, deployment: dict[str, Any], parent_otel_span: Span | None, response: object | None = None
):
"""
Tracks successful requests rpm usage.
"""

View file

@ -0,0 +1,96 @@
"""
Regression tests for #45987.
router_settings.enable_pre_call_checks filters deployments that reached their
per-minute rpm budget on /v1/chat/completions (_acompletion tracks selection
usage via _track_deployment_metrics), but passthrough routes such as
/v1/messages go through _ageneric_api_call_with_fallbacks_helper, which never
recorded selection usage, so a deployment with `rpm: 1` was never filtered out
within the same UTC minute.
"""
from unittest.mock import AsyncMock, patch
import pytest
from litellm import Router
def _make_deployment(dep_id: str, rpm: int | None = None) -> dict:
params: dict = {"model": "gpt-4o", "api_key": "key"}
if rpm is not None:
params["rpm"] = rpm
return {
"model_name": "m-group",
"litellm_params": params,
"model_info": {"id": dep_id},
}
@pytest.mark.asyncio
async def test_passthrough_helper_tracks_selection_usage():
"""
_ageneric_api_call_with_fallbacks_helper must record the picked deployment
against its per-minute rpm budget at selection time (mirroring _acompletion).
Before the fix the helper never called _track_deployment_metrics, so the
deployment's local request count stayed at 0 and _pre_call_checks never
enforced `rpm` on /v1/messages.
"""
router = Router(model_list=[_make_deployment("dep-a", rpm=1)])
async def fake_original(**kwargs):
return {"choices": [{"message": {"role": "assistant", "content": "hi"}}]}
with patch.object(
router,
"async_get_available_deployment",
AsyncMock(return_value=_make_deployment("dep-a", rpm=1)),
):
await router._ageneric_api_call_with_fallbacks_helper(
model="m-group",
original_generic_function=fake_original,
messages=[{"role": "user", "content": "hi"}],
)
count = router.cache.get_cache(key="dep-a", local_only=True)
assert count == 1
@pytest.mark.asyncio
async def test_pre_call_checks_filters_deployment_at_rpm_limit_after_tracking():
"""
End-to-end rpm accounting: once selection usage has been recorded against a
deployment (as the fixed helper now does), a follow-up _pre_call_checks pass
must drop that deployment when its count reached its rpm budget, so repeated
/v1/messages requests stop hitting an `rpm: 1` deployment within the same
UTC minute.
"""
router = Router(
model_list=[
_make_deployment("dep-a", rpm=1),
_make_deployment("dep-b"),
]
)
healthy = [_make_deployment("dep-a", rpm=1), _make_deployment("dep-b")]
# First pass: dep-a has not been used yet -> kept.
kept = router._pre_call_checks(
model="m-group",
healthy_deployments=healthy,
messages=[{"role": "user", "content": "hi"}],
)
assert any(d["model_info"]["id"] == "dep-a" for d in kept)
# Record one request against dep-a (what the fixed helper does on selection).
router._track_deployment_metrics(deployment=_make_deployment("dep-a", rpm=1), parent_otel_span=None)
# Second pass: dep-a is at its rpm budget -> filtered out, dep-b remains.
kept2 = router._pre_call_checks(
model="m-group",
healthy_deployments=healthy,
messages=[{"role": "user", "content": "hi"}],
)
ids2 = [d["model_info"]["id"] for d in kept2]
assert "dep-a" not in ids2
assert "dep-b" in ids2