diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 17d73aae6a..23f444b1ce 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -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", diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index f0b181c9a6..2cb0294206 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -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 diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 38e5229eee..160e9c23f0 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -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 diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py index 1aab74ddc2..7310c68b4e 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py @@ -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"