diff --git a/litellm/constants.py b/litellm/constants.py index 78bfc6501e82..306b8513371d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1356,6 +1356,13 @@ 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 with +# general_settings.batch_input_file_read_timeout. +DEFAULT_BATCH_INPUT_FILE_READ_TIMEOUT_SECONDS = 10.0 + 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 23fe7af99942..d4e9a18cb30d 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2408,6 +2408,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: Optional[float] = 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: Optional[str] = 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 3477af362853..45eeb567cb67 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -33,10 +33,12 @@ from fastapi import HTTPException from pydantic import BaseModel +import asyncio import json 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, @@ -96,6 +98,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. @@ -292,6 +312,28 @@ def _warn_if_unsupported_model_skip_configured(self, general_settings: Dict) -> "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, @@ -334,6 +376,12 @@ def _resolve_batch_input_file_fetch_params( Model-embedded IDs (``file-``) are not unified managed-file IDs; without decoding them, ``afile_content`` is called with the encoded ID and the upstream provider returns 404. + + The returned kwargs may carry a ``timeout`` from the deployment's + credentials (one of ``_extract_file_access_credentials``' keys). Callers + that need their own deadline must set it on the result rather than passing + it alongside ``**fetch_kwargs``, which would raise TypeError for a + duplicate kwarg. """ from litellm.proxy.openai_files_endpoints.common_utils import ( decode_model_from_file_id, @@ -518,8 +566,11 @@ async def count_input_file_usage( # 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, ) @@ -529,13 +580,35 @@ async def count_input_file_usage( 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( + # Set on the resolved kwargs, not passed as a separate keyword: + # they may already carry the deployment's own `timeout`, and + # passing both raises TypeError for a duplicate kwarg. This + # read's budget wins on purpose; a deployment timeout is sized + # for serving traffic, not for a read that blocks POST /v1/batches. + fetch_kwargs["timeout"] = timeout_seconds + # `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, **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( @@ -604,6 +677,10 @@ async def count_input_file_usage( f"Batch input file rejected for {file_id}: status={e.status_code} detail={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(f"Error counting input file usage for {file_id}: {str(e)}") raise @@ -854,6 +931,38 @@ async def async_pre_call_hook( 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(f"Error in batch rate limiting: {str(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 1376bdbed38d..5c25b7f2a939 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 7ff1bc11d811..eb11cd2679b3 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 @@ 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,290 @@ async def _deny_restricted(model, **kwargs): ) 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" + + +@pytest.mark.asyncio +async def test_read_deadline_overrides_a_deployment_configured_timeout(): + """`timeout` is one of _extract_file_access_credentials' credential keys, so a + deployment's litellm_params can already put it in fetch_kwargs. Passing the + limiter's deadline as a separate keyword alongside **fetch_kwargs then raises + TypeError for a duplicate kwarg, turning the hang into a 500. The limiter's + budget must win: a deployment timeout is sized for serving traffic, not for an + inline pre-call hook.""" + 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=["*"]) + + # A deployment whose credentials carry their own `timeout`, which is what + # _extract_file_access_credentials merges into fetch_kwargs. Driven through the + # real resolver so the merge itself is under test, not a stubbed return value. + router = MagicMock( + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": { + "model": "gpt-4o-mini", + "api_key": "sk-deployment", + "timeout": 600, + "custom_llm_provider": "openai", + }, + } + ] + ) + + with ( + patch("litellm.afile_content", new=_capture), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model", + return_value={"api_key": "sk-deployment", "timeout": 600, "custom_llm_provider": "openai"}, + ), + patch("litellm.proxy.proxy_server.llm_router", router), + patch("litellm.proxy.proxy_server.general_settings", {"batch_input_file_read_timeout": 4.0}), + ): + usage = await rate_limiter.count_input_file_usage( + file_id="file-plain", + custom_llm_provider="openai", + user_api_key_dict=user, + data={"model": "gpt-4o-mini"}, + ) + + assert usage.request_count == 1 + assert captured.get("timeout") == 4.0, "deployment timeout must not override the limiter's deadline" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 9a301f1c4740..4715ae21ca99 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -22704,6 +22704,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