fix(vertex_ai): address batch review findings

Derive the GCS batch object path from the deployment's configured model
when present, so a user-crafted JSONL body.model cannot redirect an
authorized deployment's credentials to a different endpoint; the JSONL
value remains the fallback for direct SDK calls with no deployment
config.

Route the fine-tuned endpoint resolution GET through _check_custom_proxy
so custom api_base deployments do not contact Google directly.

Prefer the publisher model path over an endpoints/ segment when parsing
GCS uris, and use the last endpoints/ occurrence, so a bucket prefix
containing endpoints/<digits> cannot shadow the real model path.

Move the custom_endpoint rejection from the batches dispatcher into the
Vertex batch handler so the provider policy lives in the provider
module.

litellm/batches/main.py

@@ -301,21 +301,6 @@ def create_batch(
                 litellm_params=litellm_params,
             )
         elif custom_llm_provider == "vertex_ai":
-            if optional_params.get("custom_endpoint"):
-                raise litellm.exceptions.BadRequestError(
-                    message=(
-                        "Vertex AI batch prediction is not supported for `custom_endpoint` deployments. "
-                        "The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; "
-                        "use a publisher model or fine-tuned Gemini endpoint deployment instead."
-                    ),
-                    model=model or "n/a",
-                    llm_provider=custom_llm_provider,
-                    response=httpx.Response(
-                        status_code=400,
-                        content="custom_endpoint deployments do not support vertex_ai batches",
-                        request=httpx.Request(method="create_batch", url="https://github.com/BerriAI/litellm"),
-                    ),
-                )
             api_base = optional_params.api_base or ""
             vertex_ai_project: Final = (
                 optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT")
@@ -334,6 +319,7 @@ def create_batch(
                 timeout=timeout,
                 max_retries=optional_params.max_retries,
                 create_batch_data=_create_batch_request,
+                custom_endpoint=optional_params.get("custom_endpoint"),
             )
         else:
             raise litellm.exceptions.BadRequestError(

litellm/llms/vertex_ai/batches/handler.py

@@ -93,7 +93,17 @@ class VertexAIBatchPrediction(VertexLLM):
         vertex_location: str | None,
         timeout: float | httpx.Timeout,
         max_retries: int | None,
+        custom_endpoint: bool | None = None,
     ) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]:
+        if custom_endpoint:
+            raise VertexAIError(
+                status_code=400,
+                message=(
+                    "Vertex AI batch prediction is not supported for `custom_endpoint` deployments. "
+                    "The OpenAI-compatible custom endpoint path has no batch surface in LiteLLM; "
+                    "use a publisher model or fine-tuned Gemini endpoint deployment instead."
+                ),
+            )
         sync_handler: Final = _get_httpx_client()
 
         access_token, project_id = self._ensure_access_token(
@@ -102,6 +112,26 @@ class VertexAIBatchPrediction(VertexLLM):
             custom_llm_provider="vertex_ai",
         )
 
+        headers: Final = {
+            "Content-Type": "application/json; charset=utf-8",
+            "Authorization": f"Bearer {access_token}",
+        }
+
+        transformed_batch_request: Final[VertexAIBatchPredictionJob] = (
+            VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request(
+                request=create_batch_data,
+                vertex_project=vertex_project or project_id,
+                vertex_location=vertex_location or "us-central1",
+            )
+        )
+        vertex_batch_request: Final = self._resolve_fine_tuned_endpoint_model(
+            vertex_batch_request=transformed_batch_request,
+            headers=headers,
+            sync_handler=sync_handler,
+            api_base=api_base,
+            vertex_location=vertex_location or "us-central1",
+        )
+
         default_api_base: Final = self.create_vertex_batch_url(
             vertex_location=vertex_location or "us-central1",
             vertex_project=vertex_project or project_id,
@@ -126,25 +156,6 @@ class VertexAIBatchPrediction(VertexLLM):
             vertex_api_version="v1",
         )
 
-        headers: Final = {
-            "Content-Type": "application/json; charset=utf-8",
-            "Authorization": f"Bearer {access_token}",
-        }
-
-        transformed_batch_request: Final[VertexAIBatchPredictionJob] = (
-            VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request(
-                request=create_batch_data,
-                vertex_project=vertex_project or project_id,
-                vertex_location=vertex_location or "us-central1",
-            )
-        )
-        vertex_batch_request: Final = self._resolve_fine_tuned_endpoint_model(
-            vertex_batch_request=transformed_batch_request,
-            headers=headers,
-            sync_handler=sync_handler,
-            vertex_location=vertex_location or "us-central1",
-        )
-
         if _is_async is True:
             return self._async_create_batch(
                 vertex_batch_request=vertex_batch_request,
@@ -170,6 +181,7 @@ class VertexAIBatchPrediction(VertexLLM):
         vertex_batch_request: VertexAIBatchPredictionJob,
         headers: dict[str, str],
         sync_handler: HTTPHandler,
+        api_base: str | None,
         vertex_location: str,
     ) -> VertexAIBatchPredictionJob:
         """
@@ -181,7 +193,19 @@ class VertexAIBatchPrediction(VertexLLM):
         if "/endpoints/" not in model:
             return vertex_batch_request
 
-        endpoint_url: Final = f"{get_vertex_base_url(vertex_location)}/v1/{model}"
+        default_endpoint_url: Final = f"{get_vertex_base_url(vertex_location)}/v1/{model}"
+        _, endpoint_url = self._check_custom_proxy(
+            api_base=api_base,
+            custom_llm_provider="vertex_ai",
+            gemini_api_key=None,
+            endpoint=(default_endpoint_url.split(":")[-1] if len(default_endpoint_url.split(":")) > 1 else ""),
+            stream=None,
+            auth_header=None,
+            url=default_endpoint_url,
+            model=None,
+            vertex_location=vertex_location,
+            vertex_api_version="v1",
+        )
         response: Final = sync_handler.get(url=endpoint_url, headers=headers)
         if response.status_code != 200:
             raise VertexAIError(

litellm/llms/vertex_ai/batches/transformation.py

@@ -262,22 +262,24 @@ class VertexAIBatchTransformation:
         """
         Returns the `publishers/<publisher>/models/<model>` or `endpoints/<numeric id>` path from a
         gcs uri, or None if the uri does not contain one.
+
+        A publisher path wins over an `endpoints/` segment, and the last `endpoints/` occurrence is
+        used, so a user-configured bucket prefix that happens to contain `endpoints/<digits>` cannot
+        override the model path LiteLLM appended after it.
         """
         unquoted_uri: Final = unquote(gcs_file_uri)
-        _, endpoint_separator, endpoint_path = unquoted_uri.partition("endpoints/")
+        _, separator, model_path = unquoted_uri.partition("publishers/")
+        if separator:
+            parts: Final = model_path.split("/")
+            if len(parts) >= 3 and parts[1] == "models" and parts[2]:
+                return f"publishers/{'/'.join(parts[:3])}"
+
+        _, endpoint_separator, endpoint_path = unquoted_uri.rpartition("endpoints/")
         endpoint_id: Final = endpoint_path.split("/")[0] if endpoint_separator else ""
         if endpoint_id.isdigit():
             return f"endpoints/{endpoint_id}"
 
-        _, separator, model_path = unquoted_uri.partition("publishers/")
-        if not separator:
-            return None
-
-        parts: Final = model_path.split("/")
-        if len(parts) < 3 or parts[1] != "models" or not parts[2]:
-            return None
-
-        return f"publishers/{'/'.join(parts[:3])}"
+        return None
 
     @classmethod
     def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: str | None) -> bool:

litellm/llms/vertex_ai/files/transformation.py

@@ -708,18 +708,28 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
     def _get_gcs_object_name_from_batch_jsonl(
         self,
         openai_jsonl_content: list[dict[str, Any]],
+        deployment_model: str | None = None,
     ) -> str:
         """
         Gets a unique GCS object name for the VertexAI batch prediction job
 
         named as: litellm-vertex-{model}-{uuid}
 
+        The stored model path decides which Vertex model the batch job later executes against, so
+        `deployment_model` (the deployment's own configured model) wins over the user-supplied
+        JSONL `body.model`; the JSONL value is only a fallback for direct SDK calls that carry no
+        deployment config.
+
         Fine-tuned Gemini deployments (numeric endpoint ids) are stored under
         `endpoints/<id>` so the batch transformation can round-trip them into a
         `projects/../locations/../endpoints/<id>` batch job model instead of a
         nonexistent publisher model.
         """
-        raw_model: Final = openai_jsonl_content[0].get("body", {}).get("model", "")
+        raw_model: Final = (
+            deployment_model.removeprefix("vertex_ai/")
+            if deployment_model
+            else openai_jsonl_content[0].get("body", {}).get("model", "")
+        )
         endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model)
         model_path: Final = (
             f"endpoints/{endpoint_id}"
@@ -730,7 +740,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
         object_name: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}"
         return object_name
 
-    def get_object_name(self, file_data: FileTypes, purpose: str) -> str:
+    def get_object_name(self, file_data: FileTypes, purpose: str, deployment_model: str | None = None) -> str:
         """
         Get the object name for the request.
 
@@ -738,10 +748,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
         upload is never materialized just to derive the GCS object name.
         """
         if purpose == "batch":
-            ## 1. If jsonl, derive the object name from the first entry's model
+            ## 1. If jsonl, derive the object name from the deployment model (or the first entry's)
             first_entry: Final = next(_iter_openai_jsonl_entries(file_data), None)
             if first_entry is not None:
-                return self._get_gcs_object_name_from_batch_jsonl([first_entry])
+                return self._get_gcs_object_name_from_batch_jsonl([first_entry], deployment_model=deployment_model)
 
         ## 2. If not jsonl, store under a server-generated managed object name
         filename, _ = extract_file_metadata(file_data)
@@ -789,7 +799,12 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
             raise ValueError("file is required")
         if purpose is None:
             raise ValueError("purpose is required")
-        object_name = self.get_object_name(file_data, purpose)
+        configured_model: Final = litellm_params.get("model")
+        object_name = self.get_object_name(
+            file_data,
+            purpose,
+            deployment_model=configured_model if isinstance(configured_model, str) else None,
+        )
         if object_prefix:
             object_name = f"{object_prefix}/{object_name}"
         encoded_object_name: Final = encode_gcs_object_name_for_url(object_name)

tests/test_litellm/batches/test_main.py

@@ -158,14 +158,12 @@ def test_create__vertex_ai_dispatch(seams):
     _assert_only(seams.vertex.create_batch, seams, "create_batch")
 
 
-def test_create__vertex_ai_custom_endpoint_raises_badrequest(seams):
-    """custom_endpoint deployments have no Vertex batch surface; creating a job would target a
-    nonexistent publisher model, so the SDK must 400 before dispatching (LIT-6899)."""
-    with pytest.raises(litellm.exceptions.BadRequestError, match="custom_endpoint"):
-        bm.create_batch(**CREATE_KW, custom_llm_provider="vertex_ai", custom_endpoint=True)
+def test_create__vertex_ai_forwards_custom_endpoint(seams):
+    """The vertex handler owns the custom_endpoint batch rejection (LIT-6899), so the dispatcher
+    must forward the flag for the handler to act on."""
+    bm.create_batch(**CREATE_KW, custom_llm_provider="vertex_ai", custom_endpoint=True)
 
-    for m in _all_seam_methods(seams, "create_batch"):
-        m.assert_not_called()
+    assert seams.vertex.create_batch.call_args.kwargs["custom_endpoint"] is True
 
 
 def test_create__provider_config_routes_to_base_http_handler(seams):

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

@@ -274,6 +274,32 @@ def test_create_batch_sync_endpoint_resolution_error_raises():
     client.post.assert_not_called()
 
 
+def test_create_batch_custom_endpoint_raises_400_without_io():
+    """custom_endpoint deployments have no Vertex batch surface; creating a job would target a
+    nonexistent publisher model, so the handler must 400 before any auth or HTTP work (LIT-6899)."""
+    h = _make_handler()
+    client = MagicMock()
+
+    with patch(f"{HMOD}._get_httpx_client", return_value=client):
+        with pytest.raises(VertexAIError) as exc_info:
+            h.create_batch(
+                _is_async=False,
+                create_batch_data=CREATE_DATA,
+                api_base=None,
+                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 "custom_endpoint" in str(exc_info.value)
+    h._ensure_access_token.assert_not_called()
+    client.post.assert_not_called()
+
+
 def test_create_batch_sync_endpoint_without_deployed_model_raises_400():
     h = _make_handler()
     client = MagicMock()

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

@@ -367,6 +367,20 @@ def test_get_model_from_gcs_file_fine_tuned_endpoint():
     assert T._get_model_from_gcs_file(ENDPOINT_INPUT_FILE) == f"endpoints/{ENDPOINT_ID}"
 
 
+def test_get_model_from_gcs_file_publisher_path_wins_over_endpoints_prefix():
+    """A bucket prefix containing endpoints/<digits> must not override the publisher model path
+    LiteLLM appended after it."""
+    uri = "gs://bucket/team-endpoints/999/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/uuid"
+    assert T._get_model_from_gcs_file(uri) == "publishers/google/models/gemini-1.5-flash-001"
+
+
+def test_get_model_from_gcs_file_last_endpoints_segment_wins():
+    """With no publisher path, the endpoint id closest to the file (last occurrence) is the one
+    LiteLLM stored; an earlier prefix segment must not shadow it."""
+    uri = f"gs://bucket/endpoints/999/litellm-vertex-files/endpoints/{ENDPOINT_ID}/uuid"
+    assert T._get_model_from_gcs_file(uri) == f"endpoints/{ENDPOINT_ID}"
+
+
 def test_get_model_from_gcs_file_non_numeric_endpoints_segment_raises_400():
     with pytest.raises(VertexAIError) as exc_info:
         T._get_model_from_gcs_file("gs://bucket/endpoints/not-a-number/file-uuid")

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

@@ -177,6 +177,36 @@ class TestBatchObjectNaming:
         object_name = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "7768560373388541952"}}])
         assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/")
 
+    def test_deployment_model_overrides_jsonl_body_model(self, config):
+        """The stored path decides which Vertex model the batch later runs against with the
+        deployment's credentials, so a user-crafted JSONL body.model must not be able to redirect
+        an authorized deployment to a different endpoint."""
+        object_name = config._get_gcs_object_name_from_batch_jsonl(
+            [{"body": {"model": "9999999999999999999"}}],
+            deployment_model="vertex_ai/gemini/7768560373388541952",
+        )
+        assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/")
+        assert "9999999999999999999" not in object_name
+
+    def test_url_derives_object_path_from_configured_model(self, config):
+        url = config.get_complete_file_url(
+            api_base=None,
+            api_key=None,
+            model="",
+            optional_params={},
+            litellm_params={
+                "gcs_bucket_name": "my-bucket",
+                "model": "vertex_ai/gemini/7768560373388541952",
+            },
+            data={
+                "file": ("batch.jsonl", b'{"body": {"model": "9999999999999999999"}}', "application/jsonl"),
+                "purpose": "batch",
+            },
+        )
+        object_name = parse_qs(urlparse(url).query)["name"][0]
+        assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/")
+        assert "9999999999999999999" not in object_name
+
 
 class TestCustomEndpointBatchUpload:
     def test_should_reject_batch_upload_for_custom_endpoint_deployment(self, config):