[diffusion] CI: improve ci reliability (#21763)

python/sglang/multimodal_gen/test/run_suite.py

@@ -234,7 +234,9 @@ def run_pytest(files, filter_expr=None):
         )
 
         is_flaky_ci_assertion = (
-            "SafetensorError" in full_output or "FileNotFoundError" in full_output
+            "SafetensorError" in full_output
+            or "FileNotFoundError" in full_output
+            or "TimeoutError" in full_output
         )
 
         is_oom_error = (

python/sglang/multimodal_gen/test/server/test_server_utils.py

@@ -55,6 +55,28 @@ globally_suppress_loggers()
 MESH_OUTPUT_PATHS: dict[str, str] = {}
 
 
+def _urlopen_with_retry(url: str, timeout: int = 30, max_retries: int = 3) -> bytes:
+    """Download content from a URL with retry on transient failures."""
+    for attempt in range(max_retries + 1):
+        try:
+            with urlopen(url, timeout=timeout) as response:
+                return response.read()
+        except (TimeoutError, OSError) as e:
+            if attempt < max_retries:
+                wait = 2**attempt
+                logger.warning(
+                    f"Download attempt {attempt + 1}/{max_retries + 1} failed "
+                    f"for {url}: {e}. Retrying in {wait}s..."
+                )
+                time.sleep(wait)
+            else:
+                logger.error(
+                    f"Failed to download from {url} after "
+                    f"{max_retries + 1} attempts: {e}"
+                )
+                raise
+
+
 def download_image_from_url(url: str) -> Path:
     """Download an image from a URL to a temporary file.
 
@@ -76,14 +98,10 @@ def download_image_from_url(url: str) -> Path:
         Path(tempfile.gettempdir()) / f"diffusion_test_image_{int(time.time())}{ext}"
     )
 
-    try:
-        with urlopen(url, timeout=30) as response:
-            temp_file.write_bytes(response.read())
-        logger.info(f"Downloaded image to: {temp_file}")
-        return temp_file
-    except Exception as e:
-        logger.error(f"Failed to download image from {url}: {e}")
-        raise
+    data = _urlopen_with_retry(url)
+    temp_file.write_bytes(data)
+    logger.info(f"Downloaded image to: {temp_file}")
+    return temp_file
 
 
 def parse_dimensions(size_string: str | None) -> tuple[int | None, int | None]:
@@ -664,8 +682,7 @@ def _download_reference_mesh(url: str) -> Path:
         return cache_path
 
     logger.info(f"Downloading reference mesh from: {url}")
-    with urlopen(url, timeout=60) as resp:
-        cache_path.write_bytes(resp.read())
+    cache_path.write_bytes(_urlopen_with_retry(url, timeout=60))
     logger.info(f"Reference mesh cached at: {cache_path}")
     return cache_path