diff --git a/vertexai/_genai/_evals_common.py b/vertexai/_genai/_evals_common.py index df2148dea0..ca45012b46 100644 --- a/vertexai/_genai/_evals_common.py +++ b/vertexai/_genai/_evals_common.py @@ -22,6 +22,7 @@ import json import logging import os +import random import threading import time from typing import Any, Callable, Literal, Optional, Union, cast @@ -29,6 +30,7 @@ from google.api_core import exceptions as api_exceptions import vertexai +from google.genai import errors as genai_errors from google.genai import types as genai_types from google.genai._api_client import BaseApiClient from google.genai.models import Models @@ -129,6 +131,45 @@ def _get_agent_engine_instance( return _thread_local_data.agent_engine_instances[agent_name] +def _is_retryable_error(e: Exception) -> bool: + """Returns True for transient errors that are safe to retry with backoff.""" + if isinstance(e, genai_errors.APIError): + return e.code in _evals_constant.RETRYABLE_STATUS_CODES + return isinstance( + e, (api_exceptions.ResourceExhausted, api_exceptions.ServiceUnavailable) + ) + + +def _get_retry_backoff( + e: Exception, attempt: int, max_retries: int, operation: str +) -> Optional[float]: + """Returns the seconds to wait before the next attempt, or None to stop. + + Retryable errors back off exponentially with jitter. Other errors fail + fast, and the last attempt never waits. + """ + if not _is_retryable_error(e) or attempt == max_retries - 1: + logger.error( + "Error during %s on attempt %d/%d: %s", + operation, + attempt + 1, + max_retries, + e, + ) + return None + backoff = 2.0**attempt + random.uniform(0, 1) + logger.warning( + "Retryable error during %s on attempt %d/%d: %s. Retrying in %.1f" + " seconds...", + operation, + attempt + 1, + max_retries, + e, + backoff, + ) + return backoff + + def _generate_content_with_retry( api_client: BaseApiClient, model: str, @@ -191,29 +232,11 @@ def _generate_content_with_retry( } else: return response - except api_exceptions.ResourceExhausted as e: - logger.warning( - "Resource Exhausted error on attempt %d/%d: %s. Retrying in %s" - " seconds...", - attempt + 1, - max_retries, - e, - 2**attempt, - ) - if attempt == max_retries - 1: - return {"error": f"Resource exhausted after retries: {e}"} - time.sleep(2**attempt) except Exception as e: # pylint: disable=broad-exception-caught - logger.error( - "Unexpected error during generate_content on attempt %d/%d: %s", - attempt + 1, - max_retries, - e, - ) - - if attempt == max_retries - 1: - return {"error": f"Failed after retries: {e}"} - time.sleep(1) + backoff = _get_retry_backoff(e, attempt, max_retries, "generate_content") + if backoff is None: + return {"error": f"Failed on attempt {attempt + 1}/{max_retries}: {e}"} + time.sleep(backoff) return {"error": f"Failed to generate content after {max_retries} retries"} @@ -2352,28 +2375,11 @@ def _execute_agent_run_with_retry( if event and CONTENT in event and PARTS in event[CONTENT]: responses.append(event) return responses - except api_exceptions.ResourceExhausted as e: - logger.warning( - "Resource Exhausted error on attempt %d/%d: %s. Retrying in %s" - " seconds...", - attempt + 1, - max_retries, - e, - 2**attempt, - ) - if attempt == max_retries - 1: - return {"error": f"Resource exhausted after retries: {e}"} - time.sleep(2**attempt) except Exception as e: # pylint: disable=broad-exception-caught - logger.error( - "Unexpected error during agent engine run on attempt %d/%d: %s", - attempt + 1, - max_retries, - e, - ) - if attempt == max_retries - 1: - return {"error": f"Failed after retries: {e}"} - time.sleep(1) + backoff = _get_retry_backoff(e, attempt, max_retries, "agent engine run") + if backoff is None: + return {"error": f"Failed on attempt {attempt + 1}/{max_retries}: {e}"} + time.sleep(backoff) return {"error": f"Failed to get agent run results after {max_retries} retries"} @@ -2475,28 +2481,13 @@ async def _execute_local_agent_run_with_retry_async( if event and CONTENT in event and PARTS in event[CONTENT]: events.append(event) return events - except api_exceptions.ResourceExhausted as e: - logger.warning( - "Resource Exhausted error on attempt %d/%d: %s. Retrying" - " in %s seconds...", - attempt + 1, - max_retries, - e, - 2**attempt, - ) - if attempt == max_retries - 1: - return {"error": f"Resource exhausted after retries: {e}"} - await asyncio.sleep(2**attempt) except Exception as e: # pylint: disable=broad-exception-caught - logger.error( - "Unexpected error during agent run on attempt %d/%d: %s", - attempt + 1, - max_retries, - e, - ) - if attempt == max_retries - 1: - return {"error": f"Failed after retries: {e}"} - await asyncio.sleep(1) + backoff = _get_retry_backoff(e, attempt, max_retries, "agent run") + if backoff is None: + return { + "error": f"Failed on attempt {attempt + 1}/{max_retries}: {e}" + } + await asyncio.sleep(backoff) return {"error": f"Failed to get agent run results after {max_retries} retries"} diff --git a/vertexai/_genai/_evals_constant.py b/vertexai/_genai/_evals_constant.py index 39d95a50a0..64c99584e2 100644 --- a/vertexai/_genai/_evals_constant.py +++ b/vertexai/_genai/_evals_constant.py @@ -116,3 +116,17 @@ AGENT_DATA, } ) + +# HTTP status codes that are safe to retry with backoff. +RETRYABLE_STATUS_CODES = frozenset( + { + 408, # RequestTimeout (DEADLINE_EXCEEDED) + 409, # Conflict / Aborted (ABORTED) + 429, # TooManyRequests / ResourceExhausted (RESOURCE_EXHAUSTED) + 499, # Client Closed Request (CANCELLED) + 500, # InternalServerError (INTERNAL) + 502, # BadGateway + 503, # ServiceUnavailable (UNAVAILABLE) + 504, # GatewayTimeout (DEADLINE_EXCEEDED) + } +) diff --git a/vertexai/_genai/_evals_metric_handlers.py b/vertexai/_genai/_evals_metric_handlers.py index 79b766d4d2..0d5f4404c0 100644 --- a/vertexai/_genai/_evals_metric_handlers.py +++ b/vertexai/_genai/_evals_metric_handlers.py @@ -38,19 +38,6 @@ logger = logging.getLogger(__name__) _MAX_RETRIES = 5 -# HTTP status codes that are safe to retry with backoff. -_RETRYABLE_STATUS_CODES = frozenset( - { - 408, # RequestTimeout (DEADLINE_EXCEEDED) - 409, # Conflict / Aborted (ABORTED) - 429, # TooManyRequests / ResourceExhausted (RESOURCE_EXHAUSTED) - 499, # Client Closed Request (CANCELLED) - 500, # InternalServerError (INTERNAL) - 502, # BadGateway - 503, # ServiceUnavailable (UNAVAILABLE) - 504, # GatewayTimeout (DEADLINE_EXCEEDED) - } -) R = TypeVar("R") T = TypeVar("T", types.Metric, types.MetricSource, types.LLMMetric) @@ -63,7 +50,7 @@ def _call_with_retry( """Calls ``fn()`` with exponential backoff + jitter on retryable errors. Retries up to ``_MAX_RETRIES`` times on errors whose HTTP status code is - in ``_RETRYABLE_STATUS_CODES`` (Aborted, DeadlineExceeded, + in ``_evals_constant.RETRYABLE_STATUS_CODES`` (Aborted, DeadlineExceeded, ResourceExhausted, ServiceUnavailable, Cancelled). Non-retryable errors are re-raised immediately. If all retries are exhausted the last exception is re-raised so the caller can decide how to handle it. @@ -83,7 +70,7 @@ def _call_with_retry( try: return fn() except genai_errors.APIError as e: - if e.code in _RETRYABLE_STATUS_CODES: + if e.code in _evals_constant.RETRYABLE_STATUS_CODES: backoff = 2**attempt + random.uniform(0, 1) logger.warning( "Retryable error (code=%s) on attempt %d/%d for metric"