From e5d51ee8be3d191b295a94366e790c1cc9b43b08 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 19:06:01 -0400 Subject: [PATCH 1/6] fix(vertex_ai): support fine-tuned Gemini endpoints in managed batches Managed batches mangled any Vertex model that is not a plain publisher model: a fine-tuned Gemini endpoint id was filed under publishers/google/models/gemini/ at upload, then the batch create parse dropped the id and targeted the nonexistent publisher model 'publishers/google/models/gemini', which Vertex rejects. Fine-tuned endpoints are now stored under endpoints/ in the GCS object path, and batch create resolves the endpoint to its deployed tuned model resource (projects/../models/) via GET endpoints/, which is the only form the v1 batch API accepts for tuned models. The cost poller's bare-model parse round-trips the endpoint id so unmanaged batch spend still maps to the configured deployment. custom_endpoint deployments have no Vertex batch surface, so batch file uploads and batch creation against them now return a clear 400 instead of creating a doomed job. Resolves LIT-6899 --- litellm/batches/main.py | 15 +++ litellm/llms/vertex_ai/batches/handler.py | 64 +++++++++- .../llms/vertex_ai/batches/transformation.py | 56 ++++++++- litellm/llms/vertex_ai/common_utils.py | 13 ++ .../llms/vertex_ai/files/transformation.py | 28 ++++- .../proxy_unit_tests/test_check_batch_cost.py | 29 +++++ tests/test_litellm/batches/test_main.py | 10 ++ .../llms/vertex_ai/batches/test_handler.py | 119 ++++++++++++++++++ .../vertex_ai/batches/test_transformation.py | 58 +++++++++ .../test_vertex_ai_files_transformation.py | 55 ++++++++ 10 files changed, 435 insertions(+), 12 deletions(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index c8360a81c7a..6cf9f070d9c 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -301,6 +301,21 @@ 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") diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 377cd9f3437..1bb8ceb747a 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -12,6 +12,7 @@ from litellm.litellm_core_utils.url_utils import ( safe_get, ) from litellm.llms.custom_httpx.http_handler import ( + HTTPHandler, _get_httpx_client, get_async_httpx_client, ) @@ -55,6 +56,20 @@ class _FetchedResponseView(TypedDict): response: ReadOnly[httpx.Response] +class _VertexEndpointDeployedModel(TypedDict, total=False): + model: ReadOnly[str] + + +class _VertexEndpointResponse(TypedDict, total=False): + deployedModels: ReadOnly[list[_VertexEndpointDeployedModel]] + + +class _VertexEndpointPayloadView(TypedDict): + """Holds one decoded GET endpoints/ response so the payload reads back typed.""" + + payload: ReadOnly[_VertexEndpointResponse] + + def _vertex_batch_payload(response: _VertexBatchJsonSource) -> VertexBatchPredictionResponse: return response.json() @@ -116,11 +131,19 @@ class VertexAIBatchPrediction(VertexLLM): "Authorization": f"Bearer {access_token}", } - vertex_batch_request: Final[VertexAIBatchPredictionJob] = ( + transformed_batch_request: Final[VertexAIBatchPredictionJob] = ( VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request( - request=create_batch_data + 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( @@ -142,6 +165,43 @@ class VertexAIBatchPrediction(VertexLLM): ) return vertex_batch_response + def _resolve_fine_tuned_endpoint_model( + self, + vertex_batch_request: VertexAIBatchPredictionJob, + headers: dict[str, str], + sync_handler: HTTPHandler, + vertex_location: str, + ) -> VertexAIBatchPredictionJob: + """ + A fine-tuned Gemini deployment is configured by its endpoint id, but the v1 batch API only + accepts Model resources, so swap the endpoint resource for its deployed tuned model + (`projects/../locations/../models/`) read from GET endpoints/. + """ + model: Final = vertex_batch_request.get("model", "") + if "/endpoints/" not in model: + return vertex_batch_request + + endpoint_url: Final = f"{get_vertex_base_url(vertex_location)}/v1/{model}" + response: Final = sync_handler.get(url=endpoint_url, headers=headers) + if response.status_code != 200: + raise VertexAIError( + status_code=response.status_code, + message=f"Failed to resolve fine-tuned Vertex endpoint '{model}': {response.text}", + ) + + payload_view: Final[_VertexEndpointPayloadView] = {"payload": response.json()} + deployed_models: Final = payload_view["payload"].get("deployedModels") or [] + deployed_model: Final = deployed_models[0].get("model", "") if deployed_models else "" + if not deployed_model: + raise VertexAIError( + status_code=400, + message=( + f"Vertex endpoint '{model}' has no deployed model, so there is no tuned model " + "resource to run batch predictions against" + ), + ) + return {**vertex_batch_request, "model": deployed_model} + async def _async_create_batch( self, vertex_batch_request: VertexAIBatchPredictionJob, diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index f284b47292b..daedf6d1959 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -22,6 +22,8 @@ class VertexAIBatchTransformation: def transform_openai_batch_request_to_vertex_ai_batch_request( cls, request: CreateBatchRequest, + vertex_project: str | None = None, + vertex_location: str | None = None, ) -> VertexAIBatchPredictionJob: """ Transforms OpenAI Batch requests to Vertex AI Batch requests @@ -31,7 +33,11 @@ class VertexAIBatchTransformation: if input_file_id is None: raise ValueError("input_file_id is required, but not provided") input_config: InputConfig = InputConfig(gcsSource=GcsSource(uris=[input_file_id]), instancesFormat="jsonl") - model: Final[str] = cls._get_model_from_gcs_file(input_file_id) + model: Final[str] = cls._get_batch_job_model( + input_file_id=input_file_id, + vertex_project=vertex_project, + vertex_location=vertex_location, + ) output_config: Final[OutputConfig] = OutputConfig( predictionsFormat="jsonl", gcsDestination=GcsDestination(outputUriPrefix=cls._get_gcs_uri_prefix_from_file(input_file_id)), @@ -188,6 +194,33 @@ class VertexAIBatchTransformation: path_parts: Final = input_file_id.rsplit("/", 1) return path_parts[0] + @classmethod + def _get_batch_job_model( + cls, + input_file_id: str, + vertex_project: str | None, + vertex_location: str | None, + ) -> str: + """ + Returns the `model` for the batchPredictionJobs request: the publisher model path as-is, or + the full `projects/../locations/../endpoints/` resource name for a fine-tuned endpoint. + + The v1 batch API only accepts Model resources, so the handler resolves an endpoint resource + to its deployed tuned model (`projects/../locations/../models/`) before sending the job. + """ + parsed_model: Final = cls._get_model_from_gcs_file(input_file_id) + if not parsed_model.startswith("endpoints/"): + return parsed_model + if not vertex_project: + raise VertexAIError( + status_code=400, + message=( + f"Vertex AI batch jobs against a fine-tuned endpoint ('{parsed_model}') require " + "`vertex_project` to build the endpoint resource name" + ), + ) + return f"projects/{vertex_project}/locations/{vertex_location or 'us-central1'}/{parsed_model}" + @classmethod def _get_model_from_gcs_file(cls, gcs_file_uri: str) -> str: """ @@ -202,6 +235,9 @@ class VertexAIBatchTransformation: gcs_file_uri format: gs://litellm-testing-bucket/litellm-vertex-files/publishers/google/models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8 returns: "publishers/google/models/gemini-1.5-flash-001" + Fine-tuned Gemini endpoints are stored as `endpoints/` in the uri and returned + in that form. + Raises a 400 `VertexAIError` when the uri carries no parseable model path. """ model: Final = cls._parse_model_from_gcs_file(gcs_file_uri) @@ -210,11 +246,13 @@ class VertexAIBatchTransformation: status_code=400, message=( "Vertex AI batch creation requires the model to be part of `input_file_id`, but " - f"'{gcs_file_uri}' contains no 'publishers//models/' path segment. " + f"'{gcs_file_uri}' contains no 'publishers//models/' or " + "'endpoints/' path segment. " "Either upload the input file through LiteLLM (POST /v1/files with " "custom_llm_provider=vertex_ai), which encodes the model into the returned file id, or " "pass a uri of the form " - "gs:////publishers//models//" + "gs:////publishers//models// " + "(or gs:////endpoints// for fine-tuned models)" ), ) return model @@ -222,10 +260,16 @@ class VertexAIBatchTransformation: @classmethod def _parse_model_from_gcs_file(cls, gcs_file_uri: str) -> str | None: """ - Returns the `publishers//models/` path from a gcs uri, or None if the uri - does not contain one. + Returns the `publishers//models/` or `endpoints/` path from a + gcs uri, or None if the uri does not contain one. """ - _, separator, model_path = unquote(gcs_file_uri).partition("publishers/") + unquoted_uri: Final = unquote(gcs_file_uri) + _, endpoint_separator, endpoint_path = unquoted_uri.partition("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 diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 970759479fe..ba654f0d851 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -310,6 +310,19 @@ def get_vertex_base_model_name(model: str) -> str: return model +def get_vertex_ai_fine_tuned_endpoint_id(model: str) -> str | None: + """ + Fine-tuned Gemini deployments are addressed by a numeric endpoint id, + configured as `vertex_ai/` or `vertex_ai/gemini/`. + + Returns the endpoint id, or None when `model` is a regular publisher model. + Mirrors the online chat path in `_get_vertex_url`, which sends numeric + models to `endpoints/{id}` instead of `publishers/google/models/{model}`. + """ + candidate: Final = model.split("/")[-1] if "gemini/" in model else model + return candidate if candidate.isdigit() else None + + def validate_vertex_location(vertex_location: str | None) -> str: """ Validate a Vertex AI location before interpolating it into a request host or diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index b6ad9fbcc04..795918f1766 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -39,6 +39,7 @@ from litellm.llms.base_llm.files.transformation import ( ) from litellm.llms.vertex_ai.common_utils import ( _convert_vertex_datetime_to_openai_datetime, + get_vertex_ai_fine_tuned_endpoint_id, ) from litellm.llms.vertex_ai.gemini.transformation import _transform_request_body from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( @@ -712,11 +713,20 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): Gets a unique GCS object name for the VertexAI batch prediction job named as: litellm-vertex-{model}-{uuid} + + Fine-tuned Gemini deployments (numeric endpoint ids) are stored under + `endpoints/` so the batch transformation can round-trip them into a + `projects/../locations/../endpoints/` batch job model instead of a + nonexistent publisher model. """ - _model = openai_jsonl_content[0].get("body", {}).get("model", "") - if "publishers/google/models" not in _model: - _model = f"publishers/google/models/{_model}" - safe_model_path: Final = sanitize_cloud_object_path(_model, fallback="model") + raw_model: Final = 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}" + if endpoint_id is not None + else (raw_model if "publishers/google/models" in raw_model else f"publishers/google/models/{raw_model}") + ) + safe_model_path: Final = sanitize_cloud_object_path(model_path, fallback="model") object_name: Final = f"{VERTEX_AI_MANAGED_GCS_PREFIX}{safe_model_path}/{uuid.uuid4()}" return object_name @@ -761,6 +771,16 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ Get the complete url for the request """ + if data.get("purpose") == "batch" and litellm_params.get("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; " + "remove this deployment from the batch request (e.g. `target_model_names`) or " + "use a publisher model / fine-tuned Gemini endpoint instead." + ), + ) bucket_name = self._get_configured_bucket_name(litellm_params) bucket_name, object_prefix = split_configured_cloud_bucket_name(bucket_name) file_data: Final = data.get("file") diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index ff5e8f89d64..b69805372ba 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1760,6 +1760,35 @@ class TestUnmanagedVertexRouting: ) router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash") + def test_flag_on_routes_fine_tuned_endpoint_to_vertex_deployment(self): + """A fine-tuned Gemini batch stores `endpoints/` in the gs:// path; the bare model + (the endpoint id) must round-trip to the deployment configured as + `vertex_ai/gemini/` (LIT-6899).""" + endpoint_id = "7768560373388541952" + router = MagicMock() + router.resolve_model_name_from_model_id.return_value = None + router.get_model_list.return_value = [ + { + "model_name": "gemini-2.5-flash-dts-usc1", + "litellm_params": { + "model": f"vertex_ai/gemini/{endpoint_id}", + "custom_llm_provider": "vertex_ai", + }, + "model_info": {"id": "deploy-ft"}, + }, + ] + instance = self._instance(track_unmanaged=True, router=router) + job = self._job( + file_object=_unmanaged_vertex_file_object( + input_file_id=f"gs://bucket/litellm-vertex-files/endpoints/{endpoint_id}/abc.jsonl" + ) + ) + + with patch(_IS_B64, return_value=False): + result = instance._resolve_job_routing(job, MagicMock()) + + assert result == ("deploy-ft", "8823717160934178816") + def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): """Flag on, but the only deployment for the model group is a non-vertex_ai provider: must not be selected, even though the model group name matches.""" diff --git a/tests/test_litellm/batches/test_main.py b/tests/test_litellm/batches/test_main.py index c3edb40c819..b10214f884d 100644 --- a/tests/test_litellm/batches/test_main.py +++ b/tests/test_litellm/batches/test_main.py @@ -158,6 +158,16 @@ 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) + + for m in _all_seam_methods(seams, "create_batch"): + m.assert_not_called() + + def test_create__provider_config_routes_to_base_http_handler(seams): """model + a provider batches config (bedrock-style) routes to the generic base_llm_http_handler, NOT the per-provider instance.""" 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 38fde3caa63..bb1c546614e 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -178,6 +178,125 @@ def test_create_batch_async_returns_coroutine_and_uses_async_client(): sync_client.post.assert_not_called() +def test_create_batch_sync_does_not_resolve_publisher_models(): + """Publisher-model jobs must not incur the endpoint-resolution GET.""" + h = _make_handler() + client = MagicMock() + client.post.return_value = _http_response() + + with patch(f"{HMOD}._get_httpx_client", return_value=client): + 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, + ) + + client.get.assert_not_called() + + +ENDPOINT_ID = "7768560373388541952" +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" + + +def _endpoint_get_response(deployed_models: list | None = None) -> MagicMock: + resp = 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}] + ), + } + return resp + + +def test_create_batch_sync_resolves_fine_tuned_endpoint_to_tuned_model(): + """A fine-tuned Gemini file id must produce a batch job against the endpoint's deployed + tuned model resource; the v1 batch API rejects endpoint resources in `model` (LIT-6899).""" + h = _make_handler() + client = MagicMock() + client.get.return_value = _endpoint_get_response() + client.post.return_value = _http_response() + + with patch(f"{HMOD}._get_httpx_client", return_value=client): + out = h.create_batch( + _is_async=False, + create_batch_data=ENDPOINT_CREATE_DATA, + api_base=None, + vertex_credentials=None, + vertex_project=PROJECT, + vertex_location=LOCATION, + timeout=600.0, + max_retries=None, + ) + + assert isinstance(out, LiteLLMBatch) + get_kwargs = client.get.call_args.kwargs + assert get_kwargs["url"] == ( + f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}" + f"/locations/{LOCATION}/endpoints/{ENDPOINT_ID}" + ) + assert get_kwargs["headers"]["Authorization"] == f"Bearer {TOKEN}" + sent = json.loads(client.post.call_args.kwargs["data"]) + assert sent["model"] == TUNED_MODEL_RESOURCE + + +def test_create_batch_sync_endpoint_resolution_error_raises(): + h = _make_handler() + client = MagicMock() + resolve_response = MagicMock() + resolve_response.status_code = 404 + resolve_response.text = "endpoint not found" + client.get.return_value = resolve_response + + 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=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 == 404 + client.post.assert_not_called() + + +def test_create_batch_sync_endpoint_without_deployed_model_raises_400(): + h = _make_handler() + client = MagicMock() + client.get.return_value = _endpoint_get_response(deployed_models=[]) + + 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=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 + assert "no deployed model" in str(exc_info.value) + client.post.assert_not_called() + + def test_create_batch_sync_httpstatuserror_propagates(): """``HTTPHandler.post`` raises for non-2xx via ``raise_for_status``; the sync create path must surface that error, not swallow it.""" 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 ccb2d7e310d..72ab87bee1a 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py @@ -34,6 +34,12 @@ INPUT_FILE = ( "models/gemini-1.5-flash-001/e9412502-2c91-42a6-8e61-f5c294cc0fc8" ) +ENDPOINT_ID = "7768560373388541952" +ENDPOINT_INPUT_FILE = ( + f"gs://litellm-testing-bucket/litellm-vertex-files/endpoints/{ENDPOINT_ID}/" + "e9412502-2c91-42a6-8e61-f5c294cc0fc8" +) + # =========================================================================== # # transform_openai_batch_request_to_vertex_ai_batch_request @@ -67,6 +73,41 @@ def test_transform_openai_request_missing_input_file_id_raises(): T.transform_openai_batch_request_to_vertex_ai_batch_request({}) +def test_transform_openai_request_fine_tuned_endpoint_builds_endpoint_resource(): + """A fine-tuned Gemini file id (endpoints/) must target the endpoint resource, + not a nonexistent publisher model (LIT-6899).""" + job = T.transform_openai_batch_request_to_vertex_ai_batch_request( + {"input_file_id": ENDPOINT_INPUT_FILE}, + vertex_project="my-project", + vertex_location="us-central1", + ) + assert job["model"] == f"projects/my-project/locations/us-central1/endpoints/{ENDPOINT_ID}" + + +def test_transform_openai_request_fine_tuned_endpoint_defaults_location(): + job = T.transform_openai_batch_request_to_vertex_ai_batch_request( + {"input_file_id": ENDPOINT_INPUT_FILE}, + vertex_project="my-project", + ) + assert job["model"] == f"projects/my-project/locations/us-central1/endpoints/{ENDPOINT_ID}" + + +def test_transform_openai_request_fine_tuned_endpoint_without_project_raises_400(): + with pytest.raises(VertexAIError) as exc_info: + T.transform_openai_batch_request_to_vertex_ai_batch_request({"input_file_id": ENDPOINT_INPUT_FILE}) + assert exc_info.value.status_code == 400 + assert "vertex_project" in str(exc_info.value) + + +def test_transform_openai_request_publisher_model_ignores_project_and_location(): + job = T.transform_openai_batch_request_to_vertex_ai_batch_request( + {"input_file_id": INPUT_FILE}, + vertex_project="my-project", + vertex_location="europe-west4", + ) + assert job["model"] == "publishers/google/models/gemini-1.5-flash-001" + + @pytest.mark.parametrize( "input_file_id", [ @@ -321,6 +362,21 @@ def test_get_model_from_gcs_file_no_publishers_raises_400(): assert exc_info.value.status_code == 400 +def test_get_model_from_gcs_file_fine_tuned_endpoint(): + """The whole endpoint id must survive parsing; the old 3-segment publishers/ parse dropped it.""" + assert T._get_model_from_gcs_file(ENDPOINT_INPUT_FILE) == 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") + assert exc_info.value.status_code == 400 + + +def test_get_bare_model_name_from_gcs_file_fine_tuned_endpoint(): + assert T.get_bare_model_name_from_gcs_file(ENDPOINT_INPUT_FILE) == ENDPOINT_ID + + # =========================================================================== # # is_unmanaged_gcs_batch_input_file_id # =========================================================================== # @@ -334,6 +390,8 @@ def test_get_model_from_gcs_file_no_publishers_raises_400(): ("file-abc123", False), ("gs://bucket/no-model-here.jsonl", False), ("gs://bucket/publishers/google/gemini-1.5-flash-001/file-uuid", False), + (ENDPOINT_INPUT_FILE, True), + ("gs://bucket/endpoints/not-a-number/file-uuid", False), ], ) def test_is_unmanaged_gcs_batch_input_file_id(input_file_id, expected): 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 3c2d56997b7..f0ea61f183c 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 @@ -159,6 +159,61 @@ class TestCreateFileUrl: assert "?" not in object_name +class TestBatchObjectNaming: + def test_should_store_publisher_model_under_publishers_path(self, config): + object_name = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "gemini-2.5-flash"}}]) + assert object_name.startswith("litellm-vertex-files/publishers/google/models/gemini-2.5-flash/") + + 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"}}] + ) + assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/") + assert "publishers" not in object_name + + def test_should_store_bare_numeric_endpoint_under_endpoints_path(self, config): + object_name = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "7768560373388541952"}}]) + assert object_name.startswith("litellm-vertex-files/endpoints/7768560373388541952/") + + +class TestCustomEndpointBatchUpload: + def test_should_reject_batch_upload_for_custom_endpoint_deployment(self, config): + """custom_endpoint deployments have no Vertex batch surface; the upload must 400 instead + of staging a file that can only produce a doomed batch job (LIT-6899).""" + from litellm.llms.vertex_ai.common_utils import VertexAIError + + with pytest.raises(VertexAIError) as exc_info: + config.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "my-bucket", "custom_endpoint": True}, + data={ + "file": ("batch.jsonl", b'{"body": {"model": "openai/gemma-2-2b-it"}}', "application/jsonl"), + "purpose": "batch", + }, + ) + assert exc_info.value.status_code == 400 + assert "custom_endpoint" in str(exc_info.value) + + def test_should_allow_non_batch_upload_for_custom_endpoint_deployment(self, config): + url = config.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={"gcs_bucket_name": "my-bucket", "custom_endpoint": True}, + data={ + "file": ("notes.txt", b"hello", "text/plain"), + "purpose": "assistants", + }, + ) + assert "/b/my-bucket/" in url + + class TestTransformRetrieveFile: def test_should_build_correct_gcs_metadata_url(self, config): file_id = "gs://my-bucket/litellm-vertex-files/path/to/file.jsonl" From 1d2ed0bdacc895b9c1e06cef2b97a4646f88f096 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 19:49:37 -0400 Subject: [PATCH 2/6] 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/ 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 | 16 +---- litellm/llms/vertex_ai/batches/handler.py | 64 +++++++++++++------ .../llms/vertex_ai/batches/transformation.py | 22 ++++--- .../llms/vertex_ai/files/transformation.py | 25 ++++++-- tests/test_litellm/batches/test_main.py | 12 ++-- .../llms/vertex_ai/batches/test_handler.py | 26 ++++++++ .../vertex_ai/batches/test_transformation.py | 14 ++++ .../test_vertex_ai_files_transformation.py | 30 +++++++++ 8 files changed, 152 insertions(+), 57 deletions(-) diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 6cf9f070d9c..77a4fdebf16 100644 --- a/litellm/batches/main.py +++ b/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( diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 1bb8ceb747a..ae99eea777c 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/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( diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index daedf6d1959..e63c80dd3cf 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -262,22 +262,24 @@ class VertexAIBatchTransformation: """ Returns the `publishers//models/` or `endpoints/` 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/` 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: diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 795918f1766..263956efc9f 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/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/` so the batch transformation can round-trip them into a `projects/../locations/../endpoints/` 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) diff --git a/tests/test_litellm/batches/test_main.py b/tests/test_litellm/batches/test_main.py index b10214f884d..b87f9489250 100644 --- a/tests/test_litellm/batches/test_main.py +++ b/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): 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 bb1c546614e..8db4775a2bc 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/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() 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 72ab87bee1a..232c6413e78 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_transformation.py +++ b/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/ 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") 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 f0ea61f183c..8a249820cbd 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 @@ -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): From 14d5ff5c2190f2a89a6b1730ce71c15d2332596f Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 19:58:37 -0400 Subject: [PATCH 3/6] fix(vertex_ai): keep new batch handler code within the mutable-collection budget --- litellm/llms/vertex_ai/batches/handler.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index ae99eea777c..7f20b197056 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -1,5 +1,5 @@ import json -from collections.abc import Coroutine +from collections.abc import Coroutine, Sequence from typing import TYPE_CHECKING, Final, Protocol import httpx @@ -61,7 +61,7 @@ class _VertexEndpointDeployedModel(TypedDict, total=False): class _VertexEndpointResponse(TypedDict, total=False): - deployedModels: ReadOnly[list[_VertexEndpointDeployedModel]] + deployedModels: ReadOnly[Sequence[_VertexEndpointDeployedModel]] class _VertexEndpointPayloadView(TypedDict): @@ -179,7 +179,7 @@ class VertexAIBatchPrediction(VertexLLM): def _resolve_fine_tuned_endpoint_model( self, vertex_batch_request: VertexAIBatchPredictionJob, - headers: dict[str, str], + headers: dict[str, str], # mutable-ok: HTTPHandler.get only accepts dict headers sync_handler: HTTPHandler, api_base: str | None, vertex_location: str, @@ -214,7 +214,7 @@ class VertexAIBatchPrediction(VertexLLM): ) payload_view: Final[_VertexEndpointPayloadView] = {"payload": response.json()} - deployed_models: Final = payload_view["payload"].get("deployedModels") or [] + deployed_models: Final = payload_view["payload"].get("deployedModels") or () deployed_model: Final = deployed_models[0].get("model", "") if deployed_models else "" if not deployed_model: raise VertexAIError( @@ -224,7 +224,8 @@ class VertexAIBatchPrediction(VertexLLM): "resource to run batch predictions against" ), ) - return {**vertex_batch_request, "model": deployed_model} + resolved_request: Final[VertexAIBatchPredictionJob] = {**vertex_batch_request, "model": deployed_model} + return resolved_request async def _async_create_batch( self, From 988ae7ca80b66dd9930c1a32ca399d18086561fa Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 20:04:05 -0400 Subject: [PATCH 4/6] test(vertex_ai): assert the publisher-model batch payload instead of only mock calls --- tests/test_litellm/llms/vertex_ai/batches/test_handler.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) 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 8db4775a2bc..d65ba92afd7 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -179,13 +179,14 @@ def test_create_batch_async_returns_coroutine_and_uses_async_client(): def test_create_batch_sync_does_not_resolve_publisher_models(): - """Publisher-model jobs must not incur the endpoint-resolution GET.""" + """Publisher-model jobs must not incur the endpoint-resolution GET, and the job model must + stay the publisher path untouched.""" h = _make_handler() client = MagicMock() client.post.return_value = _http_response() with patch(f"{HMOD}._get_httpx_client", return_value=client): - h.create_batch( + out = h.create_batch( _is_async=False, create_batch_data=CREATE_DATA, api_base=None, @@ -196,6 +197,9 @@ def test_create_batch_sync_does_not_resolve_publisher_models(): max_retries=None, ) + assert isinstance(out, LiteLLMBatch) + sent = json.loads(client.post.call_args.kwargs["data"]) + assert sent["model"] == "publishers/google/models/gemini-1.5-flash-001" client.get.assert_not_called() From 54fd69beb2525878f35b622374a76c63713f5760 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 20:13:51 -0400 Subject: [PATCH 5/6] fix(vertex_ai): build a well-formed endpoint-resolution url for path-mounted custom api_base --- litellm/llms/vertex_ai/batches/handler.py | 32 +++++++++++------ .../llms/vertex_ai/batches/test_handler.py | 34 ++++++++++++++++++- 2 files changed, 55 insertions(+), 11 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 7f20b197056..2d5afd4818e 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -1,6 +1,7 @@ import json from collections.abc import Coroutine, Sequence from typing import TYPE_CHECKING, Final, Protocol +from urllib.parse import urlparse import httpx from typing_extensions import ReadOnly, TypedDict @@ -18,6 +19,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM +from litellm.llms.vertex_ai.vertex_llm_base import _graft_default_vertex_path from litellm.types.llms.openai import CreateBatchRequest from litellm.types.llms.vertex_ai import ( VERTEX_CREDENTIALS_TYPES, @@ -176,6 +178,24 @@ class VertexAIBatchPrediction(VertexLLM): ) return vertex_batch_response + @staticmethod + def _build_endpoint_resolution_url(api_base: str | None, model: str, vertex_location: str) -> str: + """ + Builds the GET url for resolving an endpoint resource (`projects/../endpoints/`). + + A custom `api_base` replaces the Google host: its `/v1`/`/v1beta1` path swallows the + version segment (matching `_check_custom_proxy`'s grafting), any other path is kept as a + mount prefix in front of the full default path. The `:operation` suffix convention from + `_check_custom_proxy` does not apply to a plain resource GET. + """ + default_endpoint_url: Final = f"{get_vertex_base_url(vertex_location)}/v1/{model}" + if not api_base: + return default_endpoint_url + api_base_path: Final = urlparse(api_base).path.rstrip("/") + if api_base_path in ("/v1", "/v1beta1"): + return _graft_default_vertex_path(api_base=api_base, default_url=default_endpoint_url) + return api_base.rstrip("/") + urlparse(default_endpoint_url).path + def _resolve_fine_tuned_endpoint_model( self, vertex_batch_request: VertexAIBatchPredictionJob, @@ -193,18 +213,10 @@ class VertexAIBatchPrediction(VertexLLM): if "/endpoints/" not in model: return vertex_batch_request - default_endpoint_url: Final = f"{get_vertex_base_url(vertex_location)}/v1/{model}" - _, endpoint_url = self._check_custom_proxy( + endpoint_url: Final = self._build_endpoint_resolution_url( 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, + model=model, vertex_location=vertex_location, - vertex_api_version="v1", ) response: Final = sync_handler.get(url=endpoint_url, headers=headers) if response.status_code != 200: 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 d65ba92afd7..bba0eb27ed4 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -35,7 +35,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest - from litellm.llms.vertex_ai.batches.handler import ( # noqa: E402 VertexAIBatchPrediction, ) @@ -253,6 +252,39 @@ def test_create_batch_sync_resolves_fine_tuned_endpoint_to_tuned_model(): assert sent["model"] == TUNED_MODEL_RESOURCE +@pytest.mark.parametrize( + "api_base, expected", + [ + ( + None, + f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}" + f"/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + ), + ( + "https://proxy.internal", + f"https://proxy.internal/v1/projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + ), + ( + "https://proxy.internal/v1", + f"https://proxy.internal/v1/projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + ), + ( + "https://proxy.internal/vertex", + f"https://proxy.internal/vertex/v1/projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + ), + ], +) +def test_build_endpoint_resolution_url(api_base, expected): + """A custom api_base must replace the Google host for the endpoint-resolution GET without + producing a malformed url (no ':' grafting, no doubled /v1).""" + url = VertexAIBatchPrediction._build_endpoint_resolution_url( + api_base=api_base, + model=f"projects/{PROJECT}/locations/{LOCATION}/endpoints/{ENDPOINT_ID}", + vertex_location=LOCATION, + ) + assert url == expected + + def test_create_batch_sync_endpoint_resolution_error_raises(): h = _make_handler() client = MagicMock() From 18373f8f51c3e283d82a0f480dee51d84dc0e9d0 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 3 Sep 2026 20:38:24 -0400 Subject: [PATCH 6/6] fix(vertex_ai): route endpoint resolution through safe_get to guard caller-supplied api_base --- litellm/llms/vertex_ai/batches/handler.py | 12 +++++++- .../llms/vertex_ai/batches/test_handler.py | 29 ++++++++++++------- 2 files changed, 30 insertions(+), 11 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 2d5afd4818e..a15ea4d845b 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -218,7 +218,17 @@ class VertexAIBatchPrediction(VertexLLM): model=model, vertex_location=vertex_location, ) - response: Final = sync_handler.get(url=endpoint_url, headers=headers) + # ``api_base`` can come from caller-supplied request kwargs, so wrap the + # fetch in ``safe_get``: it rejects DNS-rebind / private / cloud-metadata + # targets before the bearer token leaves the process (mirrors retrieve_batch). + fetched: Final[_FetchedResponseView] = { + "response": safe_get( + sync_handler, + endpoint_url, + headers=headers, + ) + } + response: Final = fetched["response"] if response.status_code != 200: raise VertexAIError( status_code=response.status_code, 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 bba0eb27ed4..24df3214da3 100644 --- a/tests/test_litellm/llms/vertex_ai/batches/test_handler.py +++ b/tests/test_litellm/llms/vertex_ai/batches/test_handler.py @@ -184,7 +184,10 @@ def test_create_batch_sync_does_not_resolve_publisher_models(): client = MagicMock() client.post.return_value = _http_response() - with patch(f"{HMOD}._get_httpx_client", return_value=client): + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get") as safe_get, + ): out = h.create_batch( _is_async=False, create_batch_data=CREATE_DATA, @@ -199,7 +202,7 @@ def test_create_batch_sync_does_not_resolve_publisher_models(): assert isinstance(out, LiteLLMBatch) sent = json.loads(client.post.call_args.kwargs["data"]) assert sent["model"] == "publishers/google/models/gemini-1.5-flash-001" - client.get.assert_not_called() + safe_get.assert_not_called() ENDPOINT_ID = "7768560373388541952" @@ -226,10 +229,12 @@ def test_create_batch_sync_resolves_fine_tuned_endpoint_to_tuned_model(): tuned model resource; the v1 batch API rejects endpoint resources in `model` (LIT-6899).""" h = _make_handler() client = MagicMock() - client.get.return_value = _endpoint_get_response() client.post.return_value = _http_response() - with patch(f"{HMOD}._get_httpx_client", return_value=client): + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get", return_value=_endpoint_get_response()) as safe_get, + ): out = h.create_batch( _is_async=False, create_batch_data=ENDPOINT_CREATE_DATA, @@ -242,8 +247,8 @@ def test_create_batch_sync_resolves_fine_tuned_endpoint_to_tuned_model(): ) assert isinstance(out, LiteLLMBatch) - get_kwargs = client.get.call_args.kwargs - assert get_kwargs["url"] == ( + get_args, get_kwargs = safe_get.call_args + assert get_args[1] == ( f"https://{LOCATION}-aiplatform.googleapis.com/v1/projects/{PROJECT}" f"/locations/{LOCATION}/endpoints/{ENDPOINT_ID}" ) @@ -291,9 +296,11 @@ def test_create_batch_sync_endpoint_resolution_error_raises(): resolve_response = MagicMock() resolve_response.status_code = 404 resolve_response.text = "endpoint not found" - client.get.return_value = resolve_response - with patch(f"{HMOD}._get_httpx_client", return_value=client): + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get", return_value=resolve_response), + ): with pytest.raises(VertexAIError) as exc_info: h.create_batch( _is_async=False, @@ -339,9 +346,11 @@ def test_create_batch_custom_endpoint_raises_400_without_io(): def test_create_batch_sync_endpoint_without_deployed_model_raises_400(): h = _make_handler() client = MagicMock() - client.get.return_value = _endpoint_get_response(deployed_models=[]) - with patch(f"{HMOD}._get_httpx_client", return_value=client): + with ( + patch(f"{HMOD}._get_httpx_client", return_value=client), + patch(f"{HMOD}.safe_get", return_value=_endpoint_get_response(deployed_models=[])), + ): with pytest.raises(VertexAIError) as exc_info: h.create_batch( _is_async=False,