download_retry_decorator() — pytorch Function Reference

Architecture documentation for the download_retry_decorator() function in common.py from the pytorch codebase.

Function python DynamoBenchmarks Configuration calls 1

Entity Profile

DynamoBenchmarks→ Configuration→ download_retry_decorator() — pytorch Function Reference

Dependency Diagram

graph TD
  008d9589_ae5b_adb6_6987_3e212c1e7325["download_retry_decorator()"]
  ce8fd365_4112_b289_9c73_7345d5e35203["RuntimeError()"]
  008d9589_ae5b_adb6_6987_3e212c1e7325 -->|calls| ce8fd365_4112_b289_9c73_7345d5e35203
  style 008d9589_ae5b_adb6_6987_3e212c1e7325 fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/common.py lines 1578–1612

def download_retry_decorator(download_fn):
    """
    Decorator function for applying retry logic to a download function.

    The wrapped function will be called up to 5 times and raises an exception if the function fails each time.
    After each unsuccessful attempt, there is a delay before the next attempt, which is increased linearly with the number of tries.

    Usage:
    @download_retry_decorator
    def download_function(model_name: str):
        # download logic goes here
    """

    @functools.wraps(download_fn)
    def wrapper(self, *args, **kwargs) -> Any:
        tries = 0
        total_allowed_tries = MAX_DOWNLOAD_ATTEMPTS
        while tries <= total_allowed_tries:
            try:
                model = download_fn(self, *args, **kwargs)
                return model
            except Exception as e:
                tries += 1
                if tries <= total_allowed_tries:
                    wait = tries * 30
                    print(
                        f"Failed to load model: {e}. Trying again ({tries}/{total_allowed_tries}) after {wait}s"
                    )
                    time.sleep(wait)
                else:
                    raise RuntimeError(  # noqa: B904
                        f"Failed to load model '{args}' with following error(s): {str(e)}."
                    )

    return wrapper