ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(batches): skip unnecessary batch input file reads (#29114)

Browse Source 
* fix(batches): skip unnecessary batch input file reads

Skip expensive pre-read of batch input files when no batch limits apply and model allowlist checks are not required, and decode model-embedded file IDs before file-content fetches to prevent upstream 404s.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(batch-rate-limiter): prevent user metadata flag from bypassing model allowlist

The skip_batch_input_file_rate_limiting flag in litellm_metadata is
user-controllable for batch requests (request-body metadata lands in
litellm_metadata via LITELLM_METADATA_ROUTES). Honoring it
unconditionally also skipped _enforce_batch_file_model_access, letting
a restricted key submit a JSONL referencing models outside its
allowlist. Only honor the metadata-based skip when the key has no
model allowlist to enforce.

Co-authored-by: Yassin Kortam <yassin@berri.ai>

* fix(batch_rate_limiter): enforce model access check before honoring skip paths

Admin-configured skips (disable_batch_input_file_rate_limiting,
skip_batch_input_file_rate_limiting_for_models/_for_providers) and the
no-applicable-rate-limits short-circuit previously bypassed
_enforce_batch_file_model_access. A key with a restricted model
allowlist could therefore submit a batch JSONL referencing models
outside its allowlist whenever any of these skip paths fired, and the
provider-skip path was attacker-controllable via the request body's
custom_llm_provider field. Hoist the model-access guard to the top so
restricted keys always have their JSONL validated regardless of which
skip would otherwise apply.

Co-authored-by: Yassin Kortam <yassin@berri.ai>

* fix(batch_rate_limiter): wildcard model bypass + fail-open embedded model creds

- _key_requires_batch_model_access_check: check '*' / all-proxy-models
  before access_group_ids so wildcard keys skip the JSONL download.
- _resolve_batch_input_file_fetch_params: wrap embedded-model
  get_credentials_for_model in try/except HTTPException, mirroring the
  request-model fallback path, and always decode the file id.

Co-authored-by: Yassin Kortam <yassin@berri.ai>

* perf(batch_rate_limiter): reuse rate-limit descriptors across skip check and counter increment

* test(batch_rate_limiter): cover skip-path and file-fetch helpers

Add unit tests for the batch rate limiter's new skip/routing helpers so
the diff's patch coverage no longer depends on the CircleCI batches job,
whose coverage upload is blocked when an unrelated Bedrock integration
test aborts the run. Covers _get_batch_routing_model, _matches_skip_list,
_key_requires_batch_model_access_check, _has_applicable_batch_rate_limits,
_should_skip_batch_input_file_processing, _resolve_batch_input_file_fetch_params,
the descriptor-reuse path of _check_and_increment_batch_counters, and the
non-bytes file content guard in count_input_file_usage.

* fix(batch_rate_limiter): resolve provider skip from trusted deployment creds

Resolve the batch provider from router deployment credentials instead of
the user-supplied custom_llm_provider request field, so an unrestricted
key cannot spoof a skip-listed provider to bypass batch rate limiting.

Strengthen the provider-skip test to assert the file download and
descriptor work were short-circuited, and add a test that a spoofed
provider still falls through to rate-limit evaluation.

* fix(batch_rate_limiter): guard model-embedded credential lookup on llm_router presence

* test(batch_rate_limiter): drive real no-skip fetch path and pin wildcard+access-group predicate

The spoofed-provider test configured empty descriptors, so the no-limits
shortcut skipped the file fetch and the assertion only proved the provider
allow-list did not short-circuit before descriptor evaluation. Give the key an
applicable rate limit so the only thing that can prevent the fetch is the
provider skip, then assert afile_content is awaited and the counters are
incremented; the spoofed custom_llm_provider must not skip processing.

Also cover the wildcard / all-proxy-models plus access_group_ids combination in
the model-access predicate so the wildcard-wins behavior is locked down.

* fix(batch_rate_limiter): drop client-controlled skip flag to close quota bypass

The litellm_metadata.skip_batch_input_file_rate_limiting flag was read
straight from the request body, so any caller whose key had unrestricted
model access could send it and skip the input-file download, token count,
and RPM/TPM reservation, bypassing their batch rate limits. Skip decisions
now derive only from server-controlled general_settings.

* fix(batch_rate_limiter): match per-model skip on file-bound model only

The per-model skip resolved its model from _get_batch_routing_model, which
prefers the client-supplied top-level model field. That field only selects
routing credentials; the models a batch actually runs are the body.model
entries in the input JSONL. An unrestricted key could therefore name a
skip-listed deployment at the top level while routing a different,
same-provider model through the file, skipping the download, token count
and rate-limit reservation to bypass batch RPM/TPM limits.

Match the per-model skip against the file-bound model only (model-embedded
file id or unified managed file target), which is fixed when the file is
created and reflects the model the batch runs. The provider skip keeps using
the routing model since an admin opting out of a whole provider already
accepts any of that provider's models.

* fix(batch_rate_limiter): drop forgeable per-model skip to close quota bypass

The per-model skip matched skip_batch_input_file_rate_limiting_for_models
against the model bound to the input file id. That model comes from
decode_model_from_file_id / the unified file id, both unsigned base64 the
caller fully controls, so a caller could re-encode an accessible provider
file id with a skip-listed model while the JSONL still routes rate-limited
body.model entries and bypass the batch RPM/TPM counters. The models a batch
actually runs are its JSONL body.model entries, which cannot be known without
reading the file, so no caller-influenced model identifier can safely gate a
skip.

Remove the per-model skip entirely. The provider skip stays because the
provider is resolved from trusted deployment credentials and the batch is
constrained to run on that provider; the global disable and
no-applicable-limits skips stay because they do not depend on caller input.

* fix(batch_rate_limiter): warn when no-op per-model skip key is configured

* test(batch_rate_limiter): patch llm_router so model-embedded credential-error test hits fallback

* 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.

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Yassin Kortam <yassin@berri.ai>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Sameer Kankute 2026-06-02 08:33:19 +05:30 committed by GitHub
parent 609e1e9763
commit c233cbbc2a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 1012 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

321
litellm/proxy/hooks/batch_rate_limiter.py
View File

@ -17,7 +17,7 @@ Quick summary:
- async_log_success_event() fires on GET /v1/batches/{id} (batch completion)
"""
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union
from fastapi import HTTPException
from pydantic import BaseModel
@ -25,12 +25,13 @@ from pydantic import BaseModel
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.batches.batch_utils import (
_extract_file_access_credentials,
_get_batch_job_input_file_usage,
_get_file_content_as_dictionary,
_get_models_from_batch_input_file_content,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -97,6 +98,276 @@ class _PROXY_BatchRateLimiter(CustomLogger):
"""
self.internal_usage_cache = internal_usage_cache
self.parallel_request_limiter = parallel_request_limiter
self._warned_unsupported_model_skip = False
def _get_file_bound_batch_model(self, data: Dict) -> Optional[str]:
"""Resolve the model bound to the batch input file ID.
``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:
return None
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
decode_model_from_file_id,
get_models_from_unified_file_id,
)
model_from_file_id = decode_model_from_file_id(input_file_id)
if model_from_file_id:
return model_from_file_id
unified_file_id = _is_base64_encoded_unified_file_id(input_file_id)
if unified_file_id:
target_model_names = get_models_from_unified_file_id(unified_file_id)
if target_model_names:
return target_model_names[0]
return None
def _get_batch_routing_model(self, data: Dict) -> Optional[str]:
"""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 None
def _resolve_batch_provider(self, batch_model: Optional[str]) -> Optional[str]:
"""Resolve the provider from the deployment that serves ``batch_model``.
The provider is read from trusted router credentials rather than the
user-supplied ``custom_llm_provider`` request field, so a caller cannot
spoof a skip-listed provider to bypass batch rate limiting.
"""
if not batch_model:
return None
from litellm.proxy.openai_files_endpoints.common_utils import (
get_credentials_for_model,
)
from litellm.proxy.proxy_server import llm_router
if llm_router is None:
return None
try:
credentials = get_credentials_for_model(
llm_router=llm_router,
model_id=batch_model,
operation_context="batch input file read (rate limiting)",
)
except HTTPException:
return None
provider = credentials.get("custom_llm_provider")
return provider if isinstance(provider, str) and provider else None
def _create_batch_rate_limit_descriptors(
self,
user_api_key_dict: UserAPIKeyAuth,
data: Dict,
) -> List["RateLimitDescriptor"]:
return self.parallel_request_limiter._create_rate_limit_descriptors(
user_api_key_dict=user_api_key_dict,
data=data,
rpm_limit_type=None,
tpm_limit_type=None,
model_has_failures=False,
)
def _should_skip_batch_input_file_processing(
self,
data: Dict,
user_api_key_dict: UserAPIKeyAuth,
) -> Tuple[bool, Optional[List["RateLimitDescriptor"]]]:
"""
Skip downloading batch input files when the operator disabled batch
input-file rate limiting, when the batch runs entirely on a skip-listed
provider, or when there is nothing to enforce (no applicable rate
limits).
A skip is only honored for keys with unrestricted model access. When
the key has a model allowlist, the JSONL must still be downloaded so
``_enforce_batch_file_model_access`` can validate every ``body.model``
entry, otherwise a restricted key could smuggle unauthorized models
into the file via an admin-configured skip.
The skip is never keyed on a specific model name. The models a batch
actually runs are its JSONL ``body.model`` entries, and any model
identifier the caller can influence (the top-level ``model`` or the
unsigned model embedded in a ``file-...`` id) can be pointed at a
skip-listed deployment while the file routes a different, rate-limited
model. The provider skip is safe because the provider is read from the
routing deployment's trusted credentials and the batch is constrained
to run on that provider.
Returns ``(should_skip, descriptors)`` where ``descriptors`` is the
rate-limit descriptor list computed for the no-limits check, so the
caller can reuse it for counter enforcement without recomputing.
"""
from litellm.proxy.proxy_server import general_settings
self._warn_if_unsupported_model_skip_configured(general_settings)
if self._key_requires_batch_model_access_check(user_api_key_dict):
return False, None
if general_settings.get("disable_batch_input_file_rate_limiting") is True:
return True, None
skip_providers = (
general_settings.get("skip_batch_input_file_rate_limiting_for_providers")
or []
)
if skip_providers:
batch_provider = self._resolve_batch_provider(
self._get_batch_routing_model(data)
)
if batch_provider and batch_provider in skip_providers:
verbose_proxy_logger.debug(
f"Skipping batch input file processing for provider={batch_provider}"
)
return True, None
descriptors = self._create_batch_rate_limit_descriptors(
user_api_key_dict=user_api_key_dict,
data=data,
)
if not self._has_applicable_batch_rate_limits(descriptors):
verbose_proxy_logger.debug(
"Skipping batch input file processing: no rate limits configured"
)
return True, None
return False, descriptors
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.
A per-model skip is intentionally not honored because the model a batch
runs on is caller-influenced and can be pointed at a skip-listed
deployment while the JSONL routes a different, rate-limited model.
"""
if self._warned_unsupported_model_skip:
return
if general_settings.get("skip_batch_input_file_rate_limiting_for_models"):
self._warned_unsupported_model_skip = True
verbose_proxy_logger.warning(
"general_settings.skip_batch_input_file_rate_limiting_for_models is not "
"supported and has no effect. Use "
"skip_batch_input_file_rate_limiting_for_providers or "
"disable_batch_input_file_rate_limiting instead."
)
@staticmethod
def _key_requires_batch_model_access_check(
user_api_key_dict: UserAPIKeyAuth,
) -> bool:
"""True when the key may only call a subset of models (JSONL must be checked)."""
models = user_api_key_dict.models or []
if "*" in models:
return False
if SpecialModelNames.all_proxy_models.value in models:
return False
if user_api_key_dict.access_group_ids:
return True
if not models:
return False
return True
@staticmethod
def _has_applicable_batch_rate_limits(
descriptors: List["RateLimitDescriptor"],
) -> bool:
for descriptor in descriptors:
rate_limit = descriptor.get("rate_limit") or {}
if (
rate_limit.get("requests_per_unit") is not None
or rate_limit.get("tokens_per_unit") is not None
or rate_limit.get("max_parallel_requests") is not None
):
return True
return False
def _resolve_batch_input_file_fetch_params(
self,
file_id: str,
custom_llm_provider: str,
data: Dict,
) -> Tuple[str, Dict[str, Any]]:
"""
Map proxy-facing file IDs to provider file IDs and credentials.
Model-embedded IDs (``file-<base64>``) are not unified managed-file IDs;
without decoding them, ``afile_content`` is called with the encoded ID
and the upstream provider returns 404.
"""
from litellm.proxy.openai_files_endpoints.common_utils import (
decode_model_from_file_id,
get_credentials_for_model,
get_original_file_id,
)
from litellm.proxy.proxy_server import llm_router
fetch_kwargs: Dict[str, Any] = {
"custom_llm_provider": custom_llm_provider,
}
model_from_file_id = decode_model_from_file_id(file_id)
if model_from_file_id:
if llm_router is not None:
try:
credentials = get_credentials_for_model(
llm_router=llm_router,
model_id=model_from_file_id,
operation_context="batch input file read (rate limiting)",
)
fetch_kwargs.update(_extract_file_access_credentials(credentials))
fetch_kwargs["model"] = model_from_file_id
provider = credentials.get("custom_llm_provider")
if provider:
fetch_kwargs["custom_llm_provider"] = provider
except HTTPException:
pass
return get_original_file_id(file_id), fetch_kwargs
request_model = data.get("model")
if isinstance(request_model, str) and request_model and llm_router is not None:
try:
credentials = get_credentials_for_model(
llm_router=llm_router,
model_id=request_model,
operation_context="batch input file read (rate limiting)",
)
fetch_kwargs.update(_extract_file_access_credentials(credentials))
fetch_kwargs["model"] = request_model
provider = credentials.get("custom_llm_provider")
if provider:
fetch_kwargs["custom_llm_provider"] = provider
except HTTPException:
pass
return file_id, fetch_kwargs
def _raise_rate_limit_error(
self,
@ -163,6 +434,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
user_api_key_dict: UserAPIKeyAuth,
data: Dict,
batch_usage: BatchFileUsage,
descriptors: Optional[List["RateLimitDescriptor"]] = None,
) -> None:
"""
Atomically check + increment rate-limit counters by the batch amounts.
@ -171,14 +443,15 @@ class _PROXY_BatchRateLimiter(CustomLogger):
case no counter is modified. Backed by `atomic_check_and_increment_by_n`
which uses a Redis Lua script when available (multi-process atomic) and
falls back to a per-process asyncio.Lock + in-memory operation.
``descriptors`` may be passed in by the pre-call hook to reuse the list
already computed when deciding whether to skip file processing.
"""
descriptors = self.parallel_request_limiter._create_rate_limit_descriptors(
user_api_key_dict=user_api_key_dict,
data=data,
rpm_limit_type=None,
tpm_limit_type=None,
model_has_failures=False,
)
if descriptors is None:
descriptors = self._create_batch_rate_limit_descriptors(
user_api_key_dict=user_api_key_dict,
data=data,
)
increment: Dict[Literal["requests", "tokens"], int] = {
"requests": batch_usage.request_count,
@ -211,6 +484,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
file_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
user_api_key_dict: Optional[UserAPIKeyAuth] = None,
data: Optional[Dict] = None,
) -> BatchFileUsage:
"""
Count number of requests and tokens in a batch input file.
@ -238,14 +512,27 @@ class _PROXY_BatchRateLimiter(CustomLogger):
user_api_key_dict=user_api_key_dict,
)
else:
provider_file_id, fetch_kwargs = (
self._resolve_batch_input_file_fetch_params(
file_id=file_id,
custom_llm_provider=custom_llm_provider,
data=data or {},
)
)
# For non-managed files, use the standard litellm.afile_content
file_content = await litellm.afile_content(
file_id=file_id,
custom_llm_provider=custom_llm_provider,
file_id=provider_file_id,
user_api_key_dict=user_api_key_dict,
**fetch_kwargs,
)
file_content_as_dict = _get_file_content_as_dictionary(file_content.content)
file_content_bytes = getattr(file_content, "content", None)
if not isinstance(file_content_bytes, bytes):
raise ValueError(
f"Expected bytes content from file retrieval for {file_id}, "
f"got {type(file_content_bytes)}"
)
file_content_as_dict = _get_file_content_as_dictionary(file_content_bytes)
# Validate every model named in the batch JSONL against the
# caller's per-key model allowlist. Without this, a caller
@ -441,6 +728,14 @@ class _PROXY_BatchRateLimiter(CustomLogger):
)
return data
should_skip, batch_rate_limit_descriptors = (
self._should_skip_batch_input_file_processing(
data=data, user_api_key_dict=user_api_key_dict
)
)
if should_skip:
return data
# Get custom_llm_provider for token counting
custom_llm_provider = data.get("custom_llm_provider", "openai")
@ -452,6 +747,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
file_id=input_file_id,
custom_llm_provider=custom_llm_provider,
user_api_key_dict=user_api_key_dict,
data=data,
)
verbose_proxy_logger.debug(
@ -469,6 +765,7 @@ class _PROXY_BatchRateLimiter(CustomLogger):
user_api_key_dict=user_api_key_dict,
data=data,
batch_usage=batch_usage,
descriptors=batch_rate_limit_descriptors,
)
verbose_proxy_logger.debug(

703
tests/test_litellm/proxy/hooks/test_batch_file_validation.py
View File

@ -259,6 +259,188 @@ async def test_pre_call_allows_authorized_model_in_batch_file():
)
@pytest.mark.asyncio
async def test_pre_call_skips_file_fetch_when_disabled_in_general_settings():
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
user = UserAPIKeyAuth(api_key="sk-ok", user_id="alice", models=["*"])
with patch(
"litellm.proxy.proxy_server.general_settings",
{"disable_batch_input_file_rate_limiting": True},
):
result = await rate_limiter.async_pre_call_hook(
user_api_key_dict=user,
cache=MagicMock(),
data={"input_file_id": "file-abc123"},
call_type="acreate_batch",
)
assert result == {"input_file_id": "file-abc123"}
rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.assert_not_called()
@pytest.mark.asyncio
async def test_pre_call_skips_file_fetch_for_configured_provider():
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
user = UserAPIKeyAuth(api_key="sk-ok", user_id="alice", models=["*"])
data = {"input_file_id": "file-abc123", "model": "my-vllm-model"}
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",
return_value={"custom_llm_provider": "hosted_vllm"},
),
patch("litellm.afile_content", new=AsyncMock()) as mock_afile_content,
):
result = await rate_limiter.async_pre_call_hook(
user_api_key_dict=user,
cache=MagicMock(),
data=data,
call_type="acreate_batch",
)
assert result == data
# A real skip must short-circuit before any file download or rate-limit
# work — assert the skip happened rather than the hook's error-recovery
# path (which also returns data unchanged).
mock_afile_content.assert_not_awaited()
rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.assert_not_called()
@pytest.mark.asyncio
async def test_pre_call_does_not_skip_for_spoofed_provider():
"""The provider skip is resolved from trusted deployment credentials, so a
user-supplied ``custom_llm_provider`` that is not backed by the routing
deployment must not trigger a skip: the input file must still be fetched
and the rate-limit counters incremented."""
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
# An applicable rate limit keeps the no-limits shortcut from firing, so the
# only thing that could prevent the fetch below is the provider skip. If the
# spoofed ``custom_llm_provider`` were honored, afile_content would never be
# awaited.
rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.return_value = [
{"rate_limit": {"requests_per_unit": 100}}
]
rate_limiter.parallel_request_limiter.atomic_check_and_increment_by_n = AsyncMock(
return_value={"overall_code": "OK", "statuses": []}
)
user = UserAPIKeyAuth(api_key="sk-ok", user_id="alice", models=["*"])
mock_router = MagicMock()
mock_router.model_list = []
mock_router.resolve_model_name_from_model_id.return_value = "my-openai-model"
mock_content = MagicMock()
mock_content.content = (
b'{"body": {"model": "my-openai-model", '
b'"messages": [{"role": "user", "content": "hi"}]}}\n'
)
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", mock_router),
patch(
"litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model",
return_value={"custom_llm_provider": "openai"},
),
patch(
"litellm.afile_content", new=AsyncMock(return_value=mock_content)
) as mock_afile_content,
):
await rate_limiter.async_pre_call_hook(
user_api_key_dict=user,
cache=MagicMock(),
data={
"input_file_id": "file-abc123",
"model": "my-openai-model",
"custom_llm_provider": "hosted_vllm",
},
call_type="acreate_batch",
)
# The spoofed provider did not short-circuit the skip decision: the file was
# fetched and the counters were incremented.
mock_afile_content.assert_awaited_once()
rate_limiter.parallel_request_limiter.atomic_check_and_increment_by_n.assert_awaited_once()
@pytest.mark.asyncio
async def test_count_input_file_usage_decodes_model_embedded_file_id():
import base64
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
original_file_id = "file-provider-xyz"
encoded_payload = (
base64.urlsafe_b64encode(
f"litellm:{original_file_id};model,my-vllm-batch".encode()
)
.decode()
.rstrip("=")
)
encoded_file_id = f"file-{encoded_payload}"
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
mock_content = MagicMock()
mock_content.content = b'{"custom_id": "1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "my-vllm-batch", "messages": [{"role": "user", "content": "hi"}]}}\n'
with (
patch(
"litellm.afile_content",
new=AsyncMock(return_value=mock_content),
) as mock_afile_content,
patch(
"litellm.proxy.proxy_server.llm_router",
MagicMock(),
),
patch(
"litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model",
return_value={
"api_key": "test-key",
"api_base": "http://vllm:8000/v1",
"custom_llm_provider": "hosted_vllm",
},
),
):
await rate_limiter.count_input_file_usage(
file_id=encoded_file_id,
custom_llm_provider="openai",
user_api_key_dict=UserAPIKeyAuth(api_key="sk-ok", user_id="alice"),
data={},
)
mock_afile_content.assert_awaited_once()
assert mock_afile_content.await_args.kwargs["file_id"] == original_file_id
assert mock_afile_content.await_args.kwargs["custom_llm_provider"] == "hosted_vllm"
@pytest.mark.asyncio
async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias():
"""After replace_model_in_jsonl, body.model is the provider id (e.g. gpt-5.5).
@ -323,3 +505,524 @@ async def test_pre_call_skips_check_when_no_models_present():
user_api_key_dict=user,
file_content_as_dict=[{"body": {}}],
)
# ---------------------------------------------------------------------------
# Skip-path helpers
# ---------------------------------------------------------------------------
def _make_rate_limiter():
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
return _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
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
assert rate_limiter._get_batch_routing_model({"input_file_id": ""}) is None
def test_get_batch_routing_model_decodes_model_embedded_file_id():
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}"})
== "vllm-batch"
)
def test_get_batch_routing_model_uses_unified_file_id_target():
rate_limiter = _make_rate_limiter()
with (
patch(
"litellm.proxy.openai_files_endpoints.common_utils.decode_model_from_file_id",
return_value=None,
),
patch(
"litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id",
return_value="unified-id",
),
patch(
"litellm.proxy.openai_files_endpoints.common_utils.get_models_from_unified_file_id",
return_value=["model-a", "model-b"],
),
):
assert (
rate_limiter._get_batch_routing_model({"input_file_id": "file-managed"})
== "model-a"
)
def test_key_requires_batch_model_access_check_branches():
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
check = _PROXY_BatchRateLimiter._key_requires_batch_model_access_check
assert check(UserAPIKeyAuth(api_key="sk", models=["*"])) is False
assert check(UserAPIKeyAuth(api_key="sk", models=["all-proxy-models"])) is False
assert (
check(UserAPIKeyAuth(api_key="sk", models=[], access_group_ids=["grp"])) is True
)
assert check(UserAPIKeyAuth(api_key="sk", models=[])) is False
assert check(UserAPIKeyAuth(api_key="sk", models=["gpt-4o-mini"])) is True
# Wildcard / all-proxy-models grant access to every model, so
# can_key_call_model passes any model regardless of access groups (which
# only ever widen access). Such keys must not be forced to download and
# validate the JSONL even when access_group_ids are also present.
assert (
check(UserAPIKeyAuth(api_key="sk", models=["*"], access_group_ids=["grp"]))
is False
)
assert (
check(
UserAPIKeyAuth(
api_key="sk", models=["all-proxy-models"], access_group_ids=["grp"]
)
)
is False
)
# A concrete model allowlist is still a subset even with access groups.
assert (
check(
UserAPIKeyAuth(
api_key="sk", models=["gpt-4o-mini"], access_group_ids=["grp"]
)
)
is True
)
def test_has_applicable_batch_rate_limits():
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
has_limits = _PROXY_BatchRateLimiter._has_applicable_batch_rate_limits
assert has_limits([{"rate_limit": {"tokens_per_unit": 100}}]) is True
assert has_limits([{"rate_limit": {"requests_per_unit": 5}}]) is True
assert has_limits([{"rate_limit": {"max_parallel_requests": 2}}]) is True
assert has_limits([{"rate_limit": {}}, {}]) is False
def test_should_skip_returns_false_when_key_needs_model_access_check():
rate_limiter = _make_rate_limiter()
user = UserAPIKeyAuth(api_key="sk", models=["gpt-4o-mini"])
should_skip, descriptors = rate_limiter._should_skip_batch_input_file_processing(
data={"input_file_id": "file-abc"}, user_api_key_dict=user
)
assert should_skip is False
assert descriptors is None
def test_should_skip_ignores_client_supplied_metadata_flag():
"""A caller must not be able to bypass batch rate limits by setting
``litellm_metadata.skip_batch_input_file_rate_limiting`` in the request
body. The skip decision is server-controlled only, so with applicable rate
limits the JSONL is still processed despite the client flag."""
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=["*"])
with patch("litellm.proxy.proxy_server.general_settings", {}):
should_skip, descriptors = (
rate_limiter._should_skip_batch_input_file_processing(
data={
"input_file_id": "file-abc",
"litellm_metadata": {"skip_batch_input_file_rate_limiting": True},
},
user_api_key_dict=user,
)
)
assert should_skip is False
def test_should_not_skip_for_forged_model_embedded_file_id():
"""A ``file-<base64>`` id embeds an unsigned model name the caller fully
controls, so a caller can re-encode any accessible provider file id with a
skip-listed model while the JSONL still routes rate-limited ``body.model``
entries. The per-model skip must therefore never fire: with applicable rate
limits, a forged skip-listed file-bound model still falls through to file
processing and counter enforcement."""
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-xyz;model,gpt-4o-mini")
.decode()
.rstrip("=")
)
with patch(
"litellm.proxy.proxy_server.general_settings",
{"skip_batch_input_file_rate_limiting_for_models": ["gpt-4o-mini"]},
):
should_skip, descriptors = (
rate_limiter._should_skip_batch_input_file_processing(
data={"input_file_id": f"file-{encoded}"},
user_api_key_dict=user,
)
)
assert should_skip is False
assert descriptors is not None
def test_should_not_skip_for_skip_listed_top_level_model():
"""A caller must not bypass batch rate limits by naming a skip-listed model
in the top-level ``model`` while routing a different model through the JSONL
``body.model`` entries. No per-model skip exists, so a skip-listed model over
a plain file still gets processed."""
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=["*"])
with patch(
"litellm.proxy.proxy_server.general_settings",
{"skip_batch_input_file_rate_limiting_for_models": ["gpt-4o-mini"]},
):
should_skip, descriptors = (
rate_limiter._should_skip_batch_input_file_processing(
data={"model": "gpt-4o-mini", "input_file_id": "file-abc"},
user_api_key_dict=user,
)
)
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."""
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=["*"])
with (
patch(
"litellm.proxy.proxy_server.general_settings",
{"skip_batch_input_file_rate_limiting_for_models": ["gpt-4o-mini"]},
),
patch(
"litellm.proxy.hooks.batch_rate_limiter.verbose_proxy_logger"
) as mock_logger,
):
for _ in range(3):
rate_limiter._should_skip_batch_input_file_processing(
data={"model": "gpt-4o-mini", "input_file_id": "file-abc"},
user_api_key_dict=user,
)
assert mock_logger.warning.call_count == 1
assert (
"skip_batch_input_file_rate_limiting_for_models"
in mock_logger.warning.call_args[0][0]
)
def test_no_warning_when_model_skip_setting_absent():
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=["*"])
with (
patch(
"litellm.proxy.proxy_server.general_settings",
{"skip_batch_input_file_rate_limiting_for_providers": ["openai"]},
),
patch(
"litellm.proxy.hooks.batch_rate_limiter.verbose_proxy_logger"
) as mock_logger,
):
rate_limiter._should_skip_batch_input_file_processing(
data={"model": "gpt-4o-mini", "input_file_id": "file-abc"},
user_api_key_dict=user,
)
mock_logger.warning.assert_not_called()
def test_should_skip_when_no_rate_limits_configured():
rate_limiter = _make_rate_limiter()
rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.return_value = [
{"rate_limit": {}}
]
user = UserAPIKeyAuth(api_key="sk", models=["*"])
with patch("litellm.proxy.proxy_server.general_settings", {}):
should_skip, descriptors = (
rate_limiter._should_skip_batch_input_file_processing(
data={"model": "gpt-4o-mini", "input_file_id": "file-abc"},
user_api_key_dict=user,
)
)
assert should_skip is True
assert descriptors is None
def test_should_not_skip_and_reuses_descriptors_when_limits_present():
rate_limiter = _make_rate_limiter()
descriptors = [{"rate_limit": {"tokens_per_unit": 100}}]
rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.return_value = (
descriptors
)
user = UserAPIKeyAuth(api_key="sk", models=["*"])
with patch("litellm.proxy.proxy_server.general_settings", {}):
should_skip, returned = rate_limiter._should_skip_batch_input_file_processing(
data={"model": "gpt-4o-mini", "input_file_id": "file-abc"},
user_api_key_dict=user,
)
assert should_skip is False
assert returned is descriptors
def test_resolve_fetch_params_uses_request_model_credentials():
rate_limiter = _make_rate_limiter()
with (
patch("litellm.proxy.proxy_server.llm_router", MagicMock()),
patch(
"litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model",
return_value={
"api_key": "k",
"api_base": "http://vllm:8000/v1",
"custom_llm_provider": "hosted_vllm",
},
),
):
provider_file_id, fetch_kwargs = (
rate_limiter._resolve_batch_input_file_fetch_params(
file_id="file-plain-openai",
custom_llm_provider="openai",
data={"model": "my-vllm-batch"},
)
)
assert provider_file_id == "file-plain-openai"
assert fetch_kwargs["model"] == "my-vllm-batch"
assert fetch_kwargs["custom_llm_provider"] == "hosted_vllm"
assert fetch_kwargs["api_base"] == "http://vllm:8000/v1"
def test_resolve_fetch_params_fails_open_on_credential_lookup_error():
rate_limiter = _make_rate_limiter()
with (
patch("litellm.proxy.proxy_server.llm_router", MagicMock()),
patch(
"litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model",
side_effect=HTTPException(status_code=404, detail="no creds"),
),
):
provider_file_id, fetch_kwargs = (
rate_limiter._resolve_batch_input_file_fetch_params(
file_id="file-plain-openai",
custom_llm_provider="openai",
data={"model": "my-vllm-batch"},
)
)
assert provider_file_id == "file-plain-openai"
assert fetch_kwargs == {"custom_llm_provider": "openai"}
def test_resolve_fetch_params_model_embedded_fails_open_on_credential_error():
import base64
rate_limiter = _make_rate_limiter()
encoded = (
base64.urlsafe_b64encode(b"litellm:file-orig;model,vllm-batch")
.decode()
.rstrip("=")
)
encoded_file_id = f"file-{encoded}"
get_credentials = MagicMock(
side_effect=HTTPException(status_code=404, detail="no creds")
)
with (
patch("litellm.proxy.proxy_server.llm_router", MagicMock()),
patch(
"litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model",
get_credentials,
),
):
provider_file_id, fetch_kwargs = (
rate_limiter._resolve_batch_input_file_fetch_params(
file_id=encoded_file_id,
custom_llm_provider="openai",
data={},
)
)
get_credentials.assert_called_once()
assert provider_file_id == "file-orig"
assert fetch_kwargs == {"custom_llm_provider": "openai"}
@pytest.mark.asyncio
async def test_check_and_increment_computes_descriptors_when_not_passed():
from litellm.proxy.hooks.batch_rate_limiter import (
BatchFileUsage,
_PROXY_BatchRateLimiter,
)
parallel_request_limiter = MagicMock()
parallel_request_limiter._create_rate_limit_descriptors.return_value = [
{"rate_limit": {"tokens_per_unit": 100}}
]
parallel_request_limiter.atomic_check_and_increment_by_n = AsyncMock(
return_value={"overall_code": "OK", "statuses": []}
)
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=parallel_request_limiter,
)
await rate_limiter._check_and_increment_batch_counters(
user_api_key_dict=UserAPIKeyAuth(api_key="sk", models=["*"]),
data={"model": "gpt-4o-mini"},
batch_usage=BatchFileUsage(total_tokens=10, request_count=1),
descriptors=None,
)
parallel_request_limiter._create_rate_limit_descriptors.assert_called_once()
@pytest.mark.asyncio
async def test_count_input_file_usage_raises_on_non_bytes_content():
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
bad_content = MagicMock()
bad_content.content = "not-bytes"
with patch("litellm.afile_content", new=AsyncMock(return_value=bad_content)):
with pytest.raises(ValueError, match="Expected bytes content"):
await rate_limiter.count_input_file_usage(
file_id="file-plain",
custom_llm_provider="openai",
user_api_key_dict=UserAPIKeyAuth(api_key="sk", models=["*"]),
data={},
)