From 783978b39c9587cab25127e76b3e43a8ea6c7a92 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 9 Sep 2026 17:01:20 -0400 Subject: [PATCH] fix(vertex_ai): strip the batch custom_id tag via excludedFields, not keyField Probed live: instanceConfig.keyField does not remove the field from the instances the container receives, so vLLM 400s every row; excludedFields performs the strip and attaches the value to the output row. --- litellm/llms/vertex_ai/batches/handler.py | 9 +++++---- litellm/types/llms/vertex_ai.py | 1 + .../test_litellm/llms/vertex_ai/batches/test_handler.py | 2 +- 3 files changed, 7 insertions(+), 5 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index e505979c5e7..dd37e104044 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -396,12 +396,13 @@ class VertexAIBatchPrediction(VertexLLM): "outputConfig": vertex_batch_request["outputConfig"], "unmanagedContainerModel": unmanaged, "dedicatedResources": batch_resources, - # keyField strips the custom_id tag from each instance before it reaches the - # container (vLLM rejects unknown fields) and echoes it back as `key` in the output - # row; it only takes effect alongside an explicit instanceType. + # 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). "instanceConfig": { "instanceType": "object", - "keyField": VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, + "excludedFields": [VERTEX_CUSTOM_ENDPOINT_KEY_FIELD], }, } return resolved diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 1fda5d727eb..e559b039b79 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -700,6 +700,7 @@ class UnmanagedContainerModel(TypedDict, total=False): class BatchInstanceConfig(TypedDict, total=False): instanceType: ReadOnly[str] keyField: ReadOnly[str] + excludedFields: ReadOnly[list[str]] class VertexAIBatchPredictionJob(TypedDict, total=False): 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 8691d14d392..3a5b4db04bf 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -342,7 +342,7 @@ def test_create_batch_sync_custom_endpoint_builds_unmanaged_container_job(): "startingReplicaCount": 1, "maxReplicaCount": 2, } - assert sent["instanceConfig"] == {"instanceType": "object", "keyField": "litellm_custom_id"} + assert sent["instanceConfig"] == {"instanceType": "object", "excludedFields": ["litellm_custom_id"]} def test_create_batch_sync_custom_endpoint_without_container_spec_raises_400():