fix(vertex_ai): read all batch output shards and reject multi-model endpoints

Sharded prediction.results-N-of-M outputs are concatenated instead of
returning only shard zero, and an endpoint serving several deployed models
behind a traffic split is rejected at batch create rather than silently
running deployedModels[0].

litellm/llms/vertex_ai/batches/handler.py

@@ -157,11 +157,8 @@ class VertexAIBatchPrediction(VertexLLM):
                 vertex_location=vertex_location or "us-central1",
             )
         )
-        # The endpoint id in a custom-endpoints/ file path is caller-controlled (raw gs:// file
-        # ids are accepted), so it must never pick which container runs: batch jobs execute the
-        # container with the deployment's project credentials. The id has to match the endpoint
-        # the routed deployment's own api_base names, and container resolution only happens for
-        # deployments the admin marked custom_endpoint.
+        # The file-path endpoint id is caller-controlled and the job runs with the deployment's
+        # credentials, so it must match the endpoint the deployment's own api_base names.
         job_model: Final = transformed_batch_request.get("model", "")
         deployment_endpoint_id: Final = get_custom_endpoint_id_from_api_base(api_base)
         if custom_endpoint and not job_model.endswith(f"/custom-endpoints/{deployment_endpoint_id}"):
@@ -358,6 +355,15 @@ class VertexAIBatchPrediction(VertexLLM):
             )
         endpoint_view: Final[_VertexEndpointPayloadView] = {"payload": endpoint_response.json()}
         deployed_models: Final = endpoint_view["payload"].get("deployedModels") or ()
+        if len(deployed_models) > 1:
+            raise VertexAIError(
+                status_code=400,
+                message=(
+                    f"Vertex endpoint '{endpoint_resource}' serves {len(deployed_models)} deployed "
+                    "models behind a traffic split, so there is no single container to replicate "
+                    "for batch prediction; use an endpoint with exactly one deployed model"
+                ),
+            )
         deployed: Final = deployed_models[0] if deployed_models else _VertexEndpointDeployedModel()
         deployed_model_resource: Final = deployed.get("model", "")
         if not deployed_model_resource:
@@ -416,10 +422,8 @@ class VertexAIBatchPrediction(VertexLLM):
             "outputConfig": vertex_batch_request["outputConfig"],
             "unmanagedContainerModel": unmanaged,
             "dedicatedResources": batch_resources,
-            # excludedFields strips the custom_id tag from each instance before it reaches the
-            # container (vLLM rejects unknown fields) and attaches it to the output row's
-            # instance echo; keyField does NOT strip (probed live: the container still received
-            # the tag and 400'd every row).
+            # excludedFields (not keyField, which does not actually strip and 400s vLLM) removes
+            # the custom_id tag before the container sees it and echoes it in the output row.
             "instanceConfig": {
                 "instanceType": "object",
                 "excludedFields": (VERTEX_CUSTOM_ENDPOINT_KEY_FIELD,),

litellm/llms/vertex_ai/files/handler.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 import os
+import re
 import time
 from collections.abc import Coroutine, Mapping
 from typing import Any, Final
@@ -96,6 +97,41 @@ class VertexAIFilesHandler(GCSBucketBase):
             allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params),
         )
 
+    _SHARDED_RESULTS_PATTERN: Final = re.compile(r"^(?P<stem>.*-)(?P<index>\d{5})(?P<sep>-of-)(?P<total>\d{5})$")
+
+    async def _download_all_result_shards(
+        self,
+        object_path: str,
+        standard_callback_dynamic_params: StandardCallbackDynamicParams,
+    ) -> bytes | None:
+        """
+        An unmanaged-container batch writes its output as `prediction.results-000NN-of-000NN`
+        shards; a file id names shard zero, so when the shard count is above one the remaining
+        shards are fetched and concatenated (each shard is newline-delimited JSONL).
+        """
+        first_shard: Final = await self.download_gcs_object(
+            object_name=object_path,
+            standard_callback_dynamic_params=standard_callback_dynamic_params,
+        )
+        shard_match: Final = self._SHARDED_RESULTS_PATTERN.match(object_path)
+        if first_shard is None or shard_match is None:
+            return first_shard
+        total_shards: Final = int(shard_match["total"])
+        if total_shards <= 1:
+            return first_shard
+        remaining: Final = await asyncio.gather(
+            *(
+                self.download_gcs_object(
+                    object_name=f"{shard_match['stem']}{index:05d}{shard_match['sep']}{shard_match['total']}",
+                    standard_callback_dynamic_params=standard_callback_dynamic_params,
+                )
+                for index in range(1, total_shards)
+            )
+        )
+        if any(shard is None for shard in remaining):
+            return None
+        return b"\n".join((first_shard.rstrip(b"\n"), *(shard.rstrip(b"\n") for shard in remaining if shard)))
+
     async def afile_content(
         self,
         file_content_request: FileContentRequest,
@@ -141,14 +177,13 @@ class VertexAIFilesHandler(GCSBucketBase):
             litellm_params=litellm_params,
         )
 
-        download_kwargs: Final = {
-            "standard_callback_dynamic_params": {
-                "gcs_bucket_name": bucket_name,
-                "gcs_path_service_account": gcs_logging_config["path_service_account"],
-            }
-        }
-
-        file_content: Final = await self.download_gcs_object(object_name=object_path, **download_kwargs)
+        file_content: Final = await self._download_all_result_shards(
+            object_path=object_path,
+            standard_callback_dynamic_params=StandardCallbackDynamicParams(
+                gcs_bucket_name=bucket_name,
+                gcs_path_service_account=gcs_logging_config["path_service_account"],
+            ),
+        )
         decoded_file_id: Final = unquote(file_id)
 
         if file_content is None:

litellm/llms/vertex_ai/files/transformation.py

@@ -744,11 +744,9 @@ class _OpenAIToCustomEndpointBatchUploadStream(BaseFileUploadStream):
         self._openai_file_content = openai_file_content
 
     def iter_bytes(self) -> Iterator[bytes]:
-        first = True
-        for entry in _iter_openai_jsonl_entries(self._openai_file_content):
+        for index, entry in enumerate(_iter_openai_jsonl_entries(self._openai_file_content)):
             row = _openai_batch_jsonl_entry_to_custom_endpoint_row(entry)
-            prefix = b"" if first else b"\n"
-            first = False
+            prefix = b"" if index == 0 else b"\n"
             yield prefix + json.dumps(row, default=dict).encode("utf-8")
 
 

tests/test_litellm/llms/vertex_ai/batches/test_handler.py

@@ -378,6 +378,42 @@ def test_create_batch_sync_custom_endpoint_without_container_spec_raises_400():
     client.post.assert_not_called()
 
 
+def test_create_batch_sync_custom_endpoint_rejects_multi_deployment_endpoint():
+    """An endpoint behind a traffic split has no single container to replicate; index-zero
+    selection could run a different container than online traffic."""
+    h = _make_handler()
+    client = MagicMock()
+    multi = MagicMock()
+    multi.status_code = 200
+    multi.json.return_value = {
+        "deployedModels": [
+            {"model": CONTAINER_MODEL_RESOURCE},
+            {"model": f"projects/{PROJECT}/locations/{LOCATION}/models/other"},
+        ]
+    }
+
+    with (
+        patch(f"{HMOD}._get_httpx_client", return_value=client),
+        patch(f"{HMOD}.safe_get", return_value=multi),
+    ):
+        with pytest.raises(VertexAIError) as exc_info:
+            h.create_batch(
+                _is_async=False,
+                create_batch_data=CUSTOM_ENDPOINT_CREATE_DATA,
+                api_base=CUSTOM_ENDPOINT_API_BASE,
+                vertex_credentials=None,
+                vertex_project=PROJECT,
+                vertex_location=LOCATION,
+                timeout=600.0,
+                max_retries=None,
+                custom_endpoint=True,
+            )
+
+    assert exc_info.value.status_code == 400
+    assert "traffic split" in str(exc_info.value)
+    client.post.assert_not_called()
+
+
 def test_create_batch_sync_custom_endpoint_rejects_file_for_other_endpoint():
     """The endpoint id in the file path is caller-controlled (raw gs:// ids are accepted), so a
     file staged for a different endpoint must not make the deployment run that endpoint's

tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py

@@ -143,6 +143,71 @@ class TestVertexAIFilesHandler:
             assert "standard_callback_dynamic_params" in call_args.kwargs
             assert call_args.kwargs["standard_callback_dynamic_params"]["gcs_bucket_name"] == "test-bucket"
 
+    @pytest.mark.asyncio
+    async def test_afile_content_fetches_all_result_shards(self):
+        """A sharded unmanaged-container batch output names shard zero in the file id; reading
+        only that shard silently drops the rest of the batch (LIT-7387)."""
+        file_id = (
+            "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fcustom-endpoints%2F123%2F"
+            "out%2Fprediction.results-00000-of-00003"
+        )
+        shards = {
+            "litellm-vertex-files/custom-endpoints/123/out/prediction.results-00000-of-00003": b'{"a": 1}\n',
+            "litellm-vertex-files/custom-endpoints/123/out/prediction.results-00001-of-00003": b'{"b": 2}\n',
+            "litellm-vertex-files/custom-endpoints/123/out/prediction.results-00002-of-00003": b'{"c": 3}',
+        }
+
+        async def fake_download(object_name: str, **kwargs):
+            return shards[object_name]
+
+        with (
+            patch.object(self.handler, "download_gcs_object", side_effect=fake_download),
+            patch.object(
+                self.handler,
+                "get_gcs_logging_config",
+                new_callable=AsyncMock,
+                return_value=_mock_gcs_logging_config(),
+            ),
+        ):
+            result = await self.handler.afile_content(
+                file_content_request=FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None),
+                vertex_credentials=None,
+                vertex_project="test-project",
+                vertex_location="us-central1",
+                timeout=60.0,
+                max_retries=3,
+            )
+
+        assert result.response.content == b'{"a": 1}\n{"b": 2}\n{"c": 3}'
+
+    @pytest.mark.asyncio
+    async def test_afile_content_single_shard_unchanged(self):
+        file_id = (
+            "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fcustom-endpoints%2F123%2F"
+            "out%2Fprediction.results-00000-of-00001"
+        )
+        with (
+            patch.object(self.handler, "download_gcs_object", new_callable=AsyncMock) as mock_download,
+            patch.object(
+                self.handler,
+                "get_gcs_logging_config",
+                new_callable=AsyncMock,
+                return_value=_mock_gcs_logging_config(),
+            ),
+        ):
+            mock_download.return_value = b'{"a": 1}\n'
+            result = await self.handler.afile_content(
+                file_content_request=FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None),
+                vertex_credentials=None,
+                vertex_project="test-project",
+                vertex_location="us-central1",
+                timeout=60.0,
+                max_retries=3,
+            )
+
+        assert result.response.content == b'{"a": 1}\n'
+        mock_download.assert_called_once()
+
     @pytest.mark.asyncio
     async def test_afile_content_missing_file_id(self):
         """Test async file content retrieval with missing file_id"""