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):