mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-02 12:21:10 +00:00
Merge pull request #23957 from Sameerlite/litellm_vertex-cancel-batch
fix(vertex-ai): support batch cancel via Vertex API
This commit is contained in:
Sameer Kankute
+29
-3
@@ -885,7 +885,7 @@ def list_batches(
|
||||
async def acancel_batch(
|
||||
batch_id: str,
|
||||
model: Optional[str] = None,
|
||||
custom_llm_provider: Literal["openai", "azure"] = "openai",
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
|
||||
metadata: Optional[Dict[str, str]] = None,
|
||||
extra_headers: Optional[Dict[str, str]] = None,
|
||||
extra_body: Optional[Dict[str, str]] = None,
|
||||
@@ -931,7 +931,7 @@ async def acancel_batch(
|
||||
def cancel_batch(
|
||||
batch_id: str,
|
||||
model: Optional[str] = None,
|
||||
custom_llm_provider: Union[Literal["openai", "azure"], str] = "openai",
|
||||
custom_llm_provider: Union[Literal["openai", "azure", "vertex_ai"], str] = "openai",
|
||||
metadata: Optional[Dict[str, str]] = None,
|
||||
extra_headers: Optional[Dict[str, str]] = None,
|
||||
extra_body: Optional[Dict[str, str]] = None,
|
||||
@@ -1048,9 +1048,35 @@ def cancel_batch(
|
||||
cancel_batch_data=_cancel_batch_request,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
elif custom_llm_provider == "vertex_ai":
|
||||
api_base = optional_params.api_base or None
|
||||
vertex_ai_project = (
|
||||
optional_params.vertex_project
|
||||
or litellm.vertex_project
|
||||
or get_secret_str("VERTEXAI_PROJECT")
|
||||
)
|
||||
vertex_ai_location = (
|
||||
optional_params.vertex_location
|
||||
or litellm.vertex_location
|
||||
or get_secret_str("VERTEXAI_LOCATION")
|
||||
)
|
||||
vertex_credentials = optional_params.vertex_credentials or get_secret_str(
|
||||
"VERTEXAI_CREDENTIALS"
|
||||
)
|
||||
|
||||
response = vertex_ai_batches_instance.cancel_batch(
|
||||
_is_async=_is_async,
|
||||
batch_id=batch_id,
|
||||
api_base=api_base,
|
||||
vertex_project=vertex_ai_project,
|
||||
vertex_location=vertex_ai_location,
|
||||
vertex_credentials=vertex_credentials,
|
||||
timeout=timeout,
|
||||
max_retries=optional_params.max_retries,
|
||||
)
|
||||
else:
|
||||
raise litellm.exceptions.BadRequestError(
|
||||
message="LiteLLM doesn't support {} for 'cancel_batch'. Only 'openai' and 'azure' are supported.".format(
|
||||
message="LiteLLM doesn't support {} for 'cancel_batch'. Only 'openai', 'azure', and 'vertex_ai' are supported.".format(
|
||||
custom_llm_provider
|
||||
),
|
||||
model="n/a",
|
||||
|
||||
@@ -376,3 +376,148 @@ class VertexAIBatchPrediction(VertexLLM):
|
||||
response=_json_response
|
||||
)
|
||||
return vertex_batch_response
|
||||
|
||||
def cancel_batch(
|
||||
self,
|
||||
_is_async: bool,
|
||||
batch_id: str,
|
||||
api_base: Optional[str],
|
||||
vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES],
|
||||
vertex_project: Optional[str],
|
||||
vertex_location: Optional[str],
|
||||
timeout: Union[float, httpx.Timeout],
|
||||
max_retries: Optional[int],
|
||||
) -> Union[LiteLLMBatch, Coroutine[Any, Any, LiteLLMBatch]]:
|
||||
access_token, project_id = self._ensure_access_token(
|
||||
credentials=vertex_credentials,
|
||||
project_id=vertex_project,
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
|
||||
default_api_base = self.create_vertex_batch_url(
|
||||
vertex_location=vertex_location or "us-central1",
|
||||
vertex_project=vertex_project or project_id,
|
||||
)
|
||||
|
||||
retrieve_api_base_default = f"{default_api_base}/{batch_id}"
|
||||
cancel_api_base_default = f"{retrieve_api_base_default}:cancel"
|
||||
|
||||
_, api_base = self._check_custom_proxy(
|
||||
api_base=api_base,
|
||||
custom_llm_provider="vertex_ai",
|
||||
gemini_api_key=None,
|
||||
endpoint="cancel",
|
||||
stream=None,
|
||||
auth_header=None,
|
||||
url=cancel_api_base_default,
|
||||
model=None,
|
||||
vertex_project=vertex_project or project_id,
|
||||
vertex_location=vertex_location or "us-central1",
|
||||
vertex_api_version="v1",
|
||||
)
|
||||
|
||||
if api_base.endswith(":cancel"):
|
||||
retrieve_api_base = api_base.removesuffix(":cancel")
|
||||
else:
|
||||
retrieve_api_base = api_base.rsplit(":cancel", 1)[0].rstrip("/")
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json; charset=utf-8",
|
||||
"Authorization": f"Bearer {access_token}",
|
||||
}
|
||||
|
||||
if _is_async is True:
|
||||
return self._async_cancel_batch(
|
||||
api_base=api_base,
|
||||
retrieve_api_base=retrieve_api_base,
|
||||
headers=headers,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
sync_handler = _get_httpx_client()
|
||||
try:
|
||||
response = sync_handler.post(
|
||||
url=api_base,
|
||||
headers=headers,
|
||||
data=json.dumps({}),
|
||||
timeout=timeout,
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
litellm.verbose_logger.error(
|
||||
"Vertex AI batch cancel failed: status=%s, body=%s",
|
||||
e.response.status_code,
|
||||
e.response.text[:1000],
|
||||
)
|
||||
raise
|
||||
|
||||
if response.status_code != 200:
|
||||
raise Exception(f"Error: {response.status_code} {response.text}")
|
||||
|
||||
# HTTPHandler.get() does not accept a timeout parameter
|
||||
retrieve_response = sync_handler.get(
|
||||
url=retrieve_api_base,
|
||||
headers=headers,
|
||||
)
|
||||
if retrieve_response.status_code != 200:
|
||||
litellm.verbose_logger.error(
|
||||
"Vertex AI batch retrieve-after-cancel failed: status=%s, body=%s",
|
||||
retrieve_response.status_code,
|
||||
retrieve_response.text[:1000],
|
||||
)
|
||||
raise Exception(
|
||||
f"Error: {retrieve_response.status_code} {retrieve_response.text}"
|
||||
)
|
||||
|
||||
_json_response = retrieve_response.json()
|
||||
vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response(
|
||||
response=_json_response
|
||||
)
|
||||
return vertex_batch_response
|
||||
|
||||
async def _async_cancel_batch(
|
||||
self,
|
||||
api_base: str,
|
||||
retrieve_api_base: str,
|
||||
headers: Dict[str, str],
|
||||
timeout: Union[float, httpx.Timeout] = 600.0,
|
||||
) -> LiteLLMBatch:
|
||||
client = get_async_httpx_client(
|
||||
llm_provider=litellm.LlmProviders.VERTEX_AI,
|
||||
)
|
||||
try:
|
||||
response = await client.post(
|
||||
url=api_base,
|
||||
headers=headers,
|
||||
data=json.dumps({}),
|
||||
timeout=timeout,
|
||||
)
|
||||
except httpx.HTTPStatusError as e:
|
||||
litellm.verbose_logger.error(
|
||||
"Vertex AI batch cancel failed: status=%s, body=%s",
|
||||
e.response.status_code,
|
||||
e.response.text[:1000],
|
||||
)
|
||||
raise
|
||||
if response.status_code != 200:
|
||||
raise Exception(f"Error: {response.status_code} {response.text}")
|
||||
|
||||
# AsyncHTTPHandler.get() does not accept a timeout parameter
|
||||
retrieve_response = await client.get(
|
||||
url=retrieve_api_base,
|
||||
headers=headers,
|
||||
)
|
||||
if retrieve_response.status_code != 200:
|
||||
litellm.verbose_logger.error(
|
||||
"Vertex AI batch retrieve-after-cancel failed: status=%s, body=%s",
|
||||
retrieve_response.status_code,
|
||||
retrieve_response.text[:1000],
|
||||
)
|
||||
raise Exception(
|
||||
f"Error: {retrieve_response.status_code} {retrieve_response.text}"
|
||||
)
|
||||
|
||||
_json_response = retrieve_response.json()
|
||||
vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response(
|
||||
response=_json_response
|
||||
)
|
||||
return vertex_batch_response
|
||||
|
||||
@@ -5,7 +5,7 @@
|
||||
|
||||
######################################################################
|
||||
import asyncio
|
||||
from typing import Dict, Optional, cast
|
||||
from typing import Any, Dict, Optional, cast
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
|
||||
|
||||
@@ -655,7 +655,7 @@ async def list_batches(
|
||||
managed_files_obj, "list_user_batches"
|
||||
):
|
||||
verbose_proxy_logger.debug("Using managed objects table for batch listing")
|
||||
response = await managed_files_obj.list_user_batches(
|
||||
response = await cast(Any, managed_files_obj).list_user_batches(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
limit=limit,
|
||||
after=after,
|
||||
@@ -686,8 +686,9 @@ async def list_batches(
|
||||
|
||||
# Encode batch IDs in the list response so clients can use
|
||||
# them for retrieve/cancel/file downloads through the proxy.
|
||||
if response and hasattr(response, "data") and response.data:
|
||||
for batch in response.data:
|
||||
response_data = getattr(response, "data", None)
|
||||
if response_data:
|
||||
for batch in response_data:
|
||||
encode_batch_response_ids(batch, model=model_param)
|
||||
|
||||
verbose_proxy_logger.debug(f"Listed batches using model: {model_param}")
|
||||
@@ -897,7 +898,11 @@ async def cancel_batch(
|
||||
# SCENARIO 3: Fallback to custom_llm_provider (uses env variables)
|
||||
else:
|
||||
custom_llm_provider = (
|
||||
provider or data.pop("custom_llm_provider", None) or "openai"
|
||||
provider
|
||||
or data.pop("custom_llm_provider", None)
|
||||
or get_custom_llm_provider_from_request_headers(request=request)
|
||||
or get_custom_llm_provider_from_request_query(request=request)
|
||||
or "openai"
|
||||
)
|
||||
# Extract batch_id from data to avoid "multiple values for keyword argument" error
|
||||
# data was cast from CancelBatchRequest which already contains batch_id
|
||||
|
||||
@@ -1,3 +1,9 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.vertex_ai.batches.handler import VertexAIBatchPrediction
|
||||
from litellm.llms.vertex_ai.batches.transformation import VertexAIBatchTransformation
|
||||
|
||||
|
||||
@@ -36,3 +42,124 @@ def test_output_file_id_falls_back_to_output_uri_prefix_with_predictions_jsonl()
|
||||
output_file_id
|
||||
== "gs://test-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-pro/prediction-model-456/predictions.jsonl"
|
||||
)
|
||||
|
||||
|
||||
def test_vertex_ai_cancel_batch():
|
||||
"""Test that vertex_ai cancel_batch calls the correct API endpoint"""
|
||||
handler = VertexAIBatchPrediction(gcs_bucket_name="test-bucket")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"name": "projects/test-project/locations/us-central1/batchPredictionJobs/123456",
|
||||
"state": "JOB_STATE_CANCELLING",
|
||||
"createTime": "2024-03-17T10:00:00.000000Z",
|
||||
"inputConfig": {
|
||||
"gcsSource": {
|
||||
"uris": ["gs://test-bucket/input.jsonl"]
|
||||
}
|
||||
},
|
||||
"outputConfig": {
|
||||
"gcsDestination": {
|
||||
"outputUriPrefix": "gs://test-bucket/output"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
with patch("litellm.llms.vertex_ai.batches.handler._get_httpx_client") as mock_client:
|
||||
mock_client.return_value.post.return_value = mock_response
|
||||
mock_client.return_value.get.return_value = mock_response
|
||||
|
||||
with patch.object(handler, "_ensure_access_token") as mock_auth:
|
||||
mock_auth.return_value = ("fake-token", "test-project")
|
||||
|
||||
response = handler.cancel_batch(
|
||||
_is_async=False,
|
||||
batch_id="123456",
|
||||
api_base=None,
|
||||
vertex_credentials=None,
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
timeout=600.0,
|
||||
max_retries=None,
|
||||
)
|
||||
|
||||
assert response.id == "123456"
|
||||
assert response.status == "cancelling"
|
||||
|
||||
mock_client.return_value.post.assert_called_once()
|
||||
mock_client.return_value.get.assert_called_once()
|
||||
call_args = mock_client.return_value.post.call_args
|
||||
assert ":cancel" in call_args.kwargs["url"]
|
||||
|
||||
|
||||
def test_vertex_ai_cancel_batch_forwards_timeout():
|
||||
"""Test that timeout is forwarded to the POST (cancel) HTTP call.
|
||||
|
||||
Note: the follow-up GET (retrieve) call does not accept a timeout
|
||||
parameter in the underlying HTTP handler, so it is intentionally omitted.
|
||||
"""
|
||||
|
||||
|
||||
def test_vertex_ai_cancel_batch_custom_proxy_retrieve_url():
|
||||
"""Retrieve URL should go through the custom proxy, not bypass it"""
|
||||
handler = VertexAIBatchPrediction(gcs_bucket_name="test-bucket")
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"name": "projects/test-project/locations/us-central1/batchPredictionJobs/123456",
|
||||
"state": "JOB_STATE_CANCELLING",
|
||||
"createTime": "2024-03-17T10:00:00.000000Z",
|
||||
"inputConfig": {"gcsSource": {"uris": ["gs://test-bucket/input.jsonl"]}},
|
||||
"outputConfig": {"gcsDestination": {"outputUriPrefix": "gs://test-bucket/output"}},
|
||||
}
|
||||
|
||||
with patch("litellm.llms.vertex_ai.batches.handler._get_httpx_client") as mock_client:
|
||||
mock_client.return_value.post.return_value = mock_response
|
||||
mock_client.return_value.get.return_value = mock_response
|
||||
|
||||
with patch.object(handler, "_ensure_access_token") as mock_auth:
|
||||
mock_auth.return_value = ("fake-token", "test-project")
|
||||
|
||||
handler.cancel_batch(
|
||||
_is_async=False,
|
||||
batch_id="123456",
|
||||
api_base="https://my-proxy.example.com",
|
||||
vertex_credentials=None,
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
timeout=600.0,
|
||||
max_retries=None,
|
||||
)
|
||||
|
||||
post_url = mock_client.return_value.post.call_args.kwargs["url"]
|
||||
get_url = mock_client.return_value.get.call_args.kwargs["url"]
|
||||
|
||||
assert "my-proxy.example.com" in post_url
|
||||
assert ":cancel" in post_url
|
||||
assert "my-proxy.example.com" in get_url
|
||||
assert ":cancel" not in get_url
|
||||
assert "googleapis.com" not in get_url
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_litellm_cancel_batch_vertex_ai():
|
||||
"""Test that litellm.cancel_batch works with vertex_ai provider"""
|
||||
mock_response = MagicMock()
|
||||
mock_response.id = "batch_123"
|
||||
mock_response.status = "cancelling"
|
||||
|
||||
with patch("litellm.batches.main.vertex_ai_batches_instance") as mock_instance:
|
||||
mock_instance.cancel_batch.return_value = mock_response
|
||||
|
||||
response = litellm.cancel_batch(
|
||||
batch_id="batch_123",
|
||||
custom_llm_provider="vertex_ai",
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
)
|
||||
|
||||
assert mock_instance.cancel_batch.called
|
||||
assert response.id == "batch_123"
|
||||
assert response.status == "cancelling"
|
||||
|
||||
Reference in New Issue
Block a user