Merge 5d0c552694 into 0c98afa780

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/<id> 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/<id>: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/<endpoint id> 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",

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-<timestamp>`.
         """
 
         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/<id>`) 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/<publisher>/models/<model>` or `endpoints/<numeric id>` path from a
-        gcs uri, or None if the uri does not contain one.
+        Returns the `publishers/<publisher>/models/<model>`, `endpoints/<numeric id>`, or
+        `custom-endpoints/<numeric id>` 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/<digits>` 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/<digits>` 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():

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/<id>: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,

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<stem>.*/prediction-custom-unmanaged-model-[^/]+/prediction\.results-)00000(?P<sep>-of-)(?P<total>\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:

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/<endpoint id>: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.

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):

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/<id>: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()),

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
 # =========================================================================== #

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(

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/<id>: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/<id>,
         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"