sumo-wrapper-python/src/sumo/wrapper/_retry_strategy.py at 58a97ee1f0a08de6a63a1246e1b4993deeef66d2 · equinor/sumo-wrapper-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
import httpx
import tenacity as tn


def _log_retry_info(retry_state):
    # logger.log(
    #     logging.INFO,
    #     f"Attempts: {retry_state.attempt_number}; "
    #     f"Elapsed: {retry_state.seconds_since_start}",
    # )
    print(
        f"Attempts: {retry_state.attempt_number}; "
        f"Elapsed:  {retry_state.seconds_since_start}"
    )
    return


# Define the conditions for retrying based on exception types
def _is_retryable_exception(exception):
    return isinstance(
        exception,
        (
            httpx.TimeoutException,
            httpx.NetworkError,
            httpx.ProtocolError,
            httpx.ProxyError,
        ),
    )


# Define the conditions for retrying based on HTTP status codes
def _is_retryable_status_code(response):
    return response.status_code in [502, 503]


def _return_last_value(retry_state):
    return retry_state.outcome.result()


class RetryStrategy:
    def __init__(self, stop_after=6, multiplier=0.5, exp_base=2):
        self._stop_after = stop_after
        self._multiplier = multiplier
        self._exp_base = exp_base
        return

    def make_retryer(self) -> tn.Retrying:
        return tn.Retrying(
            stop=tn.stop_after_attempt(self._stop_after),
            retry=(
                tn.retry_if_exception(_is_retryable_exception)
                | tn.retry_if_result(_is_retryable_status_code)
            ),
            wait=(
                tn.wait_exponential(
                    multiplier=self._multiplier, exp_base=self._exp_base
                )
                + tn.wait_random_exponential(
                    multiplier=self._multiplier, exp_base=self._exp_base
                )
            ),
            retry_error_callback=_return_last_value,
            before_sleep=_log_retry_info,
        )

    def make_retryer_async(self) -> tn.AsyncRetrying:
        return tn.AsyncRetrying(
            stop=tn.stop_after_attempt(self._stop_after),
            retry=(
                tn.retry_if_exception(_is_retryable_exception)
                | tn.retry_if_result(_is_retryable_status_code)
            ),
            wait=(
                tn.wait_exponential(
                    multiplier=self._multiplier, exp_base=self._exp_base
                )
                + tn.wait_random_exponential(
                    multiplier=self._multiplier, exp_base=self._exp_base
                )
            ),
            retry_error_callback=_return_last_value,
            before_sleep=_log_retry_info,
        )