scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 133

Merge pull request #39668 from BerriAI/litellm_lit6899_vertex_batch_tuned_endpoints

Browse source 
fix(vertex_ai): support fine-tuned Gemini endpoints in managed batches
This commit is contained in:
Mateo Wang 2026-09-07 20:28:02 -07:00 • committed by GitHub
commit 9dbfb060bd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 619 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/batches/main.py
View file

@ -319,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(

131
litellm/llms/vertex_ai/batches/handler.py
View file

@ -1,6 +1,7 @@
import json
from collections.abc import Coroutine
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
@ -12,11 +13,13 @@ 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,
)
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,
@ -55,6 +58,20 @@ class _FetchedResponseView(TypedDict):
response: ReadOnly[httpx.Response]
class _VertexEndpointDeployedModel(TypedDict, total=False):
model: ReadOnly[str]
class _VertexEndpointResponse(TypedDict, total=False):
deployedModels: ReadOnly[Sequence[_VertexEndpointDeployedModel]]
class _VertexEndpointPayloadView(TypedDict):
"""Holds one decoded GET endpoints/<id> response so the payload reads back typed."""
payload: ReadOnly[_VertexEndpointResponse]
def _vertex_batch_payload(response: _VertexBatchJsonSource) -> VertexBatchPredictionResponse:
return response.json()
@ -78,7 +95,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(
@ -87,6 +114,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,
@ -111,17 +158,6 @@ class VertexAIBatchPrediction(VertexLLM):
vertex_api_version="v1",
)
headers: Final = {
"Content-Type": "application/json; charset=utf-8",
"Authorization": f"Bearer {access_token}",
}
vertex_batch_request: Final[VertexAIBatchPredictionJob] = (
VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request(
request=create_batch_data
)
)
if _is_async is True:
return self._async_create_batch(
vertex_batch_request=vertex_batch_request,
@ -142,6 +178,77 @@ 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/<id>`).
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,
headers: dict[str, str], # mutable-ok: HTTPHandler.get only accepts dict headers
sync_handler: HTTPHandler,
api_base: str | None,
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/<id>`) read from GET endpoints/<id>.
"""
model: Final = vertex_batch_request.get("model", "")
if "/endpoints/" not in model:
return vertex_batch_request
endpoint_url: Final = self._build_endpoint_resolution_url(
api_base=api_base,
model=model,
vertex_location=vertex_location,
)
# ``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,
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"
),
)
resolved_request: Final[VertexAIBatchPredictionJob] = {**vertex_batch_request, "model": deployed_model}
return resolved_request
async def _async_create_batch(
self,
vertex_batch_request: VertexAIBatchPredictionJob,

70
litellm/llms/vertex_ai/batches/transformation.py
View file

@ -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/<id>` 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/<id>`) 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/<numeric id>` 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/<publisher>/models/<model>' path segment. "
f"'{gcs_file_uri}' contains no 'publishers/<publisher>/models/<model>' or "
"'endpoints/<numeric endpoint id>' 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://<bucket>/<prefix>/publishers/<publisher>/models/<model>/<file>"
"gs://<bucket>/<prefix>/publishers/<publisher>/models/<model>/<file> "
"(or gs://<bucket>/<prefix>/endpoints/<numeric endpoint id>/<file> for fine-tuned models)"
),
)
return model
@ -222,18 +260,26 @@ class VertexAIBatchTransformation:
@classmethod
def _parse_model_from_gcs_file(cls, gcs_file_uri: str) -> str | None:
"""
Returns the `publishers/<publisher>/models/<model>` path from a gcs uri, or None if the uri
does not contain one.
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.
"""
_, separator, model_path = unquote(gcs_file_uri).partition("publishers/")
if not separator:
return None
unquoted_uri: Final = unquote(gcs_file_uri)
_, 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])}"
parts: Final = model_path.split("/")
if len(parts) < 3 or parts[1] != "models" or not parts[2]:
return None
_, 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}"
return f"publishers/{'/'.join(parts[:3])}"
return None
@classmethod
def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: str | None) -> bool:

13
litellm/llms/vertex_ai/common_utils.py
View file

@ -370,6 +370,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/<id>` or `vertex_ai/gemini/<id>`.
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

51
litellm/llms/vertex_ai/files/transformation.py
View file

@ -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 (
@ -707,20 +708,39 @@ 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.
"""
_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 = (
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}"
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
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.
@ -728,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)
@ -761,6 +781,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")
@ -769,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)

29
tests/proxy_unit_tests/test_check_batch_cost.py
View file

@ -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/<id>` in the gs:// path; the bare model
(the endpoint id) must round-trip to the deployment configured as
`vertex_ai/gemini/<id>` (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."""

8
tests/test_litellm/batches/test_main.py
View file

@ -158,6 +158,14 @@ def test_create__vertex_ai_dispatch(seams):
_assert_only(seams.vertex.create_batch, seams, "create_batch")
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)
assert seams.vertex.create_batch.call_args.kwargs["custom_endpoint"] is True
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."""

192
tests/test_litellm/llms/vertex_ai/batches/test_handler.py
View file

@ -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,
)
@ -178,6 +177,197 @@ 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, 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),
patch(f"{HMOD}.safe_get") as safe_get,
):
out = 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,
)
assert isinstance(out, LiteLLMBatch)
sent = json.loads(client.post.call_args.kwargs["data"])
assert sent["model"] == "publishers/google/models/gemini-1.5-flash-001"
safe_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.post.return_value = _http_response()
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,
api_base=None,
vertex_credentials=None,
vertex_project=PROJECT,
vertex_location=LOCATION,
timeout=600.0,
max_retries=None,
)
assert isinstance(out, LiteLLMBatch)
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}"
)
assert get_kwargs["headers"]["Authorization"] == f"Bearer {TOKEN}"
sent = json.loads(client.post.call_args.kwargs["data"])
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()
resolve_response = MagicMock()
resolve_response.status_code = 404
resolve_response.text = "endpoint not found"
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,
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_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()
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,
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."""

72
tests/test_litellm/llms/vertex_ai/batches/test_transformation.py
View file

@ -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/<numeric id>) 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,35 @@ 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_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")
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 +404,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):

85
tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py
View file

@ -159,6 +159,91 @@ 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/<id>,
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/")
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):
"""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"