From 697da95bf4de8f11d4374a428773e51cd9f5d9ff Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 9 Sep 2026 14:33:29 -0400 Subject: [PATCH 1/7] fix(vertex_ai): stop grafting batch storage and job urls onto the deployment api_base A deployment api_base points at online inference, commonly a full .../endpoints/:rawPredict resource url. The batch file upload used it as the GCS storage host and batch create/retrieve/list/cancel grafted job paths onto it, both producing urls Google answers with an HTML 404. Uploads now always target storage.googleapis.com, and batch operations ignore a resource-shaped api_base (path containing /projects/) while still honoring host-level or /v1 gateway mounts. --- litellm/llms/vertex_ai/batches/handler.py | 24 +++++++++--- .../llms/vertex_ai/files/transformation.py | 6 +-- .../llms/vertex_ai/batches/test_handler.py | 39 +++++++++++++++++++ .../test_vertex_ai_files_transformation.py | 26 +++++++++++++ 4 files changed, 85 insertions(+), 10 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index a15ea4d845b..8e26f6f5be5 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -72,6 +72,19 @@ class _VertexEndpointPayloadView(TypedDict): payload: ReadOnly[_VertexEndpointResponse] +def _gateway_api_base_or_none(api_base: str | None) -> str | None: + """ + A deployment `api_base` whose path names a concrete Vertex resource (contains `/projects/`, + e.g. the `.../endpoints/:rawPredict` url configured for online inference) is not a Vertex + API gateway; grafting `batchPredictionJobs` or resource-GET paths onto it can only produce + urls Google answers with an HTML 404 (LIT-7386). Batch operations ignore it and use the real + Vertex host; only a host-level or `/v1`-style gateway mount passes through. + """ + if api_base and "/projects/" in urlparse(api_base).path: + return None + return api_base + + def _vertex_batch_payload(response: _VertexBatchJsonSource) -> VertexBatchPredictionResponse: return response.json() @@ -126,11 +139,12 @@ class VertexAIBatchPrediction(VertexLLM): vertex_location=vertex_location or "us-central1", ) ) + gateway_api_base: Final = _gateway_api_base_or_none(api_base) vertex_batch_request: Final = self._resolve_fine_tuned_endpoint_model( vertex_batch_request=transformed_batch_request, headers=headers, sync_handler=sync_handler, - api_base=api_base, + api_base=gateway_api_base, vertex_location=vertex_location or "us-central1", ) @@ -145,7 +159,7 @@ class VertexAIBatchPrediction(VertexLLM): endpoint = "" _, api_base = self._check_custom_proxy( - api_base=api_base, + api_base=gateway_api_base, custom_llm_provider="vertex_ai", gemini_api_key=None, endpoint=endpoint, @@ -325,7 +339,7 @@ class VertexAIBatchPrediction(VertexLLM): endpoint = "" _, api_base = self._check_custom_proxy( - api_base=api_base, + api_base=_gateway_api_base_or_none(api_base), custom_llm_provider="vertex_ai", gemini_api_key=None, endpoint=endpoint, @@ -481,7 +495,7 @@ class VertexAIBatchPrediction(VertexLLM): endpoint = "" _, api_base = self._check_custom_proxy( - api_base=api_base, + api_base=_gateway_api_base_or_none(api_base), custom_llm_provider="vertex_ai", gemini_api_key=None, endpoint=endpoint, @@ -579,7 +593,7 @@ class VertexAIBatchPrediction(VertexLLM): cancel_api_base_default: Final = f"{retrieve_api_base_default}:cancel" _, api_base = self._check_custom_proxy( - api_base=api_base, + api_base=_gateway_api_base_or_none(api_base), custom_llm_provider="vertex_ai", gemini_api_key=None, endpoint="cancel", diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 263956efc9f..f9eadee2b5c 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -809,11 +809,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): object_name = f"{object_prefix}/{object_name}" encoded_object_name: Final = encode_gcs_object_name_for_url(object_name) endpoint: Final = f"upload/storage/v1/b/{bucket_name}/o?uploadType=media&name={encoded_object_name}" - api_base = api_base or "https://storage.googleapis.com" - if not api_base: - raise ValueError("api_base is required") - - return f"{api_base}/{endpoint}" + return f"https://storage.googleapis.com/{endpoint}" def get_supported_openai_params(self, model: str) -> list[OpenAICreateFileRequestOptionalParams]: return [] diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index 24df3214da3..6c93881bcf0 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -257,6 +257,45 @@ def test_create_batch_sync_resolves_fine_tuned_endpoint_to_tuned_model(): assert sent["model"] == TUNED_MODEL_RESOURCE +def test_create_batch_sync_ignores_resource_shaped_api_base(): + """A deployment api_base like `.../endpoints/:rawPredict` targets online inference, not + the Vertex API root; grafting batch urls onto it yields guaranteed 404s, so batch operations + must fall back to the default Vertex host (LIT-7386).""" + h = _make_handler() + client = MagicMock() + client.post.return_value = _http_response() + raw_predict_api_base = ( + f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}" + f"/locations/{LOCATION}/endpoints/{ENDPOINT_ID}:rawPredict" + ) + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get", return_value=_endpoint_get_response()) as safe_get, + ): + out = h.create_batch( + _is_async=False, + create_batch_data=ENDPOINT_CREATE_DATA, + api_base=raw_predict_api_base, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert isinstance(out, LiteLLMBatch) + resolution_url = safe_get.call_args.args[1] + assert ":rawPredict" not in resolution_url + assert resolution_url == ( + f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}" + f"/locations/{LOCATION}/endpoints/{ENDPOINT_ID}" + ) + assert h._check_custom_proxy.call_args.kwargs["api_base"] is None + sent = json.loads(client.post.call_args.kwargs["data"]) + assert sent["model"] == TUNED_MODEL_RESOURCE + + @pytest.mark.parametrize( "api_base, expected", [ diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 8a249820cbd..8df86a22664 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -158,6 +158,32 @@ class TestCreateFileUrl: assert ".." not in object_name assert "?" not in object_name + def test_should_upload_to_gcs_host_even_when_deployment_sets_api_base(self, config): + """The deployment api_base points at the inference endpoint (often a full + `.../endpoints/:rawPredict` URL); grafting the GCS upload onto it produces a + guaranteed 404 from Google, so the storage host must stay storage.googleapis.com + (LIT-7386).""" + url = config.get_complete_file_url( + api_base=( + "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project" + "/locations/us-central1/endpoints/6335039103326748672:rawPredict" + ), + api_key=None, + model="", + optional_params={}, + litellm_params={ + "gcs_bucket_name": "my-bucket", + "model": "vertex_ai/gemini/6335039103326748672", + }, + data={ + "file": ("batch.jsonl", b'{"body": {"model": "gemini-2.5-flash"}}', "application/jsonl"), + "purpose": "batch", + }, + ) + assert url.startswith("https://storage.googleapis.com/upload/storage/v1/b/my-bucket/o?") + assert "aiplatform" not in url + assert "rawPredict" not in url + class TestBatchObjectNaming: def test_should_store_publisher_model_under_publishers_path(self, config): From 7c3531265762cec82017b925ef7d92475ea8a393 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 9 Sep 2026 16:44:04 -0400 Subject: [PATCH 2/7] feat(vertex_ai): support managed batches for custom_endpoint deployments The Vertex batch API refuses both the v1beta1 BYOE endpoint field and Model-Garden-sourced model resources, so a custom_endpoint batch job instead runs batch-owned replicas of the deployment's own serving container: unmanagedContainerModel with the containerSpec read verbatim from the endpoint's deployed model, dedicatedResources copied from the online deployment, and instanceConfig.keyField round-tripping the custom_id. Uploads for these deployments stage rows as @requestFormat chatCompletions instances (the vLLM adapter's native OpenAI mode) under a custom-endpoints/ GCS path derived from the deployment api_base, and the output read unwraps prediction.predictions (already a full OpenAI chat.completion) into OpenAI batch output rows. --- litellm/llms/vertex_ai/batches/handler.py | 172 +++++++++++++++-- .../llms/vertex_ai/batches/transformation.py | 44 +++-- litellm/llms/vertex_ai/common_utils.py | 3 + .../llms/vertex_ai/files/transformation.py | 174 +++++++++++++++--- litellm/types/llms/vertex_ai.py | 37 +++- .../llms/vertex_ai/batches/test_handler.py | 125 ++++++++++++- .../vertex_ai/batches/test_transformation.py | 23 +++ .../test_vertex_ai_files_transformation.py | 102 +++++++++- 8 files changed, 617 insertions(+), 63 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 8e26f6f5be5..e505979c5e7 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -1,5 +1,5 @@ import json -from collections.abc import Coroutine, Sequence +from collections.abc import Coroutine, Mapping, Sequence from typing import TYPE_CHECKING, Final, Protocol from urllib.parse import urlparse @@ -17,12 +17,18 @@ from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, ) -from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url +from litellm.llms.vertex_ai.common_utils import ( + VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, + VertexAIError, + get_vertex_base_url, +) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.llms.vertex_ai.vertex_llm_base import _graft_default_vertex_path from litellm.types.llms.openai import CreateBatchRequest from litellm.types.llms.vertex_ai import ( VERTEX_CREDENTIALS_TYPES, + BatchDedicatedResources, + UnmanagedContainerModel, VertexAIBatchPredictionJob, VertexBatchPredictionResponse, ) @@ -58,8 +64,18 @@ class _FetchedResponseView(TypedDict): response: ReadOnly[httpx.Response] +class _VertexOnlineDedicatedResources(TypedDict, total=False): + """The dedicatedResources block on an online endpoint deployment; replica bounds are named + min/max there, unlike the batch job's starting/max.""" + + machineSpec: ReadOnly[Mapping[str, object]] + minReplicaCount: ReadOnly[int] + maxReplicaCount: ReadOnly[int] + + class _VertexEndpointDeployedModel(TypedDict, total=False): model: ReadOnly[str] + dedicatedResources: ReadOnly[_VertexOnlineDedicatedResources] class _VertexEndpointResponse(TypedDict, total=False): @@ -72,6 +88,16 @@ class _VertexEndpointPayloadView(TypedDict): payload: ReadOnly[_VertexEndpointResponse] +class _VertexModelResourceResponse(TypedDict, total=False): + containerSpec: ReadOnly[Mapping[str, object]] + + +class _VertexModelResourcePayloadView(TypedDict): + """Holds one decoded GET models/ response so the payload reads back typed.""" + + payload: ReadOnly[_VertexModelResourceResponse] + + def _gateway_api_base_or_none(api_base: str | None) -> str | None: """ A deployment `api_base` whose path names a concrete Vertex resource (contains `/projects/`, @@ -110,15 +136,6 @@ class VertexAIBatchPrediction(VertexLLM): max_retries: int | None, custom_endpoint: bool | None = None, ) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]: - if custom_endpoint: - raise VertexAIError( - status_code=400, - message=( - "Vertex AI batch prediction is not supported for `custom_endpoint` deployments. " - "The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; " - "use a publisher model or fine-tuned Gemini endpoint deployment instead." - ), - ) sync_handler: Final = _get_httpx_client() access_token, project_id = self._ensure_access_token( @@ -139,18 +156,37 @@ class VertexAIBatchPrediction(VertexLLM): vertex_location=vertex_location or "us-central1", ) ) + if custom_endpoint and "/custom-endpoints/" not in transformed_batch_request.get("model", ""): + raise VertexAIError( + status_code=400, + message=( + "Vertex AI batch prediction on a `custom_endpoint` deployment requires an input " + "file uploaded through LiteLLM against that deployment (its file id carries a " + "custom-endpoints/ path); this input file targets a publisher or " + "fine-tuned Gemini model instead." + ), + ) gateway_api_base: Final = _gateway_api_base_or_none(api_base) - vertex_batch_request: Final = self._resolve_fine_tuned_endpoint_model( + resolved_batch_request: Final = self._resolve_fine_tuned_endpoint_model( vertex_batch_request=transformed_batch_request, headers=headers, sync_handler=sync_handler, api_base=gateway_api_base, vertex_location=vertex_location or "us-central1", ) + vertex_batch_request: Final = self._resolve_custom_endpoint_container( + vertex_batch_request=resolved_batch_request, + headers=headers, + sync_handler=sync_handler, + api_base=gateway_api_base, + vertex_location=vertex_location or "us-central1", + ) + is_unmanaged_container_job: Final = "unmanagedContainerModel" in vertex_batch_request default_api_base: Final = self.create_vertex_batch_url( vertex_location=vertex_location or "us-central1", vertex_project=vertex_project or project_id, + vertex_api_version="v1beta1" if is_unmanaged_container_job else "v1", ) if len(default_api_base.split(":")) > 1: @@ -169,7 +205,7 @@ class VertexAIBatchPrediction(VertexLLM): model=None, vertex_project=vertex_project or project_id, vertex_location=vertex_location or "us-central1", - vertex_api_version="v1", + vertex_api_version="v1beta1" if is_unmanaged_container_job else "v1", ) if _is_async is True: @@ -263,6 +299,113 @@ class VertexAIBatchPrediction(VertexLLM): resolved_request: Final[VertexAIBatchPredictionJob] = {**vertex_batch_request, "model": deployed_model} return resolved_request + def _resolve_custom_endpoint_container( + self, + vertex_batch_request: VertexAIBatchPredictionJob, + headers: dict[str, str], # mutable-ok: HTTPHandler.get only accepts dict headers + sync_handler: HTTPHandler, + api_base: str | None, + vertex_location: str, + ) -> VertexAIBatchPredictionJob: + """ + A `custom_endpoint` deployment serves an OpenAI-compatible container on a Vertex endpoint. + The batch API accepts neither that endpoint (the v1beta1 BYOE `endpoint` field is refused + with "specify model or unmanaged_container_model") nor its Model-Garden-sourced model + resource ("Unknown ModelSource source_type: MODEL_GARDEN"), so the job instead runs + batch-owned replicas of the same container: `unmanagedContainerModel` with the + containerSpec read verbatim from the endpoint's deployed model (a hand-built spec loses + model-source args and crash-loops) plus `dedicatedResources` copied from the endpoint's + own deployment. + """ + model: Final = vertex_batch_request.get("model", "") + if "/custom-endpoints/" not in model: + return vertex_batch_request + endpoint_resource: Final = model.replace("/custom-endpoints/", "/endpoints/") + + endpoint_url: Final = self._build_endpoint_resolution_url( + api_base=api_base, + model=endpoint_resource, + vertex_location=vertex_location, + ) + endpoint_fetched: Final[_FetchedResponseView] = { + "response": safe_get(sync_handler, endpoint_url, headers=headers) + } + endpoint_response: Final = endpoint_fetched["response"] + if endpoint_response.status_code != 200: + raise VertexAIError( + status_code=endpoint_response.status_code, + message=f"Failed to resolve custom Vertex endpoint '{endpoint_resource}': {endpoint_response.text}", + ) + endpoint_view: Final[_VertexEndpointPayloadView] = {"payload": endpoint_response.json()} + deployed_models: Final = endpoint_view["payload"].get("deployedModels") or () + deployed: Final = deployed_models[0] if deployed_models else _VertexEndpointDeployedModel() + deployed_model_resource: Final = deployed.get("model", "") + if not deployed_model_resource: + raise VertexAIError( + status_code=400, + message=( + f"Vertex endpoint '{endpoint_resource}' has no deployed model, so there is no " + "serving container to run batch predictions with" + ), + ) + + model_url: Final = self._build_endpoint_resolution_url( + api_base=api_base, + model=deployed_model_resource, + vertex_location=vertex_location, + ) + model_fetched: Final[_FetchedResponseView] = { + "response": safe_get(sync_handler, model_url, headers=headers) + } + model_response: Final = model_fetched["response"] + if model_response.status_code != 200: + raise VertexAIError( + status_code=model_response.status_code, + message=f"Failed to read model resource '{deployed_model_resource}': {model_response.text}", + ) + model_view: Final[_VertexModelResourcePayloadView] = {"payload": model_response.json()} + container_spec: Final = model_view["payload"].get("containerSpec") + if not container_spec: + raise VertexAIError( + status_code=400, + message=( + f"Model resource '{deployed_model_resource}' carries no containerSpec, so its " + "serving container cannot be replicated for batch prediction" + ), + ) + + online_resources: Final = deployed.get("dedicatedResources") or _VertexOnlineDedicatedResources() + machine_spec: Final = online_resources.get("machineSpec") + if not machine_spec: + raise VertexAIError( + status_code=400, + message=( + f"Vertex endpoint '{endpoint_resource}' exposes no dedicatedResources machine " + "spec to size the batch replicas from" + ), + ) + batch_resources: Final[BatchDedicatedResources] = { + "machineSpec": machine_spec, + "startingReplicaCount": online_resources.get("minReplicaCount", 1), + "maxReplicaCount": online_resources.get("maxReplicaCount", 1), + } + unmanaged: Final[UnmanagedContainerModel] = {"containerSpec": container_spec} + resolved: Final[VertexAIBatchPredictionJob] = { + "displayName": vertex_batch_request["displayName"], + "inputConfig": vertex_batch_request["inputConfig"], + "outputConfig": vertex_batch_request["outputConfig"], + "unmanagedContainerModel": unmanaged, + "dedicatedResources": batch_resources, + # keyField strips the custom_id tag from each instance before it reaches the + # container (vLLM rejects unknown fields) and echoes it back as `key` in the output + # row; it only takes effect alongside an explicit instanceType. + "instanceConfig": { + "instanceType": "object", + "keyField": VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, + }, + } + return resolved + async def _async_create_batch( self, vertex_batch_request: VertexAIBatchPredictionJob, @@ -298,11 +441,12 @@ class VertexAIBatchPrediction(VertexLLM): self, vertex_location: str, vertex_project: str, + vertex_api_version: str = "v1", ) -> str: """Return the base url for the vertex garden models""" # POST https://LOCATION-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION/batchPredictionJobs base_url: Final = get_vertex_base_url(vertex_location) - return f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/batchPredictionJobs" + return f"{base_url}/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/batchPredictionJobs" def retrieve_batch( self, diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index e63c80dd3cf..6f115f6a46c 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -129,13 +129,23 @@ class VertexAIBatchTransformation: def _get_output_file_id_from_vertex_ai_batch_response(cls, response: VertexBatchPredictionResponse) -> str: """ Gets the output file id from the Vertex AI Batch response + + Gemini jobs write `predictions.jsonl`; unmanaged-container jobs (custom_endpoint + deployments) write sharded `prediction.results-*` files into a directory Vertex names + `prediction-custom-unmanaged-model-`. """ output_info: Final = response.get("outputInfo") or OutputInfo() - output_file_id: str = output_info.get("gcsOutputDirectory", "") + output_directory: Final = output_info.get("gcsOutputDirectory", "") + results_filename: Final = ( + "prediction.results-00000-of-00001" + if "prediction-custom-unmanaged-model" in output_directory + else "predictions.jsonl" + ) + output_file_id: str = output_directory if output_file_id: - output_file_id = output_file_id.rstrip("/") + "/predictions.jsonl" - if output_file_id and output_file_id != "/predictions.jsonl": + output_file_id = output_file_id.rstrip("/") + f"/{results_filename}" + if output_file_id and output_file_id != f"/{results_filename}": return output_file_id output_config: Final = response.get("outputConfig") @@ -209,17 +219,21 @@ class VertexAIBatchTransformation: to its deployed tuned model (`projects/../locations/../models/`) before sending the job. """ parsed_model: Final = cls._get_model_from_gcs_file(input_file_id) - if not parsed_model.startswith("endpoints/"): + if not parsed_model.startswith(("endpoints/", "custom-endpoints/")): return parsed_model if not vertex_project: raise VertexAIError( status_code=400, message=( - f"Vertex AI batch jobs against a fine-tuned endpoint ('{parsed_model}') require " + f"Vertex AI batch jobs against an endpoint ('{parsed_model}') require " "`vertex_project` to build the endpoint resource name" ), ) - return f"projects/{vertex_project}/locations/{vertex_location or 'us-central1'}/{parsed_model}" + location_segment: Final = f"projects/{vertex_project}/locations/{vertex_location or 'us-central1'}" + if parsed_model.startswith("custom-endpoints/"): + endpoint_id: Final = parsed_model.removeprefix("custom-endpoints/") + return f"{location_segment}/custom-endpoints/{endpoint_id}" + return f"{location_segment}/{parsed_model}" @classmethod def _get_model_from_gcs_file(cls, gcs_file_uri: str) -> str: @@ -260,12 +274,15 @@ class VertexAIBatchTransformation: @classmethod def _parse_model_from_gcs_file(cls, gcs_file_uri: str) -> str | None: """ - Returns the `publishers//models/` or `endpoints/` path from a - gcs uri, or None if the uri does not contain one. + Returns the `publishers//models/`, `endpoints/`, or + `custom-endpoints/` path from a gcs uri, or None if the uri does not contain + one. - A publisher path wins over an `endpoints/` segment, and the last `endpoints/` occurrence is - used, so a user-configured bucket prefix that happens to contain `endpoints/` cannot - override the model path LiteLLM appended after it. + A publisher path wins over an endpoints segment, `custom-endpoints/` (a custom_endpoint + deployment's serving container run as an unmanaged-container batch) wins over a plain + `endpoints/` (a fine-tuned Gemini endpoint), and the last occurrence of each is used, so a + user-configured bucket prefix that happens to contain `endpoints/` cannot override + the model path LiteLLM appended after it. """ unquoted_uri: Final = unquote(gcs_file_uri) _, separator, model_path = unquoted_uri.partition("publishers/") @@ -274,6 +291,11 @@ class VertexAIBatchTransformation: if len(parts) >= 3 and parts[1] == "models" and parts[2]: return f"publishers/{'/'.join(parts[:3])}" + _, custom_separator, custom_path = unquoted_uri.rpartition("custom-endpoints/") + custom_endpoint_id: Final = custom_path.split("/")[0] if custom_separator else "" + if custom_endpoint_id.isdigit(): + return f"custom-endpoints/{custom_endpoint_id}" + _, endpoint_separator, endpoint_path = unquoted_uri.rpartition("endpoints/") endpoint_id: Final = endpoint_path.split("/")[0] if endpoint_separator else "" if endpoint_id.isdigit(): diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 14aebcaabaf..7bda17d4334 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -370,6 +370,9 @@ def get_vertex_base_model_name(model: str) -> str: return model +VERTEX_CUSTOM_ENDPOINT_KEY_FIELD: Final = "litellm_custom_id" + + def get_vertex_ai_fine_tuned_endpoint_id(model: str) -> str | None: """ Fine-tuned Gemini deployments are addressed by a numeric endpoint id, diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index f9eadee2b5c..765931bed8d 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -7,7 +7,7 @@ import re import time from collections.abc import Callable, Iterable, Iterator, Mapping from typing import Any, Final, TypedDict -from urllib.parse import quote, unquote +from urllib.parse import quote, unquote, urlparse import httpx from httpx import Headers, Response @@ -38,6 +38,7 @@ from litellm.llms.base_llm.files.transformation import ( LiteLLMLoggingObj, ) from litellm.llms.vertex_ai.common_utils import ( + VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, _convert_vertex_datetime_to_openai_datetime, get_vertex_ai_fine_tuned_endpoint_id, ) @@ -645,6 +646,40 @@ def _parse_vertex_batch_output_row(line: str) -> _VertexBatchRow: return row +def _is_custom_endpoint_batch_output_row(row: Mapping[str, object]) -> bool: + """ + An unmanaged-container (custom_endpoint) batch output row: Vertex echoes the instance (or the + `key` extracted from it) alongside a `prediction` wrapper, unlike Gemini rows which pair + `request`/`response`/`processed_time`. + """ + return "prediction" in row and ("instance" in row or "key" in row) + + +def _custom_endpoint_row_to_openai_batch_output_row(row: Mapping[str, object]) -> _OpenAIBatchOutputRow: + """ + Unwraps one unmanaged-container batch output row. The vLLM `@requestFormat: chatCompletions` + mode already produces a full OpenAI chat.completion under `prediction.predictions`, so the + transform is: recover the custom_id (the `key` field when `instanceConfig.keyField` was + honored, else the echoed instance's tag) and re-wrap in the OpenAI batch output envelope. + """ + key: Final = row.get("key") + instance: Final = row.get("instance") + instance_map: Final = instance if isinstance(instance, Mapping) else {} + custom_id: Final = str(key if key is not None else instance_map.get(VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, "")) + + prediction: Final = row.get("prediction") + prediction_map: Final = prediction if isinstance(prediction, Mapping) else {} + body: Final = prediction_map.get("predictions") + if not isinstance(body, Mapping): + error_text: Final = str(row.get("status") or prediction or "prediction carries no response body") + return _openai_batch_output_row( + custom_id=custom_id, + error_code="vertex_ai_error", + error_message=error_text, + ) + return _openai_batch_output_row(custom_id=custom_id, body=body) + + class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): """Streams an OpenAI batch JSONL upload as Vertex-wrapped JSONL one row at a time, so the transformed payload is never held in full. @@ -673,6 +708,59 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): return self._iter_vertex_jsonl_chunks() +VERTEX_CUSTOM_ENDPOINT_GCS_SEGMENT: Final = "custom-endpoints" +_VERTEX_CHAT_COMPLETIONS_REQUEST_FORMAT: Final = "chatCompletions" + + +def get_custom_endpoint_id_from_api_base(api_base: str | None) -> str | None: + """ + The Vertex endpoint a `custom_endpoint` deployment serves from is only recorded in its + api_base (`.../endpoints/:rawPredict` or a dedicated-domain equivalent); batch jobs need + that id to read the endpoint's containerSpec, so extract it (verb suffix stripped). + """ + if not api_base: + return None + path_segments: Final = urlparse(api_base).path.split("/") + after_endpoints: Final = tuple( + segment for prior, segment in zip(path_segments, path_segments[1:]) if prior == "endpoints" + ) + if not after_endpoints: + return None + return after_endpoints[-1].split(":")[0] or None + + +def _openai_batch_jsonl_entry_to_custom_endpoint_row(openai_entry: dict[str, Any]) -> Mapping[str, object]: + """ + One OpenAI batch JSONL line as the instance a vLLM-serving Vertex container consumes: + the OpenAI request body itself tagged `@requestFormat: chatCompletions` (the container + speaks OpenAI natively, so no Gemini translation), minus `model` (the batch replica + serves exactly one model) plus the custom_id under the job's `instanceConfig.keyField`. + """ + body: Final = openai_entry.get("body") or {} + row: Final = {k: v for k, v in body.items() if k != "model"} + return { + "@requestFormat": _VERTEX_CHAT_COMPLETIONS_REQUEST_FORMAT, + **row, + VERTEX_CUSTOM_ENDPOINT_KEY_FIELD: str(openai_entry.get("custom_id", "")), + } + + +class _OpenAIToCustomEndpointBatchUploadStream(BaseFileUploadStream): + """Streams an OpenAI batch JSONL upload as `@requestFormat: chatCompletions` instances + for a custom_endpoint (OpenAI-compatible container) batch job, one row at a time.""" + + def __init__(self, openai_file_content: FileTypes) -> None: + self._openai_file_content = openai_file_content + + def iter_bytes(self) -> Iterator[bytes]: + first = True + for entry in _iter_openai_jsonl_entries(self._openai_file_content): + row = _openai_batch_jsonl_entry_to_custom_endpoint_row(entry) + prefix = b"" if first else b"\n" + first = False + yield prefix + json.dumps(row).encode("utf-8") + + class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ Config for VertexAI Files @@ -740,13 +828,25 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): object_name: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}" return object_name - def get_object_name(self, file_data: FileTypes, purpose: str, deployment_model: str | None = None) -> str: + def get_object_name( + self, + file_data: FileTypes, + purpose: str, + deployment_model: str | None = None, + custom_endpoint_id: str | None = None, + ) -> str: """ Get the object name for the request. Reads only the first JSONL entry (streamed) for batch files, so a large upload is never materialized just to derive the GCS object name. """ + if purpose == "batch" and custom_endpoint_id is not None: + safe_endpoint_id: Final = sanitize_cloud_object_path(custom_endpoint_id, fallback="endpoint") + return ( + f"{VERTEX_AI_MANAGED_GCS_PREFIX}{VERTEX_CUSTOM_ENDPOINT_GCS_SEGMENT}/" + f"{safe_endpoint_id}/{uuid.uuid4()}" + ) if purpose == "batch": ## 1. If jsonl, derive the object name from the deployment model (or the first entry's) first_entry: Final = next(_iter_openai_jsonl_entries(file_data), None) @@ -781,16 +881,6 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ Get the complete url for the request """ - if data.get("purpose") == "batch" and litellm_params.get("custom_endpoint"): - raise VertexAIError( - status_code=400, - message=( - "Vertex AI batch prediction is not supported for `custom_endpoint` deployments. " - "The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; " - "remove this deployment from the batch request (e.g. `target_model_names`) or " - "use a publisher model / fine-tuned Gemini endpoint instead." - ), - ) bucket_name = self._get_configured_bucket_name(litellm_params) bucket_name, object_prefix = split_configured_cloud_bucket_name(bucket_name) file_data: Final = data.get("file") @@ -800,10 +890,29 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): if purpose is None: raise ValueError("purpose is required") configured_model: Final = litellm_params.get("model") + deployment_api_base: Final = litellm_params.get("api_base") + custom_endpoint_id: Final = ( + get_custom_endpoint_id_from_api_base( + deployment_api_base if isinstance(deployment_api_base, str) else None + ) + if litellm_params.get("custom_endpoint") + else None + ) + if litellm_params.get("custom_endpoint") and purpose == "batch" and custom_endpoint_id is None: + raise VertexAIError( + status_code=400, + message=( + "Vertex AI batch prediction on a `custom_endpoint` deployment requires the " + "deployment's `api_base` to name its Vertex endpoint " + "(e.g. https://.../endpoints/:rawPredict), so the batch job can " + "run replicas of that endpoint's serving container." + ), + ) object_name = self.get_object_name( file_data, purpose, deployment_model=configured_model if isinstance(configured_model, str) else None, + custom_endpoint_id=custom_endpoint_id, ) if object_prefix: object_name = f"{object_prefix}/{object_name}" @@ -867,12 +976,17 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): create_file_data=create_file_data, content_type=content_type, ): + body_stream: Final[BaseFileUploadStream] = ( + _OpenAIToCustomEndpointBatchUploadStream(file_data) + if litellm_params.get("custom_endpoint") + else _OpenAIToVertexBatchUploadStream( + file_data, + self._map_openai_to_vertex_params, + ) + ) return { "streaming_media_upload": StreamingMediaUploadConfig( - body_stream=_OpenAIToVertexBatchUploadStream( - file_data, - self._map_openai_to_vertex_params, - ), + body_stream=body_stream, content_type="application/json", ) } @@ -1116,14 +1230,19 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): # first line is not valid UTF-8/JSON) raises and falls through to the # passthrough below, leaving the content untouched. first_row: Final = _parse_vertex_batch_output_row(first_line) - is_vertex_batch_output: Final = _is_vertex_embeddings_batch_output_row(first_row) or ( - "request" in first_row - and "response" in first_row - and "processed_time" in first_row - and ( - "candidates" in first_row.get("response", {}) - or "promptFeedback" in first_row.get("response", {}) - or bool(first_row.get("status")) + is_custom_endpoint_output: Final = _is_custom_endpoint_batch_output_row(first_row) + is_vertex_batch_output: Final = ( + is_custom_endpoint_output + or _is_vertex_embeddings_batch_output_row(first_row) + or ( + "request" in first_row + and "response" in first_row + and "processed_time" in first_row + and ( + "candidates" in first_row.get("response", {}) + or "promptFeedback" in first_row.get("response", {}) + or bool(first_row.get("status")) + ) ) ) if not is_vertex_batch_output: @@ -1151,6 +1270,13 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): all_lines = itertools.chain((first_line,), lines) + if is_custom_endpoint_output: + return b"\n".join( + json.dumps(_custom_endpoint_row_to_openai_batch_output_row(json.loads(line))).encode("utf-8") + for line in all_lines + if line.strip() + ) + # Embedding rows are grouped by `custom_id` rather than transformed one at a # time, since an entry that asked for several embeddings comes back as # several rows, in arbitrary order. diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 3b95b786631..1fda5d727eb 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -1,7 +1,9 @@ +from collections.abc import Mapping from enum import Enum from typing import Any, Final, Literal, Protocol from typing_extensions import ( + ReadOnly, Required, TypedDict, ) @@ -678,11 +680,36 @@ class GcsBucketResponse(TypedDict): timeFinalized: str -class VertexAIBatchPredictionJob(TypedDict): - displayName: str - model: str - inputConfig: InputConfig - outputConfig: OutputConfig +class BatchDedicatedResources(TypedDict, total=False): + """Sizing for batch-owned replicas; machineSpec is copied verbatim from the online + deployment's dedicatedResources, hence the loose Mapping.""" + + machineSpec: ReadOnly[Mapping[str, object]] + startingReplicaCount: ReadOnly[int] + maxReplicaCount: ReadOnly[int] + + +class UnmanagedContainerModel(TypedDict, total=False): + """The v1beta1 batch shape for running batch-owned replicas of a serving container. The + containerSpec is copied verbatim from the deployed model resource (hence the loose Mapping): + hand-building one loses model-source args/env and crash-loops the batch container.""" + + containerSpec: ReadOnly[Mapping[str, object]] + + +class BatchInstanceConfig(TypedDict, total=False): + instanceType: ReadOnly[str] + keyField: ReadOnly[str] + + +class VertexAIBatchPredictionJob(TypedDict, total=False): + displayName: ReadOnly[Required[str]] + model: ReadOnly[str] + unmanagedContainerModel: ReadOnly[UnmanagedContainerModel] + dedicatedResources: ReadOnly[BatchDedicatedResources] + instanceConfig: ReadOnly[BatchInstanceConfig] + inputConfig: ReadOnly[Required[InputConfig]] + outputConfig: ReadOnly[Required[OutputConfig]] class VertexBatchPredictionResponse(TypedDict, total=False): diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index 6c93881bcf0..8691d14d392 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -257,6 +257,123 @@ def test_create_batch_sync_resolves_fine_tuned_endpoint_to_tuned_model(): assert sent["model"] == TUNED_MODEL_RESOURCE +CUSTOM_ENDPOINT_ID = "4980511146650894336" +CUSTOM_ENDPOINT_CREATE_DATA = { + "input_file_id": (f"gs://bucket/litellm-vertex-files/custom-endpoints/{CUSTOM_ENDPOINT_ID}/file-uuid") +} +CONTAINER_MODEL_RESOURCE = f"projects/{PROJECT}/locations/{LOCATION}/models/google-gemma2-123" +CONTAINER_SPEC = { + "imageUri": "us-docker.pkg.dev/vertex-ai/pytorch-vllm-serve:x", + "args": ["python", "-m", "vllm.entrypoints.api_server"], + "predictRoute": "/generate", + "healthRoute": "/ping", +} +MACHINE_SPEC = {"machineType": "g2-standard-12", "acceleratorType": "NVIDIA_L4", "acceleratorCount": 1} + + +def _custom_endpoint_get_response() -> MagicMock: + resp = MagicMock() + resp.status_code = 200 + resp.json.return_value = { + "name": f"projects/{PROJECT}/locations/{LOCATION}/endpoints/{CUSTOM_ENDPOINT_ID}", + "deployedModels": [ + { + "model": CONTAINER_MODEL_RESOURCE, + "dedicatedResources": {"machineSpec": MACHINE_SPEC, "minReplicaCount": 1, "maxReplicaCount": 2}, + } + ], + } + return resp + + +def _container_model_get_response(container_spec: dict | None = CONTAINER_SPEC) -> MagicMock: + resp = MagicMock() + resp.status_code = 200 + resp.json.return_value = ( + {"name": CONTAINER_MODEL_RESOURCE, "containerSpec": container_spec} + if container_spec is not None + else {"name": CONTAINER_MODEL_RESOURCE} + ) + return resp + + +def test_create_batch_sync_custom_endpoint_builds_unmanaged_container_job(): + """A custom_endpoint batch must run batch-owned replicas of the endpoint's own serving + container: the live API refuses both the v1beta1 BYOE `endpoint` field and Model-Garden model + resources, and a hand-built containerSpec crash-loops, so the job carries the deployed + model's containerSpec verbatim under `unmanagedContainerModel` on the v1beta1 route with the + custom_id extracted server-side via instanceConfig.keyField (LIT-7387).""" + h = _make_handler() + client = MagicMock() + client.post.return_value = _http_response() + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch( + f"{HMOD}.safe_get", + side_effect=[_custom_endpoint_get_response(), _container_model_get_response()], + ) as safe_get, + ): + out = h.create_batch( + _is_async=False, + create_batch_data=CUSTOM_ENDPOINT_CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + custom_endpoint=True, + ) + + assert isinstance(out, LiteLLMBatch) + endpoint_get_url = safe_get.call_args_list[0].args[1] + assert endpoint_get_url.endswith(f"/endpoints/{CUSTOM_ENDPOINT_ID}") + model_get_url = safe_get.call_args_list[1].args[1] + assert model_get_url.endswith(CONTAINER_MODEL_RESOURCE) + + post_url = client.post.call_args.kwargs["url"] + assert "/v1beta1/" in post_url + sent = json.loads(client.post.call_args.kwargs["data"]) + assert "model" not in sent + assert sent["unmanagedContainerModel"] == {"containerSpec": CONTAINER_SPEC} + assert sent["dedicatedResources"] == { + "machineSpec": MACHINE_SPEC, + "startingReplicaCount": 1, + "maxReplicaCount": 2, + } + assert sent["instanceConfig"] == {"instanceType": "object", "keyField": "litellm_custom_id"} + + +def test_create_batch_sync_custom_endpoint_without_container_spec_raises_400(): + h = _make_handler() + client = MagicMock() + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch( + f"{HMOD}.safe_get", + side_effect=[_custom_endpoint_get_response(), _container_model_get_response(container_spec=None)], + ), + ): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data=CUSTOM_ENDPOINT_CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + custom_endpoint=True, + ) + + assert exc_info.value.status_code == 400 + assert "containerSpec" in str(exc_info.value) + client.post.assert_not_called() + + def test_create_batch_sync_ignores_resource_shaped_api_base(): """A deployment api_base like `.../endpoints/:rawPredict` targets online inference, not the Vertex API root; grafting batch urls onto it yields guaranteed 404s, so batch operations @@ -356,9 +473,9 @@ def test_create_batch_sync_endpoint_resolution_error_raises(): client.post.assert_not_called() -def test_create_batch_custom_endpoint_raises_400_without_io(): - """custom_endpoint deployments have no Vertex batch surface; creating a job would target a - nonexistent publisher model, so the handler must 400 before any auth or HTTP work (LIT-6899).""" +def test_create_batch_custom_endpoint_rejects_non_custom_endpoint_file(): + """A custom_endpoint batch create over a file staged for a publisher model would run the wrong + workload on batch replicas of the container; the handler must 400 before any HTTP work.""" h = _make_handler() client = MagicMock() @@ -378,8 +495,8 @@ def test_create_batch_custom_endpoint_raises_400_without_io(): assert exc_info.value.status_code == 400 assert "custom_endpoint" in str(exc_info.value) - h._ensure_access_token.assert_not_called() client.post.assert_not_called() + client.get.assert_not_called() def test_create_batch_sync_endpoint_without_deployed_model_raises_400(): diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index 232c6413e78..c80f244a659 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -391,6 +391,29 @@ def test_get_bare_model_name_from_gcs_file_fine_tuned_endpoint(): assert T.get_bare_model_name_from_gcs_file(ENDPOINT_INPUT_FILE) == ENDPOINT_ID +CUSTOM_ENDPOINT_ID = "4980511146650894336" +CUSTOM_ENDPOINT_INPUT_FILE = ( + f"gs://litellm-testing-bucket/litellm-vertex-files/custom-endpoints/{CUSTOM_ENDPOINT_ID}/" + "e9412502-2c91-42a6-8e61-f5c294cc0fc8" +) + + +def test_get_model_from_gcs_file_custom_endpoint(): + """`custom-endpoints/` contains `endpoints/` as a substring, so the custom marker must be + matched first or the id would be misread as a fine-tuned Gemini endpoint and the batch job + would target a nonexistent tuned model (LIT-7387).""" + assert T._get_model_from_gcs_file(CUSTOM_ENDPOINT_INPUT_FILE) == f"custom-endpoints/{CUSTOM_ENDPOINT_ID}" + + +def test_batch_job_model_custom_endpoint_builds_resource_path(): + job = T.transform_openai_batch_request_to_vertex_ai_batch_request( + {"input_file_id": CUSTOM_ENDPOINT_INPUT_FILE}, + vertex_project="my-project", + vertex_location="us-central1", + ) + assert job["model"] == f"projects/my-project/locations/us-central1/custom-endpoints/{CUSTOM_ENDPOINT_ID}" + + # =========================================================================== # # is_unmanaged_gcs_batch_input_file_id # =========================================================================== # diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 8df86a22664..be77ece8dfe 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -234,10 +234,41 @@ class TestBatchObjectNaming: assert "9999999999999999999" not in object_name +CUSTOM_ENDPOINT_ID = "4980511146650894336" +CUSTOM_ENDPOINT_API_BASE = ( + "https://us-central1-aiplatform.googleapis.com/v1/projects/my-project" + f"/locations/us-central1/endpoints/{CUSTOM_ENDPOINT_ID}:rawPredict" +) + + class TestCustomEndpointBatchUpload: - def test_should_reject_batch_upload_for_custom_endpoint_deployment(self, config): - """custom_endpoint deployments have no Vertex batch surface; the upload must 400 instead - of staging a file that can only produce a doomed batch job (LIT-6899).""" + def test_should_stage_batch_upload_under_custom_endpoints_path(self, config): + """The GCS path is how the later batch create learns which serving container to + replicate, so a custom_endpoint upload must record the endpoint id from the api_base + under the custom-endpoints/ marker (LIT-7387).""" + url = config.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={ + "gcs_bucket_name": "my-bucket", + "custom_endpoint": True, + "api_base": CUSTOM_ENDPOINT_API_BASE, + "model": "vertex_ai/openai/gemma-2-2b-it", + }, + data={ + "file": ("batch.jsonl", b'{"body": {"model": "openai/gemma-2-2b-it"}}', "application/jsonl"), + "purpose": "batch", + }, + ) + assert url.startswith("https://storage.googleapis.com/") + object_name = parse_qs(urlparse(url).query)["name"][0] + assert object_name.startswith(f"litellm-vertex-files/custom-endpoints/{CUSTOM_ENDPOINT_ID}/") + + def test_should_reject_batch_upload_when_api_base_names_no_endpoint(self, config): + """Without an endpoint id in the api_base there is no container to run the batch with, so + the upload must fail with a clear 400 instead of staging a doomed file.""" from litellm.llms.vertex_ai.common_utils import VertexAIError with pytest.raises(VertexAIError) as exc_info: @@ -246,14 +277,18 @@ class TestCustomEndpointBatchUpload: api_key=None, model="", optional_params={}, - litellm_params={"gcs_bucket_name": "my-bucket", "custom_endpoint": True}, + litellm_params={ + "gcs_bucket_name": "my-bucket", + "custom_endpoint": True, + "api_base": "https://my-gateway.internal/v1", + }, data={ "file": ("batch.jsonl", b'{"body": {"model": "openai/gemma-2-2b-it"}}', "application/jsonl"), "purpose": "batch", }, ) assert exc_info.value.status_code == 400 - assert "custom_endpoint" in str(exc_info.value) + assert "api_base" in str(exc_info.value) def test_should_allow_non_batch_upload_for_custom_endpoint_deployment(self, config): url = config.get_complete_file_url( @@ -270,6 +305,63 @@ class TestCustomEndpointBatchUpload: assert "/b/my-bucket/" in url +class TestCustomEndpointBatchRows: + def test_upload_stream_emits_chat_completions_instances(self): + """Each OpenAI batch line must become a `@requestFormat: chatCompletions` instance the + vLLM container accepts natively, with `model` dropped (the batch replica serves exactly + one model) and the custom_id under the keyField name the batch job strips server-side.""" + from litellm.llms.vertex_ai.files.transformation import ( + _OpenAIToCustomEndpointBatchUploadStream, + ) + + openai_jsonl = ( + b'{"custom_id": "req-1", "method": "POST", "url": "/v1/chat/completions",' + b' "body": {"model": "gemma", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 5}}\n' + b'{"custom_id": "req-2", "method": "POST", "url": "/v1/chat/completions",' + b' "body": {"model": "gemma", "messages": [{"role": "user", "content": "yo"}]}}' + ) + stream = _OpenAIToCustomEndpointBatchUploadStream(("batch.jsonl", openai_jsonl, "application/jsonl")) + rows = [json.loads(line) for line in b"".join(stream.iter_bytes()).split(b"\n")] + assert rows == [ + { + "@requestFormat": "chatCompletions", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 5, + "litellm_custom_id": "req-1", + }, + { + "@requestFormat": "chatCompletions", + "messages": [{"role": "user", "content": "yo"}], + "litellm_custom_id": "req-2", + }, + ] + + def test_output_rows_unwrap_to_openai_batch_format(self, config): + """An unmanaged-container output row already carries a full OpenAI chat.completion under + prediction.predictions; the transform must unwrap it and recover the custom_id from the + keyField echo, and a failed row must become an OpenAI batch error row.""" + vertex_output = ( + b'{"key": "req-1", "prediction": {"predictions": {"id": "chatcmpl-1", "object": "chat.completion",' + b' "model": "google/gemma2-2b-it", "choices": [{"index": 0, "message": {"role": "assistant",' + b' "content": "Hello"}, "finish_reason": "stop"}], "usage": {"prompt_tokens": 5,' + b' "completion_tokens": 2, "total_tokens": 7}}}}\n' + b'{"key": "req-2", "prediction": "Post request fails.", "status": "Post request fails."}' + ) + transformed = config._try_transform_vertex_batch_output_to_openai(content=vertex_output) + rows = [json.loads(line) for line in transformed.split(b"\n")] + + assert rows[0]["custom_id"] == "req-1" + assert rows[0]["error"] is None + assert rows[0]["response"]["status_code"] == 200 + assert rows[0]["response"]["body"]["choices"][0]["message"]["content"] == "Hello" + assert rows[0]["response"]["body"]["usage"]["total_tokens"] == 7 + + assert rows[1]["custom_id"] == "req-2" + assert rows[1]["response"] is None + assert rows[1]["error"]["code"] == "vertex_ai_error" + assert "Post request fails." in rows[1]["error"]["message"] + + class TestTransformRetrieveFile: def test_should_build_correct_gcs_metadata_url(self, config): file_id = "gs://my-bucket/litellm-vertex-files/path/to/file.jsonl" From 783978b39c9587cab25127e76b3e43a8ea6c7a92 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 9 Sep 2026 17:01:20 -0400 Subject: [PATCH 3/7] fix(vertex_ai): strip the batch custom_id tag via excludedFields, not keyField Probed live: instanceConfig.keyField does not remove the field from the instances the container receives, so vLLM 400s every row; excludedFields performs the strip and attaches the value to the output row. --- litellm/llms/vertex_ai/batches/handler.py | 9 +++++---- litellm/types/llms/vertex_ai.py | 1 + .../test_litellm/llms/vertex_ai/batches/test_handler.py | 2 +- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index e505979c5e7..dd37e104044 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -396,12 +396,13 @@ class VertexAIBatchPrediction(VertexLLM): "outputConfig": vertex_batch_request["outputConfig"], "unmanagedContainerModel": unmanaged, "dedicatedResources": batch_resources, - # keyField strips the custom_id tag from each instance before it reaches the - # container (vLLM rejects unknown fields) and echoes it back as `key` in the output - # row; it only takes effect alongside an explicit instanceType. + # excludedFields strips the custom_id tag from each instance before it reaches the + # container (vLLM rejects unknown fields) and attaches it to the output row's + # instance echo; keyField does NOT strip (probed live: the container still received + # the tag and 400'd every row). "instanceConfig": { "instanceType": "object", - "keyField": VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, + "excludedFields": [VERTEX_CUSTOM_ENDPOINT_KEY_FIELD], }, } return resolved diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 1fda5d727eb..e559b039b79 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -700,6 +700,7 @@ class UnmanagedContainerModel(TypedDict, total=False): class BatchInstanceConfig(TypedDict, total=False): instanceType: ReadOnly[str] keyField: ReadOnly[str] + excludedFields: ReadOnly[list[str]] class VertexAIBatchPredictionJob(TypedDict, total=False): diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index 8691d14d392..3a5b4db04bf 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -342,7 +342,7 @@ def test_create_batch_sync_custom_endpoint_builds_unmanaged_container_job(): "startingReplicaCount": 1, "maxReplicaCount": 2, } - assert sent["instanceConfig"] == {"instanceType": "object", "keyField": "litellm_custom_id"} + assert sent["instanceConfig"] == {"instanceType": "object", "excludedFields": ["litellm_custom_id"]} def test_create_batch_sync_custom_endpoint_without_container_spec_raises_400(): From b441e81f11c2edd4a7c1e1a58769698bd64c144d Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 9 Sep 2026 17:20:11 -0400 Subject: [PATCH 4/7] refactor(vertex_ai): build custom-endpoint batch rows immutably for the LIT gates --- litellm/llms/vertex_ai/batches/handler.py | 2 +- .../llms/vertex_ai/files/transformation.py | 40 +++++++++++-------- litellm/types/llms/vertex_ai.py | 4 +- 3 files changed, 27 insertions(+), 19 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index dd37e104044..c7019a5c69c 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -402,7 +402,7 @@ class VertexAIBatchPrediction(VertexLLM): # the tag and 400'd every row). "instanceConfig": { "instanceType": "object", - "excludedFields": [VERTEX_CUSTOM_ENDPOINT_KEY_FIELD], + "excludedFields": (VERTEX_CUSTOM_ENDPOINT_KEY_FIELD,), }, } return resolved diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 765931bed8d..ed55a97145b 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -6,6 +6,7 @@ import os import re import time from collections.abc import Callable, Iterable, Iterator, Mapping +from types import MappingProxyType from typing import Any, Final, TypedDict from urllib.parse import quote, unquote, urlparse @@ -664,12 +665,11 @@ def _custom_endpoint_row_to_openai_batch_output_row(row: Mapping[str, object]) - """ key: Final = row.get("key") instance: Final = row.get("instance") - instance_map: Final = instance if isinstance(instance, Mapping) else {} - custom_id: Final = str(key if key is not None else instance_map.get(VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, "")) + tagged_custom_id: Final = instance.get(VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, "") if isinstance(instance, Mapping) else "" + custom_id: Final = str(key if key is not None else tagged_custom_id) prediction: Final = row.get("prediction") - prediction_map: Final = prediction if isinstance(prediction, Mapping) else {} - body: Final = prediction_map.get("predictions") + body: Final = prediction.get("predictions") if isinstance(prediction, Mapping) else None if not isinstance(body, Mapping): error_text: Final = str(row.get("status") or prediction or "prediction carries no response body") return _openai_batch_output_row( @@ -710,6 +710,7 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): VERTEX_CUSTOM_ENDPOINT_GCS_SEGMENT: Final = "custom-endpoints" _VERTEX_CHAT_COMPLETIONS_REQUEST_FORMAT: Final = "chatCompletions" +_EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) def get_custom_endpoint_id_from_api_base(api_base: str | None) -> str | None: @@ -729,20 +730,26 @@ def get_custom_endpoint_id_from_api_base(api_base: str | None) -> str | None: return after_endpoints[-1].split(":")[0] or None -def _openai_batch_jsonl_entry_to_custom_endpoint_row(openai_entry: dict[str, Any]) -> Mapping[str, object]: +def _openai_batch_jsonl_entry_to_custom_endpoint_row(openai_entry: Mapping[str, object]) -> Mapping[str, object]: """ One OpenAI batch JSONL line as the instance a vLLM-serving Vertex container consumes: the OpenAI request body itself tagged `@requestFormat: chatCompletions` (the container speaks OpenAI natively, so no Gemini translation), minus `model` (the batch replica - serves exactly one model) plus the custom_id under the job's `instanceConfig.keyField`. + serves exactly one model) plus the custom_id tag the job's `instanceConfig.excludedFields` + strips back out before the container sees it. """ - body: Final = openai_entry.get("body") or {} - row: Final = {k: v for k, v in body.items() if k != "model"} - return { - "@requestFormat": _VERTEX_CHAT_COMPLETIONS_REQUEST_FORMAT, - **row, - VERTEX_CUSTOM_ENDPOINT_KEY_FIELD: str(openai_entry.get("custom_id", "")), - } + raw_body: Final = openai_entry.get("body") + body: Final = raw_body if isinstance(raw_body, Mapping) else _EMPTY_MAPPING + return MappingProxyType( + { + key: value + for key, value in itertools.chain( + (("@requestFormat", _VERTEX_CHAT_COMPLETIONS_REQUEST_FORMAT),), + ((k, v) for k, v in body.items() if k != "model"), + ((VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, str(openai_entry.get("custom_id", ""))),), + ) + } + ) class _OpenAIToCustomEndpointBatchUploadStream(BaseFileUploadStream): @@ -758,7 +765,7 @@ class _OpenAIToCustomEndpointBatchUploadStream(BaseFileUploadStream): row = _openai_batch_jsonl_entry_to_custom_endpoint_row(entry) prefix = b"" if first else b"\n" first = False - yield prefix + json.dumps(row).encode("utf-8") + yield prefix + json.dumps(row, default=dict).encode("utf-8") class VertexAIFilesConfig(VertexBase, BaseFilesConfig): @@ -1231,6 +1238,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): # passthrough below, leaving the content untouched. first_row: Final = _parse_vertex_batch_output_row(first_line) is_custom_endpoint_output: Final = _is_custom_endpoint_batch_output_row(first_row) + first_row_response: Final = first_row.get("response") or () is_vertex_batch_output: Final = ( is_custom_endpoint_output or _is_vertex_embeddings_batch_output_row(first_row) @@ -1239,8 +1247,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): and "response" in first_row and "processed_time" in first_row and ( - "candidates" in first_row.get("response", {}) - or "promptFeedback" in first_row.get("response", {}) + "candidates" in first_row_response + or "promptFeedback" in first_row_response or bool(first_row.get("status")) ) ) diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index e559b039b79..51c53bf6f6b 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -1,4 +1,4 @@ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from enum import Enum from typing import Any, Final, Literal, Protocol @@ -700,7 +700,7 @@ class UnmanagedContainerModel(TypedDict, total=False): class BatchInstanceConfig(TypedDict, total=False): instanceType: ReadOnly[str] keyField: ReadOnly[str] - excludedFields: ReadOnly[list[str]] + excludedFields: ReadOnly[Sequence[str]] class VertexAIBatchPredictionJob(TypedDict, total=False): From 13ec32269d60c3f9ceba06dcac29be8b01bc3e9d Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 9 Sep 2026 18:09:44 -0400 Subject: [PATCH 5/7] fix(vertex_ai): pin custom-endpoint batch jobs to the deployment's own endpoint The endpoint id in a custom-endpoints/ file path is caller-controlled (raw gs:// file ids are accepted at batch create), so it must match the endpoint the routed deployment's api_base names, and container resolution only runs for deployments marked custom_endpoint; otherwise an attacker-supplied file id could run another endpoint's container with the deployment's project credentials. --- litellm/llms/vertex_ai/batches/handler.py | 40 +++++++++--- litellm/llms/vertex_ai/common_utils.py | 18 +++++ .../llms/vertex_ai/files/transformation.py | 20 +----- .../llms/vertex_ai/batches/test_handler.py | 65 ++++++++++++++++++- 4 files changed, 113 insertions(+), 30 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index c7019a5c69c..9312da62eb8 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -20,6 +20,7 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.llms.vertex_ai.common_utils import ( VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, VertexAIError, + get_custom_endpoint_id_from_api_base, get_vertex_base_url, ) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM @@ -156,14 +157,29 @@ class VertexAIBatchPrediction(VertexLLM): vertex_location=vertex_location or "us-central1", ) ) - if custom_endpoint and "/custom-endpoints/" not in transformed_batch_request.get("model", ""): + # The endpoint id in a custom-endpoints/ file path is caller-controlled (raw gs:// file + # ids are accepted), so it must never pick which container runs: batch jobs execute the + # container with the deployment's project credentials. The id has to match the endpoint + # the routed deployment's own api_base names, and container resolution only happens for + # deployments the admin marked custom_endpoint. + job_model: Final = transformed_batch_request.get("model", "") + deployment_endpoint_id: Final = get_custom_endpoint_id_from_api_base(api_base) + if custom_endpoint and not job_model.endswith(f"/custom-endpoints/{deployment_endpoint_id}"): raise VertexAIError( status_code=400, message=( "Vertex AI batch prediction on a `custom_endpoint` deployment requires an input " - "file uploaded through LiteLLM against that deployment (its file id carries a " - "custom-endpoints/ path); this input file targets a publisher or " - "fine-tuned Gemini model instead." + "file uploaded through LiteLLM against that same deployment: the file id's " + "custom-endpoints/ path must name the endpoint the deployment's " + "`api_base` serves from." + ), + ) + if not custom_endpoint and "/custom-endpoints/" in job_model: + raise VertexAIError( + status_code=400, + message=( + "This input file was uploaded for a `custom_endpoint` deployment; create the " + "batch against that deployment instead." ), ) gateway_api_base: Final = _gateway_api_base_or_none(api_base) @@ -174,12 +190,16 @@ class VertexAIBatchPrediction(VertexLLM): api_base=gateway_api_base, vertex_location=vertex_location or "us-central1", ) - vertex_batch_request: Final = self._resolve_custom_endpoint_container( - vertex_batch_request=resolved_batch_request, - headers=headers, - sync_handler=sync_handler, - api_base=gateway_api_base, - vertex_location=vertex_location or "us-central1", + vertex_batch_request: Final = ( + self._resolve_custom_endpoint_container( + vertex_batch_request=resolved_batch_request, + headers=headers, + sync_handler=sync_handler, + api_base=gateway_api_base, + vertex_location=vertex_location or "us-central1", + ) + if custom_endpoint + else resolved_batch_request ) is_unmanaged_container_job: Final = "unmanagedContainerModel" in vertex_batch_request diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 7bda17d4334..4c14f736bd3 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -5,6 +5,7 @@ from enum import Enum from functools import lru_cache from types import MappingProxyType from typing import Any, Final, Literal, cast, get_type_hints +from urllib.parse import urlparse import httpx from pydantic import TypeAdapter, ValidationError @@ -373,6 +374,23 @@ def get_vertex_base_model_name(model: str) -> str: VERTEX_CUSTOM_ENDPOINT_KEY_FIELD: Final = "litellm_custom_id" +def get_custom_endpoint_id_from_api_base(api_base: str | None) -> str | None: + """ + The Vertex endpoint a `custom_endpoint` deployment serves from is only recorded in its + api_base (`.../endpoints/:rawPredict` or a dedicated-domain equivalent); batch jobs need + that id to read the endpoint's containerSpec, so extract it (verb suffix stripped). + """ + if not api_base: + return None + path_segments: Final = urlparse(api_base).path.split("/") + after_endpoints: Final = tuple( + segment for prior, segment in zip(path_segments, path_segments[1:]) if prior == "endpoints" + ) + if not after_endpoints: + return None + return after_endpoints[-1].split(":")[0] or None + + def get_vertex_ai_fine_tuned_endpoint_id(model: str) -> str | None: """ Fine-tuned Gemini deployments are addressed by a numeric endpoint id, diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index ed55a97145b..006342c5a4f 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -8,7 +8,7 @@ import time from collections.abc import Callable, Iterable, Iterator, Mapping from types import MappingProxyType from typing import Any, Final, TypedDict -from urllib.parse import quote, unquote, urlparse +from urllib.parse import quote, unquote import httpx from httpx import Headers, Response @@ -41,6 +41,7 @@ from litellm.llms.base_llm.files.transformation import ( from litellm.llms.vertex_ai.common_utils import ( VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, _convert_vertex_datetime_to_openai_datetime, + get_custom_endpoint_id_from_api_base, get_vertex_ai_fine_tuned_endpoint_id, ) from litellm.llms.vertex_ai.gemini.transformation import _transform_request_body @@ -713,23 +714,6 @@ _VERTEX_CHAT_COMPLETIONS_REQUEST_FORMAT: Final = "chatCompletions" _EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) -def get_custom_endpoint_id_from_api_base(api_base: str | None) -> str | None: - """ - The Vertex endpoint a `custom_endpoint` deployment serves from is only recorded in its - api_base (`.../endpoints/:rawPredict` or a dedicated-domain equivalent); batch jobs need - that id to read the endpoint's containerSpec, so extract it (verb suffix stripped). - """ - if not api_base: - return None - path_segments: Final = urlparse(api_base).path.split("/") - after_endpoints: Final = tuple( - segment for prior, segment in zip(path_segments, path_segments[1:]) if prior == "endpoints" - ) - if not after_endpoints: - return None - return after_endpoints[-1].split(":")[0] or None - - def _openai_batch_jsonl_entry_to_custom_endpoint_row(openai_entry: Mapping[str, object]) -> Mapping[str, object]: """ One OpenAI batch JSONL line as the instance a vLLM-serving Vertex container consumes: diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index 3a5b4db04bf..9bcbd723060 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -269,6 +269,10 @@ CONTAINER_SPEC = { "healthRoute": "/ping", } MACHINE_SPEC = {"machineType": "g2-standard-12", "acceleratorType": "NVIDIA_L4", "acceleratorCount": 1} +CUSTOM_ENDPOINT_API_BASE = ( + f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}" + f"/locations/{LOCATION}/endpoints/{CUSTOM_ENDPOINT_ID}:rawPredict" +) def _custom_endpoint_get_response() -> MagicMock: @@ -317,7 +321,7 @@ def test_create_batch_sync_custom_endpoint_builds_unmanaged_container_job(): out = h.create_batch( _is_async=False, create_batch_data=CUSTOM_ENDPOINT_CREATE_DATA, - api_base=None, + api_base=CUSTOM_ENDPOINT_API_BASE, vertex_credentials=None, vertex_project=PROJECT, vertex_location=LOCATION, @@ -360,7 +364,7 @@ def test_create_batch_sync_custom_endpoint_without_container_spec_raises_400(): h.create_batch( _is_async=False, create_batch_data=CUSTOM_ENDPOINT_CREATE_DATA, - api_base=None, + api_base=CUSTOM_ENDPOINT_API_BASE, vertex_credentials=None, vertex_project=PROJECT, vertex_location=LOCATION, @@ -374,6 +378,63 @@ def test_create_batch_sync_custom_endpoint_without_container_spec_raises_400(): client.post.assert_not_called() +def test_create_batch_sync_custom_endpoint_rejects_file_for_other_endpoint(): + """The endpoint id in the file path is caller-controlled (raw gs:// ids are accepted), so a + file staged for a different endpoint must not make the deployment run that endpoint's + container with its own project credentials.""" + h = _make_handler() + client = MagicMock() + foreign_file = {"input_file_id": "gs://bucket/litellm-vertex-files/custom-endpoints/999999/file-uuid"} + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get") as safe_get, + ): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data=foreign_file, + api_base=CUSTOM_ENDPOINT_API_BASE, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + custom_endpoint=True, + ) + + assert exc_info.value.status_code == 400 + safe_get.assert_not_called() + client.post.assert_not_called() + + +def test_create_batch_sync_non_custom_deployment_rejects_custom_endpoint_file(): + """A custom-endpoints file id sent through an ordinary Vertex deployment must not trigger + container resolution at all.""" + h = _make_handler() + client = MagicMock() + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get") as safe_get, + ): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data=CUSTOM_ENDPOINT_CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert exc_info.value.status_code == 400 + safe_get.assert_not_called() + client.post.assert_not_called() + + def test_create_batch_sync_ignores_resource_shaped_api_base(): """A deployment api_base like `.../endpoints/:rawPredict` targets online inference, not the Vertex API root; grafting batch urls onto it yields guaranteed 404s, so batch operations From 1818c39fed84f255e9feaeae6239691dafff5602 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 9 Sep 2026 18:39:46 -0400 Subject: [PATCH 6/7] fix(vertex_ai): read all batch output shards and reject multi-model endpoints Sharded prediction.results-N-of-M outputs are concatenated instead of returning only shard zero, and an endpoint serving several deployed models behind a traffic split is rejected at batch create rather than silently running deployedModels[0]. --- litellm/llms/vertex_ai/batches/handler.py | 22 ++++--- litellm/llms/vertex_ai/files/handler.py | 51 ++++++++++++--- .../llms/vertex_ai/files/transformation.py | 6 +- .../llms/vertex_ai/batches/test_handler.py | 36 ++++++++++ .../files/test_vertex_ai_files_handler.py | 65 +++++++++++++++++++ 5 files changed, 159 insertions(+), 21 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 9312da62eb8..7d56e60a483 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -157,11 +157,8 @@ class VertexAIBatchPrediction(VertexLLM): vertex_location=vertex_location or "us-central1", ) ) - # The endpoint id in a custom-endpoints/ file path is caller-controlled (raw gs:// file - # ids are accepted), so it must never pick which container runs: batch jobs execute the - # container with the deployment's project credentials. The id has to match the endpoint - # the routed deployment's own api_base names, and container resolution only happens for - # deployments the admin marked custom_endpoint. + # The file-path endpoint id is caller-controlled and the job runs with the deployment's + # credentials, so it must match the endpoint the deployment's own api_base names. job_model: Final = transformed_batch_request.get("model", "") deployment_endpoint_id: Final = get_custom_endpoint_id_from_api_base(api_base) if custom_endpoint and not job_model.endswith(f"/custom-endpoints/{deployment_endpoint_id}"): @@ -358,6 +355,15 @@ class VertexAIBatchPrediction(VertexLLM): ) endpoint_view: Final[_VertexEndpointPayloadView] = {"payload": endpoint_response.json()} deployed_models: Final = endpoint_view["payload"].get("deployedModels") or () + if len(deployed_models) > 1: + raise VertexAIError( + status_code=400, + message=( + f"Vertex endpoint '{endpoint_resource}' serves {len(deployed_models)} deployed " + "models behind a traffic split, so there is no single container to replicate " + "for batch prediction; use an endpoint with exactly one deployed model" + ), + ) deployed: Final = deployed_models[0] if deployed_models else _VertexEndpointDeployedModel() deployed_model_resource: Final = deployed.get("model", "") if not deployed_model_resource: @@ -416,10 +422,8 @@ class VertexAIBatchPrediction(VertexLLM): "outputConfig": vertex_batch_request["outputConfig"], "unmanagedContainerModel": unmanaged, "dedicatedResources": batch_resources, - # excludedFields strips the custom_id tag from each instance before it reaches the - # container (vLLM rejects unknown fields) and attaches it to the output row's - # instance echo; keyField does NOT strip (probed live: the container still received - # the tag and 400'd every row). + # excludedFields (not keyField, which does not actually strip and 400s vLLM) removes + # the custom_id tag before the container sees it and echoes it in the output row. "instanceConfig": { "instanceType": "object", "excludedFields": (VERTEX_CUSTOM_ENDPOINT_KEY_FIELD,), diff --git a/litellm/llms/vertex_ai/files/handler.py b/litellm/llms/vertex_ai/files/handler.py index ac95d1348f9..ee2d93814e5 100644 --- a/litellm/llms/vertex_ai/files/handler.py +++ b/litellm/llms/vertex_ai/files/handler.py @@ -1,6 +1,7 @@ import asyncio import json import os +import re import time from collections.abc import Coroutine, Mapping from typing import Any, Final @@ -96,6 +97,41 @@ class VertexAIFilesHandler(GCSBucketBase): allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params), ) + _SHARDED_RESULTS_PATTERN: Final = re.compile(r"^(?P.*-)(?P\d{5})(?P-of-)(?P\d{5})$") + + async def _download_all_result_shards( + self, + object_path: str, + standard_callback_dynamic_params: StandardCallbackDynamicParams, + ) -> bytes | None: + """ + An unmanaged-container batch writes its output as `prediction.results-000NN-of-000NN` + shards; a file id names shard zero, so when the shard count is above one the remaining + shards are fetched and concatenated (each shard is newline-delimited JSONL). + """ + first_shard: Final = await self.download_gcs_object( + object_name=object_path, + standard_callback_dynamic_params=standard_callback_dynamic_params, + ) + shard_match: Final = self._SHARDED_RESULTS_PATTERN.match(object_path) + if first_shard is None or shard_match is None: + return first_shard + total_shards: Final = int(shard_match["total"]) + if total_shards <= 1: + return first_shard + remaining: Final = await asyncio.gather( + *( + self.download_gcs_object( + object_name=f"{shard_match['stem']}{index:05d}{shard_match['sep']}{shard_match['total']}", + standard_callback_dynamic_params=standard_callback_dynamic_params, + ) + for index in range(1, total_shards) + ) + ) + if any(shard is None for shard in remaining): + return None + return b"\n".join((first_shard.rstrip(b"\n"), *(shard.rstrip(b"\n") for shard in remaining if shard))) + async def afile_content( self, file_content_request: FileContentRequest, @@ -141,14 +177,13 @@ class VertexAIFilesHandler(GCSBucketBase): litellm_params=litellm_params, ) - download_kwargs: Final = { - "standard_callback_dynamic_params": { - "gcs_bucket_name": bucket_name, - "gcs_path_service_account": gcs_logging_config["path_service_account"], - } - } - - file_content: Final = await self.download_gcs_object(object_name=object_path, **download_kwargs) + file_content: Final = await self._download_all_result_shards( + object_path=object_path, + standard_callback_dynamic_params=StandardCallbackDynamicParams( + gcs_bucket_name=bucket_name, + gcs_path_service_account=gcs_logging_config["path_service_account"], + ), + ) decoded_file_id: Final = unquote(file_id) if file_content is None: diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 006342c5a4f..22bd5f36b86 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -744,11 +744,9 @@ class _OpenAIToCustomEndpointBatchUploadStream(BaseFileUploadStream): self._openai_file_content = openai_file_content def iter_bytes(self) -> Iterator[bytes]: - first = True - for entry in _iter_openai_jsonl_entries(self._openai_file_content): + for index, entry in enumerate(_iter_openai_jsonl_entries(self._openai_file_content)): row = _openai_batch_jsonl_entry_to_custom_endpoint_row(entry) - prefix = b"" if first else b"\n" - first = False + prefix = b"" if index == 0 else b"\n" yield prefix + json.dumps(row, default=dict).encode("utf-8") diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index 9bcbd723060..6adc6869201 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -378,6 +378,42 @@ def test_create_batch_sync_custom_endpoint_without_container_spec_raises_400(): client.post.assert_not_called() +def test_create_batch_sync_custom_endpoint_rejects_multi_deployment_endpoint(): + """An endpoint behind a traffic split has no single container to replicate; index-zero + selection could run a different container than online traffic.""" + h = _make_handler() + client = MagicMock() + multi = MagicMock() + multi.status_code = 200 + multi.json.return_value = { + "deployedModels": [ + {"model": CONTAINER_MODEL_RESOURCE}, + {"model": f"projects/{PROJECT}/locations/{LOCATION}/models/other"}, + ] + } + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get", return_value=multi), + ): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data=CUSTOM_ENDPOINT_CREATE_DATA, + api_base=CUSTOM_ENDPOINT_API_BASE, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + custom_endpoint=True, + ) + + assert exc_info.value.status_code == 400 + assert "traffic split" in str(exc_info.value) + client.post.assert_not_called() + + def test_create_batch_sync_custom_endpoint_rejects_file_for_other_endpoint(): """The endpoint id in the file path is caller-controlled (raw gs:// ids are accepted), so a file staged for a different endpoint must not make the deployment run that endpoint's diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py index 0a44f0a9a74..4b937289fc3 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py @@ -143,6 +143,71 @@ class TestVertexAIFilesHandler: assert "standard_callback_dynamic_params" in call_args.kwargs assert call_args.kwargs["standard_callback_dynamic_params"]["gcs_bucket_name"] == "test-bucket" + @pytest.mark.asyncio + async def test_afile_content_fetches_all_result_shards(self): + """A sharded unmanaged-container batch output names shard zero in the file id; reading + only that shard silently drops the rest of the batch (LIT-7387).""" + file_id = ( + "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fcustom-endpoints%2F123%2F" + "out%2Fprediction.results-00000-of-00003" + ) + shards = { + "litellm-vertex-files/custom-endpoints/123/out/prediction.results-00000-of-00003": b'{"a": 1}\n', + "litellm-vertex-files/custom-endpoints/123/out/prediction.results-00001-of-00003": b'{"b": 2}\n', + "litellm-vertex-files/custom-endpoints/123/out/prediction.results-00002-of-00003": b'{"c": 3}', + } + + async def fake_download(object_name: str, **kwargs): + return shards[object_name] + + with ( + patch.object(self.handler, "download_gcs_object", side_effect=fake_download), + patch.object( + self.handler, + "get_gcs_logging_config", + new_callable=AsyncMock, + return_value=_mock_gcs_logging_config(), + ), + ): + result = await self.handler.afile_content( + file_content_request=FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None), + vertex_credentials=None, + vertex_project="test-project", + vertex_location="us-central1", + timeout=60.0, + max_retries=3, + ) + + assert result.response.content == b'{"a": 1}\n{"b": 2}\n{"c": 3}' + + @pytest.mark.asyncio + async def test_afile_content_single_shard_unchanged(self): + file_id = ( + "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fcustom-endpoints%2F123%2F" + "out%2Fprediction.results-00000-of-00001" + ) + with ( + patch.object(self.handler, "download_gcs_object", new_callable=AsyncMock) as mock_download, + patch.object( + self.handler, + "get_gcs_logging_config", + new_callable=AsyncMock, + return_value=_mock_gcs_logging_config(), + ), + ): + mock_download.return_value = b'{"a": 1}\n' + result = await self.handler.afile_content( + file_content_request=FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None), + vertex_credentials=None, + vertex_project="test-project", + vertex_location="us-central1", + timeout=60.0, + max_retries=3, + ) + + assert result.response.content == b'{"a": 1}\n' + mock_download.assert_called_once() + @pytest.mark.asyncio async def test_afile_content_missing_file_id(self): """Test async file content retrieval with missing file_id""" From 5d0c5526942d8394fb31034eac400ff4e2df92ae Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 9 Sep 2026 19:03:53 -0400 Subject: [PATCH 7/7] fix(vertex_ai): bound and pin batch output shard reads, clamp zero starting replicas The shard fan-out now only triggers on the exact Vertex unmanaged-batch output layout (prediction-custom-unmanaged-model-/prediction.results- 00000-of-NNNNN), is capped at 512 shards, and fetches sequentially, so a crafted upload filename cannot fan out GCS requests. A scale-to-zero endpoint's minReplicaCount 0 is clamped to 1 for the batch job. ruff format applied to the touched files. --- litellm/llms/vertex_ai/batches/handler.py | 14 +-- litellm/llms/vertex_ai/files/handler.py | 22 ++++- .../llms/vertex_ai/files/transformation.py | 11 +-- .../llms/vertex_ai/batches/test_handler.py | 94 ++++++++++--------- .../vertex_ai/batches/test_transformation.py | 3 +- .../files/test_vertex_ai_files_handler.py | 69 ++++++++++++-- .../test_vertex_ai_files_transformation.py | 4 +- 7 files changed, 144 insertions(+), 73 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 7d56e60a483..84c808e1f21 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -380,9 +380,7 @@ class VertexAIBatchPrediction(VertexLLM): model=deployed_model_resource, vertex_location=vertex_location, ) - model_fetched: Final[_FetchedResponseView] = { - "response": safe_get(sync_handler, model_url, headers=headers) - } + model_fetched: Final[_FetchedResponseView] = {"response": safe_get(sync_handler, model_url, headers=headers)} model_response: Final = model_fetched["response"] if model_response.status_code != 200: raise VertexAIError( @@ -410,10 +408,12 @@ class VertexAIBatchPrediction(VertexLLM): "spec to size the batch replicas from" ), ) + # A scale-to-zero online endpoint reports minReplicaCount 0, but a batch job must start + # at least one replica. batch_resources: Final[BatchDedicatedResources] = { "machineSpec": machine_spec, - "startingReplicaCount": online_resources.get("minReplicaCount", 1), - "maxReplicaCount": online_resources.get("maxReplicaCount", 1), + "startingReplicaCount": max(online_resources.get("minReplicaCount", 1), 1), + "maxReplicaCount": max(online_resources.get("maxReplicaCount", 1), 1), } unmanaged: Final[UnmanagedContainerModel] = {"containerSpec": container_spec} resolved: Final[VertexAIBatchPredictionJob] = { @@ -471,7 +471,9 @@ class VertexAIBatchPrediction(VertexLLM): """Return the base url for the vertex garden models""" # POST https://LOCATION-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION/batchPredictionJobs base_url: Final = get_vertex_base_url(vertex_location) - return f"{base_url}/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/batchPredictionJobs" + return ( + f"{base_url}/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/batchPredictionJobs" + ) def retrieve_batch( self, diff --git a/litellm/llms/vertex_ai/files/handler.py b/litellm/llms/vertex_ai/files/handler.py index ee2d93814e5..6d0ca7b0374 100644 --- a/litellm/llms/vertex_ai/files/handler.py +++ b/litellm/llms/vertex_ai/files/handler.py @@ -97,7 +97,13 @@ class VertexAIFilesHandler(GCSBucketBase): allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params), ) - _SHARDED_RESULTS_PATTERN: Final = re.compile(r"^(?P.*-)(?P\d{5})(?P-of-)(?P\d{5})$") + # Only the exact directory/file layout Vertex writes for unmanaged-container batch outputs, + # pinned to shard zero: the object path is derived from a caller-controlled file id, so a + # looser pattern would let a crafted upload filename trigger shard fan-out. + _SHARDED_RESULTS_PATTERN: Final = re.compile( + r"^(?P.*/prediction-custom-unmanaged-model-[^/]+/prediction\.results-)00000(?P-of-)(?P\d{5})$" + ) + _MAX_RESULT_SHARDS: Final = 512 async def _download_all_result_shards( self, @@ -119,14 +125,20 @@ class VertexAIFilesHandler(GCSBucketBase): total_shards: Final = int(shard_match["total"]) if total_shards <= 1: return first_shard - remaining: Final = await asyncio.gather( - *( - self.download_gcs_object( + if total_shards > self._MAX_RESULT_SHARDS: + raise ValueError( + f"Vertex batch output claims {total_shards} shards, above the supported maximum " + f"of {self._MAX_RESULT_SHARDS}" + ) + # Sequential fetch keeps memory and connection use bounded by one shard at a time. + remaining: Final = tuple( + [ + await self.download_gcs_object( object_name=f"{shard_match['stem']}{index:05d}{shard_match['sep']}{shard_match['total']}", standard_callback_dynamic_params=standard_callback_dynamic_params, ) for index in range(1, total_shards) - ) + ] ) if any(shard is None for shard in remaining): return None diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 22bd5f36b86..b96b5afb4bd 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -666,7 +666,9 @@ def _custom_endpoint_row_to_openai_batch_output_row(row: Mapping[str, object]) - """ key: Final = row.get("key") instance: Final = row.get("instance") - tagged_custom_id: Final = instance.get(VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, "") if isinstance(instance, Mapping) else "" + tagged_custom_id: Final = ( + instance.get(VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, "") if isinstance(instance, Mapping) else "" + ) custom_id: Final = str(key if key is not None else tagged_custom_id) prediction: Final = row.get("prediction") @@ -833,8 +835,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): if purpose == "batch" and custom_endpoint_id is not None: safe_endpoint_id: Final = sanitize_cloud_object_path(custom_endpoint_id, fallback="endpoint") return ( - f"{VERTEX_AI_MANAGED_GCS_PREFIX}{VERTEX_CUSTOM_ENDPOINT_GCS_SEGMENT}/" - f"{safe_endpoint_id}/{uuid.uuid4()}" + f"{VERTEX_AI_MANAGED_GCS_PREFIX}{VERTEX_CUSTOM_ENDPOINT_GCS_SEGMENT}/{safe_endpoint_id}/{uuid.uuid4()}" ) if purpose == "batch": ## 1. If jsonl, derive the object name from the deployment model (or the first entry's) @@ -881,9 +882,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): configured_model: Final = litellm_params.get("model") deployment_api_base: Final = litellm_params.get("api_base") custom_endpoint_id: Final = ( - get_custom_endpoint_id_from_api_base( - deployment_api_base if isinstance(deployment_api_base, str) else None - ) + get_custom_endpoint_id_from_api_base(deployment_api_base if isinstance(deployment_api_base, str) else None) if litellm_params.get("custom_endpoint") else None ) diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index 6adc6869201..08b6b7171da 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -47,11 +47,7 @@ PROJECT = "my-project" LOCATION = "us-central1" BATCH_ID = "3814889423749775360" -CREATE_DATA = { - "input_file_id": ( - "gs://bucket/publishers/google/models/gemini-1.5-flash-001/file-uuid" - ) -} +CREATE_DATA = {"input_file_id": ("gs://bucket/publishers/google/models/gemini-1.5-flash-001/file-uuid")} def _vertex_job_response(state: str = "JOB_STATE_SUCCEEDED") -> dict: @@ -104,8 +100,7 @@ def test_create_vertex_batch_url(): h = _make_handler() url = h.create_vertex_batch_url(vertex_location=LOCATION, vertex_project=PROJECT) assert url == ( - f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}" - f"/locations/{LOCATION}/batchPredictionJobs" + f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}/locations/{LOCATION}/batchPredictionJobs" ) @@ -206,9 +201,7 @@ def test_create_batch_sync_does_not_resolve_publisher_models(): ENDPOINT_ID = "7768560373388541952" -ENDPOINT_CREATE_DATA = { - "input_file_id": f"gs://bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/file-uuid" -} +ENDPOINT_CREATE_DATA = {"input_file_id": f"gs://bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/file-uuid"} TUNED_MODEL_RESOURCE = f"projects/{PROJECT}/locations/{LOCATION}/models/1234509876" @@ -217,9 +210,7 @@ def _endpoint_get_response(deployed_models: list | None = None) -> MagicMock: resp.status_code = 200 resp.json.return_value = { "name": f"projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", - "deployedModels": ( - deployed_models if deployed_models is not None else [{"model": TUNED_MODEL_RESOURCE}] - ), + "deployedModels": (deployed_models if deployed_models is not None else [{"model": TUNED_MODEL_RESOURCE}]), } return resp @@ -275,7 +266,7 @@ CUSTOM_ENDPOINT_API_BASE = ( ) -def _custom_endpoint_get_response() -> MagicMock: +def _custom_endpoint_get_response(min_replica_count: int = 1) -> MagicMock: resp = MagicMock() resp.status_code = 200 resp.json.return_value = { @@ -283,7 +274,11 @@ def _custom_endpoint_get_response() -> MagicMock: "deployedModels": [ { "model": CONTAINER_MODEL_RESOURCE, - "dedicatedResources": {"machineSpec": MACHINE_SPEC, "minReplicaCount": 1, "maxReplicaCount": 2}, + "dedicatedResources": { + "machineSpec": MACHINE_SPEC, + "minReplicaCount": min_replica_count, + "maxReplicaCount": 2, + }, } ], } @@ -378,6 +373,37 @@ def test_create_batch_sync_custom_endpoint_without_container_spec_raises_400(): client.post.assert_not_called() +def test_create_batch_sync_custom_endpoint_scale_to_zero_starts_one_replica(): + """A scale-to-zero online endpoint reports minReplicaCount 0, which Vertex rejects as a batch + startingReplicaCount; the job must clamp to at least one replica.""" + h = _make_handler() + client = MagicMock() + client.post.return_value = _http_response() + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch( + f"{HMOD}.safe_get", + side_effect=[_custom_endpoint_get_response(min_replica_count=0), _container_model_get_response()], + ), + ): + h.create_batch( + _is_async=False, + create_batch_data=CUSTOM_ENDPOINT_CREATE_DATA, + api_base=CUSTOM_ENDPOINT_API_BASE, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + custom_endpoint=True, + ) + + sent = json.loads(client.post.call_args.kwargs["data"]) + assert sent["dedicatedResources"]["startingReplicaCount"] == 1 + assert sent["dedicatedResources"]["maxReplicaCount"] == 2 + + def test_create_batch_sync_custom_endpoint_rejects_multi_deployment_endpoint(): """An endpoint behind a traffic split has no single container to replicate; index-zero selection could run a different container than online traffic.""" @@ -628,9 +654,7 @@ def test_create_batch_sync_httpstatuserror_propagates(): client = MagicMock() request = httpx.Request("POST", "https://x/batchPredictionJobs") err_response = httpx.Response(status_code=500, request=request, text="boom") - client.post.side_effect = httpx.HTTPStatusError( - "boom", request=request, response=err_response - ) + client.post.side_effect = httpx.HTTPStatusError("boom", request=request, response=err_response) with patch(f"{HMOD}._get_httpx_client", return_value=client): with pytest.raises(httpx.HTTPStatusError): @@ -780,9 +804,7 @@ def test_retrieve_batch_sync_invokes_logging_pre_call(): logging_obj.pre_call.assert_called_once() _, kwargs = logging_obj.pre_call.call_args - assert kwargs["additional_args"]["api_base"].endswith( - f"/batchPredictionJobs/{BATCH_ID}" - ) + assert kwargs["additional_args"]["api_base"].endswith(f"/batchPredictionJobs/{BATCH_ID}") # =========================================================================== # @@ -853,9 +875,7 @@ def test_list_batches_sync_omits_unset_pagination_params(): def test_list_batches_async_returns_coroutine(): h = _make_handler() async_client = MagicMock() - async_client.get = AsyncMock( - return_value=_http_response(json_body=_list_response()) - ) + async_client.get = AsyncMock(return_value=_http_response(json_body=_list_response())) sync_client = MagicMock() with ( @@ -910,9 +930,7 @@ def test_cancel_batch_sync_posts_cancel_then_retrieves(): h = _make_handler() client = MagicMock() client.post.return_value = _http_response(json_body={}) - client.get.return_value = _http_response( - json_body=_vertex_job_response(state="JOB_STATE_CANCELLED") - ) + client.get.return_value = _http_response(json_body=_vertex_job_response(state="JOB_STATE_CANCELLED")) with patch(f"{HMOD}._get_httpx_client", return_value=client): out = h.cancel_batch( @@ -944,9 +962,7 @@ def test_cancel_batch_async_returns_coroutine_posts_then_retrieves(): async_client = MagicMock() async_client.post = AsyncMock(return_value=_http_response(json_body={})) async_client.get = AsyncMock( - return_value=_http_response( - json_body=_vertex_job_response(state="JOB_STATE_CANCELLED") - ) + return_value=_http_response(json_body=_vertex_job_response(state="JOB_STATE_CANCELLED")) ) with ( @@ -1004,9 +1020,7 @@ def test_cancel_batch_sync_proxy_url_without_cancel_suffix_uses_rsplit_branch(): ) client = MagicMock() client.post.return_value = _http_response(json_body={}) - client.get.return_value = _http_response( - json_body=_vertex_job_response(state="JOB_STATE_CANCELLED") - ) + client.get.return_value = _http_response(json_body=_vertex_job_response(state="JOB_STATE_CANCELLED")) with patch(f"{HMOD}._get_httpx_client", return_value=client): out = h.cancel_batch( @@ -1032,9 +1046,7 @@ def test_cancel_batch_sync_httpstatuserror_logged_and_reraised(): client = MagicMock() request = httpx.Request("POST", "https://x/batchPredictionJobs/1:cancel") err_response = httpx.Response(status_code=502, request=request, text="bad gw") - client.post.side_effect = httpx.HTTPStatusError( - "boom", request=request, response=err_response - ) + client.post.side_effect = httpx.HTTPStatusError("boom", request=request, response=err_response) with patch(f"{HMOD}._get_httpx_client", return_value=client): with pytest.raises(httpx.HTTPStatusError): @@ -1056,11 +1068,7 @@ def test_create_batch_async_httpstatuserror_logged_and_reraised(): async_client = MagicMock() request = httpx.Request("POST", "https://x/batchPredictionJobs") err_response = httpx.Response(status_code=500, request=request, text="boom") - async_client.post = AsyncMock( - side_effect=httpx.HTTPStatusError( - "boom", request=request, response=err_response - ) - ) + async_client.post = AsyncMock(side_effect=httpx.HTTPStatusError("boom", request=request, response=err_response)) with ( patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()), @@ -1167,9 +1175,7 @@ def test_async_cancel_batch_httpstatuserror_and_retrieve_non_200(): async_client = MagicMock() request = httpx.Request("POST", "https://x/batchPredictionJobs/1:cancel") err_response = httpx.Response(status_code=502, request=request, text="bad") - async_client.post = AsyncMock( - side_effect=httpx.HTTPStatusError("boom", request=request, response=err_response) - ) + async_client.post = AsyncMock(side_effect=httpx.HTTPStatusError("boom", request=request, response=err_response)) async_client.get = AsyncMock() with ( patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()), diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index c80f244a659..8726ccdea9e 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -36,8 +36,7 @@ INPUT_FILE = ( ENDPOINT_ID = "7768560373388541952" ENDPOINT_INPUT_FILE = ( - f"gs://litellm-testing-bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/" - "e9412502-2c91-42a6-8e61-f5c294cc0fc8" + f"gs://litellm-testing-bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/e9412502-2c91-42a6-8e61-f5c294cc0fc8" ) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py index 4b937289fc3..442928e55c2 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py @@ -149,12 +149,12 @@ class TestVertexAIFilesHandler: only that shard silently drops the rest of the batch (LIT-7387).""" file_id = ( "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fcustom-endpoints%2F123%2F" - "out%2Fprediction.results-00000-of-00003" + "prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z%2Fprediction.results-00000-of-00003" ) shards = { - "litellm-vertex-files/custom-endpoints/123/out/prediction.results-00000-of-00003": b'{"a": 1}\n', - "litellm-vertex-files/custom-endpoints/123/out/prediction.results-00001-of-00003": b'{"b": 2}\n', - "litellm-vertex-files/custom-endpoints/123/out/prediction.results-00002-of-00003": b'{"c": 3}', + "litellm-vertex-files/custom-endpoints/123/prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z/prediction.results-00000-of-00003": b'{"a": 1}\n', + "litellm-vertex-files/custom-endpoints/123/prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z/prediction.results-00001-of-00003": b'{"b": 2}\n', + "litellm-vertex-files/custom-endpoints/123/prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z/prediction.results-00002-of-00003": b'{"c": 3}', } async def fake_download(object_name: str, **kwargs): @@ -184,7 +184,7 @@ class TestVertexAIFilesHandler: async def test_afile_content_single_shard_unchanged(self): file_id = ( "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fcustom-endpoints%2F123%2F" - "out%2Fprediction.results-00000-of-00001" + "prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z%2Fprediction.results-00000-of-00001" ) with ( patch.object(self.handler, "download_gcs_object", new_callable=AsyncMock) as mock_download, @@ -208,6 +208,62 @@ class TestVertexAIFilesHandler: assert result.response.content == b'{"a": 1}\n' mock_download.assert_called_once() + @pytest.mark.asyncio + async def test_afile_content_crafted_upload_filename_does_not_fan_out(self): + """The object path comes from a caller-controlled file id, so an ordinary upload whose + name mimics the shard suffix must not trigger shard fetches (99,998 GCS requests from one + crafted 'prediction.results-00000-of-99999' filename).""" + file_id = "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fuploads%2Fabc-prediction.results-00000-of-99999" + with ( + patch.object(self.handler, "download_gcs_object", new_callable=AsyncMock) as mock_download, + patch.object( + self.handler, + "get_gcs_logging_config", + new_callable=AsyncMock, + return_value=_mock_gcs_logging_config(), + ), + ): + mock_download.return_value = b"tiny" + result = await self.handler.afile_content( + file_content_request=FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None), + vertex_credentials=None, + vertex_project="test-project", + vertex_location="us-central1", + timeout=60.0, + max_retries=3, + ) + + assert result.response.content == b"tiny" + mock_download.assert_called_once() + + @pytest.mark.asyncio + async def test_afile_content_shard_count_above_cap_raises(self): + file_id = ( + "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fcustom-endpoints%2F123%2F" + "prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z%2Fprediction.results-00000-of-99999" + ) + with ( + patch.object(self.handler, "download_gcs_object", new_callable=AsyncMock) as mock_download, + patch.object( + self.handler, + "get_gcs_logging_config", + new_callable=AsyncMock, + return_value=_mock_gcs_logging_config(), + ), + ): + mock_download.return_value = b"tiny" + with pytest.raises(ValueError, match="shards"): + await self.handler.afile_content( + file_content_request=FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None), + vertex_credentials=None, + vertex_project="test-project", + vertex_location="us-central1", + timeout=60.0, + max_retries=3, + ) + + mock_download.assert_called_once() + @pytest.mark.asyncio async def test_afile_content_missing_file_id(self): """Test async file content retrieval with missing file_id""" @@ -247,8 +303,7 @@ class TestVertexAIFilesHandler: with pytest.raises( ValueError, match=re.escape( - "Failed to download file from GCS: " - "gs://test-bucket/litellm-vertex-files/uploads/abc-test-file.txt" + "Failed to download file from GCS: gs://test-bucket/litellm-vertex-files/uploads/abc-test-file.txt" ), ): await self.handler.afile_content( diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index be77ece8dfe..2fe8836712e 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -193,9 +193,7 @@ class TestBatchObjectNaming: def test_should_store_fine_tuned_endpoint_under_endpoints_path(self, config): """A numeric endpoint id must not be filed under publishers/google/models/gemini/, which the batch transformation later mangles into a nonexistent publisher model (LIT-6899).""" - object_name = config._get_gcs_object_name_from_batch_jsonl( - [{"body": {"model": "gemini/7768560373388541952"}}] - ) + object_name = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "gemini/7768560373388541952"}}]) assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/") assert "publishers" not in object_name