From fc29fb513cf7b7cffb67018ceb968b4c20001414 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 24 Sep 2026 14:05:31 -0700 Subject: [PATCH] fix(vertex_ai): keep batch output_file_id null until Vertex reports outputInfo (#43030) * fix(vertex_ai): keep batch output_file_id null until Vertex reports outputInfo Vertex only sets outputInfo.gcsOutputDirectory once a batch job has written output. Falling back to outputConfig's outputUriPrefix named the per-model directory shared by every batch of the deployment, an object that never exists, so the proxy minted a managed file for it under the first key and every other key's file calls on that id were 403s * fix(vertex_ai): treat a null gcsOutputDirectory as no output file yet --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../llms/vertex_ai/batches/transformation.py | 27 +--- .../test_vertex_batch_output_info_wire.py | 2 +- .../vertex_ai/batches/test_transformation.py | 120 +++++++++++------- .../test_vertex_ai_batch_transformation.py | 9 +- 4 files changed, 84 insertions(+), 74 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index a7dbb058465..11c99d130be 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -253,30 +253,15 @@ class VertexAIBatchTransformation: return uris[0] @classmethod - def _get_output_file_id_from_vertex_ai_batch_response(cls, response: VertexBatchPredictionResponse) -> str: + def _get_output_file_id_from_vertex_ai_batch_response(cls, response: VertexBatchPredictionResponse) -> str | None: """ - Gets the output file id from the Vertex AI Batch response + Gets the output file id from the Vertex AI Batch response, None until Vertex reports outputInfo """ - output_info: Final = response.get("outputInfo") or OutputInfo() - output_file_id: str = output_info.get("gcsOutputDirectory", "") - if output_file_id: - output_file_id = output_file_id.rstrip("/") + "/predictions.jsonl" - if output_file_id and output_file_id != "/predictions.jsonl": - return output_file_id - - output_config: Final = response.get("outputConfig") - if output_config is None: - return output_file_id - - gcs_destination: Final = output_config.get("gcsDestination") - if gcs_destination is None: - return output_file_id - - output_uri_prefix: Final = gcs_destination.get("outputUriPrefix", "") - if output_uri_prefix.endswith("/predictions.jsonl"): - return output_uri_prefix - return output_uri_prefix.rstrip("/") + "/predictions.jsonl" + gcs_output_directory: Final = (output_info.get("gcsOutputDirectory") or "").rstrip("/") + if not gcs_output_directory: + return None + return f"{gcs_output_directory}/predictions.jsonl" @classmethod def _get_batch_job_status_from_vertex_ai_batch_response( diff --git a/tests/integration/providers/test_vertex_batch_output_info_wire.py b/tests/integration/providers/test_vertex_batch_output_info_wire.py index a7ac896076c..de41e0790d1 100644 --- a/tests/integration/providers/test_vertex_batch_output_info_wire.py +++ b/tests/integration/providers/test_vertex_batch_output_info_wire.py @@ -116,7 +116,7 @@ def test_vertex_batch_create_survives_explicit_null_output_info(gateway: Gateway "batch", "validating", _encoded(INPUT_FILE_ID, model, "file-"), - _encoded(f"{OUTPUT_PREFIX}/predictions.jsonl", model, "file-"), + None, None, "24h", ), response.text 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 ae045edbec5..5f969620a80 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -13,6 +13,8 @@ There are no real I/O seams here; ``uuid.uuid4`` is the only nondeterministic dependency and is patched where the displayName is asserted. """ +from collections.abc import Mapping +from typing import Final from unittest.mock import patch import pytest @@ -35,8 +37,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" ) @@ -248,9 +249,78 @@ def test_get_input_file_id_empty_uris(): # =========================================================================== # -# _get_output_file_id_from_vertex_ai_batch_response +# _get_output_file_id_from_vertex_ai_batch_response: None until Vertex reports outputInfo # =========================================================================== # +SHARED_OUTPUT_PREFIX: Final = "gs://bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-flash" +SUCCEEDED_OUTPUT_DIRECTORY: Final = f"{SHARED_OUTPUT_PREFIX}/prediction-model-2026-09-24T19:41:00.000000Z" + + +def _vertex_job(state: str) -> dict[str, object]: + return { + "name": "projects/510528649030/locations/us-central1/batchPredictionJobs/3814889423749775360", + "state": state, + "createTime": "2026-09-24T19:37:25.775603Z", + "inputConfig": { + "instancesFormat": "jsonl", + "gcsSource": {"uris": [f"{SHARED_OUTPUT_PREFIX}/0586ba52-4f8b-4988-aa8d-3573550a4b0f"]}, + }, + "outputConfig": { + "predictionsFormat": "jsonl", + "gcsDestination": {"outputUriPrefix": SHARED_OUTPUT_PREFIX}, + }, + } + + +@pytest.mark.parametrize( + "vertex_state,output_info_field,expected_status,expected_output_file_id", + [ + ("JOB_STATE_PENDING", {}, "validating", None), + ("JOB_STATE_RUNNING", {"outputInfo": {}}, "in_progress", None), + ("JOB_STATE_CANCELLED", {"outputInfo": None}, "cancelled", None), + ( + "JOB_STATE_SUCCEEDED", + {"outputInfo": {"gcsOutputDirectory": SUCCEEDED_OUTPUT_DIRECTORY}}, + "completed", + f"{SUCCEEDED_OUTPUT_DIRECTORY}/predictions.jsonl", + ), + ], + ids=["create_or_pending", "running", "cancelled", "succeeded"], +) +def test_transform_vertex_response_output_file_id_is_none_until_output_info( + vertex_state: str, + output_info_field: Mapping[str, object], + expected_status: str, + expected_output_file_id: str | None, +) -> None: + batch: Final = T.transform_vertex_ai_batch_response_to_openai_batch_response( + {**_vertex_job(vertex_state), **output_info_field} + ) + + assert batch.status == expected_status + assert batch.output_file_id == expected_output_file_id + + +@pytest.mark.parametrize( + "response", + [ + {}, + {"outputConfig": {}}, + {"outputInfo": None}, + {"outputInfo": {"gcsOutputDirectory": ""}}, + {"outputInfo": {"gcsOutputDirectory": None}}, + ], + ids=[ + "no_fields", + "output_config_without_destination", + "null_output_info", + "empty_output_directory", + "null_output_directory", + ], +) +def test_get_output_file_id_is_none_without_output_directory(response: Mapping[str, object]) -> None: + assert T._get_output_file_id_from_vertex_ai_batch_response(response) is None + def test_get_output_file_id_from_output_info(): # outputInfo branch: rstrip trailing slash, append predictions.jsonl @@ -267,49 +337,7 @@ def test_get_output_file_id_output_info_no_trailing_slash(): ) -def test_get_output_file_id_empty_output_info_falls_through_to_output_config(): - # gcsOutputDirectory missing -> "" -> the "/predictions.jsonl" guard skips - # the outputInfo branch, falls through to outputConfig - resp = { - "outputInfo": {}, - "outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://b/cfg"}}, - } - assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl" - - -def test_get_output_file_id_output_info_explicit_none_falls_through_to_output_config(): - resp = { - "outputInfo": None, - "outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://b/cfg"}}, - } - assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl" - - -def test_get_output_file_id_output_info_explicit_none_and_no_output_config(): - assert T._get_output_file_id_from_vertex_ai_batch_response({"outputInfo": None}) == "" - - -def test_get_output_file_id_no_output_info_and_no_output_config(): - assert T._get_output_file_id_from_vertex_ai_batch_response({}) == "" - - -def test_get_output_file_id_output_config_missing_gcs_destination(): - # outputConfig present but no gcsDestination -> returns the running "" value - assert T._get_output_file_id_from_vertex_ai_batch_response({"outputConfig": {}}) == "" - - -def test_get_output_file_id_output_config_already_has_suffix(): - # outputUriPrefix already ends in /predictions.jsonl -> returned as-is (no double append) - resp = {"outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://b/cfg/predictions.jsonl"}}} - assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl" - - -def test_get_output_file_id_output_config_strips_trailing_slash(): - resp = {"outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://b/cfg/"}}} - assert T._get_output_file_id_from_vertex_ai_batch_response(resp) == "gs://b/cfg/predictions.jsonl" - - -def test_get_output_file_id_output_info_takes_precedence_over_output_config(): +def test_get_output_file_id_output_info_ignores_output_uri_prefix(): resp = { "outputInfo": {"gcsOutputDirectory": "gs://from-info"}, "outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://from-config"}}, diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py index 8a135dac3bb..3bc66d017a0 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py @@ -26,12 +26,12 @@ def test_output_file_id_uses_predictions_jsonl_with_output_info(): ) -def test_output_file_id_falls_back_to_output_uri_prefix_with_predictions_jsonl(): +def test_output_file_id_is_none_until_output_info(): response = { "outputInfo": {}, "outputConfig": { "gcsDestination": { - "outputUriPrefix": "gs://test-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-pro/prediction-model-456" + "outputUriPrefix": "gs://test-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-pro" } }, } @@ -42,10 +42,7 @@ def test_output_file_id_falls_back_to_output_uri_prefix_with_predictions_jsonl() ) ) - assert ( - output_file_id - == "gs://test-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-pro/prediction-model-456/predictions.jsonl" - ) + assert output_file_id is None def test_vertex_ai_cancel_batch():