From c44fafaa246a6ba07c0412101148e1bd15812e95 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 1 Jun 2026 21:38:42 +0000 Subject: [PATCH] fix(batch_rate_limiter): resolve provider skip from file-bound model create_batch routes a model-embedded or unified file id on the model bound to that file and ignores the top-level model, so deriving the provider skip from the top-level model first let a caller point model at a skip-listed provider while the file routed a rate-limited one, skipping counter enforcement. Resolve the routing model from the file binding first, matching the batch endpoint. --- litellm/proxy/hooks/batch_rate_limiter.py | 23 +++- .../proxy/hooks/test_batch_file_validation.py | 108 +++++++++++++++++- 2 files changed, 125 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index b4abd86102a..435b6eea45b 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -103,9 +103,11 @@ class _PROXY_BatchRateLimiter(CustomLogger): def _get_file_bound_batch_model(self, data: Dict) -> Optional[str]: """Resolve the model bound to the batch input file ID. - Used only as a fallback routing model when the request omits a - top-level ``model``; the provider is then read from that deployment's - trusted credentials for the provider-level skip decision. + ``create_batch`` routes a file-bound id (model-embedded ``file-...`` or + unified managed file) on that bound model and ignores the top-level + ``model``, so this is the authoritative routing model whenever the file + binds one. The provider is then read from that deployment's trusted + credentials for the provider-level skip decision. """ input_file_id = data.get("input_file_id") if not isinstance(input_file_id, str) or not input_file_id: @@ -130,12 +132,23 @@ class _PROXY_BatchRateLimiter(CustomLogger): return None def _get_batch_routing_model(self, data: Dict) -> Optional[str]: - """Resolve the deployment/model used for this batch from request data.""" + """Resolve the deployment/model used for this batch from request data. + + Mirrors ``create_batch`` routing precedence: a model bound to the input + file id wins over the top-level ``model``, because the batch endpoint + ignores the top-level model for file-bound ids. Resolving the provider + skip from the top-level model first would let a caller point ``model`` + at a skip-listed provider while the file routes a rate-limited one. + """ + file_bound_model = self._get_file_bound_batch_model(data) + if file_bound_model: + return file_bound_model + model = data.get("model") if isinstance(model, str) and model: return model - return self._get_file_bound_batch_model(data) + return None def _resolve_batch_provider(self, batch_model: Optional[str]) -> Optional[str]: """Resolve the provider from the deployment that serves ``batch_model``. diff --git a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py index e9bfaf7205e..af5a5cde8ba 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -521,13 +521,34 @@ def _make_rate_limiter(): ) -def test_get_batch_routing_model_prefers_request_model(): +def test_get_batch_routing_model_uses_request_model_for_plain_file(): rate_limiter = _make_rate_limiter() assert ( rate_limiter._get_batch_routing_model({"model": "gpt-4o-mini"}) == "gpt-4o-mini" ) +def test_get_batch_routing_model_prefers_file_bound_over_request_model(): + """``create_batch`` routes a model-embedded file id on its bound model and + ignores the top-level ``model``. The skip decision must use the same + precedence, otherwise a caller could point ``model`` at a skip-listed + provider while the file routes a rate-limited one.""" + import base64 + + rate_limiter = _make_rate_limiter() + encoded = ( + base64.urlsafe_b64encode(b"litellm:file-xyz;model,vllm-batch") + .decode() + .rstrip("=") + ) + assert ( + rate_limiter._get_batch_routing_model( + {"input_file_id": f"file-{encoded}", "model": "gpt-4o-mini"} + ) + == "vllm-batch" + ) + + def test_get_batch_routing_model_returns_none_without_model_or_file(): rate_limiter = _make_rate_limiter() assert rate_limiter._get_batch_routing_model({}) is None @@ -708,6 +729,91 @@ def test_should_not_skip_for_skip_listed_top_level_model(): assert should_skip is False +def test_should_not_skip_when_file_bound_provider_is_rate_limited(): + """A caller must not bypass batch rate limits by pointing the top-level + ``model`` at a skip-listed provider while the model-embedded ``input_file_id`` + routes to a rate-limited provider. ``create_batch`` runs the batch on the + file-bound model, so the skip decision must resolve the provider from that + model and still process the file when its provider is not skip-listed.""" + import base64 + + rate_limiter = _make_rate_limiter() + rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.return_value = [ + {"rate_limit": {"requests_per_unit": 5}} + ] + user = UserAPIKeyAuth(api_key="sk", models=["*"]) + encoded = ( + base64.urlsafe_b64encode(b"litellm:file-orig;model,vllm-batch") + .decode() + .rstrip("=") + ) + + def _creds(model_id, **kwargs): + provider = "hosted_vllm" if model_id == "vllm-batch" else "openai" + return {"custom_llm_provider": provider} + + with ( + patch( + "litellm.proxy.proxy_server.general_settings", + {"skip_batch_input_file_rate_limiting_for_providers": ["openai"]}, + ), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model", + side_effect=_creds, + ), + ): + should_skip, descriptors = ( + rate_limiter._should_skip_batch_input_file_processing( + data={"input_file_id": f"file-{encoded}", "model": "gpt-skip"}, + user_api_key_dict=user, + ) + ) + assert should_skip is False + assert descriptors is not None + + +def test_should_skip_when_file_bound_provider_is_skip_listed(): + """The provider skip must still fire when the model the batch actually runs + on (the file-bound model) resolves to a skip-listed provider, even if the + top-level ``model`` resolves to a different, non-skipped provider.""" + import base64 + + rate_limiter = _make_rate_limiter() + rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.return_value = [ + {"rate_limit": {"requests_per_unit": 5}} + ] + user = UserAPIKeyAuth(api_key="sk", models=["*"]) + encoded = ( + base64.urlsafe_b64encode(b"litellm:file-orig;model,vllm-batch") + .decode() + .rstrip("=") + ) + + def _creds(model_id, **kwargs): + provider = "hosted_vllm" if model_id == "vllm-batch" else "openai" + return {"custom_llm_provider": provider} + + with ( + patch( + "litellm.proxy.proxy_server.general_settings", + {"skip_batch_input_file_rate_limiting_for_providers": ["hosted_vllm"]}, + ), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model", + side_effect=_creds, + ), + ): + should_skip, descriptors = ( + rate_limiter._should_skip_batch_input_file_processing( + data={"input_file_id": f"file-{encoded}", "model": "gpt-skip"}, + user_api_key_dict=user, + ) + ) + assert should_skip is True + + def test_warns_once_for_unsupported_model_skip_setting(): """Operators who set the no-op per-model skip key get a single warning so a misconfigured deployment does not silently leave batch limits unenforced."""