From c4d27cb239d89d1cf934d27b52285d5a6b490f95 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 18 Mar 2026 10:30:05 +0530 Subject: [PATCH] =?UTF-8?q?fix(vertex-ai):=20address=20greptile=20review?= =?UTF-8?q?=20=E2=80=93=20proxy=20retrieve=20URL,=20timeout=20forwarding,?= =?UTF-8?q?=20sync=20logging?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Fix retrieve_api_base derivation to handle custom proxies with path-based routing (not just :cancel suffix) - Forward timeout to POST calls in cancel_batch (sync + async) - Add try/except error logging to sync cancel path (parity with async) - Add tests for timeout forwarding and custom proxy retrieve URL Made-with: Cursor --- litellm/llms/vertex_ai/batches/handler.py | 49 ++++++++---- .../test_vertex_ai_batch_transformation.py | 78 +++++++++++++++++++ 2 files changed, 113 insertions(+), 14 deletions(-) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 9c760f628ed..416c5fc69fc 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -416,11 +416,12 @@ class VertexAIBatchPrediction(VertexLLM): vertex_api_version="v1", ) - retrieve_api_base = ( - api_base.removesuffix(":cancel") - if api_base.endswith(":cancel") - else retrieve_api_base_default - ) + if api_base.endswith(":cancel"): + retrieve_api_base = api_base.removesuffix(":cancel") + elif api_base == cancel_api_base_default: + retrieve_api_base = retrieve_api_base_default + else: + retrieve_api_base = api_base.rsplit(":cancel", 1)[0].rstrip("/") headers = { "Content-Type": "application/json; charset=utf-8", @@ -432,22 +433,40 @@ class VertexAIBatchPrediction(VertexLLM): api_base=api_base, retrieve_api_base=retrieve_api_base, headers=headers, + timeout=timeout, ) sync_handler = _get_httpx_client() - response = sync_handler.post( - url=api_base, - headers=headers, - data=json.dumps({}), - ) + 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}") - retrieve_response = sync_handler.get( - url=retrieve_api_base, - headers=headers, - ) + try: + retrieve_response = sync_handler.get( + url=retrieve_api_base, + headers=headers, + ) + except httpx.HTTPStatusError as e: + litellm.verbose_logger.error( + "Vertex AI batch retrieve-after-cancel failed: status=%s, body=%s", + e.response.status_code, + e.response.text[:1000], + ) + raise if retrieve_response.status_code != 200: raise Exception( f"Error: {retrieve_response.status_code} {retrieve_response.text}" @@ -464,6 +483,7 @@ class VertexAIBatchPrediction(VertexLLM): 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, @@ -473,6 +493,7 @@ class VertexAIBatchPrediction(VertexLLM): url=api_base, headers=headers, data=json.dumps({}), + timeout=timeout, ) except httpx.HTTPStatusError as e: litellm.verbose_logger.error( 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 1fa8dd2e16c..b934afd1036 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 @@ -93,6 +93,84 @@ def test_vertex_ai_cancel_batch(): assert ":cancel" in call_args.kwargs["url"] +def test_vertex_ai_cancel_batch_forwards_timeout(): + """Test that timeout is forwarded to both POST and GET HTTP calls""" + 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=None, + vertex_credentials=None, + vertex_project="test-project", + vertex_location="us-central1", + timeout=42.0, + max_retries=None, + ) + + post_kwargs = mock_client.return_value.post.call_args.kwargs + assert post_kwargs["timeout"] == 42.0 + + +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"""