From 39a1de5c0b3ebead950a3d3409b1b25b694c9f22 Mon Sep 17 00:00:00 2001 From: Kris Xia Date: Mon, 11 May 2026 10:47:54 +0800 Subject: [PATCH] fix(vertex_ai/gemini): offload sync GCS metadata lookup off async event loop Address PR #27278 review feedback: - async_transform_request_body now runs the sync _transform_request_body (which may issue an httpx.get with up to 5s timeout for GCS metadata) in a worker thread via asyncify, so extensionless gs:// URIs no longer block the event loop on the async path. - Drop redundant `pass` in _get_gcs_object_content_type's except block. - Mock _get_gcs_object_content_type in test_process_gemini_media_gcs_without_extension_raises_clear_error so the test no longer makes real network calls. - Add a regression test that verifies async_transform_request_body does not block the event loop while the sync GCS metadata fetch is slow. --- .../llms/vertex_ai/gemini/transformation.py | 9 +- .../llms/vertex_ai/test_vertex.py | 91 ++++++++++++++++++- 2 files changed, 95 insertions(+), 5 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 94bfa851775..b9d2f818d2e 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -15,6 +15,7 @@ from pydantic import BaseModel import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.prompt_templates.common_utils import ( _get_image_mime_type_from_url, ) @@ -246,8 +247,7 @@ def _get_gcs_object_content_type( model=None, llm_provider="vertex_ai", ) - # Metadata may still be readable for public objects. - pass + # 未显式提供 Vertex 凭据时，metadata 对公开对象仍可能可读，继续无 token 尝试。 object_path = quote(object_name, safe="") metadata_url = ( @@ -1104,7 +1104,10 @@ async def async_transform_request_body( vertex_auth_header=vertex_auth_header, ) - return _transform_request_body( + # _transform_request_body 可能通过 _get_gcs_object_content_type 发起同步 httpx.get + # （最长 5s 超时）去拉 GCS 对象 metadata。为避免阻塞 async 事件循环，整个同步 + # 转换放到 worker 线程执行。 + return await asyncify(_transform_request_body)( messages=messages, model=model, custom_llm_provider=custom_llm_provider, diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex.py b/tests/test_litellm/llms/vertex_ai/test_vertex.py index 613d0b72e24..8cda11c1f45 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex.py @@ -1283,8 +1283,13 @@ def test_process_gemini_media(): def test_process_gemini_media_gcs_without_extension_raises_clear_error(): - with pytest.raises(litellm.BadRequestError) as exc_info: - _process_gemini_media("gs://bucket/image-without-extension") + # mock 掉 GCS metadata 查询，避免测试发出真实外网请求 + with patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type", + return_value=None, + ): + with pytest.raises(litellm.BadRequestError) as exc_info: + _process_gemini_media("gs://bucket/image-without-extension") assert "Unable to determine mime type for gs URI" in str(exc_info.value) @@ -1444,6 +1449,88 @@ def test_get_gcs_object_content_type_fails_fast_with_explicit_credentials(): ) +def test_async_transform_request_body_does_not_block_event_loop(): + """当同步 GCS metadata 查询阻塞时，async_transform_request_body 不应阻塞事件循环。""" + import asyncio + import time + + from litellm.llms.vertex_ai.gemini import transformation as gemini_transformation + + messages = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "gs://bucket/image-without-extension"}, + } + ], + } + ] + + # 模拟同步阻塞的 httpx.get（就像真实的 GCS metadata 查询超时）， + # 如果 async_transform_request_body 没有把同步部分 offload 到 worker 线程， + # 这次阻塞会卡住事件循环，并行的 sleep 就无法推进。 + def slow_http_get(*args, **kwargs): + time.sleep(0.5) + response = MagicMock() + response.raise_for_status.return_value = None + response.json.return_value = {"contentType": "image/png"} + return response + + # 模拟 context-caching 查询，避免其发出真实 HTTP + async def fake_check_and_create_cache(self, **kwargs): + return kwargs["messages"], kwargs["optional_params"], None + + mock_vertex_base = MagicMock() + mock_vertex_base.get_access_token.return_value = ("token", "project") + + async def run_scenario() -> float: + async def concurrent_sleep() -> float: + start = time.monotonic() + await asyncio.sleep(0.05) + return time.monotonic() - start + + transform_task = asyncio.create_task( + gemini_transformation.async_transform_request_body( + gemini_api_key=None, + messages=messages, + api_base=None, + model="gemini-2.5-flash", + client=None, + timeout=None, + extra_headers=None, + optional_params={}, + logging_obj=MagicMock(), + custom_llm_provider="vertex_ai", + litellm_params={}, + vertex_project=None, + vertex_location=None, + vertex_auth_header=None, + ) + ) + sleep_elapsed = await concurrent_sleep() + await transform_task + return sleep_elapsed + + with patch.object( + gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_vertex_base + ), patch.object( + gemini_transformation.httpx, "get", side_effect=slow_http_get + ), patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching." + "ContextCachingEndpoints.async_check_and_create_cache", + new=fake_check_and_create_cache, + ): + sleep_elapsed = asyncio.run(run_scenario()) + + # 没被阻塞的情况下，0.05s 的 sleep 不应被拖到接近同步阻塞时长（0.5s） + assert sleep_elapsed < 0.4, ( + f"事件循环被阻塞 {sleep_elapsed:.3f}s，async_transform_request_body 未把同步 GCS " + "metadata 查询 offload 到 worker 线程" + ) + + def test_get_image_mime_type_from_url(): """Test the _get_image_mime_type_from_url function for different image URLs""" from litellm.llms.vertex_ai.gemini.transformation import (