Skip to content

Commit c7d50c8

Browse files
speedstorm1copybara-github
authored andcommitted
fix: Await asyncio.sleep in async HTTPX upload retry backoff
fixes #2938 PiperOrigin-RevId: 998864079
1 parent e7f9ad3 commit c7d50c8

2 files changed

Lines changed: 44 additions & 1 deletion

File tree

‎google/genai/_api_client.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2853,7 +2853,7 @@ async def _async_upload_fd(
28532853
break
28542854
delay_seconds = INITIAL_RETRY_DELAY * (DELAY_MULTIPLIER**retry_count)
28552855
retry_count += 1
2856-
time.sleep(delay_seconds)
2856+
await asyncio.sleep(delay_seconds)
28572857

28582858
offset += chunk_size
28592859
if (

‎google/genai/tests/client/test_upload_errors.py‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515

1616
"""Tests for error handling in the client api client layer for uploads."""
1717

18+
import asyncio
19+
import contextlib
1820
import io
1921
import json
2022
from unittest import mock
@@ -132,6 +134,47 @@ async def test_async_upload_url_rewrite_httpx(client: api_client.BaseApiClient):
132134
)
133135

134136

137+
@pytest.mark.asyncio
138+
async def test_async_upload_retry_does_not_block_event_loop_httpx(
139+
client: api_client.BaseApiClient,
140+
):
141+
"""The HTTPX retry backoff yields so other tasks keep running."""
142+
mock_async_httpx_client = mock.MagicMock(spec=httpx.AsyncClient)
143+
mock_async_httpx_client.request = mock.AsyncMock(
144+
side_effect=[
145+
_httpx_response(200), # No X-Goog-Upload-Status, so it retries.
146+
_httpx_response(200, headers={"X-Goog-Upload-Status": "final"}),
147+
]
148+
)
149+
client._async_httpx_client = mock_async_httpx_client
150+
ticks = 0
151+
152+
async def ticker():
153+
nonlocal ticks
154+
while True:
155+
await asyncio.sleep(0.01)
156+
ticks += 1
157+
158+
ticker_task = asyncio.create_task(ticker())
159+
try:
160+
with mock.patch.object(
161+
client, "_use_aiohttp", return_value=False
162+
), mock.patch.object(api_client, "INITIAL_RETRY_DELAY", 0.2), io.BytesIO(
163+
b"test"
164+
) as f:
165+
await client._async_upload_fd(
166+
f, "http://fake/upload", 4, http_options=types.HttpOptions()
167+
)
168+
ticks_during_upload = ticks
169+
finally:
170+
ticker_task.cancel()
171+
with contextlib.suppress(asyncio.CancelledError):
172+
await ticker_task
173+
174+
assert mock_async_httpx_client.request.call_count == 2
175+
assert ticks_during_upload > 0
176+
177+
135178
@pytest.mark.asyncio
136179
async def test_async_upload_fd_error_httpx(client: api_client.BaseApiClient):
137180
error_content = json.dumps({

0 commit comments

Comments
 (0)