fix: GenAI Client(evals) - Mirror the inference retry fix to agentplatform/private and vertexai

vertexai/_genai/_evals_common.py

@@ -22,13 +22,15 @@
 import json
 import logging
 import os
+import random
 import threading
 import time
 from typing import Any, Callable, Literal, Optional, Union, cast
 import uuid
 
 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"}
 
 

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)
+    }
+)

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"