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,75 @@ 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+ # retry on bad gateway, service unavailable, gateway timeout, too many requests
40+ retryable_status_codes = (502 , 503 , 504 , 429 ),
41+ retryable_exceptions = (RequestException ,),
42+ ):
43+ """
44+ A decorator for retrying functions that might fail due to transient network issues.
45+
46+ Args:
47+ max_retries: Maximum number of retry attempts
48+ initial_delay: Initial delay between retries in seconds
49+ backoff_factor: Factor by which the delay increases with each retry
50+ retryable_status_codes: HTTP status codes that trigger a retry
51+ retryable_exceptions: Exception types that trigger a retry
52+
53+ Returns:
54+ Decorated function with retry logic
55+ """
56+
57+ def decorator (func ):
58+ @functools .wraps (func )
59+ def wrapper (* args , ** kwargs ):
60+ delay = initial_delay
61+ last_exception = None
62+
63+ for retry_count in range (max_retries + 1 ):
64+ try :
65+ if retry_count > 0 :
66+ logger .info (
67+ f"Retry attempt { retry_count } /{ max_retries } for { func .__name__ } "
68+ )
69+
70+ response = func (* args , ** kwargs )
71+
72+ # Check for retryable status codes in the response
73+ if (
74+ hasattr (response , "status_code" )
75+ and response .status_code in retryable_status_codes
76+ ):
77+ status_code = response .status_code
78+ logger .warning (
79+ f"Received status code { status_code } from { func .__name__ } , retrying..."
80+ )
81+ last_exception = HTTPError (f"HTTP Error { status_code } " )
82+ else :
83+ # Success, return the response
84+ return response
85+
86+ except retryable_exceptions as e :
87+ logger .warning (f"Request failed in { func .__name__ } : { str (e )} " )
88+ last_exception = e
89+
90+ # Don't sleep if this was the last attempt
91+ if retry_count < max_retries :
92+ sleep_time = delay * (backoff_factor ** retry_count )
93+ logger .info (f"Waiting { sleep_time :.2f} seconds before retry" )
94+ time .sleep (sleep_time )
95+
96+ # If we got here, all retries failed
97+ logger .error (f"All { max_retries } retries failed for { func .__name__ } " )
98+ if last_exception :
99+ raise last_exception
100+ return None
101+
102+ return wrapper
103+
104+ return decorator
0 commit comments