diff --git a/src/deadline/job_attachments/_utils.py b/src/deadline/job_attachments/_utils.py index 5a148a810..b9694eb76 100644 --- a/src/deadline/job_attachments/_utils.py +++ b/src/deadline/job_attachments/_utils.py @@ -1,9 +1,12 @@ # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. import datetime +from functools import wraps from hashlib import shake_256 from pathlib import Path -from typing import Tuple, Union +import random +import time +from typing import Any, Callable, Optional, Tuple, Type, Union import uuid import ctypes import sys @@ -98,3 +101,56 @@ def _is_windows_file_path_limit() -> bool: return bool(ntdll.RtlAreLongPathsEnabled()) return True + + +def _retry( + ExceptionToCheck: Union[Type[Exception], Tuple[Type[Exception], ...]] = AssertionError, + tries: int = 4, + delay: Union[int, float, Tuple[Union[int, float], Union[int, float]]] = 3.0, + backoff: float = 2.0, + logger: Optional[Callable] = print, +) -> Callable: + """Retry calling the decorated function using an exponential backoff. + + http://www.saltycrane.com/blog/2009/11/trying-out-retry-decorator-python/ + original from: http://wiki.python.org/moin/PythonDecoratorLibrary#Retry + + :param ExceptionToCheck: the exception to check. may be a tuple of + exceptions to check + :type ExceptionToCheck: Exception or tuple + :param tries: number of times to try (not retry) before giving up + :type tries: int + :param delay: initial delay between retries in seconds + :type delay: float or tuple + :param backoff: backoff multiplier e.g. value of 2 will double the delay + each retry + :type backoff: float + :param logger: logging function to use. If None, won't log + :type logger: logging.Logger instance + """ + + def deco_retry(f: Callable) -> Callable: + @wraps(f) + def f_retry(*args: Any, **kwargs: Any) -> Callable: + mtries: int = tries + if isinstance(delay, (float, int)): + mdelay = delay + elif isinstance(delay, tuple): + mdelay = random.uniform(delay[0], delay[1]) + else: + raise ValueError(f"Provided delay {delay} isn't supported") + + while mtries > 1: + try: + return f(*args, **kwargs) + except ExceptionToCheck as e: + if logger: + logger(f"{str(e)}, Retrying in {mdelay} seconds...") + time.sleep(mdelay) + mtries -= 1 + mdelay *= backoff + return f(*args, **kwargs) + + return f_retry # true decorator + + return deco_retry diff --git a/test/unit/deadline_job_attachments/test_utils.py b/test/unit/deadline_job_attachments/test_utils.py index b33fb3389..808b29ef9 100644 --- a/test/unit/deadline_job_attachments/test_utils.py +++ b/test/unit/deadline_job_attachments/test_utils.py @@ -7,6 +7,7 @@ from deadline.job_attachments._utils import ( _is_relative_to, + _retry, ) @@ -60,3 +61,23 @@ def test_is_relative_to_on_windows(self, path1, path2, expected): Tests if the is_relative_to() works correctly when using Windows paths. """ assert _is_relative_to(path1, path2) == expected + + def test_retry(self): + """ + Test a function that throws an exception is retried. + """ + call_count = 0 + + # Given + @_retry(ExceptionToCheck=NotImplementedError, tries=2, delay=0.1, backoff=0.1) + def test_bad_function(): + nonlocal call_count + call_count = call_count + 1 + if call_count == 1: + raise NotImplementedError() + + # When + test_bad_function() + + # Then + assert call_count == 2