add backoff for download cubin files, and add number of retries

Author: yzh119

flashinfer/artifacts.py

@@ -26,7 +26,7 @@
 from .jit.core import logger
 from .jit.cubin_loader import (
     FLASHINFER_CUBINS_REPOSITORY,
-    get_cubin,
+    download_file,
     safe_urljoin,
     FLASHINFER_CUBIN_DIR,
 )
@@ -125,26 +125,31 @@ def download_artifacts() -> None:
     # HTTPS connections.
     session = requests.Session()
 
-    with temp_env_var("FLASHINFER_CUBIN_CHECKSUM_DISABLED", "1"):
-        cubin_files = list(get_cubin_file_list())
-        num_threads = int(os.environ.get("FLASHINFER_CUBIN_DOWNLOAD_THREADS", "4"))
-        with tqdm_logging_redirect(
-            total=len(cubin_files), desc="Downloading cubins"
-        ) as pbar:
-
-            def update_pbar_cb(_) -> None:
-                pbar.update(1)
-
-            with ThreadPoolExecutor(num_threads) as pool:
-                futures = []
-                for name in cubin_files:
-                    fut = pool.submit(get_cubin, name, "", session)
-                    fut.add_done_callback(update_pbar_cb)
-                    futures.append(fut)
-
-                results = [fut.result() for fut in as_completed(futures)]
-
-        all_success = all(results)
+    cubin_files = list(get_cubin_file_list())
+    num_threads = int(os.environ.get("FLASHINFER_CUBIN_DOWNLOAD_THREADS", "4"))
+    with tqdm_logging_redirect(
+        total=len(cubin_files), desc="Downloading cubins"
+    ) as pbar:
+
+        def update_pbar_cb(_) -> None:
+            pbar.update(1)
+
+        with ThreadPoolExecutor(num_threads) as pool:
+            futures = []
+            for name in cubin_files:
+                source = safe_urljoin(FLASHINFER_CUBINS_REPOSITORY, name)
+                local_path = FLASHINFER_CUBIN_DIR / name
+                # Ensure parent directory exists
+                local_path.parent.mkdir(parents=True, exist_ok=True)
+                fut = pool.submit(
+                    download_file, source, str(local_path), session=session
+                )
+                fut.add_done_callback(update_pbar_cb)
+                futures.append(fut)
+
+            results = [fut.result() for fut in as_completed(futures)]
+
+    all_success = all(results)
     if not all_success:
         raise RuntimeError("Failed to download cubins")
 

flashinfer/jit/cubin_loader.py

@@ -44,7 +44,7 @@ def safe_urljoin(base, path):
 def download_file(
     source: str,
     local_path: str,
-    retries: int = 3,
+    retries: int = 4,
     delay: int = 5,
     timeout: int = 10,
     lock_timeout: int = 30,
@@ -57,7 +57,7 @@ def download_file(
     - source (str): The URL or local file path of the file to download.
     - local_path (str): The local file path to save the downloaded/copied file.
     - retries (int): Number of retry attempts for URL downloads (default: 3).
-    - delay (int): Delay in seconds between retries (default: 5).
+    - delay (int): Initial delay in seconds for exponential backoff (default: 5).
     - timeout (int): Timeout for the HTTP request in seconds (default: 10).
     - lock_timeout (int): Timeout in seconds for the file lock (default: 30).
 
@@ -87,7 +87,7 @@ def download_file(
                     logger.error(f"Failed to copy local file: {e}")
                     return False
 
-            # Handle URL downloads
+            # Handle URL downloads with exponential backoff
             for attempt in range(1, retries + 1):
                 try:
                     response = session.get(source, timeout=timeout)
@@ -107,8 +107,9 @@ def download_file(
                     )
 
                     if attempt < retries:
-                        logger.info(f"Retrying in {delay} seconds...")
-                        time.sleep(delay)
+                        backoff_delay = delay * (2 ** (attempt - 1))
+                        logger.info(f"Retrying in {backoff_delay} seconds...")
+                        time.sleep(backoff_delay)
                     else:
                         logger.error("Max retries reached. Download failed.")
                         return False