Merge 55f77ac709 into 159b5b2f72

litellm/router.py

@@ -6213,6 +6213,12 @@ class Router:
                 splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS,
             )
 
+            # 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
@@ -15089,7 +15095,7 @@ class Router:
         messages: list[dict[str, str]] | None = None,
         input: str | list | None = None,
         specific_deployment: bool | None = False,
-    ):
+    ) -> dict[str, Any]:  # mutable-ok: returns pass-through deployment map
         """
         Async implementation of 'get_available_deployments'.
 
@@ -16579,7 +16585,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: Mapping[str, Any], parent_otel_span: Span | None, response: object | None = None
+    ):
         """
         Tracks successful requests rpm usage.
         """

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