scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

fix(vertex_ai): honor timeout on batch create, retrieve, list and cancel

Browse source 
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.
This commit is contained in:
Abhishek Acharya 2026-10-19 23:44:39 +11:00
parent 748c973f42
commit a655b3960d
3 changed files with 173 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
litellm/batches/main.py
View file

@ -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":

29
litellm/llms/vertex_ai/batches/handler.py
View file

@ -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(

128
tests/unit/llms/vertex_ai/batches/test_handler.py
View file

@ -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]