diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index c7019a5c69c..9312da62eb8 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -20,6 +20,7 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.llms.vertex_ai.common_utils import ( VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, VertexAIError, + get_custom_endpoint_id_from_api_base, get_vertex_base_url, ) from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM @@ -156,14 +157,29 @@ class VertexAIBatchPrediction(VertexLLM): vertex_location=vertex_location or "us-central1", ) ) - if custom_endpoint and "/custom-endpoints/" not in transformed_batch_request.get("model", ""): + # 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. + 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}"): raise VertexAIError( status_code=400, message=( "Vertex AI batch prediction on a `custom_endpoint` deployment requires an input " - "file uploaded through LiteLLM against that deployment (its file id carries a " - "custom-endpoints/ path); this input file targets a publisher or " - "fine-tuned Gemini model instead." + "file uploaded through LiteLLM against that same deployment: the file id's " + "custom-endpoints/ path must name the endpoint the deployment's " + "`api_base` serves from." + ), + ) + if not custom_endpoint and "/custom-endpoints/" in job_model: + raise VertexAIError( + status_code=400, + message=( + "This input file was uploaded for a `custom_endpoint` deployment; create the " + "batch against that deployment instead." ), ) gateway_api_base: Final = _gateway_api_base_or_none(api_base) @@ -174,12 +190,16 @@ class VertexAIBatchPrediction(VertexLLM): api_base=gateway_api_base, vertex_location=vertex_location or "us-central1", ) - vertex_batch_request: Final = self._resolve_custom_endpoint_container( - vertex_batch_request=resolved_batch_request, - headers=headers, - sync_handler=sync_handler, - api_base=gateway_api_base, - vertex_location=vertex_location or "us-central1", + vertex_batch_request: Final = ( + self._resolve_custom_endpoint_container( + vertex_batch_request=resolved_batch_request, + headers=headers, + sync_handler=sync_handler, + api_base=gateway_api_base, + vertex_location=vertex_location or "us-central1", + ) + if custom_endpoint + else resolved_batch_request ) is_unmanaged_container_job: Final = "unmanagedContainerModel" in vertex_batch_request diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 7bda17d4334..4c14f736bd3 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -5,6 +5,7 @@ from enum import Enum from functools import lru_cache from types import MappingProxyType from typing import Any, Final, Literal, cast, get_type_hints +from urllib.parse import urlparse import httpx from pydantic import TypeAdapter, ValidationError @@ -373,6 +374,23 @@ def get_vertex_base_model_name(model: str) -> str: VERTEX_CUSTOM_ENDPOINT_KEY_FIELD: Final = "litellm_custom_id" +def get_custom_endpoint_id_from_api_base(api_base: str | None) -> str | None: + """ + The Vertex endpoint a `custom_endpoint` deployment serves from is only recorded in its + api_base (`.../endpoints/:rawPredict` or a dedicated-domain equivalent); batch jobs need + that id to read the endpoint's containerSpec, so extract it (verb suffix stripped). + """ + if not api_base: + return None + path_segments: Final = urlparse(api_base).path.split("/") + after_endpoints: Final = tuple( + segment for prior, segment in zip(path_segments, path_segments[1:]) if prior == "endpoints" + ) + if not after_endpoints: + return None + return after_endpoints[-1].split(":")[0] or None + + def get_vertex_ai_fine_tuned_endpoint_id(model: str) -> str | None: """ Fine-tuned Gemini deployments are addressed by a numeric endpoint id, diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index ed55a97145b..006342c5a4f 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -8,7 +8,7 @@ import time from collections.abc import Callable, Iterable, Iterator, Mapping from types import MappingProxyType from typing import Any, Final, TypedDict -from urllib.parse import quote, unquote, urlparse +from urllib.parse import quote, unquote import httpx from httpx import Headers, Response @@ -41,6 +41,7 @@ from litellm.llms.base_llm.files.transformation import ( from litellm.llms.vertex_ai.common_utils import ( VERTEX_CUSTOM_ENDPOINT_KEY_FIELD, _convert_vertex_datetime_to_openai_datetime, + get_custom_endpoint_id_from_api_base, get_vertex_ai_fine_tuned_endpoint_id, ) from litellm.llms.vertex_ai.gemini.transformation import _transform_request_body @@ -713,23 +714,6 @@ _VERTEX_CHAT_COMPLETIONS_REQUEST_FORMAT: Final = "chatCompletions" _EMPTY_MAPPING: Final[Mapping[str, object]] = MappingProxyType({}) -def get_custom_endpoint_id_from_api_base(api_base: str | None) -> str | None: - """ - The Vertex endpoint a `custom_endpoint` deployment serves from is only recorded in its - api_base (`.../endpoints/:rawPredict` or a dedicated-domain equivalent); batch jobs need - that id to read the endpoint's containerSpec, so extract it (verb suffix stripped). - """ - if not api_base: - return None - path_segments: Final = urlparse(api_base).path.split("/") - after_endpoints: Final = tuple( - segment for prior, segment in zip(path_segments, path_segments[1:]) if prior == "endpoints" - ) - if not after_endpoints: - return None - return after_endpoints[-1].split(":")[0] or None - - def _openai_batch_jsonl_entry_to_custom_endpoint_row(openai_entry: Mapping[str, object]) -> Mapping[str, object]: """ One OpenAI batch JSONL line as the instance a vLLM-serving Vertex container consumes: 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 3a5b4db04bf..9bcbd723060 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -269,6 +269,10 @@ CONTAINER_SPEC = { "healthRoute": "/ping", } MACHINE_SPEC = {"machineType": "g2-standard-12", "acceleratorType": "NVIDIA_L4", "acceleratorCount": 1} +CUSTOM_ENDPOINT_API_BASE = ( + f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}" + f"/locations/{LOCATION}/endpoints/{CUSTOM_ENDPOINT_ID}:rawPredict" +) def _custom_endpoint_get_response() -> MagicMock: @@ -317,7 +321,7 @@ def test_create_batch_sync_custom_endpoint_builds_unmanaged_container_job(): out = h.create_batch( _is_async=False, create_batch_data=CUSTOM_ENDPOINT_CREATE_DATA, - api_base=None, + api_base=CUSTOM_ENDPOINT_API_BASE, vertex_credentials=None, vertex_project=PROJECT, vertex_location=LOCATION, @@ -360,7 +364,7 @@ def test_create_batch_sync_custom_endpoint_without_container_spec_raises_400(): h.create_batch( _is_async=False, create_batch_data=CUSTOM_ENDPOINT_CREATE_DATA, - api_base=None, + api_base=CUSTOM_ENDPOINT_API_BASE, vertex_credentials=None, vertex_project=PROJECT, vertex_location=LOCATION, @@ -374,6 +378,63 @@ 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_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 + container with its own project credentials.""" + h = _make_handler() + client = MagicMock() + foreign_file = {"input_file_id": "gs://bucket/litellm-vertex-files/custom-endpoints/999999/file-uuid"} + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get") as safe_get, + ): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data=foreign_file, + 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 + safe_get.assert_not_called() + client.post.assert_not_called() + + +def test_create_batch_sync_non_custom_deployment_rejects_custom_endpoint_file(): + """A custom-endpoints file id sent through an ordinary Vertex deployment must not trigger + container resolution at all.""" + h = _make_handler() + client = MagicMock() + + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get") as safe_get, + ): + with pytest.raises(VertexAIError) as exc_info: + h.create_batch( + _is_async=False, + create_batch_data=CUSTOM_ENDPOINT_CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert exc_info.value.status_code == 400 + safe_get.assert_not_called() + client.post.assert_not_called() + + def test_create_batch_sync_ignores_resource_shaped_api_base(): """A deployment api_base like `.../endpoints/:rawPredict` targets online inference, not the Vertex API root; grafting batch urls onto it yields guaranteed 404s, so batch operations