Skip to content

Commit efa3980

Browse files
jsondaicopybara-github
jsondai
authored and
copybara-github
committed
fix: GenAI Client(evals) - Mirror the inference retry fix to agentplatform/private and vertexai
PiperOrigin-RevId: 995455628
1 parent e631a16 commit efa3980

3 files changed

Lines changed: 71 additions & 79 deletions

File tree

‎vertexai/_genai/_evals_common.py‎

Lines changed: 55 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -22,13 +22,15 @@
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
2829
import uuid
2930

3031
from google.api_core import exceptions as api_exceptions
3132
import vertexai
33+
from google.genai import errors as genai_errors
3234
from google.genai import types as genai_types
3335
from google.genai._api_client import BaseApiClient
3436
from google.genai.models import Models
@@ -129,6 +131,45 @@ def _get_agent_engine_instance(
129131
return _thread_local_data.agent_engine_instances[agent_name]
130132

131133

134+
def _is_retryable_error(e: Exception) -> bool:
135+
"""Returns True for transient errors that are safe to retry with backoff."""
136+
if isinstance(e, genai_errors.APIError):
137+
return e.code in _evals_constant.RETRYABLE_STATUS_CODES
138+
return isinstance(
139+
e, (api_exceptions.ResourceExhausted, api_exceptions.ServiceUnavailable)
140+
)
141+
142+
143+
def _get_retry_backoff(
144+
e: Exception, attempt: int, max_retries: int, operation: str
145+
) -> Optional[float]:
146+
"""Returns the seconds to wait before the next attempt, or None to stop.
147+
148+
Retryable errors back off exponentially with jitter. Other errors fail
149+
fast, and the last attempt never waits.
150+
"""
151+
if not _is_retryable_error(e) or attempt == max_retries - 1:
152+
logger.error(
153+
"Error during %s on attempt %d/%d: %s",
154+
operation,
155+
attempt + 1,
156+
max_retries,
157+
e,
158+
)
159+
return None
160+
backoff = 2.0**attempt + random.uniform(0, 1)
161+
logger.warning(
162+
"Retryable error during %s on attempt %d/%d: %s. Retrying in %.1f"
163+
" seconds...",
164+
operation,
165+
attempt + 1,
166+
max_retries,
167+
e,
168+
backoff,
169+
)
170+
return backoff
171+
172+
132173
def _generate_content_with_retry(
133174
api_client: BaseApiClient,
134175
model: str,
@@ -191,29 +232,11 @@ def _generate_content_with_retry(
191232
}
192233
else:
193234
return response
194-
except api_exceptions.ResourceExhausted as e:
195-
logger.warning(
196-
"Resource Exhausted error on attempt %d/%d: %s. Retrying in %s"
197-
" seconds...",
198-
attempt + 1,
199-
max_retries,
200-
e,
201-
2**attempt,
202-
)
203-
if attempt == max_retries - 1:
204-
return {"error": f"Resource exhausted after retries: {e}"}
205-
time.sleep(2**attempt)
206235
except Exception as e: # pylint: disable=broad-exception-caught
207-
logger.error(
208-
"Unexpected error during generate_content on attempt %d/%d: %s",
209-
attempt + 1,
210-
max_retries,
211-
e,
212-
)
213-
214-
if attempt == max_retries - 1:
215-
return {"error": f"Failed after retries: {e}"}
216-
time.sleep(1)
236+
backoff = _get_retry_backoff(e, attempt, max_retries, "generate_content")
237+
if backoff is None:
238+
return {"error": f"Failed on attempt {attempt + 1}/{max_retries}: {e}"}
239+
time.sleep(backoff)
217240
return {"error": f"Failed to generate content after {max_retries} retries"}
218241

219242

@@ -2352,28 +2375,11 @@ def _execute_agent_run_with_retry(
23522375
if event and CONTENT in event and PARTS in event[CONTENT]:
23532376
responses.append(event)
23542377
return responses
2355-
except api_exceptions.ResourceExhausted as e:
2356-
logger.warning(
2357-
"Resource Exhausted error on attempt %d/%d: %s. Retrying in %s"
2358-
" seconds...",
2359-
attempt + 1,
2360-
max_retries,
2361-
e,
2362-
2**attempt,
2363-
)
2364-
if attempt == max_retries - 1:
2365-
return {"error": f"Resource exhausted after retries: {e}"}
2366-
time.sleep(2**attempt)
23672378
except Exception as e: # pylint: disable=broad-exception-caught
2368-
logger.error(
2369-
"Unexpected error during agent engine run on attempt %d/%d: %s",
2370-
attempt + 1,
2371-
max_retries,
2372-
e,
2373-
)
2374-
if attempt == max_retries - 1:
2375-
return {"error": f"Failed after retries: {e}"}
2376-
time.sleep(1)
2379+
backoff = _get_retry_backoff(e, attempt, max_retries, "agent engine run")
2380+
if backoff is None:
2381+
return {"error": f"Failed on attempt {attempt + 1}/{max_retries}: {e}"}
2382+
time.sleep(backoff)
23772383
return {"error": f"Failed to get agent run results after {max_retries} retries"}
23782384

23792385

@@ -2475,28 +2481,13 @@ async def _execute_local_agent_run_with_retry_async(
24752481
if event and CONTENT in event and PARTS in event[CONTENT]:
24762482
events.append(event)
24772483
return events
2478-
except api_exceptions.ResourceExhausted as e:
2479-
logger.warning(
2480-
"Resource Exhausted error on attempt %d/%d: %s. Retrying"
2481-
" in %s seconds...",
2482-
attempt + 1,
2483-
max_retries,
2484-
e,
2485-
2**attempt,
2486-
)
2487-
if attempt == max_retries - 1:
2488-
return {"error": f"Resource exhausted after retries: {e}"}
2489-
await asyncio.sleep(2**attempt)
24902484
except Exception as e: # pylint: disable=broad-exception-caught
2491-
logger.error(
2492-
"Unexpected error during agent run on attempt %d/%d: %s",
2493-
attempt + 1,
2494-
max_retries,
2495-
e,
2496-
)
2497-
if attempt == max_retries - 1:
2498-
return {"error": f"Failed after retries: {e}"}
2499-
await asyncio.sleep(1)
2485+
backoff = _get_retry_backoff(e, attempt, max_retries, "agent run")
2486+
if backoff is None:
2487+
return {
2488+
"error": f"Failed on attempt {attempt + 1}/{max_retries}: {e}"
2489+
}
2490+
await asyncio.sleep(backoff)
25002491
return {"error": f"Failed to get agent run results after {max_retries} retries"}
25012492

25022493

‎vertexai/_genai/_evals_constant.py‎

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

‎vertexai/_genai/_evals_metric_handlers.py‎

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

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

5542
R = TypeVar("R")
5643
T = TypeVar("T", types.Metric, types.MetricSource, types.LLMMetric)
@@ -63,7 +50,7 @@ def _call_with_retry(
6350
"""Calls ``fn()`` with exponential backoff + jitter on retryable errors.
6451
6552
Retries up to ``_MAX_RETRIES`` times on errors whose HTTP status code is
66-
in ``_RETRYABLE_STATUS_CODES`` (Aborted, DeadlineExceeded,
53+
in ``_evals_constant.RETRYABLE_STATUS_CODES`` (Aborted, DeadlineExceeded,
6754
ResourceExhausted, ServiceUnavailable, Cancelled). Non-retryable errors
6855
are re-raised immediately. If all retries are exhausted the last
6956
exception is re-raised so the caller can decide how to handle it.
@@ -83,7 +70,7 @@ def _call_with_retry(
8370
try:
8471
return fn()
8572
except genai_errors.APIError as e:
86-
if e.code in _RETRYABLE_STATUS_CODES:
73+
if e.code in _evals_constant.RETRYABLE_STATUS_CODES:
8774
backoff = 2**attempt + random.uniform(0, 1)
8875
logger.warning(
8976
"Retryable error (code=%s) on attempt %d/%d for metric"

0 commit comments

Comments
 (0)