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.
This commit is contained in:
Kris Xia 2026-05-11 10:47:54 +08:00 • committed by S0ngRu1
parent 29020e9316
commit 39a1de5c0b
2 changed files with 95 additions and 5 deletions

View file

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

View file

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