265 lines
7.7 KiB
Python
265 lines
7.7 KiB
Python
"""Safe retry helpers for transient cloud API failures."""
|
|
|
|
import math
|
|
import random
|
|
import time
|
|
from collections.abc import Mapping
|
|
from typing import Any, Callable, Optional
|
|
|
|
|
|
MAX_API_ATTEMPTS = 10
|
|
RETRYABLE_STATUS_CODES = {408, 425, 429, 500, 502, 503, 504}
|
|
_TRANSIENT_EXCEPTION_NAMES = {
|
|
"ConnectTimeout",
|
|
"ConnectionError",
|
|
"ConnectionResetError",
|
|
"ProtocolError",
|
|
"ReadTimeout",
|
|
"Timeout",
|
|
"TimeoutError",
|
|
}
|
|
|
|
|
|
class APIRequestError(RuntimeError):
|
|
"""A sanitized cloud API failure with retry classification metadata."""
|
|
|
|
def __init__(
|
|
self,
|
|
operation: str,
|
|
status_code: Optional[int] = None,
|
|
code: Optional[str] = None,
|
|
message: Optional[str] = None,
|
|
request_id: Optional[str] = None,
|
|
retry_after: Optional[float] = None,
|
|
) -> None:
|
|
self.operation = operation
|
|
self.status_code = status_code
|
|
self.code = code
|
|
self.service_message = message
|
|
self.request_id = request_id
|
|
self.retry_after = retry_after
|
|
|
|
details = []
|
|
if status_code is not None:
|
|
details.append(f"status={status_code}")
|
|
if code:
|
|
details.append(f"code={code}")
|
|
if request_id:
|
|
details.append(f"request_id={request_id}")
|
|
suffix = f" ({', '.join(details)})" if details else ""
|
|
super().__init__(f"{operation} failed{suffix}")
|
|
|
|
|
|
def _optional_text(value: Any) -> Optional[str]:
|
|
if value is None:
|
|
return None
|
|
text = str(value).strip()
|
|
return text or None
|
|
|
|
|
|
def _optional_status(value: Any) -> Optional[int]:
|
|
try:
|
|
return int(value)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
|
|
def _optional_delay(value: Any) -> Optional[float]:
|
|
try:
|
|
delay = float(value)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
if not math.isfinite(delay):
|
|
return None
|
|
return max(0.0, delay)
|
|
|
|
|
|
def _response_value(
|
|
response: Any,
|
|
name: str,
|
|
default: Any = None,
|
|
) -> Any:
|
|
"""Read mapping-style and object-style API responses safely."""
|
|
if isinstance(response, Mapping):
|
|
return response.get(name, default)
|
|
try:
|
|
return getattr(response, name, default)
|
|
except (AttributeError, KeyError):
|
|
return default
|
|
|
|
|
|
def _header_value(headers: Any, name: str) -> Any:
|
|
if isinstance(headers, Mapping):
|
|
for key, value in headers.items():
|
|
if str(key).lower() == name.lower():
|
|
return value
|
|
return None
|
|
|
|
try:
|
|
getter = getattr(headers, "get")
|
|
except (AttributeError, KeyError):
|
|
return None
|
|
|
|
for key in (name, name.lower()):
|
|
try:
|
|
value = getter(key)
|
|
except (AttributeError, KeyError, TypeError):
|
|
continue
|
|
if value is not None:
|
|
return value
|
|
return None
|
|
|
|
|
|
def api_error_from_response(response: Any, operation: str) -> APIRequestError:
|
|
"""Build a sanitized error from a DashScope-style response object."""
|
|
headers = _response_value(response, "headers")
|
|
retry_after = _optional_delay(
|
|
_response_value(response, "retry_after")
|
|
)
|
|
if retry_after is None:
|
|
retry_after = _optional_delay(
|
|
_header_value(headers, "Retry-After")
|
|
)
|
|
|
|
return APIRequestError(
|
|
operation=operation,
|
|
status_code=_optional_status(
|
|
_response_value(response, "status_code")
|
|
),
|
|
code=_optional_text(_response_value(response, "code")),
|
|
message=_optional_text(_response_value(response, "message")),
|
|
request_id=_optional_text(
|
|
_response_value(response, "request_id")
|
|
),
|
|
retry_after=retry_after,
|
|
)
|
|
|
|
|
|
def is_retryable_api_error(error: Exception) -> bool:
|
|
"""Return whether an error is likely to succeed on a later attempt."""
|
|
if isinstance(error, APIRequestError):
|
|
return error.status_code in RETRYABLE_STATUS_CODES
|
|
if isinstance(error, (TimeoutError, ConnectionError)):
|
|
return True
|
|
|
|
current: Optional[BaseException] = error
|
|
visited = set()
|
|
while current is not None and id(current) not in visited:
|
|
visited.add(id(current))
|
|
if type(current).__name__ in _TRANSIENT_EXCEPTION_NAMES:
|
|
return True
|
|
current = current.__cause__ or current.__context__
|
|
return False
|
|
|
|
|
|
def throw_if_processing_interrupted() -> None:
|
|
"""Raise ComfyUI's interruption exception when processing was cancelled."""
|
|
try:
|
|
import comfy.model_management # type: ignore
|
|
except ImportError:
|
|
return
|
|
comfy.model_management.throw_exception_if_processing_interrupted()
|
|
|
|
|
|
def is_processing_interrupt(error: Exception) -> bool:
|
|
try:
|
|
import comfy.model_management # type: ignore
|
|
return isinstance(
|
|
error,
|
|
comfy.model_management.InterruptProcessingException,
|
|
)
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def normalize_max_attempts(
|
|
max_attempts: int,
|
|
logger: Callable[[str], None] = print,
|
|
) -> int:
|
|
try:
|
|
attempts = int(max_attempts)
|
|
except (TypeError, ValueError) as error:
|
|
raise ValueError("max_retries must be an integer") from error
|
|
if attempts < 1:
|
|
raise ValueError("max_retries must be at least 1")
|
|
if attempts > MAX_API_ATTEMPTS:
|
|
logger(
|
|
f"[API Retry] Requested {attempts} attempts; "
|
|
f"capping at {MAX_API_ATTEMPTS}"
|
|
)
|
|
return MAX_API_ATTEMPTS
|
|
return attempts
|
|
|
|
|
|
def compute_retry_delay(
|
|
failed_attempt: int,
|
|
retry_after: Optional[float] = None,
|
|
random_fn: Callable[[], float] = random.random,
|
|
max_delay: float = 30.0,
|
|
) -> float:
|
|
if retry_after is not None:
|
|
return min(max(0.0, retry_after), max_delay)
|
|
base_delay = min(2 ** max(0, failed_attempt - 1), max_delay)
|
|
jitter = max(0.0, min(1.0, float(random_fn()))) * base_delay * 0.25
|
|
return min(base_delay + jitter, max_delay)
|
|
|
|
|
|
def _interruptible_sleep(
|
|
delay: float,
|
|
sleep_fn: Callable[[float], None],
|
|
interrupt_check: Callable[[], None],
|
|
) -> None:
|
|
remaining = max(0.0, delay)
|
|
while remaining > 0:
|
|
interrupt_check()
|
|
interval = min(0.25, remaining)
|
|
sleep_fn(interval)
|
|
remaining -= interval
|
|
interrupt_check()
|
|
|
|
|
|
def call_with_retry(
|
|
operation: Callable[[], Any],
|
|
max_attempts: int,
|
|
operation_name: str,
|
|
*,
|
|
sleep_fn: Callable[[float], None] = time.sleep,
|
|
random_fn: Callable[[], float] = random.random,
|
|
interrupt_check: Callable[[], None] = throw_if_processing_interrupted,
|
|
interrupt_error_check: Callable[[Exception], bool] = is_processing_interrupt,
|
|
logger: Callable[[str], None] = print,
|
|
) -> Any:
|
|
"""Run one API operation, retrying only classified transient failures."""
|
|
attempts = normalize_max_attempts(max_attempts, logger=logger)
|
|
|
|
for attempt in range(1, attempts + 1):
|
|
interrupt_check()
|
|
try:
|
|
return operation()
|
|
except Exception as error:
|
|
if interrupt_error_check(error):
|
|
raise
|
|
if not is_retryable_api_error(error):
|
|
raise
|
|
if attempt >= attempts:
|
|
raise
|
|
|
|
retry_after = (
|
|
error.retry_after
|
|
if isinstance(error, APIRequestError)
|
|
else None
|
|
)
|
|
delay = compute_retry_delay(
|
|
failed_attempt=attempt,
|
|
retry_after=retry_after,
|
|
random_fn=random_fn,
|
|
)
|
|
logger(
|
|
f"[API Retry] {operation_name} failed with "
|
|
f"{type(error).__name__}; retrying attempt "
|
|
f"{attempt + 1}/{attempts} in {delay:.2f}s"
|
|
)
|
|
_interruptible_sleep(delay, sleep_fn, interrupt_check)
|
|
|
|
raise RuntimeError("API retry loop exited unexpectedly")
|