diff --git a/litellm/router.py b/litellm/router.py index cdede46ef5c..0644005423c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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. """ diff --git a/tests/unit/test_router_pre_call_checks_passthrough_rpm.py b/tests/unit/test_router_pre_call_checks_passthrough_rpm.py new file mode 100644 index 00000000000..a3904b1e4bf --- /dev/null +++ b/tests/unit/test_router_pre_call_checks_passthrough_rpm.py @@ -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