diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 9b5630d81ad..6305b3ac068 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -13,7 +13,7 @@ https://platform.openai.com/docs/api-reference/batch import asyncio import contextvars import os -from collections.abc import Coroutine +from collections.abc import Coroutine, Mapping from functools import partial from typing import Any, Final, Literal, cast @@ -104,6 +104,24 @@ def _resolve_timeout( return float(timeout) +def _caller_timeout( + optional_params: GenericLiteLLMParams, + resolved_timeout: float | httpx.Timeout, +) -> float | httpx.Timeout | None: + """ + ``resolved_timeout`` when the caller asked for a timeout, else None. + + The resolution above falls back to 600s when nothing was set. Sending that on every + request would replace the HTTP client's own default (a 5s connect timeout, and an + explicitly configured ``litellm.request_timeout``), so None keeps that default. + """ + # ``request_timeout`` is the legacy spelling; GenericLiteLLMParams keeps it as an extra field + extra_params: Final[Mapping[str, object]] = optional_params.model_extra or {} + if optional_params.timeout or extra_params.get("request_timeout"): + return resolved_timeout + return None + + @client async def acreate_batch( completion_window: Literal["24h"], @@ -355,7 +373,7 @@ def create_batch( vertex_project=vertex_ai_project, vertex_location=vertex_ai_location, vertex_credentials=vertex_credentials, - timeout=timeout, + timeout=_caller_timeout(optional_params, timeout), max_retries=optional_params.max_retries, create_batch_data=_create_batch_request, custom_endpoint=optional_params.get("custom_endpoint"), @@ -519,7 +537,7 @@ def _handle_retrieve_batch_providers_without_provider_config( vertex_project=vertex_ai_project, vertex_location=vertex_ai_location, vertex_credentials=vertex_credentials, - timeout=timeout, + timeout=_caller_timeout(optional_params, timeout), max_retries=optional_params.max_retries, logging_obj=logging_obj, ) @@ -868,7 +886,7 @@ def list_batches( vertex_project=vertex_ai_project, vertex_location=vertex_ai_location, vertex_credentials=vertex_credentials, - timeout=timeout, + timeout=_caller_timeout(optional_params, timeout), max_retries=optional_params.max_retries, ) else: @@ -1066,7 +1084,7 @@ def cancel_batch( vertex_project=vertex_ai_project, vertex_location=vertex_ai_location, vertex_credentials=vertex_credentials, - timeout=timeout, + timeout=_caller_timeout(optional_params, timeout), max_retries=optional_params.max_retries, ) elif custom_llm_provider == "bedrock": diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index de5fe39cfc9..2dc0a20fbf7 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -93,7 +93,7 @@ class VertexAIBatchPrediction(VertexLLM): vertex_credentials: VERTEX_CREDENTIALS_TYPES | None, vertex_project: str | None, vertex_location: str | None, - timeout: float | httpx.Timeout, + timeout: float | httpx.Timeout | None, max_retries: int | None, custom_endpoint: bool | None = None, ) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]: @@ -132,6 +132,7 @@ class VertexAIBatchPrediction(VertexLLM): sync_handler=sync_handler, api_base=api_base, vertex_location=vertex_location or "us-central1", + timeout=timeout, ) default_api_base: Final = self.create_vertex_batch_url( @@ -163,12 +164,14 @@ class VertexAIBatchPrediction(VertexLLM): vertex_batch_request=vertex_batch_request, api_base=api_base, headers=headers, + timeout=timeout, ) response: Final = sync_handler.post( url=api_base, headers=headers, data=json.dumps(vertex_batch_request), + timeout=timeout, ) payload_view: Final[_VertexBatchPayloadView] = {"payload": response.json()} @@ -203,6 +206,7 @@ class VertexAIBatchPrediction(VertexLLM): sync_handler: HTTPHandler, api_base: str | None, vertex_location: str, + timeout: float | httpx.Timeout | None = None, ) -> VertexAIBatchPredictionJob: """ A fine-tuned Gemini deployment is configured by its endpoint id, but the v1 batch API only @@ -226,6 +230,7 @@ class VertexAIBatchPrediction(VertexLLM): sync_handler, endpoint_url, headers=headers, + timeout=timeout, ) } response: Final = fetched["response"] @@ -254,6 +259,7 @@ class VertexAIBatchPrediction(VertexLLM): vertex_batch_request: VertexAIBatchPredictionJob, api_base: str, headers: dict[str, str], + timeout: float | httpx.Timeout | None = None, ) -> LiteLLMBatch: client: Final = get_async_httpx_client( llm_provider=litellm.LlmProviders.VERTEX_AI, @@ -263,6 +269,7 @@ class VertexAIBatchPrediction(VertexLLM): url=api_base, headers=headers, data=json.dumps(vertex_batch_request), + timeout=timeout, ) except httpx.HTTPStatusError as e: error_body: Final = e.response.text @@ -298,7 +305,7 @@ class VertexAIBatchPrediction(VertexLLM): vertex_credentials: VERTEX_CREDENTIALS_TYPES | None, vertex_project: str | None, vertex_location: str | None, - timeout: float | httpx.Timeout, + timeout: float | httpx.Timeout | None, max_retries: int | None, logging_obj: "LiteLLMLoggingObj | None" = None, ) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]: @@ -348,6 +355,7 @@ class VertexAIBatchPrediction(VertexLLM): api_base=api_base, headers=headers, logging_obj=logging_obj, + timeout=timeout, ) # Log the request using logging_obj if available @@ -382,6 +390,7 @@ class VertexAIBatchPrediction(VertexLLM): sync_handler, api_base, headers=headers, + timeout=timeout, ) } response: Final = fetched["response"] @@ -401,6 +410,7 @@ class VertexAIBatchPrediction(VertexLLM): api_base: str, headers: dict[str, str], logging_obj: "LiteLLMLoggingObj | None" = None, + timeout: float | httpx.Timeout | None = None, ) -> LiteLLMBatch: client: Final = get_async_httpx_client( llm_provider=litellm.LlmProviders.VERTEX_AI, @@ -437,6 +447,7 @@ class VertexAIBatchPrediction(VertexLLM): client, api_base, headers=headers, + timeout=timeout, ) } response: Final = fetched["response"] @@ -459,7 +470,7 @@ class VertexAIBatchPrediction(VertexLLM): vertex_credentials: VERTEX_CREDENTIALS_TYPES | None, vertex_project: str | None, vertex_location: str | None, - timeout: float | httpx.Timeout, + timeout: float | httpx.Timeout | None, max_retries: int | None, ): sync_handler: Final = get_httpx_client() @@ -504,12 +515,14 @@ class VertexAIBatchPrediction(VertexLLM): api_base=api_base, headers=headers, params=params, + timeout=timeout, ) response: Final = sync_handler.get( url=api_base, headers=headers, params=params, + timeout=timeout, ) if response.status_code != 200: @@ -530,6 +543,7 @@ class VertexAIBatchPrediction(VertexLLM): api_base: str, headers: dict[str, str], params: dict[str, str], + timeout: float | httpx.Timeout | None = None, ): client: Final = get_async_httpx_client( llm_provider=litellm.LlmProviders.VERTEX_AI, @@ -538,6 +552,7 @@ class VertexAIBatchPrediction(VertexLLM): url=api_base, headers=headers, params=params, + timeout=timeout, ) if response.status_code != 200: raise VertexAIError( @@ -560,7 +575,7 @@ class VertexAIBatchPrediction(VertexLLM): vertex_credentials: VERTEX_CREDENTIALS_TYPES | None, vertex_project: str | None, vertex_location: str | None, - timeout: float | httpx.Timeout, + timeout: float | httpx.Timeout | None, max_retries: int | None, ) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]: access_token, project_id = self._ensure_access_token( @@ -626,10 +641,10 @@ class VertexAIBatchPrediction(VertexLLM): ) raise - # HTTPHandler.get() does not accept a timeout parameter retrieve_response: Final = sync_handler.get( url=retrieve_api_base, headers=headers, + timeout=timeout, ) if retrieve_response.status_code != 200: litellm.verbose_logger.error( @@ -652,7 +667,7 @@ class VertexAIBatchPrediction(VertexLLM): api_base: str, retrieve_api_base: str, headers: dict[str, str], - timeout: float | httpx.Timeout = 600.0, + timeout: float | httpx.Timeout | None = None, ) -> LiteLLMBatch: client: Final = get_async_httpx_client( llm_provider=litellm.LlmProviders.VERTEX_AI, @@ -672,10 +687,10 @@ class VertexAIBatchPrediction(VertexLLM): ) raise - # AsyncHTTPHandler.get() does not accept a timeout parameter retrieve_response: Final = await client.get( url=retrieve_api_base, headers=headers, + timeout=timeout, ) if retrieve_response.status_code != 200: litellm.verbose_logger.error( diff --git a/tests/unit/llms/vertex_ai/batches/test_handler.py b/tests/unit/llms/vertex_ai/batches/test_handler.py index 91a2adb9d19..78f72f7ac37 100644 --- a/tests/unit/llms/vertex_ai/batches/test_handler.py +++ b/tests/unit/llms/vertex_ai/batches/test_handler.py @@ -37,6 +37,7 @@ import pytest import litellm +from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS # noqa: E402 from litellm.llms.vertex_ai.batches.handler import ( # noqa: E402 VertexAIBatchPrediction, ) @@ -1082,3 +1083,130 @@ async def test_vertex_async_create_batch_logs_error_body_on_http_error(): assert exc_info.value.response.status_code == 400 assert "gemini-2.0-flash" in exc_info.value.response.text + + +VERTEX_BATCH_JOBS_URL = ( + f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}/locations/{LOCATION}/batchPredictionJobs" +) +SEVEN_SECONDS = {"connect": 7.0, "read": 7.0, "write": 7.0, "pool": 7.0} +# The cached client's own default: 600s, with a 5s connect timeout +CLIENT_DEFAULT = {"connect": 5.0, "read": 600.0, "write": 600.0, "pool": 600.0} + +BATCH_OPERATIONS = [ + pytest.param( + "create_batch", + {"completion_window": "24h", "endpoint": "/v1/chat/completions", **CREATE_DATA}, + [("POST", "/batchPredictionJobs")], + id="create", + ), + pytest.param( + "create_batch", + {"completion_window": "24h", "endpoint": "/v1/chat/completions", **ENDPOINT_CREATE_DATA}, + [("GET", f"/endpoints/{ENDPOINT_ID}"), ("POST", "/batchPredictionJobs")], + id="create_from_fine_tuned_endpoint", + ), + pytest.param( + "retrieve_batch", {"batch_id": BATCH_ID}, [("GET", f"/batchPredictionJobs/{BATCH_ID}")], id="retrieve" + ), + pytest.param("list_batches", {}, [("GET", "/batchPredictionJobs")], id="list"), + pytest.param( + "cancel_batch", + {"batch_id": BATCH_ID}, + [("POST", f"/batchPredictionJobs/{BATCH_ID}:cancel"), ("GET", f"/batchPredictionJobs/{BATCH_ID}")], + id="cancel", + ), +] + + +@pytest.fixture +def vertex_batch_api(respx_mock, monkeypatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + # respx answers for the Google host, so skip the DNS based SSRF rewrite to an IP + monkeypatch.setattr(litellm, "user_url_validation", False) + monkeypatch.setattr( + VertexAIBatchPrediction, + "_ensure_access_token", + lambda self, credentials, project_id, custom_llm_provider: (TOKEN, PROJECT), + ) + # Running, so the success logger has no output file to fetch + running = _vertex_job_response(state="JOB_STATE_RUNNING") + respx_mock.post(VERTEX_BATCH_JOBS_URL).mock(return_value=httpx.Response(200, json=running)) + respx_mock.post(f"{VERTEX_BATCH_JOBS_URL}/{BATCH_ID}:cancel").mock(return_value=httpx.Response(200, json={})) + respx_mock.get(f"{VERTEX_BATCH_JOBS_URL}/{BATCH_ID}").mock(return_value=httpx.Response(200, json=running)) + respx_mock.get( + f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}" + ).mock(return_value=httpx.Response(200, json={"deployedModels": [{"model": TUNED_MODEL_RESOURCE}]})) + respx_mock.get(VERTEX_BATCH_JOBS_URL).mock( + return_value=httpx.Response(200, json={"batchPredictionJobs": [running], "nextPageToken": ""}) + ) + return respx_mock + + +async def _call_batch_api(operation, kwargs, sync_mode): + call_kwargs = {**kwargs, "custom_llm_provider": "vertex_ai", "vertex_project": PROJECT, "vertex_location": LOCATION} + if sync_mode: + getattr(litellm, operation)(**call_kwargs) + else: + await getattr(litellm, f"a{operation}")(**call_kwargs) + + +def _sent_requests(router): + return [(call.request.method, call.request.url.path) for call in router.calls] + + +def _sent_timeouts(router): + return [call.request.extensions["timeout"] for call in router.calls] + + +@pytest.mark.parametrize("operation, kwargs, expected_requests", BATCH_OPERATIONS) +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +@pytest.mark.asyncio +@pytest.mark.respx(assert_all_called=False) +async def test_batch_requests_send_the_callers_timeout( + vertex_batch_api, operation, kwargs, expected_requests, sync_mode +): + await _call_batch_api(operation, {**kwargs, "timeout": 7.0}, sync_mode) + + assert _sent_requests(vertex_batch_api) == [ + (method, f"/v1/projects/{PROJECT}/locations/{LOCATION}{path}") for method, path in expected_requests + ] + assert _sent_timeouts(vertex_batch_api) == [SEVEN_SECONDS] * len(expected_requests) + + +@pytest.mark.parametrize("operation, kwargs, expected_requests", BATCH_OPERATIONS) +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +@pytest.mark.asyncio +@pytest.mark.respx(assert_all_called=False) +async def test_batch_requests_without_a_timeout_keep_the_client_default( + vertex_batch_api, monkeypatch, operation, kwargs, expected_requests, sync_mode +): + # No global timeout configured, so a REQUEST_TIMEOUT in the environment cannot move the default + monkeypatch.setattr(litellm, "request_timeout", DEFAULT_REQUEST_TIMEOUT_SECONDS) + monkeypatch.setattr(litellm, "request_timeout_explicitly_set", False) + + await _call_batch_api(operation, kwargs, sync_mode) + + assert len(_sent_requests(vertex_batch_api)) == len(expected_requests) + assert _sent_timeouts(vertex_batch_api) == [CLIENT_DEFAULT] * len(expected_requests) + + +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +@pytest.mark.asyncio +@pytest.mark.respx(assert_all_called=False) +async def test_batch_requests_without_a_timeout_use_the_configured_request_timeout( + vertex_batch_api, monkeypatch, sync_mode +): + monkeypatch.setattr(litellm, "request_timeout", 30.0) + + await _call_batch_api("retrieve_batch", {"batch_id": BATCH_ID}, sync_mode) + + assert _sent_timeouts(vertex_batch_api) == [{"connect": 5.0, "read": 30.0, "write": 30.0, "pool": 30.0}] + + +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +@pytest.mark.asyncio +@pytest.mark.respx(assert_all_called=False) +async def test_batch_requests_send_the_legacy_request_timeout(vertex_batch_api, sync_mode): + await _call_batch_api("retrieve_batch", {"batch_id": BATCH_ID, "request_timeout": 7.0}, sync_mode) + + assert _sent_timeouts(vertex_batch_api) == [SEVEN_SECONDS]