Coverage for llm_dataset_engine/utils/retry_handler.py: 96%

1"""

2 Retry handling with exponential backoff.

3

4 Provides robust retry logic for transient failures following clean code

5 principles with single responsibility.

6"""

7

8 import time

9 from typing import Callable, Optional, Type, TypeVar

10

11 from tenacity import (

12 Retrying,

13 retry_if_exception_type,

14 stop_after_attempt,

15 wait_exponential,

16)

17

18 T = TypeVar("T")

19

20

21 class RetryableError(Exception):

22 """Base class for errors that should be retried."""

23

24 pass

25

26

27 class RateLimitError(RetryableError):

28 """Rate limit exceeded error."""

29

30 pass

31

32

33 class NetworkError(RetryableError):

34 """Network-related error."""

35

36 pass

37

38

39 class RetryHandler:

40 """

41 Handles retry logic with exponential backoff.

42

43 Follows SOLID principles with single responsibility for retry logic.

44 """

45

46 def __init__(

47 self,

48 max_attempts: int = 3,

49 initial_delay: float = 1.0,

50 max_delay: float = 60.0,

51 exponential_base: int = 2,

52 retryable_exceptions: Optional[tuple[Type[Exception], ...]] = None,

53 ):

54 """

55 Initialize retry handler.

56

57 Args:

58 max_attempts: Maximum retry attempts

59 initial_delay: Initial delay in seconds

60 max_delay: Maximum delay in seconds

61 exponential_base: Base for exponential backoff

62 retryable_exceptions: Exception types to retry

63 """

64 self.max_attempts = max_attempts

65 self.initial_delay = initial_delay

66 self.max_delay = max_delay

67 self.exponential_base = exponential_base

68

69 if retryable_exceptions is None:

70 self.retryable_exceptions = (

71 RetryableError,

72 RateLimitError,

73 NetworkError,

74 )

75 else:

76 self.retryable_exceptions = retryable_exceptions

77

78 def execute(self, func: Callable[[], T]) -> T:

79 """

80 Execute function with retry logic.

81

82 Args:

83 func: Function to execute

84

85 Returns:

86 Result from function

87

88 Raises:

89 Exception: If all retries exhausted

90 """

91 retryer = Retrying(

92 stop=stop_after_attempt(self.max_attempts),

93 wait=wait_exponential(

94 multiplier=self.initial_delay,

95 max=self.max_delay,

96 exp_base=self.exponential_base,

97 ),

98 retry=retry_if_exception_type(self.retryable_exceptions),

99 reraise=True,

100 )

101

102 return retryer(func)

103

104 def calculate_delay(self, attempt: int) -> float:

105 """

106 Calculate delay for given attempt number.

107

108 Args:

109 attempt: Attempt number (1-based)

110

111 Returns:

112 Delay in seconds

113 """

114 delay = self.initial_delay * (self.exponential_base ** (attempt - 1))

115 return min(delay, self.max_delay)

116