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