From 00228deaaa2b05376567797d08fa2e7edef1b3a0 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Mon, 6 Jul 2026 11:58:19 +0800 Subject: [PATCH] fix(download): retry on Civitai 429 rate limit instead of removing images from metadata MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When Civitai returns 429 (Too Many Requests) during example image downloads, the previous behavior treated all failures identically and permanently removed the corresponding images from model metadata — making them impossible to retry. This commit adds: - 429 detection + Retry-After header parsing in download_to_memory - Exponential backoff retry (up to 3 attempts) in download_model_images_with_tracking - Separate tracking of rate-limited vs permanently failed URLs - rate_limited_models progress tracking persisted to disk - Rate-limited models are NOT added to failed_models/processed_models so they are automatically retried on subsequent download runs - Force mode clears failed_models when rate-limited images exist --- py/services/downloader.py | 37 +++++ py/utils/example_images_download_manager.py | 105 ++++++++++---- py/utils/example_images_processor.py | 131 ++++++++++++------ ...t_example_images_download_manager_async.py | 4 +- tests/services/test_issue_760_repro.py | 2 +- 5 files changed, 213 insertions(+), 66 deletions(-) diff --git a/py/services/downloader.py b/py/services/downloader.py index 06bfabd6..b498db75 100644 --- a/py/services/downloader.py +++ b/py/services/downloader.py @@ -46,6 +46,30 @@ def is_ssl_cert_verify_error(exc: BaseException) -> bool: return "CERTIFICATE_VERIFY_FAILED" in str(exc) +def _parse_retry_after(value: str) -> int: + """Parse a Retry-After header value into seconds. + + Supports both integer seconds and HTTP-date formats. + Returns a default of 60 seconds on invalid/missing input. + """ + if not value or not value.strip(): + return 60 + + value = value.strip() + try: + return max(1, int(value)) + except ValueError: + pass + + try: + parsed = parsedate_to_datetime(value) + now = datetime.now().astimezone() + delta = (parsed - now).total_seconds() + return max(1, int(delta)) + except (ValueError, OverflowError, OSError): + return 60 + + @dataclass(frozen=True) class DownloadProgress: """Snapshot of a download transfer at a moment in time.""" @@ -911,6 +935,19 @@ class Downloader: elif response.status == 404: error_msg = "File not found" return False, error_msg, None + elif response.status == 429: + raw_retry_after = response.headers.get("Retry-After") + retry_after = _parse_retry_after(raw_retry_after or "") + if raw_retry_after: + logger.warning( + "Rate limited (429) for %s, Retry-After: %ss", url, retry_after + ) + else: + logger.warning( + "Rate limited (429) for %s, no Retry-After header; defaulting to %ss", + url, retry_after, + ) + return False, f"Rate limited (429), retry after {retry_after}s", None else: error_msg = f"Download failed with status {response.status}" return False, error_msg, None diff --git a/py/utils/example_images_download_manager.py b/py/utils/example_images_download_manager.py index 0ecae86f..994ca62d 100644 --- a/py/utils/example_images_download_manager.py +++ b/py/utils/example_images_download_manager.py @@ -72,6 +72,7 @@ class _DownloadProgress(dict): refreshed_models=set(), failed_models=set(), reprocessed_models=set(), + rate_limited_models=set(), ) def snapshot(self) -> dict: @@ -82,6 +83,7 @@ class _DownloadProgress(dict): snapshot["refreshed_models"] = list(self["refreshed_models"]) snapshot["failed_models"] = list(self["failed_models"]) snapshot["reprocessed_models"] = list(self.get("reprocessed_models", set())) + snapshot["rate_limited_models"] = list(self.get("rate_limited_models", set())) return snapshot @@ -153,13 +155,15 @@ class DownloadManager: # Step 3: Load progress file (I/O operation, done outside lock) processed_models = set() failed_models = set() + rate_limited_models = set() try: - progress_file, processed_models, failed_models = await self._load_progress_file(output_dir) + progress_file, processed_models, failed_models, rate_limited_models = await self._load_progress_file(output_dir) logger.debug( - "Loaded previous progress, %s models already processed, %s models marked as failed", + "Loaded previous progress, %s models already processed, %s models marked as failed, %s models rate-limited", len(processed_models), len(failed_models), + len(rate_limited_models), ) except Exception as e: logger.error(f"Failed to load progress file: {e}") @@ -175,6 +179,7 @@ class DownloadManager: self._progress.reset() self._progress["processed_models"] = processed_models self._progress["failed_models"] = failed_models + self._progress["rate_limited_models"] = rate_limited_models self._stop_requested = False self._progress["status"] = "running" self._progress["start_time"] = time.time() @@ -242,8 +247,8 @@ class DownloadManager: "status": self._progress.snapshot(), } - async def _load_progress_file(self, output_dir: str) -> tuple[str, set, set]: - """Load progress file from disk. Returns (progress_file_path, processed_models, failed_models). + async def _load_progress_file(self, output_dir: str) -> tuple[str, set, set, set]: + """Load progress file from disk. Returns (progress_file_path, processed_models, failed_models, rate_limited_models). This is a separate async method to allow running in executor to avoid blocking event loop. """ @@ -252,8 +257,12 @@ class DownloadManager: None, self._load_progress_file_sync, output_dir ) - def _load_progress_file_sync(self, output_dir: str) -> tuple[str, set, set]: - """Synchronous implementation of progress file loading.""" + def _load_progress_file_sync(self, output_dir: str) -> tuple[str, set, set, set]: + """Synchronous implementation of progress file loading. + + Returns: + tuple: (progress_file_path, processed_models, failed_models, rate_limited_models) + """ progress_file = os.path.join(output_dir, ".download_progress.json") progress_source = progress_file @@ -289,6 +298,7 @@ class DownloadManager: processed_models = set() failed_models = set() + rate_limited_models = set() if os.path.exists(progress_source): try: @@ -296,11 +306,11 @@ class DownloadManager: saved_progress = json.load(f) processed_models = set(saved_progress.get("processed_models", [])) failed_models = set(saved_progress.get("failed_models", [])) + rate_limited_models = set(saved_progress.get("rate_limited_models", [])) except Exception: - # Return empty sets on error pass - return progress_file, processed_models, failed_models + return progress_file, processed_models, failed_models, rate_limited_models def _load_progress_sets_sync(self, progress_file: str) -> tuple[set, set]: """Load only the processed and failed model sets from progress file. @@ -732,11 +742,13 @@ class DownloadManager: success, is_stale, failed_images, + rate_limited_images, ) = await ExampleImagesProcessor.download_model_images_with_tracking( model_hash, model_name, images, model_dir, optimize, downloader ) failed_urls: Set[str] = set(failed_images) + rate_limited_urls: Set[str] = set(rate_limited_images) # If metadata is stale, try to refresh it if is_stale and model_hash not in self._progress["refreshed_models"]: @@ -760,6 +772,7 @@ class DownloadManager: success, _, additional_failed, + additional_rate_limited, ) = await ExampleImagesProcessor.download_model_images_with_tracking( model_hash, model_name, @@ -770,29 +783,50 @@ class DownloadManager: ) failed_urls.update(additional_failed) + rate_limited_urls.update(additional_rate_limited) self._progress["refreshed_models"].add(model_hash) - if failed_urls: + # Separate permanent failures from rate-limited ones + permanent_failures = failed_urls - rate_limited_urls + + if permanent_failures: await self._remove_failed_images_from_metadata( model_hash, model_name, model_dir, - failed_urls, + permanent_failures, scanner, ) - if failed_urls: + if rate_limited_urls: + self._progress["rate_limited_models"].add(model_hash) + logger.warning( + "%d example images for %s are rate-limited (429), will retry next time", + len(rate_limited_urls), + model_name, + ) + # Clear failed_models so non-force runs can retry + if force and model_hash in self._progress["failed_models"]: + self._progress["failed_models"].discard(model_hash) + logger.info( + f"Removed {model_name} from failed_models after force retry with rate-limited images" + ) + + if rate_limited_urls: + # Don't mark as failed or fully processed — rate-limited + # images will be retried next time. + pass + elif permanent_failures: self._progress["failed_models"].add(model_hash) self._progress["processed_models"].add(model_hash) logger.info( "Removed %s failed example images for %s", - len(failed_urls), + len(permanent_failures), model_name, ) elif success: self._progress["processed_models"].add(model_hash) - # Remove from failed_models if force mode enabled and model was previously failed if force and model_hash in self._progress["failed_models"]: self._progress["failed_models"].discard(model_hash) logger.info( @@ -850,6 +884,7 @@ class DownloadManager: "processed_models": list(self._progress["processed_models"]), "refreshed_models": list(self._progress["refreshed_models"]), "failed_models": list(self._progress["failed_models"]), + "rate_limited_models": list(self._progress.get("rate_limited_models", set())), "completed": self._progress["completed"], "total": self._progress["total"], "last_update": time.time(), @@ -1155,11 +1190,13 @@ class DownloadManager: success, is_stale, failed_images, + rate_limited_images, ) = await ExampleImagesProcessor.download_model_images_with_tracking( model_hash, model_name, images, model_dir, optimize, downloader ) failed_urls: Set[str] = set(failed_images) + rate_limited_urls: Set[str] = set(rate_limited_images) # If metadata is stale, try to refresh it if is_stale and model_hash not in self._progress["refreshed_models"]: @@ -1183,6 +1220,7 @@ class DownloadManager: success, _, additional_failed_images, + additional_rate_limited, ) = await ExampleImagesProcessor.download_model_images_with_tracking( model_hash, model_name, @@ -1192,21 +1230,35 @@ class DownloadManager: downloader, ) - # Combine failed images from both attempts failed_urls.update(additional_failed_images) + rate_limited_urls.update(additional_rate_limited) self._progress["refreshed_models"].add(model_hash) - # For forced downloads, remove failed images from metadata - if failed_urls: + # Separate permanent failures from rate-limited ones + permanent_failures = failed_urls - rate_limited_urls + + # Only remove permanently failed images from metadata + if permanent_failures: await self._remove_failed_images_from_metadata( - model_hash, model_name, model_dir, failed_urls, scanner + model_hash, model_name, model_dir, permanent_failures, scanner ) - # Mark as processed - if ( - success or failed_urls - ): # Mark as processed if we successfully downloaded some images or removed failed ones + if rate_limited_urls: + self._progress["rate_limited_models"].add(model_hash) + logger.warning( + "%d example images for %s are rate-limited (429), will retry next time", + len(rate_limited_urls), + model_name, + ) + + # Mark as processed only when no rate-limited images remain + if rate_limited_urls: + pass + elif permanent_failures: + self._progress["processed_models"].add(model_hash) + self._progress["failed_models"].add(model_hash) + elif success: self._progress["processed_models"].add(model_hash) return True # Return True to indicate a remote download happened @@ -1229,15 +1281,20 @@ class DownloadManager: model_dir: str, failed_images: Iterable[str], scanner, + error_type: str = "not_found", ) -> None: - """Mark failed images in model metadata so they won't be retried.""" + """Mark failed images in model metadata so they won't be retried. + + Args: + error_type: Reason string stored in the image's ``downloadError`` field + (default ``"not_found"``). + """ failed_set: Set[str] = {url for url in failed_images if url} if not failed_set: return try: - # Get current model data model_data = await MetadataUpdater.get_updated_model(model_hash, scanner) if not model_data: logger.warning( @@ -1268,7 +1325,7 @@ class DownloadManager: continue image["downloadFailed"] = True - image.setdefault("downloadError", "not_found") + image.setdefault("downloadError", error_type) logger.debug( "Marked example image %s for %s as failed due to missing remote asset", image_url, diff --git a/py/utils/example_images_processor.py b/py/utils/example_images_processor.py index 101f239d..da617706 100644 --- a/py/utils/example_images_processor.py +++ b/py/utils/example_images_processor.py @@ -1,3 +1,4 @@ +import asyncio import logging import os import re @@ -194,16 +195,22 @@ class ExampleImagesProcessor: return model_success, False # (success, is_metadata_stale) + @staticmethod + def _extract_retry_after(error_message: str) -> int: + if not error_message: + return 60 + match = re.search(r"retry after (\d+)s", str(error_message)) + if match: + return max(1, int(match.group(1))) + return 60 + @staticmethod async def download_model_images_with_tracking(model_hash, model_name, model_images, model_dir, optimize, downloader): - """Download images for a single model with tracking of failed image URLs - - Returns: - tuple: (success, is_stale_metadata, failed_images) - whether download was successful, whether metadata is stale, list of failed image URLs - """ model_success = True failed_images = [] - + rate_limited_images = [] + any_successful_download = False + for i, image in enumerate(model_images): image_url = image.get('url') if not image_url: @@ -221,64 +228,110 @@ class ExampleImagesProcessor: original_url = image_url if optimize and 'civitai.com' in image_url: image_url = ExampleImagesProcessor.get_civitai_optimized_url(image_url) - - # Download the file first to determine the actual file type - try: - logger.debug(f"Downloading media file {i} for {model_name}") - - # Download using the unified downloader with headers - success, content, headers = await downloader.download_to_memory( + + async def _attempt_download() -> tuple: + logger.debug("Downloading media file %s for %s", i, model_name) + return await downloader.download_to_memory( image_url, - use_auth=False, # Example images don't need auth - return_headers=True + use_auth=False, + return_headers=True, ) - + + try: + success, content, headers = await _attempt_download() + if success: - # Determine file extension from content or headers media_ext = ExampleImagesProcessor._get_file_extension_from_content_or_headers( content, headers, original_url, image.get("type") ) - - # Check if the detected file type is supported is_image = media_ext in SUPPORTED_MEDIA_EXTENSIONS['images'] is_video = media_ext in SUPPORTED_MEDIA_EXTENSIONS['videos'] - + if not (is_image or is_video): - logger.debug(f"Skipping unsupported file type: {media_ext}") + logger.debug("Skipping unsupported file type: %s", media_ext) continue - - # Use 0-based indexing with the detected extension + save_filename = f"image_{i}{media_ext}" save_path = os.path.join(model_dir, save_filename) - - # Check if already downloaded + if os.path.exists(save_path): - logger.debug(f"File already exists: {save_path}") + logger.debug("File already exists: %s", save_path) continue - - # Save the file + with open(save_path, 'wb') as f: f.write(content) - + any_successful_download = True + elif ExampleImagesProcessor._is_not_found_error(content): error_msg = f"Failed to download file: {image_url}, status code: 404 - Model metadata might be stale" logger.warning(error_msg) - model_success = False # Mark the model as failed due to 404 error - failed_images.append(image_url) # Track failed URL - # Return early to trigger metadata refresh attempt - return False, True, failed_images # (success, is_metadata_stale, failed_images) + model_success = False + failed_images.append(image_url) + return False, True, failed_images, rate_limited_images + + elif "Rate limited (429)" in str(content): + max_attempts = 3 + for attempt in range(1, max_attempts + 1): + wait = ExampleImagesProcessor._extract_retry_after(str(content)) * (2 ** (attempt - 1)) + logger.warning( + "Rate limited (429) for %s, retry %d/%d after %ds", + image_url, attempt, max_attempts, wait, + ) + await asyncio.sleep(wait) + + success, content, headers = await _attempt_download() + if success: + media_ext = ExampleImagesProcessor._get_file_extension_from_content_or_headers( + content, headers, original_url, image.get("type") + ) + is_image = media_ext in SUPPORTED_MEDIA_EXTENSIONS['images'] + is_video = media_ext in SUPPORTED_MEDIA_EXTENSIONS['videos'] + + if not (is_image or is_video): + logger.debug("Skipping unsupported file type: %s", media_ext) + break + + save_filename = f"image_{i}{media_ext}" + save_path = os.path.join(model_dir, save_filename) + if os.path.exists(save_path): + logger.debug("File already exists: %s", save_path) + break + + with open(save_path, 'wb') as f: + f.write(content) + any_successful_download = True + break + elif "Rate limited (429)" in str(content): + continue + elif ExampleImagesProcessor._is_not_found_error(content): + logger.warning("Failed to download file: %s, status code: 404", image_url) + model_success = False + failed_images.append(image_url) + break + else: + logger.warning("Failed to download file: %s, error: %s", image_url, content) + model_success = False + failed_images.append(image_url) + break + else: + logger.warning( + "Giving up on %s after %d retries due to rate limiting", + image_url, max_attempts, + ) + rate_limited_images.append(image_url) + model_success = False else: error_msg = f"Failed to download file: {image_url}, error: {content}" logger.warning(error_msg) - model_success = False # Mark the model as failed - failed_images.append(image_url) # Track failed URL + model_success = False + failed_images.append(image_url) except Exception as e: error_msg = f"Error downloading file {image_url}: {str(e)}" logger.error(error_msg) - model_success = False # Mark the model as failed - failed_images.append(image_url) # Track failed URL - - return model_success, False, failed_images # (success, is_metadata_stale, failed_images) + model_success = False + failed_images.append(image_url) + + return any_successful_download or model_success, False, failed_images, rate_limited_images @staticmethod async def process_local_examples(model_file_path, model_file_name, model_name, model_dir, optimize): diff --git a/tests/services/test_example_images_download_manager_async.py b/tests/services/test_example_images_download_manager_async.py index 8ccbe48d..40de65b6 100644 --- a/tests/services/test_example_images_download_manager_async.py +++ b/tests/services/test_example_images_download_manager_async.py @@ -164,7 +164,7 @@ async def test_pause_resume_blocks_processing( await first_release.wait() else: second_call_started.set() - return True, False, [] + return True, False, [], [] async def fake_get_downloader(): class _Downloader: @@ -288,7 +288,7 @@ async def test_legacy_folder_migrated_and_skipped( async def fake_download_model_images(*_args, **_kwargs): nonlocal download_called download_called = True - return True, False, [] + return True, False, [], [] async def fake_get_downloader(): class _Downloader: diff --git a/tests/services/test_issue_760_repro.py b/tests/services/test_issue_760_repro.py index 174e25fd..be8fa4ed 100644 --- a/tests/services/test_issue_760_repro.py +++ b/tests/services/test_issue_760_repro.py @@ -77,7 +77,7 @@ async def test_reprocessing_triggered_when_folder_missing(monkeypatch, tmp_path) model_dir = args[3] Path(model_dir).mkdir(parents=True, exist_ok=True) (Path(model_dir) / "image_0.png").write_text("fixed") - return True, False, [] + return True, False, [], [] monkeypatch.setattr(download_module.ExampleImagesProcessor, "download_model_images_with_tracking", fake_download_model_images)