Merge 239b7d60f4 into 8699998c9e

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
             )

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()