Skip to content

Commit f6bf219

Browse files
jsondaicopybara-github
jsondai
authored and
copybara-github
committed
fix: GenAI Client(evals) - Back off only on retryable errors in inference and agent-run retries
PiperOrigin-RevId: 995397403
1 parent 2c2d69e commit f6bf219

4 files changed

Lines changed: 184 additions & 97 deletions

File tree

‎agentplatform/_genai/_evals_common.py‎

Lines changed: 54 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import json
2323
import logging
2424
import os
25+
import random
2526
import threading
2627
import time
2728
from typing import Any, Callable, Literal, Optional, Union, cast
@@ -160,6 +161,45 @@ def _get_runtime_instance(
160161
return _thread_local_data.runtime_instances[agent_name]
161162

162163

164+
def _is_retryable_error(e: Exception) -> bool:
165+
"""Returns True for transient errors that are safe to retry with backoff."""
166+
if isinstance(e, genai_errors.APIError):
167+
return e.code in _evals_constant.RETRYABLE_STATUS_CODES
168+
return isinstance(
169+
e, (api_exceptions.ResourceExhausted, api_exceptions.ServiceUnavailable)
170+
)
171+
172+
173+
def _get_retry_backoff(
174+
e: Exception, attempt: int, max_retries: int, operation: str
175+
) -> Optional[float]:
176+
"""Returns the seconds to wait before the next attempt, or None to stop.
177+
178+
Retryable errors back off exponentially with jitter. Other errors fail
179+
fast, and the last attempt never waits.
180+
"""
181+
if not _is_retryable_error(e) or attempt == max_retries - 1:
182+
logger.error(
183+
"Error during %s on attempt %d/%d: %s",
184+
operation,
185+
attempt + 1,
186+
max_retries,
187+
e,
188+
)
189+
return None
190+
backoff = 2**attempt + random.uniform(0, 1)
191+
logger.warning(
192+
"Retryable error during %s on attempt %d/%d: %s. Retrying in %.1f"
193+
" seconds...",
194+
operation,
195+
attempt + 1,
196+
max_retries,
197+
e,
198+
backoff,
199+
)
200+
return backoff
201+
202+
163203
def _generate_content_with_retry(
164204
api_client: BaseApiClient,
165205
model: str,
@@ -222,29 +262,11 @@ def _generate_content_with_retry(
222262
}
223263
else:
224264
return response
225-
except api_exceptions.ResourceExhausted as e:
226-
logger.warning(
227-
"Resource Exhausted error on attempt %d/%d: %s. Retrying in %s"
228-
" seconds...",
229-
attempt + 1,
230-
max_retries,
231-
e,
232-
2**attempt,
233-
)
234-
if attempt == max_retries - 1:
235-
return {"error": f"Resource exhausted after retries: {e}"}
236-
time.sleep(2**attempt)
237265
except Exception as e: # pylint: disable=broad-exception-caught
238-
logger.error(
239-
"Unexpected error during generate_content on attempt %d/%d: %s",
240-
attempt + 1,
241-
max_retries,
242-
e,
243-
)
244-
245-
if attempt == max_retries - 1:
246-
return {"error": f"Failed after retries: {e}"}
247-
time.sleep(1)
266+
backoff = _get_retry_backoff(e, attempt, max_retries, "generate_content")
267+
if backoff is None:
268+
return {"error": f"Failed on attempt {attempt + 1}/{max_retries}: {e}"}
269+
time.sleep(backoff)
248270
return {"error": f"Failed to generate content after {max_retries} retries"}
249271

250272

@@ -3570,28 +3592,11 @@ def _execute_agent_run_with_retry(
35703592
if event and CONTENT in event and PARTS in event[CONTENT]:
35713593
responses.append(event)
35723594
return responses
3573-
except api_exceptions.ResourceExhausted as e:
3574-
logger.warning(
3575-
"Resource Exhausted error on attempt %d/%d: %s. Retrying in %s"
3576-
" seconds...",
3577-
attempt + 1,
3578-
max_retries,
3579-
e,
3580-
2**attempt,
3581-
)
3582-
if attempt == max_retries - 1:
3583-
return {"error": f"Resource exhausted after retries: {e}"}
3584-
time.sleep(2**attempt)
35853595
except Exception as e: # pylint: disable=broad-exception-caught
3586-
logger.error(
3587-
"Unexpected error during agent engine run on attempt %d/%d: %s",
3588-
attempt + 1,
3589-
max_retries,
3590-
e,
3591-
)
3592-
if attempt == max_retries - 1:
3593-
return {"error": f"Failed after retries: {e}"}
3594-
time.sleep(1)
3596+
backoff = _get_retry_backoff(e, attempt, max_retries, "agent engine run")
3597+
if backoff is None:
3598+
return {"error": f"Failed on attempt {attempt + 1}/{max_retries}: {e}"}
3599+
time.sleep(backoff)
35953600
return {"error": f"Failed to get agent run results after {max_retries} retries"}
35963601

35973602

@@ -3697,28 +3702,13 @@ async def _execute_local_agent_run_with_retry_async(
36973702
if event and CONTENT in event and PARTS in event[CONTENT]:
36983703
events.append(event)
36993704
return events
3700-
except api_exceptions.ResourceExhausted as e:
3701-
logger.warning(
3702-
"Resource Exhausted error on attempt %d/%d: %s. Retrying"
3703-
" in %s seconds...",
3704-
attempt + 1,
3705-
max_retries,
3706-
e,
3707-
2**attempt,
3708-
)
3709-
if attempt == max_retries - 1:
3710-
return {"error": f"Resource exhausted after retries: {e}"}
3711-
await asyncio.sleep(2**attempt)
37123705
except Exception as e: # pylint: disable=broad-exception-caught
3713-
logger.error(
3714-
"Unexpected error during agent run on attempt %d/%d: %s",
3715-
attempt + 1,
3716-
max_retries,
3717-
e,
3718-
)
3719-
if attempt == max_retries - 1:
3720-
return {"error": f"Failed after retries: {e}"}
3721-
await asyncio.sleep(1)
3706+
backoff = _get_retry_backoff(e, attempt, max_retries, "agent run")
3707+
if backoff is None:
3708+
return {
3709+
"error": f"Failed on attempt {attempt + 1}/{max_retries}: {e}"
3710+
}
3711+
await asyncio.sleep(backoff)
37223712
return {"error": f"Failed to get agent run results after {max_retries} retries"}
37233713

37243714

‎agentplatform/_genai/_evals_constant.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,3 +120,17 @@
120120
INTERACTION_ID,
121121
}
122122
)
123+
124+
# HTTP status codes that are safe to retry with backoff.
125+
RETRYABLE_STATUS_CODES = frozenset(
126+
{
127+
408, # RequestTimeout (DEADLINE_EXCEEDED)
128+
409, # Conflict / Aborted (ABORTED)
129+
429, # TooManyRequests / ResourceExhausted (RESOURCE_EXHAUSTED)
130+
499, # Client Closed Request (CANCELLED)
131+
500, # InternalServerError (INTERNAL)
132+
502, # BadGateway
133+
503, # ServiceUnavailable (UNAVAILABLE)
134+
504, # GatewayTimeout (DEADLINE_EXCEEDED)
135+
}
136+
)

‎agentplatform/_genai/_evals_metric_handlers.py‎

Lines changed: 2 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -39,19 +39,6 @@
3939

4040
logger = logging.getLogger(__name__)
4141
_MAX_RETRIES = 5
42-
# HTTP status codes that are safe to retry with backoff.
43-
_RETRYABLE_STATUS_CODES = frozenset(
44-
{
45-
408, # RequestTimeout (DEADLINE_EXCEEDED)
46-
409, # Conflict / Aborted (ABORTED)
47-
429, # TooManyRequests / ResourceExhausted (RESOURCE_EXHAUSTED)
48-
499, # Client Closed Request (CANCELLED)
49-
500, # InternalServerError (INTERNAL)
50-
502, # BadGateway
51-
503, # ServiceUnavailable (UNAVAILABLE)
52-
504, # GatewayTimeout (DEADLINE_EXCEEDED)
53-
}
54-
)
5542

5643
R = TypeVar("R")
5744
T = TypeVar("T", types.Metric, types.MetricSource, types.LLMMetric)
@@ -64,7 +51,7 @@ def _call_with_retry(
6451
"""Calls ``fn()`` with exponential backoff + jitter on retryable errors.
6552
6653
Retries up to ``_MAX_RETRIES`` times on errors whose HTTP status code is
67-
in ``_RETRYABLE_STATUS_CODES`` (Aborted, DeadlineExceeded,
54+
in ``_evals_constant.RETRYABLE_STATUS_CODES`` (Aborted, DeadlineExceeded,
6855
ResourceExhausted, ServiceUnavailable, Cancelled). Non-retryable errors
6956
are re-raised immediately. If all retries are exhausted the last
7057
exception is re-raised so the caller can decide how to handle it.
@@ -84,7 +71,7 @@ def _call_with_retry(
8471
try:
8572
return fn()
8673
except genai_errors.APIError as e:
87-
if e.code in _RETRYABLE_STATUS_CODES:
74+
if e.code in _evals_constant.RETRYABLE_STATUS_CODES:
8875
backoff = 2**attempt + random.uniform(0, 1)
8976
logger.warning(
9077
"Retryable error (code=%s) on attempt %d/%d for metric"

‎tests/unit/agentplatform/genai/test_evals.py‎

Lines changed: 114 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
import tempfile
2727
from unittest import mock
2828

29+
from google.api_core import exceptions as api_exceptions
2930
import google.auth.credentials
3031
from google.cloud import aiplatform
3132
import agentplatform
@@ -3390,9 +3391,7 @@ def test_run_inference_with_runtime_and_session_inputs_dict(
33903391
]
33913392

33923393
mock_runtime.stream_query.return_value = iter(stream_query_return_value)
3393-
mock_agentplatform_client.return_value.runtimes.get.return_value = (
3394-
mock_runtime
3395-
)
3394+
mock_agentplatform_client.return_value.runtimes.get.return_value = mock_runtime
33963395

33973396
inference_result = self.client.evals.run_inference(
33983397
agent="projects/test-project/locations/us-central1/reasoningEngines/123",
@@ -3499,9 +3498,7 @@ def test_run_inference_with_runtime_and_session_inputs_literal_string(
34993498
]
35003499

35013500
mock_runtime.stream_query.return_value = iter(stream_query_return_value)
3502-
mock_agentplatform_client.return_value.runtimes.get.return_value = (
3503-
mock_runtime
3504-
)
3501+
mock_agentplatform_client.return_value.runtimes.get.return_value = mock_runtime
35053502

35063503
inference_result = self.client.evals.run_inference(
35073504
agent="projects/test-project/locations/us-central1/reasoningEngines/123",
@@ -3591,9 +3588,7 @@ def test_run_inference_with_runtime_with_response_column_raises_error(
35913588
)
35923589

35933590
mock_runtime = mock.Mock()
3594-
mock_agentplatform_client.return_value.runtimes.get.return_value = (
3595-
mock_runtime
3596-
)
3591+
mock_agentplatform_client.return_value.runtimes.get.return_value = mock_runtime
35973592

35983593
with pytest.raises(ValueError) as excinfo:
35993594
self.client.evals.run_inference(
@@ -3644,9 +3639,7 @@ def test_run_inference_with_runtime_falls_back_to_managed_sessions_api(
36443639
"projects/test-project/locations/us-central1"
36453640
"/reasoningEngines/123/sessions/managed-session-1"
36463641
)
3647-
mock_runtime.api_client.sessions.create.return_value = (
3648-
mock_session_operation
3649-
)
3642+
mock_runtime.api_client.sessions.create.return_value = mock_session_operation
36503643

36513644
stream_query_return_value = [
36523645
{
@@ -3663,9 +3656,7 @@ def test_run_inference_with_runtime_falls_back_to_managed_sessions_api(
36633656
},
36643657
]
36653658
mock_runtime.stream_query.return_value = iter(stream_query_return_value)
3666-
mock_agentplatform_client.return_value.runtimes.get.return_value = (
3667-
mock_runtime
3668-
)
3659+
mock_agentplatform_client.return_value.runtimes.get.return_value = mock_runtime
36693660

36703661
inference_result = self.client.evals.run_inference(
36713662
agent="projects/test-project/locations/us-central1/reasoningEngines/123",
@@ -4392,9 +4383,7 @@ def test_run_inference_non_gemini_agent_routes_to_runtime(
43924383
}
43934384
]
43944385
)
4395-
mock_agentplatform_client.return_value.runtimes.get.return_value = (
4396-
mock_runtime
4397-
)
4386+
mock_agentplatform_client.return_value.runtimes.get.return_value = mock_runtime
43984387

43994388
self.client.evals.run_inference(
44004389
src=mock_df,
@@ -11141,6 +11130,113 @@ def test_call_with_retry_no_retry_on_non_retryable(self, mock_sleep):
1114111130
assert mock_sleep.call_count == 0
1114211131

1114311132

11133+
class TestInferenceRetry:
11134+
@pytest.mark.parametrize(
11135+
"error, expected",
11136+
[
11137+
(genai_errors.ClientError(code=429, response_json={}), True),
11138+
(genai_errors.ServerError(code=503, response_json={}), True),
11139+
(genai_errors.ClientError(code=400, response_json={}), False),
11140+
(genai_errors.ClientError(code=403, response_json={}), False),
11141+
(genai_errors.ClientError(code=404, response_json={}), False),
11142+
(api_exceptions.ResourceExhausted("quota"), True),
11143+
(api_exceptions.ServiceUnavailable("unavailable"), True),
11144+
(api_exceptions.InvalidArgument("bad request"), False),
11145+
(ValueError("bad value"), False),
11146+
],
11147+
)
11148+
def test_is_retryable_error(self, error, expected):
11149+
assert _evals_common._is_retryable_error(error) is expected
11150+
11151+
@mock.patch("time.sleep", return_value=None)
11152+
@mock.patch.object(_evals_common, "Models")
11153+
def test_generate_content_retries_429_with_backoff(self, mock_models, mock_sleep):
11154+
response = genai_types.GenerateContentResponse(
11155+
candidates=[
11156+
genai_types.Candidate(
11157+
content=genai_types.Content(parts=[genai_types.Part(text="ok")]),
11158+
finish_reason=genai_types.FinishReason.STOP,
11159+
)
11160+
]
11161+
)
11162+
error = genai_errors.ClientError(code=429, response_json={})
11163+
mock_models.return_value.generate_content.side_effect = [error, error, response]
11164+
11165+
result = _evals_common._generate_content_with_retry(
11166+
api_client=mock.Mock(), model="gemini-pro", contents="prompt"
11167+
)
11168+
11169+
assert result is response
11170+
assert mock_models.return_value.generate_content.call_count == 3
11171+
backoffs = [call.args[0] for call in mock_sleep.call_args_list]
11172+
assert len(backoffs) == 2
11173+
assert 1 <= backoffs[0] <= 2
11174+
assert 2 <= backoffs[1] <= 3
11175+
11176+
@mock.patch("time.sleep", return_value=None)
11177+
@mock.patch.object(_evals_common, "Models")
11178+
def test_generate_content_does_not_retry_400(self, mock_models, mock_sleep):
11179+
mock_models.return_value.generate_content.side_effect = (
11180+
genai_errors.ClientError(code=400, response_json={})
11181+
)
11182+
11183+
result = _evals_common._generate_content_with_retry(
11184+
api_client=mock.Mock(), model="gemini-pro", contents="prompt"
11185+
)
11186+
11187+
assert result["error"].startswith("Failed on attempt 1/3: 400")
11188+
assert mock_models.return_value.generate_content.call_count == 1
11189+
mock_sleep.assert_not_called()
11190+
11191+
@mock.patch("time.sleep", return_value=None)
11192+
def test_agent_engine_run_retries_429_and_fails_fast_on_400(self, mock_sleep):
11193+
runtime = mock.Mock()
11194+
runtime.stream_query.side_effect = [
11195+
genai_errors.ClientError(code=429, response_json={}),
11196+
genai_errors.ClientError(code=400, response_json={}),
11197+
]
11198+
11199+
with mock.patch.object(
11200+
_evals_common, "_create_runtime_session", return_value="session-id"
11201+
):
11202+
result = _evals_common._execute_agent_run_with_retry(
11203+
row=pd.Series({"prompt": "prompt"}), contents="prompt", runtime=runtime
11204+
)
11205+
11206+
assert result["error"].startswith("Failed on attempt 2/3: 400")
11207+
assert runtime.stream_query.call_count == 2
11208+
mock_sleep.assert_called_once()
11209+
assert 1 <= mock_sleep.call_args.args[0] <= 2
11210+
11211+
@mock.patch("asyncio.sleep", new_callable=mock.AsyncMock)
11212+
def test_local_agent_run_retries_429_and_fails_fast_on_400(self, mock_sleep):
11213+
mock_runners = mock.MagicMock()
11214+
mock_runners.Runner.return_value.run_async.side_effect = [
11215+
genai_errors.ClientError(code=429, response_json={}),
11216+
genai_errors.ClientError(code=400, response_json={}),
11217+
]
11218+
mock_sessions = mock.MagicMock()
11219+
mock_sessions.InMemorySessionService.return_value.create_session = (
11220+
mock.AsyncMock()
11221+
)
11222+
11223+
with mock.patch.dict(
11224+
sys.modules,
11225+
{"google.adk.runners": mock_runners, "google.adk.sessions": mock_sessions},
11226+
):
11227+
result = _evals_common._execute_local_agent_run_with_retry(
11228+
row=pd.Series({"prompt": "prompt"}),
11229+
contents="prompt",
11230+
agent=mock.Mock(),
11231+
api_client=mock.Mock(),
11232+
)
11233+
11234+
assert result["error"].startswith("Failed on attempt 2/3: 400")
11235+
assert mock_runners.Runner.return_value.run_async.call_count == 2
11236+
mock_sleep.assert_awaited_once()
11237+
assert 1 <= mock_sleep.await_args.args[0] <= 2
11238+
11239+
1114411240
class TestComputationMetricRetry:
1114511241
"""Tests for retry behavior in ComputationMetricHandler."""
1114611242

0 commit comments

Comments
 (0)