feat(downloads): return structured 429 rate-limit responses with retry_after

On a CivitAI/CivArchive 429, download-model and download-model-get now
return HTTP 429 with {"reason": "rate_limited", "retry_after": N}
instead of a generic 500 string, and the queue row goes back to
"queued" rather than history as failed — so queue drivers can
auto-pause and retry later instead of burning through the queue.

- new DownloadRateLimitError carrying retry_after/host (opt-in via
  raise_on_rate_limit on Downloader; other call sites keep the legacy
  string behavior)
- fail-fast pre-flight gate in DownloadManager consults
  RateLimitCoordinator before acquiring the semaphore slot: hosts in
  cooldown get an immediate structured 429, no HTTP request attempted
- best-effort 429 detection for the aria2 backend
This commit is contained in:
Will Miao
2026-10-02 21:16:06 +08:00
parent 2a667df98c
commit 034660d8c4
8 changed files with 627 additions and 9 deletions
+33 -1
View File
@@ -31,7 +31,7 @@ from .connectivity_guard import (
OFFLINE_FRIENDLY_MESSAGE,
ConnectivityGuard,
)
from .errors import RateLimitError
from .errors import DownloadRateLimitError, RateLimitError
from .rate_limit_coordinator import RateLimitCoordinator
logger = logging.getLogger(__name__)
@@ -434,6 +434,7 @@ class Downloader:
custom_headers: Optional[Dict[str, str]] = None,
allow_resume: bool = True,
pause_event: Optional[DownloadStreamControl] = None,
raise_on_rate_limit: bool = False,
) -> Tuple[bool, str]:
"""
Download a file with resumable downloads and retry mechanism
@@ -446,6 +447,11 @@ class Downloader:
custom_headers: Additional headers to include in request
allow_resume: Whether to support resumable downloads
pause_event: Optional stream control used to pause/resume and request reconnects
raise_on_rate_limit: When True, a 429 response raises
``DownloadRateLimitError`` instead of returning a plain error
string, so callers (the model download manager) can build the
structured rate-limit result required by the download queue
contract. Defaults to the legacy tuple behavior.
Returns:
Tuple[bool, str]: (success, save_path or error message)
@@ -610,6 +616,12 @@ class Downloader:
logger.warning(
f"Rate limited (429) for {url}, retry_after={retry_after}"
)
if raise_on_rate_limit:
raise DownloadRateLimitError(
f"Download rate limited (429), retry after {retry_after}s",
retry_after=retry_after,
host=self._guard_destination(url),
)
return False, f"Download rate limited (429), retry after {retry_after}s"
else:
logger.error(
@@ -902,6 +914,11 @@ class Downloader:
f"Network error after {self.max_retries + 1} attempts: {str(e)}",
)
except DownloadRateLimitError:
# 429s are never retried in-band; the structured error must
# reach the caller (download manager) unmodified.
raise
except Exception as e:
logger.error(f"Unexpected download error: {e}")
return False, str(e)
@@ -931,6 +948,7 @@ class Downloader:
use_auth: bool = False,
custom_headers: Optional[Dict[str, str]] = None,
return_headers: bool = False,
raise_on_rate_limit: bool = False,
) -> Tuple[bool, Union[bytes, str], Optional[Dict[str, Any]]]:
"""
Download a file to memory (for small files like preview images)
@@ -940,6 +958,10 @@ class Downloader:
use_auth: Whether to include authentication headers
custom_headers: Additional headers to include in request
return_headers: Whether to return response headers along with content
raise_on_rate_limit: When True, a 429 response raises
``DownloadRateLimitError`` instead of returning a plain error
string (see ``download_file``). Defaults to the legacy tuple
behavior.
Returns:
Tuple[bool, Union[bytes, str], Optional[Dict]]: (success, content or error message, response headers if requested)
@@ -1002,11 +1024,21 @@ class Downloader:
"Rate limited (429) for %s, no Retry-After header; defaulting to %ss",
url, retry_after,
)
if raise_on_rate_limit:
raise DownloadRateLimitError(
f"Rate limited (429), retry after {retry_after}s",
retry_after=retry_after,
host=destination,
)
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
except DownloadRateLimitError:
# Structured rate-limit errors must reach the caller unmodified.
raise
except Exception as e:
if guard.is_network_unreachable_error(e):
guard.register_network_failure(e, destination)