diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index a15ea4d845b..84c808e1f21 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,19 @@ 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_custom_endpoint_id_from_api_base, + 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 +65,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 +89,29 @@ 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/`, + 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() @@ -97,15 +137,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( @@ -126,17 +157,53 @@ class VertexAIBatchPrediction(VertexLLM): vertex_location=vertex_location or "us-central1", ) ) - vertex_batch_request: Final = self._resolve_fine_tuned_endpoint_model( + # 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}"): + raise VertexAIError( + status_code=400, + message=( + "Vertex AI batch prediction on a `custom_endpoint` deployment requires an input " + "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) + resolved_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", ) + 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 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: @@ -145,7 +212,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, @@ -155,7 +222,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: @@ -249,6 +316,121 @@ 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 () + 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: + 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" + ), + ) + # 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": max(online_resources.get("minReplicaCount", 1), 1), + "maxReplicaCount": max(online_resources.get("maxReplicaCount", 1), 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, + # 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,), + }, + } + return resolved + async def _async_create_batch( self, vertex_batch_request: VertexAIBatchPredictionJob, @@ -284,11 +466,14 @@ 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, @@ -325,7 +510,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 +666,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 +764,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/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..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 @@ -370,6 +371,26 @@ def get_vertex_base_model_name(model: str) -> str: return model +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/handler.py b/litellm/llms/vertex_ai/files/handler.py index ac95d1348f9..6d0ca7b0374 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,53 @@ class VertexAIFilesHandler(GCSBucketBase): allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params), ) + # 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, + 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 + 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 + 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 +189,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 263956efc9f..b96b5afb4bd 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 @@ -38,7 +39,9 @@ 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_custom_endpoint_id_from_api_base, get_vertex_ai_fine_tuned_endpoint_id, ) from litellm.llms.vertex_ai.gemini.transformation import _transform_request_body @@ -645,6 +648,41 @@ 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") + 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") + 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( + 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 +711,47 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): return self._iter_vertex_jsonl_chunks() +VERTEX_CUSTOM_ENDPOINT_GCS_SEGMENT: Final = "custom-endpoints" +_VERTEX_CHAT_COMPLETIONS_REQUEST_FORMAT: Final = "chatCompletions" +_EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) + + +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 tag the job's `instanceConfig.excludedFields` + strips back out before the container sees it. + """ + 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): + """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]: + 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 index == 0 else b"\n" + yield prefix + json.dumps(row, default=dict).encode("utf-8") + + class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ Config for VertexAI Files @@ -740,13 +819,24 @@ 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}/{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 +871,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,20 +880,33 @@ 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}" 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 [] @@ -871,12 +964,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", ) } @@ -1120,14 +1218,20 @@ 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) + 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) + or ( + "request" in first_row + and "response" in first_row + and "processed_time" in first_row + and ( + "candidates" in first_row_response + or "promptFeedback" in first_row_response + or bool(first_row.get("status")) + ) ) ) if not is_vertex_batch_output: @@ -1155,6 +1259,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..51c53bf6f6b 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -1,7 +1,9 @@ +from collections.abc import Mapping, Sequence from enum import Enum from typing import Any, Final, Literal, Protocol from typing_extensions import ( + ReadOnly, Required, TypedDict, ) @@ -678,11 +680,37 @@ 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] + excludedFields: ReadOnly[Sequence[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 24df3214da3..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 @@ -257,6 +248,294 @@ 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} +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(min_replica_count: int = 1) -> 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": min_replica_count, + "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=CUSTOM_ENDPOINT_API_BASE, + 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", "excludedFields": ["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=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 "containerSpec" in str(exc_info.value) + 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.""" + 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 + 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 + 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", [ @@ -317,9 +596,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() @@ -339,8 +618,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(): @@ -375,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): @@ -527,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}") # =========================================================================== # @@ -600,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 ( @@ -657,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( @@ -691,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 ( @@ -751,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( @@ -779,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): @@ -803,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()), @@ -914,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 232c6413e78..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" ) @@ -391,6 +390,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_handler.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py index 0a44f0a9a74..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 @@ -143,6 +143,127 @@ 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" + "prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z%2Fprediction.results-00000-of-00003" + ) + shards = { + "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): + 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" + "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, + 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_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""" @@ -182,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 8a249820cbd..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 @@ -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): @@ -167,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 @@ -208,10 +232,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: @@ -220,14 +275,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( @@ -244,6 +303,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"