111 lines
3.9 KiB
Python
111 lines
3.9 KiB
Python
"""Retry policies for the three external boundaries the batch depends on:
|
|
LLM providers, Weaviate, and SQL Server. `tenacity` was already a pinned
|
|
dependency (see requirements.txt) but unused until now — the original
|
|
single-field CLI ran once a day and a failure just meant re-running by hand.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Callable, TypeVar
|
|
|
|
import pyodbc
|
|
from tenacity import (
|
|
RetryCallState,
|
|
Retrying,
|
|
retry_if_exception,
|
|
stop_after_attempt,
|
|
wait_exponential_jitter,
|
|
)
|
|
|
|
from pipeline.config import RetrySettings
|
|
from pipeline.errors import NoModelConfiguredError, PipelineError
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
T = TypeVar("T")
|
|
|
|
# SQL Server deadlock victim / serialization failure - safe to retry, the
|
|
# transaction was rolled back and never committed.
|
|
_DEADLOCK_SQLSTATES = {"40001"}
|
|
# Link failure / connection timeout - safe to retry a fresh statement.
|
|
_CONNECTION_SQLSTATES = {"08S01", "08001", "HYT00", "HYT01"}
|
|
|
|
|
|
def _log_retry(retry_state: RetryCallState) -> None:
|
|
exc = retry_state.outcome.exception() if retry_state.outcome else None
|
|
fn_name = getattr(retry_state.fn, "__name__", "call")
|
|
logger.warning(
|
|
"Retrying %s after attempt %d (%s: %s)",
|
|
fn_name,
|
|
retry_state.attempt_number,
|
|
type(exc).__name__ if exc else "?",
|
|
exc,
|
|
)
|
|
|
|
|
|
def is_llm_retryable(exc: BaseException) -> bool:
|
|
"""
|
|
Transient LLM failures worth a retry: HTTP 429 / 5xx, timeouts, and
|
|
transport errors from any of the three SDKs, plus a `PipelineError`
|
|
raised because the model returned empty or unparseable output (a re-roll
|
|
often produces valid output on the next attempt).
|
|
|
|
Never retried: `NoModelConfiguredError` and other pipeline errors that
|
|
are not about LLM output — those describe a configuration or data
|
|
problem that will fail identically on every attempt.
|
|
"""
|
|
if isinstance(exc, NoModelConfiguredError):
|
|
return False
|
|
if isinstance(exc, PipelineError):
|
|
return True
|
|
|
|
status_code = getattr(exc, "status_code", None)
|
|
if isinstance(status_code, int) and (status_code == 429 or status_code >= 500):
|
|
return True
|
|
|
|
name = type(exc).__name__
|
|
return any(
|
|
token in name
|
|
for token in ("RateLimit", "APIConnection", "APITimeout", "Timeout", "ServiceUnavailable")
|
|
)
|
|
|
|
|
|
def is_weaviate_retryable(exc: BaseException) -> bool:
|
|
"""Connection/timeout/availability failures from the weaviate-client SDK."""
|
|
name = type(exc).__name__
|
|
return any(token in name for token in ("Timeout", "Connection", "Unavailable", "Deadline"))
|
|
|
|
|
|
def is_sql_retryable(exc: BaseException) -> bool:
|
|
"""Deadlocks and transient connection failures only; anything else is a real bug."""
|
|
if not isinstance(exc, pyodbc.Error):
|
|
return False
|
|
sqlstate = exc.args[0] if exc.args else ""
|
|
return sqlstate in _DEADLOCK_SQLSTATES or sqlstate in _CONNECTION_SQLSTATES
|
|
|
|
|
|
def _make_retrying(retry_settings: RetrySettings, predicate: Callable[[BaseException], bool]) -> Retrying:
|
|
return Retrying(
|
|
stop=stop_after_attempt(max(1, retry_settings.attempts)),
|
|
wait=wait_exponential_jitter(
|
|
initial=retry_settings.initial_backoff_seconds,
|
|
max=retry_settings.max_backoff_seconds,
|
|
),
|
|
retry=retry_if_exception(predicate),
|
|
before_sleep=_log_retry,
|
|
reraise=True,
|
|
)
|
|
|
|
|
|
def call_with_llm_retry(retry_settings: RetrySettings, fn: Callable[..., T], *args, **kwargs) -> T:
|
|
return _make_retrying(retry_settings, is_llm_retryable)(fn, *args, **kwargs)
|
|
|
|
|
|
def call_with_weaviate_retry(retry_settings: RetrySettings, fn: Callable[..., T], *args, **kwargs) -> T:
|
|
return _make_retrying(retry_settings, is_weaviate_retryable)(fn, *args, **kwargs)
|
|
|
|
|
|
def call_with_sql_retry(retry_settings: RetrySettings, fn: Callable[..., T], *args, **kwargs) -> T:
|
|
return _make_retrying(retry_settings, is_sql_retryable)(fn, *args, **kwargs)
|