diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 5b814ad28fd..34d3027e5c0 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -361,6 +361,47 @@ class _PROXY_BatchRateLimiter(CustomLogger): return False, descriptors + def _batch_input_file_models_only( + self, + data: dict, # mutable-ok: read only here, and mirrors the sibling skip helper's signature + user_api_key_dict: UserAPIKeyAuth, + has_enqueued_scopes: bool = False, + ) -> bool: + """True when the JSONL is read only to validate ``body.model``. + + ``_should_skip_batch_input_file_processing`` returns on the model + allowlist check before it consults the operator opt-outs, so a key with a + restricted ``models`` list always takes the full path even when batch + input-file rate limiting is disabled. The download itself is still + required -- the allowlist can only be enforced by inspecting every row -- + but the per-row token counting it also performs is then discarded by the + caller. + + Reporting that case lets the read collect ``body.model`` without + tokenizing each row. Enforcement is unchanged: this returns False + whenever anything still needs the totals, including enqueued-token + scopes, whose reservation is priced from them. + """ + from litellm.proxy.proxy_server import general_settings + + if has_enqueued_scopes: + return False + + if not self._key_requires_batch_model_access_check(user_api_key_dict): + # An unrestricted key is already covered by the full-skip paths. + return False + + if general_settings.get("disable_batch_input_file_rate_limiting") is True: + return True + + skip_providers: Final = tuple(general_settings.get("skip_batch_input_file_rate_limiting_for_providers") or ()) + if skip_providers: + batch_provider: Final = self._resolve_batch_provider(self._get_batch_routing_model(data)) + if batch_provider and batch_provider in skip_providers: + return True + + return False + def _warn_if_unsupported_model_skip_configured(self, general_settings: dict) -> None: """Warn once that ``skip_batch_input_file_rate_limiting_for_models`` is a no-op. @@ -729,6 +770,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): user_api_key_dict: UserAPIKeyAuth | None = None, data: dict | None = None, descriptors: Sequence["RateLimitDescriptor"] | None = None, + models_only: bool = False, ) -> BatchFileUsage: """ Count number of requests and tokens in a batch input file. @@ -739,6 +781,10 @@ class _PROXY_BatchRateLimiter(CustomLogger): user_api_key_dict: User authentication information for file access (required for managed files) descriptors: Rate limit descriptors already computed for this batch, so the configured project OTPM limit can scale the no-``max_tokens`` output floor + models_only: Collect and validate every row's ``body.model`` but skip + per-row token counting. Set when the file is read solely for + allowlist enforcement, so the token fields come back 0 and the + caller must not charge them. Returns: BatchFileUsage with total_tokens, output_tokens, request_count, and @@ -831,6 +877,10 @@ class _PROXY_BatchRateLimiter(CustomLogger): try: entry = json.loads(raw_line) except Exception: + if models_only: + # A malformed row names no model, so skipping its size + # estimate cannot weaken the allowlist check below. + continue entry_total_tokens = _estimate_batch_entry_tokens(raw_line) entry_output_tokens = self.parallel_request_limiter.no_max_tokens_output_floor( min_configured_otpm_limit @@ -843,6 +893,12 @@ class _PROXY_BatchRateLimiter(CustomLogger): if model: models.add(model) + if models_only: + # Every row's `body.model` is collected above, which is all + # `_enforce_batch_file_model_access` needs. Tokenizing here + # would be work the caller throws away. + continue + if isinstance(entry, dict): entry_output_tokens = self._estimate_entry_output_tokens(entry, min_configured_otpm_limit) else: @@ -1109,6 +1165,12 @@ class _PROXY_BatchRateLimiter(CustomLogger): if should_skip: return data + models_only: Final = self._batch_input_file_models_only( + data=data, + user_api_key_dict=user_api_key_dict, + has_enqueued_scopes=bool(enqueued_scopes), + ) + # Get custom_llm_provider for token counting custom_llm_provider: Final = data.get("custom_llm_provider", "openai") @@ -1120,8 +1182,19 @@ class _PROXY_BatchRateLimiter(CustomLogger): user_api_key_dict=user_api_key_dict, data=data, descriptors=batch_rate_limit_descriptors, + models_only=models_only, ) + if models_only: + # The allowlist was enforced inside count_input_file_usage and + # nothing charges this batch, so there are no counters to check. + # Returning here keeps the zeroed totals out of `data`, matching + # the full-skip path above which also leaves them unset. + verbose_proxy_logger.debug( + "Batch model-access validation passed; batch input file rate limiting disabled" + ) + return data + verbose_proxy_logger.debug( "Batch input file usage - Tokens: %s, Requests: %s", batch_usage.total_tokens, batch_usage.request_count ) 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 4ef94f0965b..c585d14f732 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -2286,3 +2286,268 @@ async def test_disable_flag_still_skips_batch_processing_with_enqueued_limits(): assert result is data afile_content_mock.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# Models-only read: the allowlist forces the download, nothing charges the +# tokens, so the per-row tokenization is skipped. +# --------------------------------------------------------------------------- + +_MODELS_ONLY_ALLOWED = "gpt-4o-mini" +_MODELS_ONLY_DENIED = "gpt-4o" + + +def _models_only_file(models): + """A JSONL batch file naming ``models``, one chat row each.""" + import json as _json + + body = "\n".join( + _json.dumps({"body": {"model": m, "messages": [{"role": "user", "content": "x" * 64}]}}) + for m in models + ) + content = MagicMock() + content.content = body.encode("utf-8") + return content + + +def _restricted_user(**kwargs): + return UserAPIKeyAuth( + api_key="sk-restricted-models-only", + user_id="alice", + models=[_MODELS_ONLY_ALLOWED], + user_role=LitellmUserRoles.INTERNAL_USER.value, + **kwargs, + ) + + +@pytest.mark.asyncio +async def test_models_only_read_skips_tokenization_for_restricted_key(): + """Opt-out set + model allowlist: read the file, validate it, don't tokenize. + + ``async_pre_call_hook`` swallows unexpected exceptions and returns ``data``, + so the returned value alone proves nothing — assert on the calls that + separate the intended path from error recovery. + """ + rate_limiter = _make_rate_limiter() + data = {"input_file_id": "file-abc123"} + afile_content = AsyncMock(return_value=_models_only_file([_MODELS_ONLY_ALLOWED] * 2)) + + with ( + patch( # test-quality-ok: operator config is a module global read at call time; the hook exposes no injection point + "litellm.proxy.proxy_server.general_settings", + {"disable_batch_input_file_rate_limiting": True}, + ), + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: module global; None selects the no-router provider-resolution path + patch("litellm.afile_content", new=afile_content), # test-quality-ok: this is the file-read boundary the test fakes; nothing HTTP sits below it here + patch("litellm.proxy.hooks.batch_rate_limiter._count_entry_tokens") as count_entry_tokens, # test-quality-ok: the call under assertion -- the test's subject is whether it runs + patch.object(rate_limiter, "_enforce_batch_file_model_access", new=AsyncMock()) as enforce, + patch.object(rate_limiter, "_check_and_increment_batch_counters", new=AsyncMock()) as counters, + ): + result = await rate_limiter.async_pre_call_hook( + user_api_key_dict=_restricted_user(), + cache=MagicMock(), + data=data, + call_type="acreate_batch", + ) + + assert result == data + # Not the full-skip path: the file really was downloaded... + afile_content.assert_awaited_once() + # ...and every model in it was still checked against the allowlist. + enforce.assert_awaited_once() + assert set(enforce.await_args.kwargs["models"]) == {_MODELS_ONLY_ALLOWED} + # Nothing charges this batch, so the discarded tokenization is skipped. + count_entry_tokens.assert_not_called() + counters.assert_not_awaited() + assert "_batch_token_count" not in data + + +@pytest.mark.asyncio +async def test_models_only_read_still_rejects_unauthorized_model(): + """The opt-out must not turn into an authorization bypass.""" + rate_limiter = _make_rate_limiter() + afile_content = AsyncMock(return_value=_models_only_file([_MODELS_ONLY_DENIED])) + + async def _raise_unauthorized(**kwargs): + raise Exception( + f"Key not allowed to access model. This key only has access to " + f"models={kwargs['valid_token'].models}" + ) + + with ( + patch( # test-quality-ok: operator config is a module global read at call time; the hook exposes no injection point + "litellm.proxy.proxy_server.general_settings", + {"disable_batch_input_file_rate_limiting": True}, + ), + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: module global; None selects the no-router provider-resolution path + patch("litellm.afile_content", new=afile_content), # test-quality-ok: this is the file-read boundary the test fakes; nothing HTTP sits below it here + patch( # test-quality-ok: stands in for the authorization decision being simulated + "litellm.proxy.auth.auth_checks.can_key_call_model", + new=AsyncMock(side_effect=_raise_unauthorized), + ), + ): + with pytest.raises(HTTPException) as exc: + await rate_limiter.async_pre_call_hook( + user_api_key_dict=_restricted_user(), + cache=MagicMock(), + data={"input_file_id": "file-abc123"}, + call_type="acreate_batch", + ) + + assert exc.value.status_code == 403 + assert _MODELS_ONLY_DENIED in str(exc.value.detail) + afile_content.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_restricted_key_without_opt_out_still_counts_tokens(): + """No opt-out configured: the pre-existing path is untouched.""" + rate_limiter = _make_rate_limiter() + afile_content = AsyncMock(return_value=_models_only_file([_MODELS_ONLY_ALLOWED] * 2)) + data = {"input_file_id": "file-abc123"} + + with ( + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: operator config is a module global read at call time; the hook exposes no injection point + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: module global; None selects the no-router provider-resolution path + patch("litellm.afile_content", new=afile_content), # test-quality-ok: this is the file-read boundary the test fakes; nothing HTTP sits below it here + patch( # test-quality-ok: the call under assertion -- the test's subject is whether it runs + "litellm.proxy.hooks.batch_rate_limiter._count_entry_tokens", return_value=7 + ) as count_entry_tokens, + patch.object(rate_limiter, "_enforce_batch_file_model_access", new=AsyncMock()), + patch.object(rate_limiter, "_check_and_increment_batch_counters", new=AsyncMock()) as counters, + ): + result = await rate_limiter.async_pre_call_hook( + user_api_key_dict=_restricted_user(), + cache=MagicMock(), + data=data, + call_type="acreate_batch", + ) + + assert result == data + afile_content.assert_awaited_once() + assert count_entry_tokens.call_count == 2 + counters.assert_awaited_once() + assert data["_batch_token_count"] == 14 + assert data["_batch_request_count"] == 2 + + +@pytest.mark.asyncio +async def test_models_only_read_disabled_when_enqueued_scopes_apply(): + """Enqueued-token reservations are priced from the totals, so keep counting.""" + rate_limiter = _make_rate_limiter() + afile_content = AsyncMock(return_value=_models_only_file([_MODELS_ONLY_ALLOWED] * 2)) + + with ( + patch( # test-quality-ok: operator config is a module global read at call time; the hook exposes no injection point + "litellm.proxy.proxy_server.general_settings", + {"disable_batch_input_file_rate_limiting": True}, + ), + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: module global; None selects the no-router provider-resolution path + patch("litellm.afile_content", new=afile_content), # test-quality-ok: this is the file-read boundary the test fakes; nothing HTTP sits below it here + patch( # test-quality-ok: the call under assertion -- the test's subject is whether it runs + "litellm.proxy.hooks.batch_rate_limiter._count_entry_tokens", return_value=7 + ) as count_entry_tokens, + patch.object(rate_limiter, "_enforce_batch_file_model_access", new=AsyncMock()), + patch.object(rate_limiter, "_reserve_batch_enqueued_tokens", new=AsyncMock()) as reserve, + ): + data = {"input_file_id": "file-abc123"} + await rate_limiter.async_pre_call_hook( + user_api_key_dict=_restricted_user(metadata={"batch_enqueued_token_limit": 10}), + cache=MagicMock(), + data=data, + call_type="acreate_batch", + ) + + assert count_entry_tokens.call_count == 2 + reserve.assert_awaited_once() + assert data["_batch_token_count"] == 14 + + +def _models_only_file_with_malformed_row(model): + """One valid row plus a line that is not JSON at all.""" + import json as _json + + body = "\n".join( + [ + _json.dumps({"body": {"model": model, "messages": [{"role": "user", "content": "x" * 64}]}}), + "{ this is not json", + ] + ) + content = MagicMock() + content.content = body.encode("utf-8") + return content + + +@pytest.mark.asyncio +async def test_models_only_read_applies_to_skip_listed_provider(): + """The provider skip list reaches the models-only path too, not just the global flag.""" + rate_limiter = _make_rate_limiter() + data = {"input_file_id": "file-abc123", "model": "my-vllm-model"} + afile_content = AsyncMock(return_value=_models_only_file([_MODELS_ONLY_ALLOWED])) + + with ( + patch( # test-quality-ok: operator config is a module global read at call time; the hook exposes no injection point + "litellm.proxy.proxy_server.general_settings", + {"skip_batch_input_file_rate_limiting_for_providers": ["hosted_vllm"]}, + ), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), # test-quality-ok: module global; a router must exist for provider resolution + patch( # test-quality-ok: stands in for the deployment credentials the provider is resolved from + "litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model", + return_value={"custom_llm_provider": "hosted_vllm"}, + ), + patch("litellm.afile_content", new=afile_content), # test-quality-ok: this is the file-read boundary the test fakes; nothing HTTP sits below it here + patch("litellm.proxy.hooks.batch_rate_limiter._count_entry_tokens") as count_entry_tokens, # test-quality-ok: the call under assertion -- the test's subject is whether it runs + patch.object(rate_limiter, "_enforce_batch_file_model_access", new=AsyncMock()) as enforce, + ): + result = await rate_limiter.async_pre_call_hook( + user_api_key_dict=_restricted_user(), + cache=MagicMock(), + data=data, + call_type="acreate_batch", + ) + + assert result == data + # A restricted key still forces the download, so this is the models-only path + # rather than the full skip that an unrestricted key would take. + afile_content.assert_awaited_once() + enforce.assert_awaited_once() + count_entry_tokens.assert_not_called() + + +@pytest.mark.asyncio +async def test_models_only_read_skips_size_estimate_for_malformed_row(): + """A row that is not JSON names no model, so it needs no size estimate either.""" + rate_limiter = _make_rate_limiter() + afile_content = AsyncMock(return_value=_models_only_file_with_malformed_row(_MODELS_ONLY_ALLOWED)) + + with ( + patch( # test-quality-ok: operator config is a module global read at call time; the hook exposes no injection point + "litellm.proxy.proxy_server.general_settings", + {"disable_batch_input_file_rate_limiting": True}, + ), + patch("litellm.proxy.proxy_server.llm_router", None), # test-quality-ok: module global; None selects the no-router provider-resolution path + patch("litellm.afile_content", new=afile_content), # test-quality-ok: this is the file-read boundary the test fakes; nothing HTTP sits below it here + patch("litellm.proxy.hooks.batch_rate_limiter._count_entry_tokens") as count_entry_tokens, # test-quality-ok: the call under assertion -- the test's subject is whether it runs + patch("litellm.proxy.hooks.batch_rate_limiter._estimate_batch_entry_tokens") as estimate_tokens, # test-quality-ok: the malformed-row fallback under assertion + patch.object(rate_limiter, "_enforce_batch_file_model_access", new=AsyncMock()) as enforce, + ): + data = {"input_file_id": "file-abc123"} + result = await rate_limiter.async_pre_call_hook( + user_api_key_dict=_restricted_user(), + cache=MagicMock(), + data=data, + call_type="acreate_batch", + ) + + # What the caller observes: the batch is admitted and nothing is charged for + # it, malformed row included. + assert result == data + assert "_batch_token_count" not in data + assert "_batch_request_count" not in data + afile_content.assert_awaited_once() + # The malformed row cannot smuggle a model past the allowlist, and neither + # counter runs for it. + enforce.assert_awaited_once() + assert set(enforce.await_args.kwargs["models"]) == {_MODELS_ONLY_ALLOWED} + count_entry_tokens.assert_not_called() + estimate_tokens.assert_not_called()