1+ import functools
12import json
3+ import logging
24import os
35import re
6+ import sys
7+ import time
48
9+ from requests .exceptions import RequestException , HTTPError
10+
11+ logger = logging .getLogger (__name__ )
12+ logging .basicConfig (stream = sys .stdout , level = logging .INFO )
513
614def write_dict_to_json_file (input , output_path ):
715 if os .path .exists (output_path ):
@@ -22,3 +30,74 @@ def get_uri_for_digest(uri, digest):
2230 """
2331 base_uri = re .split (r"[@:]" , uri , maxsplit = 1 )[0 ]
2432 return f"{ base_uri } @{ digest } "
33+
34+
35+ def retry_on_error (
36+ max_retries = 3 ,
37+ initial_delay = 1 ,
38+ backoff_factor = 2 ,
39+ retryable_status_codes = (502 , 503 , 504 , 429 ),
40+ retryable_exceptions = (RequestException ,),
41+ ):
42+ """
43+ A decorator for retrying functions that might fail due to transient network issues.
44+
45+ Args:
46+ max_retries: Maximum number of retry attempts
47+ initial_delay: Initial delay between retries in seconds
48+ backoff_factor: Factor by which the delay increases with each retry
49+ retryable_status_codes: HTTP status codes that trigger a retry
50+ retryable_exceptions: Exception types that trigger a retry
51+
52+ Returns:
53+ Decorated function with retry logic
54+ """
55+
56+ def decorator (func ):
57+ @functools .wraps (func )
58+ def wrapper (* args , ** kwargs ):
59+ delay = initial_delay
60+ last_exception = None
61+
62+ for retry_count in range (max_retries + 1 ):
63+ try :
64+ if retry_count > 0 :
65+ logger .info (
66+ f"Retry attempt { retry_count } /{ max_retries } for { func .__name__ } "
67+ )
68+
69+ response = func (* args , ** kwargs )
70+
71+ # Check for retryable status codes in the response
72+ if (
73+ hasattr (response , "status_code" )
74+ and response .status_code in retryable_status_codes
75+ ):
76+ status_code = response .status_code
77+ logger .warning (
78+ f"Received status code { status_code } from { func .__name__ } , retrying..."
79+ )
80+ last_exception = HTTPError (f"HTTP Error { status_code } " )
81+ else :
82+ # Success, return the response
83+ return response
84+
85+ except retryable_exceptions as e :
86+ logger .warning (f"Request failed in { func .__name__ } : { str (e )} " )
87+ last_exception = e
88+
89+ # Don't sleep if this was the last attempt
90+ if retry_count < max_retries :
91+ sleep_time = delay * (backoff_factor ** retry_count )
92+ logger .info (f"Waiting { sleep_time :.2f} seconds before retry" )
93+ time .sleep (sleep_time )
94+
95+ # If we got here, all retries failed
96+ logger .error (f"All { max_retries } retries failed for { func .__name__ } " )
97+ if last_exception :
98+ raise last_exception
99+ return None
100+
101+ return wrapper
102+
103+ return decorator
0 commit comments