From f1c4a40faba969c2320b6149f039cd4fe2254e7c Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Fri, 31 Jul 2026 13:09:44 -0700 Subject: [PATCH] fix(batches): bound the batch rate limiter's input-file read POST /v1/batches could hang indefinitely. BatchRateLimiter.async_pre_call_hook runs inline in the request path and, for keys with applicable rpm/tpm limits, read the input file to count tokens with no deadline. With none set the OpenAI SDK default applied (600s, max_retries=2), so a slow or stalled Files API held the request open far past any client timeout; 63.6s was observed on stage against a 60s client read timeout. The read does double duty: it counts tokens for rate limiting, and it validates every body.model in the JSONL against the caller's allowlist. Those have opposite safe defaults, so the timeout policy splits on whether the key needs that check. A key restricted to a subset of models is rejected, because admitting it unchecked grants exactly the bypass _should_skip_batch_input_file_processing refuses to allow via operator config. A key with unrestricted access is admitted unmetered, matching the existing fail-open, so a degraded Files API does not become an outage. The deadline is passed to afile_content as well as to wait_for. afile_content runs the sync client via run_in_executor, and cancelling that await does not interrupt a thread already in the pool, so bounding only the await would leak the worker until the SDK's own timeout fired. Also unskips the e2e test that guards the LIT-3266 unattributed-spend-row regression, which was blocked on this hang. Defaults to 10s; override with general_settings.batch_input_file_read_timeout. --- litellm/constants.py | 7 + litellm/proxy/_types.py | 4 + litellm/proxy/hooks/batch_rate_limiter.py | 105 +++++++- tests/e2e/batches/test_batches_e2e.py | 10 - .../proxy/hooks/test_batch_file_validation.py | 237 +++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 5 + 6 files changed, 354 insertions(+), 14 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 164f5a77a76..4c2d14dae02 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1376,6 +1376,13 @@ BASE_MCP_ROUTE = "/mcp" BATCH_STATUS_POLL_INTERVAL_SECONDS = int(os.getenv("BATCH_STATUS_POLL_INTERVAL_SECONDS", 3600)) # 1 hour BATCH_STATUS_POLL_MAX_ATTEMPTS = int(os.getenv("BATCH_STATUS_POLL_MAX_ATTEMPTS", 24)) # for 24 hours +# Deadline for the batch rate limiter's input-file read. The read happens inline +# in POST /v1/batches, so it must resolve well within a client's read timeout; +# unbounded, the OpenAI SDK default (600s, max_retries=2) applies and a stalled +# Files API holds the request open indefinitely. Override per-deployment with +# general_settings.batch_input_file_read_timeout. +DEFAULT_BATCH_INPUT_FILE_READ_TIMEOUT_SECONDS = float(os.getenv("BATCH_INPUT_FILE_READ_TIMEOUT_SECONDS", 10)) + HEALTH_CHECK_TIMEOUT_SECONDS = int(os.getenv("HEALTH_CHECK_TIMEOUT_SECONDS", 60)) # 60 seconds _background_health_check_max_tokens_env = os.getenv("BACKGROUND_HEALTH_CHECK_MAX_TOKENS") try: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index fd21c5b8334..67faed2e2ea 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2417,6 +2417,10 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): None, description="By default, the user calling /team/new is automatically added to the new team as a team admin. If True, proxy admins are no longer auto-added; members explicitly listed in members_with_roles are unaffected. Default is False.", ) + batch_input_file_read_timeout: float | None = Field( + None, + description="Seconds the batch rate limiter may spend reading a batch input file to count tokens (default 10). The read runs inline in POST /v1/batches, so this must stay well inside client read timeouts. On timeout, keys whose model allowlist must be validated against the file are rejected; keys with unrestricted model access are admitted without rate limiting.", + ) maximum_spend_logs_retention_period: str | None = Field( None, description="Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted.", diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index afaf0ebf392..f628123e311 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -17,6 +17,7 @@ Quick summary: - async_log_success_event() fires on GET /v1/batches/{id} (batch completion) """ +import asyncio import json from collections.abc import Iterable from typing import ( @@ -32,6 +33,7 @@ from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import DEFAULT_BATCH_INPUT_FILE_READ_TIMEOUT_SECONDS from litellm.batches.batch_utils import ( _count_entry_tokens, _estimate_batch_entry_tokens, @@ -91,6 +93,24 @@ class BatchFileUsage(BaseModel): request_count: int +class BatchInputFileReadTimeout(Exception): + """The batch input-file read exceeded its deadline. + + Distinct from a generic failure because the read serves two purposes and the + two have opposite safe defaults: it counts tokens for rate limiting (where + admitting the batch unmetered is tolerable) and it validates every + ``body.model`` in the JSONL against the caller's allowlist (where admitting + the batch unchecked is a privilege escalation). Carrying its own type lets + ``async_pre_call_hook`` fail closed only for keys that need the allowlist + check, instead of the blanket fail-open its generic handler applies. + """ + + def __init__(self, file_id: str, timeout_seconds: float) -> None: + self.file_id = file_id + self.timeout_seconds = timeout_seconds + super().__init__(f"Timed out after {timeout_seconds}s reading batch input file {file_id}") + + class _PROXY_BatchRateLimiter(CustomLogger): """ Rate limiter for batch API requests. @@ -287,6 +307,28 @@ class _PROXY_BatchRateLimiter(CustomLogger): "disable_batch_input_file_rate_limiting instead." ) + @staticmethod + def _batch_input_file_read_timeout() -> float: + """Seconds the input-file read may take before it is abandoned. + + Falls back to the default when the operator's value is missing or not a + positive number: a zero/negative deadline would make wait_for expire + immediately and reject every batch from a restricted key. + """ + from litellm.proxy.proxy_server import general_settings + + configured = general_settings.get("batch_input_file_read_timeout") + if isinstance(configured, bool) or not isinstance(configured, (int, float)): + return DEFAULT_BATCH_INPUT_FILE_READ_TIMEOUT_SECONDS + if configured <= 0: + verbose_proxy_logger.warning( + "Ignoring general_settings.batch_input_file_read_timeout=%s: must be > 0. Using %ss.", + configured, + DEFAULT_BATCH_INPUT_FILE_READ_TIMEOUT_SECONDS, + ) + return DEFAULT_BATCH_INPUT_FILE_READ_TIMEOUT_SECONDS + return float(configured) + @staticmethod def _key_requires_batch_model_access_check( user_api_key_dict: UserAPIKeyAuth, @@ -513,8 +555,11 @@ class _PROXY_BatchRateLimiter(CustomLogger): # For managed files the unified file id encodes the proxy model # alias(es) the file was uploaded for; auth validates against those. target_model_names = get_models_from_unified_file_id(is_managed_file) if is_managed_file else [] + # Resolved before the coroutine is built so a failure here can never + # leave an un-awaited coroutine behind. + timeout_seconds = self._batch_input_file_read_timeout() if is_managed_file and user_api_key_dict is not None: - file_content = await self._fetch_managed_file_content( + fetch = self._fetch_managed_file_content( file_id=file_id, user_api_key_dict=user_api_key_dict, ) @@ -524,13 +569,31 @@ class _PROXY_BatchRateLimiter(CustomLogger): 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( + # For non-managed files, use the standard litellm.afile_content. + # `timeout` reaches file_content's TIMEOUT LOGIC via + # GenericLiteLLMParams, capping the upstream HTTP request itself + # rather than only the await, so an abandoned read stops + # occupying its executor thread too. + fetch = litellm.afile_content( file_id=provider_file_id, user_api_key_dict=user_api_key_dict, + timeout=timeout_seconds, **fetch_kwargs, ) + # Bound the read: it runs inline in POST /v1/batches, so unbounded the + # SDK default (600s x 3 attempts) outlives every client timeout and the + # request just hangs (LIT-5027). + # + # wait_for is what guarantees the handler stops waiting, and it is the + # only bound the managed-files path has (that hook takes no timeout + # argument), so an abandoned managed read may hold its executor thread + # until the SDK's own timeout fires. + try: + file_content = await asyncio.wait_for(fetch, timeout=timeout_seconds) + except asyncio.TimeoutError as exc: + raise BatchInputFileReadTimeout(file_id=file_id, timeout_seconds=timeout_seconds) from exc + file_content_bytes = getattr(file_content, "content", None) if not isinstance(file_content_bytes, bytes): raise ValueError( @@ -599,6 +662,10 @@ class _PROXY_BatchRateLimiter(CustomLogger): "Batch input file rejected for %s: status=%s detail=%s", file_id, e.status_code, e.detail ) raise + except BatchInputFileReadTimeout: + # The caller decides the policy (reject vs admit unmetered) and logs + # accordingly; a generic error line here would just duplicate it. + raise except Exception as e: verbose_proxy_logger.error("Error counting input file usage for %s: %s", file_id, e) raise @@ -845,6 +912,38 @@ class _PROXY_BatchRateLimiter(CustomLogger): except HTTPException: # Re-raise HTTP exceptions (rate limit exceeded) raise + except BatchInputFileReadTimeout as e: + # The read is both the token count and the JSONL model-allowlist + # check, so the two cases diverge. A key restricted to a subset of + # models cannot be admitted without validating the file: doing so + # would grant exactly the bypass _should_skip_batch_input_file_processing + # refuses to allow via operator config. An unrestricted key has only + # rate-limit accuracy at stake, so it is admitted unmetered, matching + # the generic fail-open below. + if self._key_requires_batch_model_access_check(user_api_key_dict): + verbose_proxy_logger.error( + "Rejecting batch: could not read input file %s within %ss to validate " + "the models it references against the key's allowlist.", + e.file_id, + e.timeout_seconds, + ) + raise ProxyException( + message=( + f"Could not read the batch input file within {e.timeout_seconds}s to " + "validate the models it references. Retry, or contact your proxy admin " + "if the files API is degraded." + ), + type=ProxyErrorTypes.internal_server_error, + param="input_file_id", + code=504, + ) from e + verbose_proxy_logger.warning( + "Batch admitted without rate limiting: reading input file %s timed out after %ss. " + "Its tokens and requests are not counted against this key's limits.", + e.file_id, + e.timeout_seconds, + ) + return data except Exception as e: verbose_proxy_logger.error("Error in batch rate limiting: %s", e, exc_info=True) # Don't block the request if rate limiting fails diff --git a/tests/e2e/batches/test_batches_e2e.py b/tests/e2e/batches/test_batches_e2e.py index 1376bdbed38..5c25b7f2a93 100644 --- a/tests/e2e/batches/test_batches_e2e.py +++ b/tests/e2e/batches/test_batches_e2e.py @@ -397,16 +397,6 @@ def unattributed_rows(rows: list[SpendLogRow]) -> list[SpendLogRow]: return [row for row in rows if not row.api_key] -@pytest.mark.skip( - reason=( - "LIT-5027: the path under test hangs. The batch rate limiter reads the input file " - "to count tokens by awaiting litellm.afile_content with no timeout, so a slow Files " - "API holds POST /v1/batches open past any client deadline (63.6s observed on stage " - "against a 60s read timeout). The unattributed-spend-row contract below is never " - "reached, so the test reports a timeout rather than the behavior it guards. Unskip " - "once the fetch is bounded." - ) -) def test_rate_limited_batch_create_leaves_no_unattributed_spend_row( client: BatchClient, resources: ResourceManager, batch_deployments: None ) -> None: 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 7ff1bc11d81..bbd17b8d46d 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -7,12 +7,14 @@ VERIA-39 regression tests: models the caller is not authorized to use. """ +import asyncio +import time from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi import HTTPException -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth def _models(file_content_as_dict): @@ -1671,3 +1673,236 @@ async def test_count_input_file_usage_collects_models_after_malformed_line(): ) assert exc.value.status_code == 403 + + +# --------------------------------------------------------------------------- +# LIT-5027: the input-file read must be bounded +# +# The read runs inline in POST /v1/batches and served two purposes with opposite +# safe defaults: counting tokens for rate limiting, and validating every +# body.model in the JSONL against the caller's allowlist. Unbounded, a stalled +# Files API held the request open past any client deadline (63.6s observed on +# stage against a 60s read timeout). +# +# These use a genuinely slow fetch against a short deadline rather than faking +# asyncio.TimeoutError, so they fail if wait_for is removed and the await goes +# back to being unbounded. +# --------------------------------------------------------------------------- + + +def _slow_fetch(delay: float = 10.0): + """A file read that outlives the test-scale deadline (0.05s) by 200x. + + Paired with `@pytest.mark.timeout` on each test so that removing the bound + surfaces as a fast failure rather than a hung CI job: unbounded, the await + runs the full `delay` and pytest-timeout kills it well before that. + """ + + async def _fetch(*args, **kwargs): + await asyncio.sleep(delay) + raise AssertionError("slow fetch should have been abandoned, not awaited to completion") + + return _fetch + + +@pytest.mark.timeout(30) +@pytest.mark.asyncio +async def test_input_file_read_is_abandoned_at_the_deadline(): + """The read must not outlive its budget. Without the bound this awaits the + full 30s sleep (in prod, the SDK's 600s x 3) instead of giving up.""" + from litellm.proxy.hooks.batch_rate_limiter import BatchInputFileReadTimeout + + rate_limiter = _make_rate_limiter() + user = UserAPIKeyAuth(api_key="sk", models=["*"]) + + with ( + patch("litellm.afile_content", new=_slow_fetch()), + patch("litellm.proxy.proxy_server.general_settings", {"batch_input_file_read_timeout": 0.05}), + patch("litellm.proxy.proxy_server.llm_router", None), + ): + started = time.monotonic() + with pytest.raises(BatchInputFileReadTimeout) as exc: + await rate_limiter.count_input_file_usage( + file_id="file-slow", + custom_llm_provider="openai", + user_api_key_dict=user, + ) + elapsed = time.monotonic() - started + + assert elapsed < 5, f"read was not abandoned at its deadline (took {elapsed:.2f}s)" + assert exc.value.file_id == "file-slow" + assert exc.value.timeout_seconds == 0.05 + + +@pytest.mark.timeout(30) +@pytest.mark.asyncio +async def test_managed_file_read_is_also_bounded(): + """Managed files take a different code path (the managed-files hook, which + accepts no timeout kwarg), so it needs the same bound.""" + from litellm.proxy.hooks.batch_rate_limiter import BatchInputFileReadTimeout + + rate_limiter = _make_rate_limiter() + user = UserAPIKeyAuth(api_key="sk", models=["*"]) + + with ( + patch.object(rate_limiter, "_fetch_managed_file_content", new=_slow_fetch()), + patch( + "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id", + return_value="litellm_proxy/gpt-4o-mini", + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_models_from_unified_file_id", + return_value=["gpt-4o-mini"], + ), + patch("litellm.proxy.proxy_server.general_settings", {"batch_input_file_read_timeout": 0.05}), + ): + started = time.monotonic() + with pytest.raises(BatchInputFileReadTimeout): + await rate_limiter.count_input_file_usage( + file_id="file-managed-slow", + custom_llm_provider="openai", + user_api_key_dict=user, + ) + elapsed = time.monotonic() - started + + assert elapsed < 5, f"managed-file read was not abandoned at its deadline (took {elapsed:.2f}s)" + + +@pytest.mark.timeout(30) +@pytest.mark.asyncio +async def test_timeout_rejects_batch_when_key_has_a_model_allowlist(): + """A restricted key's batch cannot be admitted on timeout: the file read is + the only thing that validates the models inside the JSONL, so admitting it + unchecked would let the caller run models outside its allowlist under the + proxy's shared credentials.""" + 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-restricted", models=["gpt-4o-mini"]) + + with ( + patch("litellm.afile_content", new=_slow_fetch()), + patch("litellm.proxy.proxy_server.general_settings", {"batch_input_file_read_timeout": 0.05}), + patch("litellm.proxy.proxy_server.llm_router", None), + ): + with pytest.raises(ProxyException) as exc: + await rate_limiter.async_pre_call_hook( + user_api_key_dict=user, + cache=MagicMock(), + data={"input_file_id": "file-slow", "model": "gpt-4o-mini"}, + call_type="acreate_batch", + ) + + assert exc.value.code == "504" + assert "validate the models" in exc.value.message + + +@pytest.mark.timeout(30) +@pytest.mark.asyncio +async def test_timeout_admits_batch_when_key_has_unrestricted_model_access(): + """With no allowlist to enforce, only rate-limit accuracy is at stake, so a + degraded Files API must not turn batch creation into an outage.""" + rate_limiter = _make_rate_limiter() + rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.return_value = [ + {"rate_limit": {"requests_per_unit": 5}} + ] + rate_limiter._check_and_increment_batch_counters = AsyncMock() + user = UserAPIKeyAuth(api_key="sk-open", models=["*"]) + data = {"input_file_id": "file-slow", "model": "gpt-4o-mini"} + + with ( + patch("litellm.afile_content", new=_slow_fetch()), + patch("litellm.proxy.proxy_server.general_settings", {"batch_input_file_read_timeout": 0.05}), + patch("litellm.proxy.proxy_server.llm_router", None), + ): + result = await rate_limiter.async_pre_call_hook( + user_api_key_dict=user, + cache=MagicMock(), + data=data, + call_type="acreate_batch", + ) + + assert result is data + # Admitted unmetered: no counters were reserved, and no token count was + # stamped for the completion-side reconciliation to read. + rate_limiter._check_and_increment_batch_counters.assert_not_awaited() + assert "_batch_token_count" not in data + + +@pytest.mark.timeout(30) +@pytest.mark.asyncio +async def test_access_group_key_is_rejected_on_timeout(): + """Access-group keys carry no literal model list but still require the JSONL + check (_key_requires_batch_model_access_check returns True), so they must + fail closed alongside allowlisted keys.""" + 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-group", models=[], access_group_ids=["grp-1"]) + + with ( + patch("litellm.afile_content", new=_slow_fetch()), + patch("litellm.proxy.proxy_server.general_settings", {"batch_input_file_read_timeout": 0.05}), + patch("litellm.proxy.proxy_server.llm_router", None), + ): + with pytest.raises(ProxyException) as exc: + await rate_limiter.async_pre_call_hook( + user_api_key_dict=user, + cache=MagicMock(), + data={"input_file_id": "file-slow", "model": "gpt-4o-mini"}, + call_type="acreate_batch", + ) + + assert exc.value.code == "504" + + +def test_read_timeout_defaults_and_rejects_unusable_values(): + """A zero/negative or non-numeric deadline must fall back to the default; a + 0s budget would expire instantly and reject every restricted key's batch.""" + from litellm.constants import DEFAULT_BATCH_INPUT_FILE_READ_TIMEOUT_SECONDS + + rate_limiter = _make_rate_limiter() + resolve = rate_limiter._batch_input_file_read_timeout + + for settings, expected in ( + ({}, DEFAULT_BATCH_INPUT_FILE_READ_TIMEOUT_SECONDS), + ({"batch_input_file_read_timeout": 0}, DEFAULT_BATCH_INPUT_FILE_READ_TIMEOUT_SECONDS), + ({"batch_input_file_read_timeout": -5}, DEFAULT_BATCH_INPUT_FILE_READ_TIMEOUT_SECONDS), + ({"batch_input_file_read_timeout": "20"}, DEFAULT_BATCH_INPUT_FILE_READ_TIMEOUT_SECONDS), + ({"batch_input_file_read_timeout": True}, DEFAULT_BATCH_INPUT_FILE_READ_TIMEOUT_SECONDS), + ({"batch_input_file_read_timeout": 45}, 45.0), + ({"batch_input_file_read_timeout": 2.5}, 2.5), + ): + with patch("litellm.proxy.proxy_server.general_settings", settings): + assert resolve() == expected, f"unexpected deadline for {settings}" + + +@pytest.mark.asyncio +async def test_read_deadline_is_passed_to_the_upstream_file_fetch(): + """wait_for alone only abandons the await; afile_content runs the sync client + in an executor thread that a cancelled await does not interrupt. The same + deadline must therefore reach afile_content so the HTTP request is capped and + the thread is released.""" + captured: dict = {} + + async def _capture(*args, **kwargs): + captured.update(kwargs) + return MagicMock(content=b'{"body": {"model": "gpt-4o-mini", "messages": []}}\n') + + rate_limiter = _make_rate_limiter() + user = UserAPIKeyAuth(api_key="sk", models=["*"]) + + with ( + patch("litellm.afile_content", new=_capture), + patch("litellm.proxy.proxy_server.general_settings", {"batch_input_file_read_timeout": 3.5}), + patch("litellm.proxy.proxy_server.llm_router", None), + ): + await rate_limiter.count_input_file_usage( + file_id="file-plain", + custom_llm_provider="openai", + user_api_key_dict=user, + ) + + assert captured.get("timeout") == 3.5, "read deadline never reached the upstream fetch" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 9133bfb5cf4..e20cf9630a2 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -22865,6 +22865,11 @@ export interface components { * @description run health checks in background */ background_health_checks?: boolean | null; + /** + * Batch Input File Read Timeout + * @description Seconds the batch rate limiter may spend reading a batch input file to count tokens (default 10). The read runs inline in POST /v1/batches, so this must stay well inside client read timeouts. On timeout, keys whose model allowlist must be validated against the file are rejected; keys with unrestricted model access are admitted without rate limiting. + */ + batch_input_file_read_timeout?: number | null; /** * Cancel On Disconnect * @description cancel the in-flight upstream LLM request (non-streaming) when the client disconnects, freeing backend capacity (e.g. a vLLM GPU slot); the request is logged as a 499 failure