From a655b3960dcbcd7ceb9efb05d206387bd395c54b Mon Sep 17 00:00:00 2001 From: Abhishek Acharya Date: Mon, 19 Oct 2026 23:44:39 +1100 Subject: [PATCH] fix(vertex_ai): honor timeout on batch create, retrieve, list and cancel The Vertex AI batch handler accepted a timeout but dropped it on the create POST, the fine tuned endpoint lookup, the retrieve and list GETs and the GET that follows a cancel, sync and async. Those requests always ran with the client default instead of the caller's timeout. batches/main.py resolves a missing timeout to 600s. Sending that would replace the client default (5s connect, plus an explicitly configured litellm.request_timeout), so the Vertex branches now pass the timeout only when the caller set one and None otherwise. This also stops the cancel POST from forcing 600s on calls that set no timeout. The comments saying HTTPHandler.get and AsyncHTTPHandler.get take no timeout were out of date, so they are removed. --- litellm/batches/main.py | 28 +++- litellm/llms/vertex_ai/batches/handler.py | 29 +++- .../llms/vertex_ai/batches/test_handler.py | 128 ++++++++++++++++++ 3 files changed, 173 insertions(+), 12 deletions(-) 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]