fix(docker): add retry logic to model downloads for transient 5xx/timeout errors

download_models.py used bare urllib.request.urlretrieve() with no retry
logic. CI hit a HTTP 504 Gateway Timeout mid-build, failing the Docker
Build Test. Added _urlretrieve() wrapper that retries up to 3 times with
a 10s delay on any 5xx or network error. Also adds imports for time and
urllib.error.
This commit is contained in:
Siddharth Kumar Sah
2026-04-14 22:49:49 +08:00
parent 0f156d13a3
commit fb6cff0fec
+30 -10
View File
@@ -5,6 +5,8 @@ failing the build. No silent fallbacks.
"""
import os
import sys
import time
import urllib.error
import urllib.request
# Some servers (e.g. Berkeley) block the default Python-urllib User-Agent.
@@ -12,6 +14,24 @@ _opener = urllib.request.build_opener()
_opener.addheaders = [("User-Agent", "ashim/1.0")]
urllib.request.install_opener(_opener)
def _urlretrieve(url: str, path: str, max_retries: int = 3) -> None:
"""Download url to path, retrying on transient network errors (5xx, timeout)."""
for attempt in range(1, max_retries + 1):
try:
urllib.request.urlretrieve(url, path)
return
except urllib.error.HTTPError as e:
if attempt == max_retries or e.code < 500:
raise
print(f" HTTP {e.code} on attempt {attempt}/{max_retries}, retrying in 10s...")
time.sleep(10)
except (urllib.error.URLError, OSError) as e:
if attempt == max_retries:
raise
print(f" Network error on attempt {attempt}/{max_retries}: {e}, retrying in 10s...")
time.sleep(10)
# Force CPU mode during build - no GPU driver available at build time.
# Must be set before any ML library import.
os.environ["PADDLE_DEVICE"] = "cpu"
@@ -171,7 +191,7 @@ def download_lama_model():
print("=== Downloading LaMa ONNX model ===")
os.makedirs(LAMA_MODEL_DIR, exist_ok=True)
print(f" Downloading from {LAMA_MODEL_URL}...")
urllib.request.urlretrieve(LAMA_MODEL_URL, LAMA_MODEL_PATH)
_urlretrieve(LAMA_MODEL_URL, LAMA_MODEL_PATH)
size = os.path.getsize(LAMA_MODEL_PATH)
assert size > LAMA_MIN_SIZE, (
@@ -185,7 +205,7 @@ def download_realesrgan_model():
print("=== Downloading RealESRGAN model ===")
os.makedirs(REALESRGAN_MODEL_DIR, exist_ok=True)
print(f" Downloading from {REALESRGAN_MODEL_URL}...")
urllib.request.urlretrieve(REALESRGAN_MODEL_URL, REALESRGAN_MODEL_PATH)
_urlretrieve(REALESRGAN_MODEL_URL, REALESRGAN_MODEL_PATH)
size = os.path.getsize(REALESRGAN_MODEL_PATH)
assert size > REALESRGAN_MIN_SIZE, (
@@ -199,7 +219,7 @@ def download_gfpgan_model():
print("=== Downloading GFPGAN model ===")
os.makedirs(GFPGAN_MODEL_DIR, exist_ok=True)
print(f" Downloading from {GFPGAN_MODEL_URL}...")
urllib.request.urlretrieve(GFPGAN_MODEL_URL, GFPGAN_MODEL_PATH)
_urlretrieve(GFPGAN_MODEL_URL, GFPGAN_MODEL_PATH)
size = os.path.getsize(GFPGAN_MODEL_PATH)
assert size > GFPGAN_MIN_SIZE, (
@@ -213,7 +233,7 @@ def download_codeformer_model():
print("=== Downloading CodeFormer model ===")
os.makedirs(CODEFORMER_MODEL_DIR, exist_ok=True)
print(f" Downloading from {CODEFORMER_MODEL_URL}...")
urllib.request.urlretrieve(CODEFORMER_MODEL_URL, CODEFORMER_MODEL_PATH)
_urlretrieve(CODEFORMER_MODEL_URL, CODEFORMER_MODEL_PATH)
size = os.path.getsize(CODEFORMER_MODEL_PATH)
assert size > CODEFORMER_MIN_SIZE, (
@@ -320,7 +340,7 @@ def download_scunet_model():
"""Download SCUNet real-noise denoising model."""
print(f"Downloading SCUNet model to {SCUNET_MODEL_PATH}...")
os.makedirs(SCUNET_MODEL_DIR, exist_ok=True)
urllib.request.urlretrieve(SCUNET_MODEL_URL, SCUNET_MODEL_PATH)
_urlretrieve(SCUNET_MODEL_URL, SCUNET_MODEL_PATH)
size = os.path.getsize(SCUNET_MODEL_PATH)
assert size > SCUNET_MIN_SIZE, (
f"SCUNet model too small: {size} bytes (expected >{SCUNET_MIN_SIZE})"
@@ -332,7 +352,7 @@ def download_nafnet_model():
"""Download NAFNet SIDD width-64 denoising model."""
print(f"Downloading NAFNet model to {NAFNET_MODEL_PATH}...")
os.makedirs(NAFNET_MODEL_DIR, exist_ok=True)
urllib.request.urlretrieve(NAFNET_MODEL_URL, NAFNET_MODEL_PATH)
_urlretrieve(NAFNET_MODEL_URL, NAFNET_MODEL_PATH)
size = os.path.getsize(NAFNET_MODEL_PATH)
assert size > NAFNET_MIN_SIZE, (
f"NAFNet model too small: {size} bytes (expected >{NAFNET_MIN_SIZE})"
@@ -350,7 +370,7 @@ def download_facexlib_models():
os.makedirs(FACEXLIB_MODEL_DIR, exist_ok=True)
print(f" Downloading detection_Resnet50_Final.pth...")
urllib.request.urlretrieve(FACEXLIB_DET_URL, FACEXLIB_DET_PATH)
_urlretrieve(FACEXLIB_DET_URL, FACEXLIB_DET_PATH)
size = os.path.getsize(FACEXLIB_DET_PATH)
assert size > FACEXLIB_DET_MIN_SIZE, (
f"Face detection model too small: {size} bytes (expected > {FACEXLIB_DET_MIN_SIZE})"
@@ -358,7 +378,7 @@ def download_facexlib_models():
print(f" detection_Resnet50_Final.pth downloaded ({size / 1_000_000:.1f} MB)")
print(f" Downloading parsing_parsenet.pth...")
urllib.request.urlretrieve(FACEXLIB_PARSE_URL, FACEXLIB_PARSE_PATH)
_urlretrieve(FACEXLIB_PARSE_URL, FACEXLIB_PARSE_PATH)
size = os.path.getsize(FACEXLIB_PARSE_PATH)
assert size > FACEXLIB_PARSE_MIN_SIZE, (
f"Face parsing model too small: {size} bytes (expected > {FACEXLIB_PARSE_MIN_SIZE})"
@@ -380,7 +400,7 @@ def download_opencv_colorize_models():
(OPENCV_POINTS_URL, OPENCV_POINTS_PATH, "pts_in_hull.npy"),
]:
print(f" Downloading {name}...")
urllib.request.urlretrieve(url, path)
_urlretrieve(url, path)
size = os.path.getsize(path)
print(f" {name} downloaded ({size / 1_000_000:.1f} MB)")
@@ -409,7 +429,7 @@ def download_mediapipe_task_models():
"face_landmarker", FACE_LANDMARKER_MIN_SIZE),
]:
print(f" Downloading {name}...")
urllib.request.urlretrieve(url, path)
_urlretrieve(url, path)
size = os.path.getsize(path)
assert size > min_size, (
f"{name} model too small: {size} bytes (expected > {min_size})"