| 1 | """Retry utility for model queries."""
|
| 2 |
|
| 3 | import logging
|
| 4 | import os
|
| 5 |
|
| 6 | from tenacity import Retrying, before_sleep_log, retry_if_not_exception_type, stop_after_attempt, wait_exponential
|
| 7 |
|
| 8 |
|
| 9 | def retry(*, logger: logging.Logger, abort_exceptions: list[type[Exception]]) -> Retrying:
|
| 10 | """Thin wrapper around tenacity.Retrying to make use of global config etc.
|
| 11 |
|
| 12 | Args:
|
| 13 | logger: Logger to use for reporting retries
|
| 14 | abort_exceptions: Exceptions to abort on.
|
| 15 |
|
| 16 | Returns:
|
| 17 | A tenacity.Retrying object.
|
| 18 | """
|
| 19 | return Retrying(
|
| 20 | reraise=True,
|
| 21 | stop=stop_after_attempt(int(os.getenv("MSWEA_MODEL_RETRY_STOP_AFTER_ATTEMPT", "10"))),
|
| 22 | wait=wait_exponential(multiplier=1, min=4, max=60),
|
| 23 | before_sleep=before_sleep_log(logger, logging.WARNING),
|
| 24 | retry=retry_if_not_exception_type(tuple(abort_exceptions)),
|
| 25 | )
|
| 26 |
|
