From 1818c39fed84f255e9feaeae6239691dafff5602 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 9 Sep 2026 18:39:46 -0400 Subject: [PATCH] fix(vertex_ai): read all batch output shards and reject multi-model endpoints Sharded prediction.results-N-of-M outputs are concatenated instead of returning only shard zero, and an endpoint serving several deployed models behind a traffic split is rejected at batch create rather than silently running deployedModels[0]. --- litellm/llms/vertex_ai/batches/handler.py | 22 ++++--- litellm/llms/vertex_ai/files/handler.py | 51 ++++++++++++--- .../llms/vertex_ai/files/transformation.py | 6 +- .../llms/vertex_ai/batches/test_handler.py | 36 ++++++++++ .../files/test_vertex_ai_files_handler.py | 65 +++++++++++++++++++ 5 files changed, 159 insertions(+), 21 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 9312da62eb8..7d56e60a483 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -157,11 +157,8 @@ class VertexAIBatchPrediction(VertexLLM): vertex_location=vertex_location or "us-central1", ) ) - # The endpoint id in a custom-endpoints/ file path is caller-controlled (raw gs:// file - # ids are accepted), so it must never pick which container runs: batch jobs execute the - # container with the deployment's project credentials. The id has to match the endpoint - # the routed deployment's own api_base names, and container resolution only happens for - # deployments the admin marked custom_endpoint. + # The file-path endpoint id is caller-controlled and the job runs with the deployment's + # credentials, so it must match the endpoint the deployment's own api_base names. job_model: Final = transformed_batch_request.get("model", "") deployment_endpoint_id: Final = get_custom_endpoint_id_from_api_base(api_base) if custom_endpoint and not job_model.endswith(f"/custom-endpoints/{deployment_endpoint_id}"): @@ -358,6 +355,15 @@ class VertexAIBatchPrediction(VertexLLM): ) endpoint_view: Final[_VertexEndpointPayloadView] = {"payload": endpoint_response.json()} deployed_models: Final = endpoint_view["payload"].get("deployedModels") or () + if len(deployed_models) > 1: + raise VertexAIError( + status_code=400, + message=( + f"Vertex endpoint '{endpoint_resource}' serves {len(deployed_models)} deployed " + "models behind a traffic split, so there is no single container to replicate " + "for batch prediction; use an endpoint with exactly one deployed model" + ), + ) deployed: Final = deployed_models[0] if deployed_models else _VertexEndpointDeployedModel() deployed_model_resource: Final = deployed.get("model", "") if not deployed_model_resource: @@ -416,10 +422,8 @@ class VertexAIBatchPrediction(VertexLLM): "outputConfig": vertex_batch_request["outputConfig"], "unmanagedContainerModel": unmanaged, "dedicatedResources": batch_resources, - # excludedFields strips the custom_id tag from each instance before it reaches the - # container (vLLM rejects unknown fields) and attaches it to the output row's - # instance echo; keyField does NOT strip (probed live: the container still received - # the tag and 400'd every row). + # excludedFields (not keyField, which does not actually strip and 400s vLLM) removes + # the custom_id tag before the container sees it and echoes it in the output row. "instanceConfig": { "instanceType": "object", "excludedFields": (VERTEX_CUSTOM_ENDPOINT_KEY_FIELD,), diff --git a/litellm/llms/vertex_ai/files/handler.py b/litellm/llms/vertex_ai/files/handler.py index ac95d1348f9..ee2d93814e5 100644 --- a/litellm/llms/vertex_ai/files/handler.py +++ b/litellm/llms/vertex_ai/files/handler.py @@ -1,6 +1,7 @@ import asyncio import json import os +import re import time from collections.abc import Coroutine, Mapping from typing import Any, Final @@ -96,6 +97,41 @@ class VertexAIFilesHandler(GCSBucketBase): allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params), ) + _SHARDED_RESULTS_PATTERN: Final = re.compile(r"^(?P.*-)(?P\d{5})(?P-of-)(?P\d{5})$") + + async def _download_all_result_shards( + self, + object_path: str, + standard_callback_dynamic_params: StandardCallbackDynamicParams, + ) -> bytes | None: + """ + An unmanaged-container batch writes its output as `prediction.results-000NN-of-000NN` + shards; a file id names shard zero, so when the shard count is above one the remaining + shards are fetched and concatenated (each shard is newline-delimited JSONL). + """ + first_shard: Final = await self.download_gcs_object( + object_name=object_path, + standard_callback_dynamic_params=standard_callback_dynamic_params, + ) + shard_match: Final = self._SHARDED_RESULTS_PATTERN.match(object_path) + if first_shard is None or shard_match is None: + return first_shard + total_shards: Final = int(shard_match["total"]) + if total_shards <= 1: + return first_shard + remaining: Final = await asyncio.gather( + *( + self.download_gcs_object( + object_name=f"{shard_match['stem']}{index:05d}{shard_match['sep']}{shard_match['total']}", + standard_callback_dynamic_params=standard_callback_dynamic_params, + ) + for index in range(1, total_shards) + ) + ) + if any(shard is None for shard in remaining): + return None + return b"\n".join((first_shard.rstrip(b"\n"), *(shard.rstrip(b"\n") for shard in remaining if shard))) + async def afile_content( self, file_content_request: FileContentRequest, @@ -141,14 +177,13 @@ class VertexAIFilesHandler(GCSBucketBase): litellm_params=litellm_params, ) - download_kwargs: Final = { - "standard_callback_dynamic_params": { - "gcs_bucket_name": bucket_name, - "gcs_path_service_account": gcs_logging_config["path_service_account"], - } - } - - file_content: Final = await self.download_gcs_object(object_name=object_path, **download_kwargs) + file_content: Final = await self._download_all_result_shards( + object_path=object_path, + standard_callback_dynamic_params=StandardCallbackDynamicParams( + gcs_bucket_name=bucket_name, + gcs_path_service_account=gcs_logging_config["path_service_account"], + ), + ) decoded_file_id: Final = unquote(file_id) if file_content is None: diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 006342c5a4f..22bd5f36b86 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -744,11 +744,9 @@ class _OpenAIToCustomEndpointBatchUploadStream(BaseFileUploadStream): self._openai_file_content = openai_file_content def iter_bytes(self) -> Iterator[bytes]: - first = True - for entry in _iter_openai_jsonl_entries(self._openai_file_content): + for index, entry in enumerate(_iter_openai_jsonl_entries(self._openai_file_content)): row = _openai_batch_jsonl_entry_to_custom_endpoint_row(entry) - prefix = b"" if first else b"\n" - first = False + prefix = b"" if index == 0 else b"\n" yield prefix + json.dumps(row, default=dict).encode("utf-8") diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index 9bcbd723060..6adc6869201 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -378,6 +378,42 @@ def test_create_batch_sync_custom_endpoint_without_container_spec_raises_400(): client.post.assert_not_called() +def test_create_batch_sync_custom_endpoint_rejects_multi_deployment_endpoint(): + """An endpoint behind a traffic split has no single container to replicate; index-zero + selection could run a different container than online traffic.""" + h = _make_handler() + client = MagicMock() + multi = MagicMock() + multi.status_code = 200 + multi.json.return_value = { + "deployedModels": [ + {"model": CONTAINER_MODEL_RESOURCE}, + {"model": f"projects/{PROJECT}/locations/{LOCATION}/models/other"}, + ] + } + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get", return_value=multi), + ): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data=CUSTOM_ENDPOINT_CREATE_DATA, + api_base=CUSTOM_ENDPOINT_API_BASE, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + custom_endpoint=True, + ) + + assert exc_info.value.status_code == 400 + assert "traffic split" in str(exc_info.value) + client.post.assert_not_called() + + def test_create_batch_sync_custom_endpoint_rejects_file_for_other_endpoint(): """The endpoint id in the file path is caller-controlled (raw gs:// ids are accepted), so a file staged for a different endpoint must not make the deployment run that endpoint's diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py index 0a44f0a9a74..4b937289fc3 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py @@ -143,6 +143,71 @@ class TestVertexAIFilesHandler: assert "standard_callback_dynamic_params" in call_args.kwargs assert call_args.kwargs["standard_callback_dynamic_params"]["gcs_bucket_name"] == "test-bucket" + @pytest.mark.asyncio + async def test_afile_content_fetches_all_result_shards(self): + """A sharded unmanaged-container batch output names shard zero in the file id; reading + only that shard silently drops the rest of the batch (LIT-7387).""" + file_id = ( + "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fcustom-endpoints%2F123%2F" + "out%2Fprediction.results-00000-of-00003" + ) + shards = { + "litellm-vertex-files/custom-endpoints/123/out/prediction.results-00000-of-00003": b'{"a": 1}\n', + "litellm-vertex-files/custom-endpoints/123/out/prediction.results-00001-of-00003": b'{"b": 2}\n', + "litellm-vertex-files/custom-endpoints/123/out/prediction.results-00002-of-00003": b'{"c": 3}', + } + + async def fake_download(object_name: str, **kwargs): + return shards[object_name] + + with ( + patch.object(self.handler, "download_gcs_object", side_effect=fake_download), + patch.object( + self.handler, + "get_gcs_logging_config", + new_callable=AsyncMock, + return_value=_mock_gcs_logging_config(), + ), + ): + result = await self.handler.afile_content( + file_content_request=FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None), + vertex_credentials=None, + vertex_project="test-project", + vertex_location="us-central1", + timeout=60.0, + max_retries=3, + ) + + assert result.response.content == b'{"a": 1}\n{"b": 2}\n{"c": 3}' + + @pytest.mark.asyncio + async def test_afile_content_single_shard_unchanged(self): + file_id = ( + "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fcustom-endpoints%2F123%2F" + "out%2Fprediction.results-00000-of-00001" + ) + with ( + patch.object(self.handler, "download_gcs_object", new_callable=AsyncMock) as mock_download, + patch.object( + self.handler, + "get_gcs_logging_config", + new_callable=AsyncMock, + return_value=_mock_gcs_logging_config(), + ), + ): + mock_download.return_value = b'{"a": 1}\n' + result = await self.handler.afile_content( + file_content_request=FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None), + vertex_credentials=None, + vertex_project="test-project", + vertex_location="us-central1", + timeout=60.0, + max_retries=3, + ) + + assert result.response.content == b'{"a": 1}\n' + mock_download.assert_called_once() + @pytest.mark.asyncio async def test_afile_content_missing_file_id(self): """Test async file content retrieval with missing file_id"""