diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 7d56e60a483..84c808e1f21 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -380,9 +380,7 @@ class VertexAIBatchPrediction(VertexLLM): model=deployed_model_resource, vertex_location=vertex_location, ) - model_fetched: Final[_FetchedResponseView] = { - "response": safe_get(sync_handler, model_url, headers=headers) - } + model_fetched: Final[_FetchedResponseView] = {"response": safe_get(sync_handler, model_url, headers=headers)} model_response: Final = model_fetched["response"] if model_response.status_code != 200: raise VertexAIError( @@ -410,10 +408,12 @@ class VertexAIBatchPrediction(VertexLLM): "spec to size the batch replicas from" ), ) + # A scale-to-zero online endpoint reports minReplicaCount 0, but a batch job must start + # at least one replica. batch_resources: Final[BatchDedicatedResources] = { "machineSpec": machine_spec, - "startingReplicaCount": online_resources.get("minReplicaCount", 1), - "maxReplicaCount": online_resources.get("maxReplicaCount", 1), + "startingReplicaCount": max(online_resources.get("minReplicaCount", 1), 1), + "maxReplicaCount": max(online_resources.get("maxReplicaCount", 1), 1), } unmanaged: Final[UnmanagedContainerModel] = {"containerSpec": container_spec} resolved: Final[VertexAIBatchPredictionJob] = { @@ -471,7 +471,9 @@ class VertexAIBatchPrediction(VertexLLM): """Return the base url for the vertex garden models""" # POST https://LOCATION-aiplatform.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION/batchPredictionJobs base_url: Final = get_vertex_base_url(vertex_location) - return f"{base_url}/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/batchPredictionJobs" + return ( + f"{base_url}/{vertex_api_version}/projects/{vertex_project}/locations/{vertex_location}/batchPredictionJobs" + ) def retrieve_batch( self, diff --git a/litellm/llms/vertex_ai/files/handler.py b/litellm/llms/vertex_ai/files/handler.py index ee2d93814e5..6d0ca7b0374 100644 --- a/litellm/llms/vertex_ai/files/handler.py +++ b/litellm/llms/vertex_ai/files/handler.py @@ -97,7 +97,13 @@ class VertexAIFilesHandler(GCSBucketBase): allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(litellm_params), ) - _SHARDED_RESULTS_PATTERN: Final = re.compile(r"^(?P.*-)(?P\d{5})(?P-of-)(?P\d{5})$") + # Only the exact directory/file layout Vertex writes for unmanaged-container batch outputs, + # pinned to shard zero: the object path is derived from a caller-controlled file id, so a + # looser pattern would let a crafted upload filename trigger shard fan-out. + _SHARDED_RESULTS_PATTERN: Final = re.compile( + r"^(?P.*/prediction-custom-unmanaged-model-[^/]+/prediction\.results-)00000(?P-of-)(?P\d{5})$" + ) + _MAX_RESULT_SHARDS: Final = 512 async def _download_all_result_shards( self, @@ -119,14 +125,20 @@ class VertexAIFilesHandler(GCSBucketBase): total_shards: Final = int(shard_match["total"]) if total_shards <= 1: return first_shard - remaining: Final = await asyncio.gather( - *( - self.download_gcs_object( + if total_shards > self._MAX_RESULT_SHARDS: + raise ValueError( + f"Vertex batch output claims {total_shards} shards, above the supported maximum " + f"of {self._MAX_RESULT_SHARDS}" + ) + # Sequential fetch keeps memory and connection use bounded by one shard at a time. + remaining: Final = tuple( + [ + await self.download_gcs_object( object_name=f"{shard_match['stem']}{index:05d}{shard_match['sep']}{shard_match['total']}", standard_callback_dynamic_params=standard_callback_dynamic_params, ) for index in range(1, total_shards) - ) + ] ) if any(shard is None for shard in remaining): return None diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 22bd5f36b86..b96b5afb4bd 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -666,7 +666,9 @@ def _custom_endpoint_row_to_openai_batch_output_row(row: Mapping[str, object]) - """ key: Final = row.get("key") instance: Final = row.get("instance") - tagged_custom_id: Final = instance.get(VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, "") if isinstance(instance, Mapping) else "" + tagged_custom_id: Final = ( + instance.get(VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, "") if isinstance(instance, Mapping) else "" + ) custom_id: Final = str(key if key is not None else tagged_custom_id) prediction: Final = row.get("prediction") @@ -833,8 +835,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): if purpose == "batch" and custom_endpoint_id is not None: safe_endpoint_id: Final = sanitize_cloud_object_path(custom_endpoint_id, fallback="endpoint") return ( - f"{VERTEX_AI_MANAGED_GCS_PREFIX}{VERTEX_CUSTOM_ENDPOINT_GCS_SEGMENT}/" - f"{safe_endpoint_id}/{uuid.uuid4()}" + f"{VERTEX_AI_MANAGED_GCS_PREFIX}{VERTEX_CUSTOM_ENDPOINT_GCS_SEGMENT}/{safe_endpoint_id}/{uuid.uuid4()}" ) if purpose == "batch": ## 1. If jsonl, derive the object name from the deployment model (or the first entry's) @@ -881,9 +882,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): configured_model: Final = litellm_params.get("model") deployment_api_base: Final = litellm_params.get("api_base") custom_endpoint_id: Final = ( - get_custom_endpoint_id_from_api_base( - deployment_api_base if isinstance(deployment_api_base, str) else None - ) + get_custom_endpoint_id_from_api_base(deployment_api_base if isinstance(deployment_api_base, str) else None) if litellm_params.get("custom_endpoint") else None ) diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py index 6adc6869201..08b6b7171da 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -47,11 +47,7 @@ PROJECT = "my-project" LOCATION = "us-central1" BATCH_ID = "3814889423749775360" -CREATE_DATA = { - "input_file_id": ( - "gs://bucket/publishers/google/models/gemini-1.5-flash-001/file-uuid" - ) -} +CREATE_DATA = {"input_file_id": ("gs://bucket/publishers/google/models/gemini-1.5-flash-001/file-uuid")} def _vertex_job_response(state: str = "JOB_STATE_SUCCEEDED") -> dict: @@ -104,8 +100,7 @@ def test_create_vertex_batch_url(): h = _make_handler() url = h.create_vertex_batch_url(vertex_location=LOCATION, vertex_project=PROJECT) assert url == ( - f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}" - f"/locations/{LOCATION}/batchPredictionJobs" + f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}/locations/{LOCATION}/batchPredictionJobs" ) @@ -206,9 +201,7 @@ def test_create_batch_sync_does_not_resolve_publisher_models(): ENDPOINT_ID = "7768560373388541952" -ENDPOINT_CREATE_DATA = { - "input_file_id": f"gs://bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/file-uuid" -} +ENDPOINT_CREATE_DATA = {"input_file_id": f"gs://bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/file-uuid"} TUNED_MODEL_RESOURCE = f"projects/{PROJECT}/locations/{LOCATION}/models/1234509876" @@ -217,9 +210,7 @@ def _endpoint_get_response(deployed_models: list | None = None) -> MagicMock: resp.status_code = 200 resp.json.return_value = { "name": f"projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", - "deployedModels": ( - deployed_models if deployed_models is not None else [{"model": TUNED_MODEL_RESOURCE}] - ), + "deployedModels": (deployed_models if deployed_models is not None else [{"model": TUNED_MODEL_RESOURCE}]), } return resp @@ -275,7 +266,7 @@ CUSTOM_ENDPOINT_API_BASE = ( ) -def _custom_endpoint_get_response() -> MagicMock: +def _custom_endpoint_get_response(min_replica_count: int = 1) -> MagicMock: resp = MagicMock() resp.status_code = 200 resp.json.return_value = { @@ -283,7 +274,11 @@ def _custom_endpoint_get_response() -> MagicMock: "deployedModels": [ { "model": CONTAINER_MODEL_RESOURCE, - "dedicatedResources": {"machineSpec": MACHINE_SPEC, "minReplicaCount": 1, "maxReplicaCount": 2}, + "dedicatedResources": { + "machineSpec": MACHINE_SPEC, + "minReplicaCount": min_replica_count, + "maxReplicaCount": 2, + }, } ], } @@ -378,6 +373,37 @@ def test_create_batch_sync_custom_endpoint_without_container_spec_raises_400(): client.post.assert_not_called() +def test_create_batch_sync_custom_endpoint_scale_to_zero_starts_one_replica(): + """A scale-to-zero online endpoint reports minReplicaCount 0, which Vertex rejects as a batch + startingReplicaCount; the job must clamp to at least one replica.""" + h = _make_handler() + client = MagicMock() + client.post.return_value = _http_response() + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch( + f"{HMOD}.safe_get", + side_effect=[_custom_endpoint_get_response(min_replica_count=0), _container_model_get_response()], + ), + ): + h.create_batch( + _is_async=False, + create_batch_data=CUSTOM_ENDPOINT_CREATE_DATA, + api_base=CUSTOM_ENDPOINT_API_BASE, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + custom_endpoint=True, + ) + + sent = json.loads(client.post.call_args.kwargs["data"]) + assert sent["dedicatedResources"]["startingReplicaCount"] == 1 + assert sent["dedicatedResources"]["maxReplicaCount"] == 2 + + def test_create_batch_sync_custom_endpoint_rejects_multi_deployment_endpoint(): """An endpoint behind a traffic split has no single container to replicate; index-zero selection could run a different container than online traffic.""" @@ -628,9 +654,7 @@ def test_create_batch_sync_httpstatuserror_propagates(): client = MagicMock() request = httpx.Request("POST", "https://x/batchPredictionJobs") err_response = httpx.Response(status_code=500, request=request, text="boom") - client.post.side_effect = httpx.HTTPStatusError( - "boom", request=request, response=err_response - ) + client.post.side_effect = httpx.HTTPStatusError("boom", request=request, response=err_response) with patch(f"{HMOD}._get_httpx_client", return_value=client): with pytest.raises(httpx.HTTPStatusError): @@ -780,9 +804,7 @@ def test_retrieve_batch_sync_invokes_logging_pre_call(): logging_obj.pre_call.assert_called_once() _, kwargs = logging_obj.pre_call.call_args - assert kwargs["additional_args"]["api_base"].endswith( - f"/batchPredictionJobs/{BATCH_ID}" - ) + assert kwargs["additional_args"]["api_base"].endswith(f"/batchPredictionJobs/{BATCH_ID}") # =========================================================================== # @@ -853,9 +875,7 @@ def test_list_batches_sync_omits_unset_pagination_params(): def test_list_batches_async_returns_coroutine(): h = _make_handler() async_client = MagicMock() - async_client.get = AsyncMock( - return_value=_http_response(json_body=_list_response()) - ) + async_client.get = AsyncMock(return_value=_http_response(json_body=_list_response())) sync_client = MagicMock() with ( @@ -910,9 +930,7 @@ def test_cancel_batch_sync_posts_cancel_then_retrieves(): h = _make_handler() client = MagicMock() client.post.return_value = _http_response(json_body={}) - client.get.return_value = _http_response( - json_body=_vertex_job_response(state="JOB_STATE_CANCELLED") - ) + client.get.return_value = _http_response(json_body=_vertex_job_response(state="JOB_STATE_CANCELLED")) with patch(f"{HMOD}._get_httpx_client", return_value=client): out = h.cancel_batch( @@ -944,9 +962,7 @@ def test_cancel_batch_async_returns_coroutine_posts_then_retrieves(): async_client = MagicMock() async_client.post = AsyncMock(return_value=_http_response(json_body={})) async_client.get = AsyncMock( - return_value=_http_response( - json_body=_vertex_job_response(state="JOB_STATE_CANCELLED") - ) + return_value=_http_response(json_body=_vertex_job_response(state="JOB_STATE_CANCELLED")) ) with ( @@ -1004,9 +1020,7 @@ def test_cancel_batch_sync_proxy_url_without_cancel_suffix_uses_rsplit_branch(): ) client = MagicMock() client.post.return_value = _http_response(json_body={}) - client.get.return_value = _http_response( - json_body=_vertex_job_response(state="JOB_STATE_CANCELLED") - ) + client.get.return_value = _http_response(json_body=_vertex_job_response(state="JOB_STATE_CANCELLED")) with patch(f"{HMOD}._get_httpx_client", return_value=client): out = h.cancel_batch( @@ -1032,9 +1046,7 @@ def test_cancel_batch_sync_httpstatuserror_logged_and_reraised(): client = MagicMock() request = httpx.Request("POST", "https://x/batchPredictionJobs/1:cancel") err_response = httpx.Response(status_code=502, request=request, text="bad gw") - client.post.side_effect = httpx.HTTPStatusError( - "boom", request=request, response=err_response - ) + client.post.side_effect = httpx.HTTPStatusError("boom", request=request, response=err_response) with patch(f"{HMOD}._get_httpx_client", return_value=client): with pytest.raises(httpx.HTTPStatusError): @@ -1056,11 +1068,7 @@ def test_create_batch_async_httpstatuserror_logged_and_reraised(): async_client = MagicMock() request = httpx.Request("POST", "https://x/batchPredictionJobs") err_response = httpx.Response(status_code=500, request=request, text="boom") - async_client.post = AsyncMock( - side_effect=httpx.HTTPStatusError( - "boom", request=request, response=err_response - ) - ) + async_client.post = AsyncMock(side_effect=httpx.HTTPStatusError("boom", request=request, response=err_response)) with ( patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()), @@ -1167,9 +1175,7 @@ def test_async_cancel_batch_httpstatuserror_and_retrieve_non_200(): async_client = MagicMock() request = httpx.Request("POST", "https://x/batchPredictionJobs/1:cancel") err_response = httpx.Response(status_code=502, request=request, text="bad") - async_client.post = AsyncMock( - side_effect=httpx.HTTPStatusError("boom", request=request, response=err_response) - ) + async_client.post = AsyncMock(side_effect=httpx.HTTPStatusError("boom", request=request, response=err_response)) async_client.get = AsyncMock() with ( patch(f"{HMOD}._get_httpx_client", return_value=MagicMock()), diff --git a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py index c80f244a659..8726ccdea9e 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -36,8 +36,7 @@ INPUT_FILE = ( ENDPOINT_ID = "7768560373388541952" ENDPOINT_INPUT_FILE = ( - f"gs://litellm-testing-bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/" - "e9412502-2c91-42a6-8e61-f5c294cc0fc8" + f"gs://litellm-testing-bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/e9412502-2c91-42a6-8e61-f5c294cc0fc8" ) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py index 4b937289fc3..442928e55c2 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py @@ -149,12 +149,12 @@ class TestVertexAIFilesHandler: only that shard silently drops the rest of the batch (LIT-7387).""" file_id = ( "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fcustom-endpoints%2F123%2F" - "out%2Fprediction.results-00000-of-00003" + "prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z%2Fprediction.results-00000-of-00003" ) shards = { - "litellm-vertex-files/custom-endpoints/123/out/prediction.results-00000-of-00003": b'{"a": 1}\n', - "litellm-vertex-files/custom-endpoints/123/out/prediction.results-00001-of-00003": b'{"b": 2}\n', - "litellm-vertex-files/custom-endpoints/123/out/prediction.results-00002-of-00003": b'{"c": 3}', + "litellm-vertex-files/custom-endpoints/123/prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z/prediction.results-00000-of-00003": b'{"a": 1}\n', + "litellm-vertex-files/custom-endpoints/123/prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z/prediction.results-00001-of-00003": b'{"b": 2}\n', + "litellm-vertex-files/custom-endpoints/123/prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z/prediction.results-00002-of-00003": b'{"c": 3}', } async def fake_download(object_name: str, **kwargs): @@ -184,7 +184,7 @@ class TestVertexAIFilesHandler: async def test_afile_content_single_shard_unchanged(self): file_id = ( "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fcustom-endpoints%2F123%2F" - "out%2Fprediction.results-00000-of-00001" + "prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z%2Fprediction.results-00000-of-00001" ) with ( patch.object(self.handler, "download_gcs_object", new_callable=AsyncMock) as mock_download, @@ -208,6 +208,62 @@ class TestVertexAIFilesHandler: assert result.response.content == b'{"a": 1}\n' mock_download.assert_called_once() + @pytest.mark.asyncio + async def test_afile_content_crafted_upload_filename_does_not_fan_out(self): + """The object path comes from a caller-controlled file id, so an ordinary upload whose + name mimics the shard suffix must not trigger shard fetches (99,998 GCS requests from one + crafted 'prediction.results-00000-of-99999' filename).""" + file_id = "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fuploads%2Fabc-prediction.results-00000-of-99999" + with ( + patch.object(self.handler, "download_gcs_object", new_callable=AsyncMock) as mock_download, + patch.object( + self.handler, + "get_gcs_logging_config", + new_callable=AsyncMock, + return_value=_mock_gcs_logging_config(), + ), + ): + mock_download.return_value = b"tiny" + result = await self.handler.afile_content( + file_content_request=FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None), + vertex_credentials=None, + vertex_project="test-project", + vertex_location="us-central1", + timeout=60.0, + max_retries=3, + ) + + assert result.response.content == b"tiny" + mock_download.assert_called_once() + + @pytest.mark.asyncio + async def test_afile_content_shard_count_above_cap_raises(self): + file_id = ( + "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fcustom-endpoints%2F123%2F" + "prediction-custom-unmanaged-model-2026_09_09T13_00_00_000Z%2Fprediction.results-00000-of-99999" + ) + with ( + patch.object(self.handler, "download_gcs_object", new_callable=AsyncMock) as mock_download, + patch.object( + self.handler, + "get_gcs_logging_config", + new_callable=AsyncMock, + return_value=_mock_gcs_logging_config(), + ), + ): + mock_download.return_value = b"tiny" + with pytest.raises(ValueError, match="shards"): + await self.handler.afile_content( + file_content_request=FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None), + vertex_credentials=None, + vertex_project="test-project", + vertex_location="us-central1", + timeout=60.0, + max_retries=3, + ) + + mock_download.assert_called_once() + @pytest.mark.asyncio async def test_afile_content_missing_file_id(self): """Test async file content retrieval with missing file_id""" @@ -247,8 +303,7 @@ class TestVertexAIFilesHandler: with pytest.raises( ValueError, match=re.escape( - "Failed to download file from GCS: " - "gs://test-bucket/litellm-vertex-files/uploads/abc-test-file.txt" + "Failed to download file from GCS: gs://test-bucket/litellm-vertex-files/uploads/abc-test-file.txt" ), ): await self.handler.afile_content( diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index be77ece8dfe..2fe8836712e 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -193,9 +193,7 @@ class TestBatchObjectNaming: def test_should_store_fine_tuned_endpoint_under_endpoints_path(self, config): """A numeric endpoint id must not be filed under publishers/google/models/gemini/, which the batch transformation later mangles into a nonexistent publisher model (LIT-6899).""" - object_name = config._get_gcs_object_name_from_batch_jsonl( - [{"body": {"model": "gemini/7768560373388541952"}}] - ) + object_name = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "gemini/7768560373388541952"}}]) assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/") assert "publishers" not in object_name