perf(example-images): parallelize downloads and recover transient failures
- Add a shared MediaDownloadLimiter (default 6 concurrent) with a per-run 429 cooldown and AIMD recovery; drop the empirical inter-model delay. - Request width-aware Civitai renditions (800px images, transcoded 450px videos) and map Genur signed/direct URLs to c.genur.art _small.webp or _large.webp. - Return an ImageDownloadOutcome so offline/connect/timeout/5xx failures stay transient instead of being written as downloadFailed; track retryable_models and leave those models pending. - Force retries flagged Genur images and clears downloadFailed/downloadError only after a successful write, through the normal metadata save/cache path. - Force bypasses the populated-folder guard; success is false when flagged images were skipped, so models are no longer falsely marked recovered. - Add settings, tests, and a design/measurement note.
This commit is contained in:
@@ -0,0 +1,133 @@
|
||||
# Example-image download throughput
|
||||
|
||||
## Problem
|
||||
|
||||
Downloading example media for a large library was effectively serial: one
|
||||
Civitai CDN request per image/video, one model after another, plus an
|
||||
empirical `delay` (0.2 s) between models. A model with 10–20 gallery items
|
||||
therefore paid 10–20 network round-trips before the next model started.
|
||||
|
||||
The inter-model delay is not the cost. In a 10-image probe it accounted for
|
||||
about 2 s of a 7 s run; the remaining time was the per-image round-trips.
|
||||
|
||||
## Is there a batch endpoint?
|
||||
|
||||
No. Civitai's public API exposes:
|
||||
|
||||
- `GET /api/v1/images?ids=...` — **metadata** for up to 100 images in one
|
||||
request (`IMAGE_IDS_BATCH_MAX = 100`), not the binaries.
|
||||
- One model-version response that already contains the image URLs the
|
||||
downloader uses.
|
||||
- One binary per `image.civitai.com` URL. Each URL 301-redirects to a
|
||||
`blobs-b2.civitai.com` object.
|
||||
|
||||
The image cacher also has no multi-URL delivery endpoint. The practical
|
||||
"batch" is therefore client-side pipelining: bounded concurrent requests, a
|
||||
rendition that is cheaper to transfer, and a shared cooldown when the CDN
|
||||
returns 429.
|
||||
|
||||
## Measurements (single sample, 2026-11)
|
||||
|
||||
Ten images from one model, cold process, Civitai CDN over the local link:
|
||||
|
||||
| Strategy | Wall time | Bytes |
|
||||
| --- | ---: | ---: |
|
||||
| Old: serial, `/optimized=true` | 7.07 s | 1669 KiB |
|
||||
| New: 4 concurrent, `width=800` | 3.08 s | 834 KiB |
|
||||
| New: 8 concurrent, `width=800` | 1.62 s | 834 KiB |
|
||||
| New: 6 concurrent (default), `width=800` | 3.23 s | 834 KiB |
|
||||
|
||||
A 720×924 sample video:
|
||||
|
||||
| Rendition | Size |
|
||||
| --- | ---: |
|
||||
| Old: `/optimized=true` (original MP4) | 73.4 MB |
|
||||
| New: `transcode=true,width=450,optimized=true` | 27.0 MB |
|
||||
|
||||
The video case was the largest single regression in the old path: the URL
|
||||
helper only appended `optimized=true`, which the cacher serves as the original
|
||||
video.
|
||||
|
||||
## Design
|
||||
|
||||
- `MediaDownloadLimiter` (`py/utils/example_images_processor.py`):
|
||||
- bounds concurrent requests (default 6),
|
||||
- shares a 429 cooldown across every worker instead of letting each task
|
||||
sleep and retry independently,
|
||||
- halves the limit on a 429 and adds one slot back after a run of successes
|
||||
(AIMD), so a sustained limit settles below the threshold.
|
||||
- `download_model_images_with_tracking` builds one task per pending item and
|
||||
runs them through the limiter. The task result still drives the existing
|
||||
per-model bookkeeping (`failed_models`, `rate_limited_models`, metadata
|
||||
refresh on 404).
|
||||
- URLs are built with `build_civitai_media_url`:
|
||||
- images: `width=800,optimized=true` (the filename is preserved),
|
||||
- videos: `transcode=true,width=450,optimized=true`,
|
||||
- non-Civitai URLs: unchanged,
|
||||
- `optimize=False`: the original URL.
|
||||
- The legacy inter-model `delay` defaults to 0. API callers that still send it
|
||||
keep the old behavior.
|
||||
- Failure lists now store the *original* metadata URL, which is what
|
||||
`_remove_failed_images_from_metadata` matches against.
|
||||
|
||||
## Settings (`DEFAULT_SETTINGS`)
|
||||
|
||||
| Key | Default | Meaning |
|
||||
| --- | ---: | --- |
|
||||
| `example_images_download_concurrency` | 6 | Simultaneous CDN requests per run |
|
||||
| `example_images_max_width` | 800 | Image rendition width |
|
||||
| `example_images_video_max_width` | 450 | Video transcode width |
|
||||
| `example_images_rate_limit_cooldown` | 30 | Fallback cooldown when `Retry-After` is absent |
|
||||
|
||||
Widths are chosen from the image cacher's common-size ladder
|
||||
(96/320/450/512/800/1200/1600/2200); off-ladder widths are snapped up by the
|
||||
cacher, so they do not buy a smaller transfer.
|
||||
|
||||
## Follow-up: transient failures and Genur recovery
|
||||
|
||||
The first implementation still treated *any* non-404 failure as permanent. A
|
||||
network outage (or the ConnectivityGuard's shared cooldown message) was
|
||||
written to the sidecar as `downloadFailed: true, downloadError: "not_found"`,
|
||||
which made a transient failure permanent and produced empty example folders
|
||||
that were reprocessed every run. Genur was the visible victim: its signed
|
||||
proxy host (`img.genur.art`) can be unreachable (fake-IP DNS / TLS reset)
|
||||
while the direct host serves the same asset.
|
||||
|
||||
The follow-up changes:
|
||||
|
||||
- `download_model_images_with_tracking` returns an `ImageDownloadOutcome`
|
||||
(`downloaded`, `failed`, `rate_limited`, `transient`, `recovered`,
|
||||
`skipped_failed`) instead of a tuple. Offline/connect/timeout/5xx errors go
|
||||
to `transient` and are **not** written to the sidecar; the model is left
|
||||
pending (tracked in `retryable_models`) so the next run retries it.
|
||||
- `retry_failed=True` (force / explicit targets) re-attempts Genur entries
|
||||
that carry `downloadFailed`; the flag is cleared only after a file is
|
||||
actually written, through `_clear_failed_images_from_metadata` and the
|
||||
normal `save_metadata` + cache-update path. No bulk sidecar rewrite is
|
||||
needed.
|
||||
- Genur signed URLs are resolved at download time:
|
||||
`img.genur.art/sig/.../<base64>` → `c.genur.art/<uuid>_small.webp`
|
||||
(`_large.webp` when optimization is disabled). Direct
|
||||
`c.genur.art/<uuid>` URLs get the same `_small.webp` suffix. The direct
|
||||
renditions return 200 for the sampled Organic Sauce / Better Lego assets.
|
||||
- Force now bypasses the model-level "populated folder" guard, so flagged
|
||||
images in partially populated folders can be retried; per-image existence
|
||||
checks still skip files that are already there.
|
||||
|
||||
Existing affected libraries recover by restarting the instance and running a
|
||||
force download (global or per-model). The 398 Genur entries counted in the
|
||||
user's sidecars are eligible; Civitai 404 entries stay flagged because those
|
||||
assets are genuinely gone.
|
||||
|
||||
## Future work
|
||||
|
||||
- Pipeline models, not just images within a model: start model N+1's image
|
||||
jobs while model N's metadata/failure bookkeeping finishes. This helps
|
||||
models with only a few gallery items.
|
||||
- Stream media to a temp file instead of `download_to_memory`; the current
|
||||
pool bounds memory at `concurrency × item size`, which is fine for
|
||||
transcoded videos but large if `optimize_example_images` is disabled.
|
||||
- Cache the `image.civitai.com` → `blobs-b2.civitai.com` redirect when a
|
||||
rendition is retried, to save one hop per retry.
|
||||
- Expose the four settings in the settings UI if users need to tune them
|
||||
without editing `settings.json`.
|
||||
@@ -147,6 +147,15 @@ DEFAULT_SETTINGS: Dict[str, Any] = {
|
||||
"example_images_local_root": "",
|
||||
"example_images_open_uri_template": "",
|
||||
"optimize_example_images": True,
|
||||
# Example-media downloads are one CDN request per image/video. Serial
|
||||
# requests spend most of their wall time waiting on round-trips, so a
|
||||
# small bounded pool plus a shared 429 cooldown is the throughput lever.
|
||||
# Widths follow the Civitai image-cacher's common-size ladder; videos are
|
||||
# capped lower because a video rendition is orders of magnitude larger.
|
||||
"example_images_download_concurrency": 6,
|
||||
"example_images_max_width": 800,
|
||||
"example_images_video_max_width": 450,
|
||||
"example_images_rate_limit_cooldown": 30,
|
||||
"auto_download_example_images": False,
|
||||
"blur_mature_content": True,
|
||||
"mature_blur_level": "R",
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import posixpath
|
||||
import re
|
||||
from typing import Any, Dict, Iterable, Mapping, Sequence
|
||||
from urllib.parse import parse_qs, urlparse, urlunparse
|
||||
@@ -338,14 +339,40 @@ def resolve_license_info(
|
||||
return payload, build_license_flags(payload)
|
||||
|
||||
|
||||
def _normalize_media_width(width: int | None) -> int | None:
|
||||
if width is None:
|
||||
return None
|
||||
try:
|
||||
return max(1, int(width))
|
||||
except (TypeError, ValueError):
|
||||
return 450
|
||||
|
||||
|
||||
def _media_rendition_params(media_type: str | None, width: int | None) -> str:
|
||||
"""Return the Civitai image-cacher params for a media rendition."""
|
||||
|
||||
is_video = (media_type or "").lower() == "video"
|
||||
normalized_width = _normalize_media_width(width)
|
||||
if normalized_width is None:
|
||||
return "transcode=true,optimized=true" if is_video else "optimized=true"
|
||||
if is_video:
|
||||
return f"transcode=true,width={normalized_width},optimized=true"
|
||||
return f"width={normalized_width},optimized=true"
|
||||
|
||||
|
||||
def rewrite_preview_url(
|
||||
source_url: str | None, media_type: str | None = None
|
||||
source_url: str | None,
|
||||
media_type: str | None = None,
|
||||
*,
|
||||
width: int | None = 450,
|
||||
) -> tuple[str | None, bool]:
|
||||
"""Rewrite Civitai preview URLs to use optimized renditions.
|
||||
|
||||
Args:
|
||||
source_url: Original preview URL from the Civitai API.
|
||||
media_type: Optional media type hint (e.g. ``"image"`` or ``"video"``).
|
||||
width: Rendition width. ``None`` keeps the legacy width-less
|
||||
``optimized=true`` form.
|
||||
|
||||
Returns:
|
||||
A tuple of the potentially rewritten URL and a flag indicating whether the
|
||||
@@ -368,9 +395,7 @@ def rewrite_preview_url(
|
||||
if hostname == "civitai.com" or not hostname.endswith(".civitai.com"):
|
||||
return source_url, False
|
||||
|
||||
replacement = "/width=450,optimized=true"
|
||||
if (media_type or "").lower() == "video":
|
||||
replacement = "/transcode=true,width=450,optimized=true"
|
||||
replacement = f"/{_media_rendition_params(media_type, width)}"
|
||||
|
||||
if "/original=true" not in parsed.path:
|
||||
return source_url, False
|
||||
@@ -383,8 +408,67 @@ def rewrite_preview_url(
|
||||
return rewritten, True
|
||||
|
||||
|
||||
def build_civitai_media_url(
|
||||
source_url: str | None,
|
||||
media_type: str | None = None,
|
||||
*,
|
||||
width: int | None = 450,
|
||||
optimize: bool = True,
|
||||
) -> str | None:
|
||||
"""Build the image-cacher delivery URL for a Civitai media asset.
|
||||
|
||||
Unlike :func:`rewrite_preview_url`, this handles URLs that do not carry the
|
||||
current ``/original=true`` params segment (older cached metadata, legacy
|
||||
``.../<uuid>`` URLs) and replaces an existing rendition segment instead of
|
||||
appending a second one. The filename is preserved: the image cacher uses it
|
||||
to derive the output name, and dropping it makes the requested rendition
|
||||
ambiguous.
|
||||
"""
|
||||
|
||||
if not source_url or not optimize:
|
||||
return source_url
|
||||
|
||||
try:
|
||||
parsed = urlparse(source_url)
|
||||
except ValueError:
|
||||
return source_url
|
||||
|
||||
hostname = parsed.hostname
|
||||
if hostname is None:
|
||||
return source_url
|
||||
hostname = hostname.lower()
|
||||
if hostname == "civitai.com" or not hostname.endswith(".civitai.com"):
|
||||
return source_url
|
||||
|
||||
rewritten, changed = rewrite_preview_url(source_url, media_type, width=width)
|
||||
if changed:
|
||||
return rewritten
|
||||
|
||||
segments = [segment for segment in parsed.path.split("/") if segment]
|
||||
if not segments:
|
||||
return source_url
|
||||
|
||||
params = _media_rendition_params(media_type, width)
|
||||
last = segments[-1]
|
||||
if "=" in last:
|
||||
# Legacy ``.../<uuid>/optimized=true`` form (no filename).
|
||||
segments[-1] = params
|
||||
elif len(segments) >= 2 and "=" in segments[-2]:
|
||||
# The rendition segment is immediately before the filename.
|
||||
segments[-2] = params
|
||||
elif posixpath.splitext(last)[1]:
|
||||
# ``.../<uuid>/<filename>``: insert the rendition before the filename.
|
||||
segments.insert(len(segments) - 1, params)
|
||||
else:
|
||||
# ``.../<uuid>``: no filename to preserve.
|
||||
segments.append(params)
|
||||
|
||||
return urlunparse(parsed._replace(path="/" + "/".join(segments)))
|
||||
|
||||
|
||||
__all__ = [
|
||||
"build_civitai_image_page_url",
|
||||
"build_civitai_media_url",
|
||||
"build_license_flags",
|
||||
"extract_civitai_image_id",
|
||||
"extract_civitai_image_id_from_cdn_url",
|
||||
|
||||
@@ -20,7 +20,7 @@ from ..utils.example_images_paths import (
|
||||
uses_library_scoped_folders,
|
||||
)
|
||||
from ..utils.metadata_manager import MetadataManager
|
||||
from .example_images_processor import ExampleImagesProcessor
|
||||
from .example_images_processor import ExampleImagesProcessor, MediaDownloadLimiter
|
||||
from .example_images_metadata import (
|
||||
MetadataUpdater,
|
||||
update_cache_from_metadata,
|
||||
@@ -79,6 +79,7 @@ class _DownloadProgress(dict[str, Any]):
|
||||
failed_models=set(),
|
||||
reprocessed_models=set(),
|
||||
rate_limited_models=set(),
|
||||
retryable_models=set(),
|
||||
)
|
||||
|
||||
def snapshot(self) -> Dict[str, Any]:
|
||||
@@ -90,6 +91,7 @@ class _DownloadProgress(dict[str, Any]):
|
||||
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()))
|
||||
snapshot["retryable_models"] = list(self.get("retryable_models", set()))
|
||||
return snapshot
|
||||
|
||||
|
||||
@@ -163,6 +165,38 @@ class DownloadManager:
|
||||
return ""
|
||||
return ensure_library_root_exists(library_name)
|
||||
|
||||
def _create_media_download_limiter(
|
||||
self, optimize: bool = True
|
||||
) -> MediaDownloadLimiter:
|
||||
"""Build the run-wide limiter from settings.
|
||||
|
||||
The limiter is shared by every model in a run so a 429 seen while
|
||||
downloading one model also pauses the next one instead of being
|
||||
forgotten at the model boundary.
|
||||
"""
|
||||
|
||||
settings = get_settings_manager()
|
||||
try:
|
||||
concurrency = int(
|
||||
settings.get("example_images_download_concurrency", 6)
|
||||
)
|
||||
except (TypeError, ValueError):
|
||||
concurrency = 6
|
||||
try:
|
||||
cooldown = float(
|
||||
settings.get("example_images_rate_limit_cooldown", 30)
|
||||
)
|
||||
except (TypeError, ValueError):
|
||||
cooldown = 30.0
|
||||
if not optimize:
|
||||
# Original renditions can be tens of megabytes per video; keep the
|
||||
# memory/bandwidth spike bounded when optimization is disabled.
|
||||
concurrency = min(concurrency, 2)
|
||||
return MediaDownloadLimiter(
|
||||
max_concurrency=max(1, concurrency),
|
||||
default_cooldown=max(1.0, cooldown),
|
||||
)
|
||||
|
||||
async def start_download(self, options: Dict[str, Any]):
|
||||
"""Start downloading example images for models."""
|
||||
|
||||
@@ -171,7 +205,7 @@ class DownloadManager:
|
||||
auto_mode = data.get("auto_mode", False)
|
||||
optimize = data.get("optimize", True)
|
||||
model_types = data.get("model_types", ["lora", "checkpoint"])
|
||||
delay = float(data.get("delay", 0.2))
|
||||
delay = float(data.get("delay", 0) or 0)
|
||||
force = data.get("force", False)
|
||||
model_hashes = data.get("model_hashes", [])
|
||||
|
||||
@@ -589,6 +623,7 @@ class DownloadManager:
|
||||
"""Download example images for all models (or only the given hashes)."""
|
||||
|
||||
downloader = await get_downloader()
|
||||
limiter = self._create_media_download_limiter(optimize)
|
||||
|
||||
try:
|
||||
# Get scanners
|
||||
@@ -654,6 +689,7 @@ class DownloadManager:
|
||||
library_name,
|
||||
force,
|
||||
explicit_targets,
|
||||
limiter=limiter,
|
||||
)
|
||||
|
||||
# Update progress
|
||||
@@ -671,9 +707,12 @@ class DownloadManager:
|
||||
if should_stop:
|
||||
break
|
||||
|
||||
# Only add delay after remote download of models, and not after processing the last model
|
||||
# Pacing is handled by the shared media-download limiter. The
|
||||
# legacy ``delay`` option is kept for API callers that still
|
||||
# send it, but defaults to 0.
|
||||
if (
|
||||
was_remote_download
|
||||
delay > 0
|
||||
and was_remote_download
|
||||
and i < len(all_models) - 1
|
||||
and current_status == "running"
|
||||
):
|
||||
@@ -751,6 +790,7 @@ class DownloadManager:
|
||||
library_name,
|
||||
force: bool = False,
|
||||
explicit_targets: bool = False,
|
||||
limiter: MediaDownloadLimiter | None = None,
|
||||
):
|
||||
"""Process a single model download."""
|
||||
|
||||
@@ -784,10 +824,12 @@ class DownloadManager:
|
||||
)
|
||||
existing_files = _model_directory_has_files(model_dir)
|
||||
|
||||
# Model-level guard: a populated folder counts as done. Explicitly
|
||||
# targeted models bypass it so the per-image existence pre-check can
|
||||
# fill individual gaps without re-fetching existing files.
|
||||
if not explicit_targets:
|
||||
# Model-level guard: a populated folder counts as done. Force and
|
||||
# explicitly targeted models bypass it so the per-image existence
|
||||
# pre-check can fill individual gaps, retry flagged images, or
|
||||
# clear a stale ``downloadFailed`` marker without re-fetching the
|
||||
# files that are already present.
|
||||
if not force and not explicit_targets:
|
||||
# Skip if already processed AND directory exists with files
|
||||
if model_hash in self._progress["processed_models"]:
|
||||
if existing_files:
|
||||
@@ -846,20 +888,30 @@ class DownloadManager:
|
||||
if civitai_payload.get("images"):
|
||||
images = civitai_payload.get("images", [])
|
||||
|
||||
(
|
||||
success,
|
||||
is_stale,
|
||||
failed_images,
|
||||
rate_limited_images,
|
||||
) = await ExampleImagesProcessor.download_model_images_with_tracking(
|
||||
model_hash, model_name, images, model_dir, optimize, downloader
|
||||
outcome = await ExampleImagesProcessor.download_model_images_with_tracking(
|
||||
model_hash,
|
||||
model_name,
|
||||
images,
|
||||
model_dir,
|
||||
optimize,
|
||||
downloader,
|
||||
limiter=limiter,
|
||||
retry_failed=bool(force or explicit_targets),
|
||||
)
|
||||
|
||||
failed_urls: Set[str] = set(failed_images)
|
||||
rate_limited_urls: Set[str] = set(rate_limited_images)
|
||||
failed_urls: Set[str] = set(outcome.failed)
|
||||
rate_limited_urls: Set[str] = set(outcome.rate_limited)
|
||||
transient_urls: Set[str] = set(outcome.transient)
|
||||
recovered_urls: Set[str] = set(outcome.recovered)
|
||||
final_success = outcome.success
|
||||
|
||||
if recovered_urls:
|
||||
await self._clear_failed_images_from_metadata(
|
||||
model_hash, model_name, recovered_urls, scanner
|
||||
)
|
||||
|
||||
# If metadata is stale, try to refresh it
|
||||
if is_stale and model_hash not in self._progress["refreshed_models"]:
|
||||
if outcome.is_stale and model_hash not in self._progress["refreshed_models"]:
|
||||
await MetadataUpdater.refresh_model_metadata(
|
||||
model_hash, model_name, scanner_type, scanner, self._progress
|
||||
)
|
||||
@@ -876,26 +928,35 @@ class DownloadManager:
|
||||
if updated_civitai.get("images"):
|
||||
# Retry download with updated metadata
|
||||
updated_images = updated_civitai.get("images", [])
|
||||
(
|
||||
success,
|
||||
_,
|
||||
additional_failed,
|
||||
additional_rate_limited,
|
||||
) = await ExampleImagesProcessor.download_model_images_with_tracking(
|
||||
model_hash,
|
||||
model_name,
|
||||
updated_images,
|
||||
model_dir,
|
||||
optimize,
|
||||
downloader,
|
||||
retry_outcome = (
|
||||
await ExampleImagesProcessor.download_model_images_with_tracking(
|
||||
model_hash,
|
||||
model_name,
|
||||
updated_images,
|
||||
model_dir,
|
||||
optimize,
|
||||
downloader,
|
||||
limiter=limiter,
|
||||
retry_failed=bool(force or explicit_targets),
|
||||
)
|
||||
)
|
||||
|
||||
failed_urls.update(additional_failed)
|
||||
rate_limited_urls.update(additional_rate_limited)
|
||||
# The refreshed metadata is now the source of truth:
|
||||
# failures against the pre-refresh URLs are gone (or
|
||||
# will fail again on this pass).
|
||||
failed_urls = set(retry_outcome.failed)
|
||||
rate_limited_urls = set(retry_outcome.rate_limited)
|
||||
transient_urls = set(retry_outcome.transient)
|
||||
recovered_urls.update(retry_outcome.recovered)
|
||||
final_success = retry_outcome.success
|
||||
if retry_outcome.recovered:
|
||||
await self._clear_failed_images_from_metadata(
|
||||
model_hash, model_name, retry_outcome.recovered, scanner
|
||||
)
|
||||
|
||||
self._progress["refreshed_models"].add(model_hash)
|
||||
|
||||
# Separate permanent failures from rate-limited ones
|
||||
# Separate permanent failures from retryable ones
|
||||
permanent_failures = failed_urls - rate_limited_urls
|
||||
|
||||
if permanent_failures:
|
||||
@@ -914,17 +975,25 @@ class DownloadManager:
|
||||
len(rate_limited_urls),
|
||||
model_name,
|
||||
)
|
||||
# Clear failed_models so non-force runs can retry
|
||||
|
||||
if transient_urls:
|
||||
self._progress["retryable_models"].add(model_hash)
|
||||
logger.warning(
|
||||
"%d example images for %s hit a transient network error; will retry next run",
|
||||
len(transient_urls),
|
||||
model_name,
|
||||
)
|
||||
|
||||
if rate_limited_urls or transient_urls:
|
||||
# Leave the model pending: the next run retries the
|
||||
# rate-limited/transient URLs. Permanent failures are
|
||||
# still flagged above; do not mark the model processed.
|
||||
if (force or explicit_targets) 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"
|
||||
"Removed %s from failed_models for a retryable failure",
|
||||
model_name,
|
||||
)
|
||||
|
||||
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)
|
||||
@@ -933,7 +1002,7 @@ class DownloadManager:
|
||||
len(permanent_failures),
|
||||
model_name,
|
||||
)
|
||||
elif success:
|
||||
elif final_success:
|
||||
self._progress["processed_models"].add(model_hash)
|
||||
if (force or explicit_targets) and model_hash in self._progress["failed_models"]:
|
||||
self._progress["failed_models"].discard(model_hash)
|
||||
@@ -1035,7 +1104,7 @@ class DownloadManager:
|
||||
model_hashes = data.get("model_hashes", [])
|
||||
optimize = data.get("optimize", True)
|
||||
model_types = data.get("model_types", ["lora", "checkpoint"])
|
||||
delay = float(data.get("delay", 0.2))
|
||||
delay = float(data.get("delay", 0) or 0)
|
||||
|
||||
if not model_hashes:
|
||||
raise DownloadConfigurationError("Missing model_hashes parameter")
|
||||
@@ -1106,6 +1175,7 @@ class DownloadManager:
|
||||
"""Download example images for specific models only - synchronous version."""
|
||||
|
||||
downloader = await get_downloader()
|
||||
limiter = self._create_media_download_limiter(optimize)
|
||||
|
||||
try:
|
||||
# Get scanners
|
||||
@@ -1160,6 +1230,7 @@ class DownloadManager:
|
||||
optimize,
|
||||
downloader,
|
||||
library_name,
|
||||
limiter=limiter,
|
||||
)
|
||||
|
||||
if was_successful:
|
||||
@@ -1181,9 +1252,12 @@ class DownloadManager:
|
||||
if should_stop:
|
||||
break
|
||||
|
||||
# Only add delay after remote download, and not after processing the last model
|
||||
# Pacing is handled by the shared media-download limiter. The
|
||||
# legacy ``delay`` option is kept for API callers that still
|
||||
# send it, but defaults to 0.
|
||||
if (
|
||||
was_successful
|
||||
delay > 0
|
||||
and was_successful
|
||||
and i < len(models_to_process) - 1
|
||||
and current_status == "running"
|
||||
):
|
||||
@@ -1255,6 +1329,7 @@ class DownloadManager:
|
||||
optimize,
|
||||
downloader,
|
||||
library_name,
|
||||
limiter: MediaDownloadLimiter | None = None,
|
||||
):
|
||||
"""Process a specific model for forced download, ignoring previous download status."""
|
||||
|
||||
@@ -1313,20 +1388,30 @@ class DownloadManager:
|
||||
if civitai_payload.get("images"):
|
||||
images = civitai_payload.get("images", [])
|
||||
|
||||
(
|
||||
success,
|
||||
is_stale,
|
||||
failed_images,
|
||||
rate_limited_images,
|
||||
) = await ExampleImagesProcessor.download_model_images_with_tracking(
|
||||
model_hash, model_name, images, model_dir, optimize, downloader
|
||||
outcome = await ExampleImagesProcessor.download_model_images_with_tracking(
|
||||
model_hash,
|
||||
model_name,
|
||||
images,
|
||||
model_dir,
|
||||
optimize,
|
||||
downloader,
|
||||
limiter=limiter,
|
||||
retry_failed=True,
|
||||
)
|
||||
|
||||
failed_urls: Set[str] = set(failed_images)
|
||||
rate_limited_urls: Set[str] = set(rate_limited_images)
|
||||
failed_urls: Set[str] = set(outcome.failed)
|
||||
rate_limited_urls: Set[str] = set(outcome.rate_limited)
|
||||
transient_urls: Set[str] = set(outcome.transient)
|
||||
recovered_urls: Set[str] = set(outcome.recovered)
|
||||
final_success = outcome.success
|
||||
|
||||
if recovered_urls:
|
||||
await self._clear_failed_images_from_metadata(
|
||||
model_hash, model_name, recovered_urls, scanner
|
||||
)
|
||||
|
||||
# If metadata is stale, try to refresh it
|
||||
if is_stale and model_hash not in self._progress["refreshed_models"]:
|
||||
if outcome.is_stale and model_hash not in self._progress["refreshed_models"]:
|
||||
await MetadataUpdater.refresh_model_metadata(
|
||||
model_hash, model_name, scanner_type, scanner, self._progress
|
||||
)
|
||||
@@ -1343,26 +1428,35 @@ class DownloadManager:
|
||||
if updated_civitai.get("images"):
|
||||
# Retry download with updated metadata
|
||||
updated_images = updated_civitai.get("images", [])
|
||||
(
|
||||
success,
|
||||
_,
|
||||
additional_failed_images,
|
||||
additional_rate_limited,
|
||||
) = await ExampleImagesProcessor.download_model_images_with_tracking(
|
||||
model_hash,
|
||||
model_name,
|
||||
updated_images,
|
||||
model_dir,
|
||||
optimize,
|
||||
downloader,
|
||||
retry_outcome = (
|
||||
await ExampleImagesProcessor.download_model_images_with_tracking(
|
||||
model_hash,
|
||||
model_name,
|
||||
updated_images,
|
||||
model_dir,
|
||||
optimize,
|
||||
downloader,
|
||||
limiter=limiter,
|
||||
retry_failed=True,
|
||||
)
|
||||
)
|
||||
|
||||
failed_urls.update(additional_failed_images)
|
||||
rate_limited_urls.update(additional_rate_limited)
|
||||
# The refreshed metadata is now the source of truth:
|
||||
# failures against the pre-refresh URLs are gone (or
|
||||
# will fail again on this pass).
|
||||
failed_urls = set(retry_outcome.failed)
|
||||
rate_limited_urls = set(retry_outcome.rate_limited)
|
||||
transient_urls = set(retry_outcome.transient)
|
||||
recovered_urls.update(retry_outcome.recovered)
|
||||
final_success = retry_outcome.success
|
||||
if retry_outcome.recovered:
|
||||
await self._clear_failed_images_from_metadata(
|
||||
model_hash, model_name, retry_outcome.recovered, scanner
|
||||
)
|
||||
|
||||
self._progress["refreshed_models"].add(model_hash)
|
||||
|
||||
# Separate permanent failures from rate-limited ones
|
||||
# Separate permanent failures from retryable ones
|
||||
permanent_failures = failed_urls - rate_limited_urls
|
||||
|
||||
# Only remove permanently failed images from metadata
|
||||
@@ -1379,13 +1473,21 @@ class DownloadManager:
|
||||
model_name,
|
||||
)
|
||||
|
||||
# Mark as processed only when no rate-limited images remain
|
||||
if rate_limited_urls:
|
||||
pass
|
||||
if transient_urls:
|
||||
self._progress["retryable_models"].add(model_hash)
|
||||
logger.warning(
|
||||
"%d example images for %s hit a transient network error; will retry next run",
|
||||
len(transient_urls),
|
||||
model_name,
|
||||
)
|
||||
|
||||
# Mark as processed only when nothing is left to retry.
|
||||
if rate_limited_urls or transient_urls:
|
||||
self._progress["failed_models"].discard(model_hash)
|
||||
elif permanent_failures:
|
||||
self._progress["processed_models"].add(model_hash)
|
||||
self._progress["failed_models"].add(model_hash)
|
||||
elif success:
|
||||
elif final_success:
|
||||
self._progress["processed_models"].add(model_hash)
|
||||
|
||||
return True # Return True to indicate a remote download happened
|
||||
@@ -1401,6 +1503,76 @@ class DownloadManager:
|
||||
self._progress["last_error"] = error_msg
|
||||
return False # Return False on exception
|
||||
|
||||
async def _clear_failed_images_from_metadata(
|
||||
self,
|
||||
model_hash: str,
|
||||
model_name: str,
|
||||
recovered_images: Iterable[str],
|
||||
scanner,
|
||||
) -> None:
|
||||
"""Clear stale ``downloadFailed`` flags for images that are now on disk.
|
||||
|
||||
Option B keeps the sidecar untouched until a retry actually succeeds;
|
||||
this is the write half of that contract. Only the two failure fields
|
||||
are removed; the URL and every other image field stay as they are.
|
||||
"""
|
||||
|
||||
recovered_set: Set[str] = {url for url in recovered_images if url}
|
||||
if not recovered_set:
|
||||
return
|
||||
|
||||
try:
|
||||
model_data = await MetadataUpdater.get_updated_model(model_hash, scanner)
|
||||
if not model_data:
|
||||
logger.warning(
|
||||
"Could not find model data for %s to clear recovered images",
|
||||
model_name,
|
||||
)
|
||||
return
|
||||
|
||||
civitai_payload = model_data.get("civitai") or {}
|
||||
current_images = civitai_payload.get("images") or []
|
||||
cleared = 0
|
||||
for image in current_images:
|
||||
if image.get("url") not in recovered_set:
|
||||
continue
|
||||
if not image.get("downloadFailed"):
|
||||
continue
|
||||
image.pop("downloadFailed", None)
|
||||
image.pop("downloadError", None)
|
||||
cleared += 1
|
||||
|
||||
if not cleared:
|
||||
return
|
||||
|
||||
file_path = model_data.get("file_path")
|
||||
if not file_path:
|
||||
return
|
||||
|
||||
model_copy = model_data.copy()
|
||||
model_copy.pop("folder", None)
|
||||
await MetadataManager.save_metadata(file_path, model_copy)
|
||||
try:
|
||||
await update_cache_from_metadata(scanner, file_path, model_copy)
|
||||
except Exception as cache_error:
|
||||
logger.warning(
|
||||
"Failed to update cache after clearing recovered flags for %s: %s",
|
||||
model_name,
|
||||
cache_error,
|
||||
)
|
||||
logger.info(
|
||||
"Recovered %s previously failed example images for %s",
|
||||
cleared,
|
||||
model_name,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to clear recovered example image flags for %s: %s",
|
||||
model_name,
|
||||
e,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
async def _remove_failed_images_from_metadata(
|
||||
self,
|
||||
model_hash: str,
|
||||
@@ -1439,13 +1611,7 @@ class DownloadManager:
|
||||
|
||||
for image in current_images:
|
||||
image_url = image.get("url")
|
||||
optimized_url = (
|
||||
ExampleImagesProcessor.get_civitai_optimized_url(image_url)
|
||||
if image_url and "civitai.com" in image_url
|
||||
else None
|
||||
)
|
||||
|
||||
if image_url not in failed_set and optimized_url not in failed_set:
|
||||
if image_url not in failed_set:
|
||||
continue
|
||||
|
||||
if image.get("downloadFailed"):
|
||||
|
||||
@@ -1,14 +1,25 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import logging
|
||||
import os
|
||||
import posixpath
|
||||
import re
|
||||
import random
|
||||
import string
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from aiohttp import web
|
||||
from ..utils.constants import SUPPORTED_MEDIA_EXTENSIONS
|
||||
from ..services.connectivity_guard import (
|
||||
OFFLINE_COOLDOWN_ERROR,
|
||||
OFFLINE_FRIENDLY_MESSAGE,
|
||||
)
|
||||
from ..services.service_registry import ServiceRegistry
|
||||
from ..services.settings_manager import get_settings_manager
|
||||
from ..utils.civitai_utils import build_civitai_media_url
|
||||
from ..utils.example_images_paths import get_model_folder, get_model_relative_path
|
||||
from .example_images_metadata import MetadataUpdater, update_cache_from_metadata
|
||||
from .file_utils import detect_media_extension
|
||||
@@ -16,6 +27,236 @@ from ..utils.metadata_manager import MetadataManager
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MAX_DOWNLOAD_ATTEMPTS = 3
|
||||
|
||||
_GENUR_HOSTS = frozenset({"img.genur.art", "c.genur.art"})
|
||||
_TRANSIENT_ERROR_MARKERS = (
|
||||
"network offline",
|
||||
"offline_cooldown",
|
||||
"cannot connect to host",
|
||||
"server disconnected",
|
||||
"connection reset",
|
||||
"connection refused",
|
||||
"temporary failure",
|
||||
"name or service not known",
|
||||
"timed out",
|
||||
"timeout",
|
||||
)
|
||||
|
||||
|
||||
def _write_media_file(path: str, content: bytes) -> None:
|
||||
with open(path, "wb") as handle:
|
||||
handle.write(content)
|
||||
|
||||
|
||||
def is_genur_media_url(url: str) -> bool:
|
||||
"""Return True for Genur image URLs handled by the fallback resolver."""
|
||||
|
||||
try:
|
||||
hostname = urlparse(url or "").hostname
|
||||
except ValueError:
|
||||
return False
|
||||
return bool(hostname) and hostname.lower() in _GENUR_HOSTS
|
||||
|
||||
|
||||
def resolve_genur_media_url(
|
||||
url: str,
|
||||
media_type: str | None = None,
|
||||
*,
|
||||
size: str = "small",
|
||||
) -> str:
|
||||
"""Map Genur image URLs to a direct rendition that is currently reachable.
|
||||
|
||||
Genur's signed proxy (``img.genur.art/sig/.../<base64>``) embeds the direct
|
||||
asset URL in the last path segment. The proxy host is not reachable from
|
||||
every network (fake-IP DNS / TLS resets), while the direct host serves the
|
||||
same image as ``<uuid>_small.webp`` and ``<uuid>_large.webp``. Returning
|
||||
the direct rendition avoids a guaranteed failed request and the
|
||||
ConnectivityGuard cooldown that follows it.
|
||||
|
||||
Non-Genur URLs, videos, and already-suffixed renditions are returned
|
||||
unchanged.
|
||||
"""
|
||||
|
||||
if not url or (media_type or "").lower() == "video":
|
||||
return url
|
||||
|
||||
try:
|
||||
parsed = urlparse(url)
|
||||
except ValueError:
|
||||
return url
|
||||
|
||||
hostname = (parsed.hostname or "").lower()
|
||||
if hostname == "img.genur.art":
|
||||
payload = parsed.path.rsplit("/", 1)[-1]
|
||||
try:
|
||||
padded = payload + "=" * (-len(payload) % 4)
|
||||
decoded = base64.b64decode(padded).decode("utf-8")
|
||||
except (ValueError, UnicodeDecodeError):
|
||||
return url
|
||||
try:
|
||||
decoded_parsed = urlparse(decoded)
|
||||
except ValueError:
|
||||
return url
|
||||
if (decoded_parsed.hostname or "").lower() != "c.genur.art":
|
||||
return url
|
||||
base = f"https://c.genur.art{decoded_parsed.path}"
|
||||
elif hostname == "c.genur.art":
|
||||
base = f"https://c.genur.art{parsed.path}"
|
||||
else:
|
||||
return url
|
||||
|
||||
base = base.rstrip("/")
|
||||
if base.endswith(("_small.webp", "_large.webp")):
|
||||
return base
|
||||
if posixpath.splitext(base)[1].lower() in SUPPORTED_MEDIA_EXTENSIONS["images"]:
|
||||
# Already a direct image path; do not invent a suffix.
|
||||
return base
|
||||
suffix = "_large.webp" if size == "large" else "_small.webp"
|
||||
return f"{base}{suffix}"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImageDownloadOutcome:
|
||||
"""Result of one model's example-media download pass.
|
||||
|
||||
``downloaded`` and ``recovered`` contain the *original* metadata URLs so
|
||||
callers can match them against ``civitai.images``. ``success`` is false
|
||||
when anything is still missing: a permanent failure, a rate-limit retry,
|
||||
a transient network failure, or a previously flagged image that this pass
|
||||
did not retry.
|
||||
"""
|
||||
|
||||
downloaded: list[str] = field(default_factory=list)
|
||||
failed: list[str] = field(default_factory=list)
|
||||
rate_limited: list[str] = field(default_factory=list)
|
||||
transient: list[str] = field(default_factory=list)
|
||||
recovered: list[str] = field(default_factory=list)
|
||||
skipped_failed: int = 0
|
||||
is_stale: bool = False
|
||||
|
||||
@property
|
||||
def success(self) -> bool:
|
||||
return not (
|
||||
self.failed
|
||||
or self.rate_limited
|
||||
or self.transient
|
||||
or self.skipped_failed
|
||||
)
|
||||
|
||||
|
||||
class MediaDownloadLimiter:
|
||||
"""Bound concurrent media requests and share 429 cooldowns across workers.
|
||||
|
||||
Civitai serves one binary per request and each request costs a CDN
|
||||
round-trip, so serial downloads spend most of their wall clock waiting.
|
||||
Concurrency is the useful lever, but it has to stay below the CDN's
|
||||
threshold: a 429 is a signal for the whole destination, not for the one
|
||||
task that happened to receive it. A cooldown armed here therefore parks
|
||||
every worker, and the limit is halved on a 429 and grows back one slot at
|
||||
a time after a run of successes (AIMD) instead of oscillating around the
|
||||
limit.
|
||||
|
||||
All waits are bounded by ``max_cooldown``. The per-request retry budget
|
||||
lives in the caller; this class only decides when a request may start.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_concurrency: int = 6,
|
||||
*,
|
||||
min_cooldown: float = 1.0,
|
||||
default_cooldown: float = 30.0,
|
||||
max_cooldown: float = 300.0,
|
||||
recovery_successes: int | None = None,
|
||||
) -> None:
|
||||
self.max_concurrency = max(1, int(max_concurrency))
|
||||
self._limit = self.max_concurrency
|
||||
self._in_flight = 0
|
||||
self._cooldown_until = 0.0
|
||||
self._min_cooldown = max(0.0, float(min_cooldown))
|
||||
self._default_cooldown = max(self._min_cooldown, float(default_cooldown))
|
||||
self._max_cooldown = max(self._default_cooldown, float(max_cooldown))
|
||||
self._recovery_successes = max(
|
||||
1, int(recovery_successes or max(8, self.max_concurrency * 2))
|
||||
)
|
||||
self._successes = 0
|
||||
self._condition = asyncio.Condition()
|
||||
|
||||
@property
|
||||
def limit(self) -> int:
|
||||
"""Current concurrency limit (may be below ``max_concurrency``)."""
|
||||
|
||||
return self._limit
|
||||
|
||||
async def acquire(self) -> None:
|
||||
"""Wait for the shared cooldown and a concurrency slot."""
|
||||
|
||||
async with self._condition:
|
||||
while True:
|
||||
remaining = self._cooldown_until - time.monotonic()
|
||||
if remaining > 0:
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self._condition.wait(), timeout=remaining
|
||||
)
|
||||
except asyncio.TimeoutError:
|
||||
pass
|
||||
continue
|
||||
if self._in_flight < self._limit:
|
||||
self._in_flight += 1
|
||||
return
|
||||
await self._condition.wait()
|
||||
|
||||
async def release(self) -> None:
|
||||
async with self._condition:
|
||||
self._in_flight = max(0, self._in_flight - 1)
|
||||
self._condition.notify()
|
||||
|
||||
async def report_rate_limit(self, retry_after: float | None = None) -> float:
|
||||
"""Arm the shared cooldown and shrink the concurrency limit.
|
||||
|
||||
Returns the cooldown actually applied, after clamping.
|
||||
"""
|
||||
|
||||
async with self._condition:
|
||||
requested = (
|
||||
retry_after
|
||||
if retry_after is not None and retry_after > 0
|
||||
else self._default_cooldown
|
||||
)
|
||||
cooldown = min(
|
||||
max(self._min_cooldown, float(requested)), self._max_cooldown
|
||||
)
|
||||
self._cooldown_until = max(
|
||||
self._cooldown_until, time.monotonic() + cooldown
|
||||
)
|
||||
self._limit = max(1, self._limit // 2)
|
||||
self._successes = 0
|
||||
self._condition.notify_all()
|
||||
return cooldown
|
||||
|
||||
async def report_success(self) -> None:
|
||||
"""Slowly restore concurrency after a run of successful requests."""
|
||||
|
||||
async with self._condition:
|
||||
if self._limit >= self.max_concurrency:
|
||||
return
|
||||
self._successes += 1
|
||||
if self._successes >= self._recovery_successes:
|
||||
self._limit += 1
|
||||
self._successes = 0
|
||||
self._condition.notify()
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Test seam: restore the initial state."""
|
||||
|
||||
self._limit = self.max_concurrency
|
||||
self._in_flight = 0
|
||||
self._cooldown_until = 0.0
|
||||
self._successes = 0
|
||||
|
||||
|
||||
|
||||
class ExampleImagesImportError(RuntimeError):
|
||||
"""Base error for example image import operations."""
|
||||
@@ -70,16 +311,64 @@ class ExampleImagesProcessor:
|
||||
return ''.join(random.choice(chars) for _ in range(length))
|
||||
|
||||
@staticmethod
|
||||
def get_civitai_optimized_url(media_url):
|
||||
"""Convert Civitai media URL (image or video) to its optimized version"""
|
||||
base_pattern = r'(https://image\.civitai\.com/[^/]+/[^/]+)'
|
||||
match = re.match(base_pattern, media_url)
|
||||
|
||||
if match:
|
||||
base_url = match.group(1)
|
||||
return f"{base_url}/optimized=true"
|
||||
|
||||
return media_url
|
||||
def _is_video_media(media_url: str, media_type_hint: str | None = None) -> bool:
|
||||
if (media_type_hint or "").lower() == "video":
|
||||
return True
|
||||
path = urlparse(media_url or "").path
|
||||
extension = posixpath.splitext(path)[1].lower()
|
||||
return extension in SUPPORTED_MEDIA_EXTENSIONS["videos"]
|
||||
|
||||
@staticmethod
|
||||
def _media_width(media_type: str | None) -> int:
|
||||
"""Return the configured rendition width for example media."""
|
||||
|
||||
is_video = (media_type or "").lower() == "video"
|
||||
key = "example_images_video_max_width" if is_video else "example_images_max_width"
|
||||
default = 450 if is_video else 800
|
||||
try:
|
||||
width = int(get_settings_manager().get(key, default))
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
# The image cacher's ladder starts at 96px; a smaller value is a
|
||||
# configuration mistake and would request an unusable thumbnail.
|
||||
return width if width >= 96 else default
|
||||
|
||||
@staticmethod
|
||||
def build_example_media_url(
|
||||
media_url: str,
|
||||
media_type_hint: str | None,
|
||||
optimize: bool,
|
||||
) -> str:
|
||||
"""Return the URL to download for an example image or video.
|
||||
|
||||
Civitai URLs are rewritten to a resized rendition (images default to
|
||||
800px, videos to 450px; both adjustable in settings). Genur signed
|
||||
proxy URLs are mapped to the direct ``c.genur.art`` rendition, which
|
||||
the broken proxy host never reaches. Other URLs are returned as-is.
|
||||
"""
|
||||
|
||||
media_url = (media_url or "").strip()
|
||||
if not media_url:
|
||||
return media_url
|
||||
|
||||
media_type = (
|
||||
"video"
|
||||
if ExampleImagesProcessor._is_video_media(media_url, media_type_hint)
|
||||
else "image"
|
||||
)
|
||||
|
||||
if is_genur_media_url(media_url):
|
||||
return resolve_genur_media_url(
|
||||
media_url, media_type, size="small" if optimize else "large"
|
||||
)
|
||||
|
||||
if not optimize:
|
||||
return media_url
|
||||
|
||||
width = ExampleImagesProcessor._media_width(media_type)
|
||||
return build_civitai_media_url(
|
||||
media_url, media_type, width=width, optimize=True
|
||||
) or media_url
|
||||
|
||||
@staticmethod
|
||||
def _get_file_extension_from_content_or_headers(content, headers, fallback_url=None, media_type_hint=None):
|
||||
@@ -120,234 +409,238 @@ class ExampleImagesProcessor:
|
||||
|
||||
@staticmethod
|
||||
async def download_model_images(model_hash, model_name, model_images, model_dir, optimize, downloader):
|
||||
"""Download images for a single model
|
||||
|
||||
"""Download example media for a single model.
|
||||
|
||||
Returns:
|
||||
tuple: (success, is_stale_metadata) - whether download was successful, whether metadata is stale
|
||||
tuple: (success, is_stale_metadata) - whether download was successful,
|
||||
whether metadata is stale.
|
||||
"""
|
||||
model_success = True
|
||||
|
||||
for i, image in enumerate(model_images):
|
||||
image_url = image.get('url')
|
||||
if not image_url:
|
||||
continue
|
||||
|
||||
if image.get('downloadFailed'):
|
||||
logger.debug(
|
||||
"Skipping example image %s for %s because it previously failed to download",
|
||||
image_url,
|
||||
model_name,
|
||||
)
|
||||
continue
|
||||
outcome = await ExampleImagesProcessor.download_model_images_with_tracking(
|
||||
model_hash,
|
||||
model_name,
|
||||
model_images,
|
||||
model_dir,
|
||||
optimize,
|
||||
downloader,
|
||||
)
|
||||
return outcome.success, outcome.is_stale
|
||||
|
||||
# Apply optimization for Civitai URLs if enabled
|
||||
original_url = image_url
|
||||
if optimize and 'civitai.com' in image_url:
|
||||
image_url = ExampleImagesProcessor.get_civitai_optimized_url(image_url)
|
||||
|
||||
# Skip the download when the file already exists on disk
|
||||
if ExampleImagesProcessor._example_image_file_exists(model_dir, i, image.get("type")):
|
||||
logger.debug("File already exists, skipping download for %s", image_url)
|
||||
continue
|
||||
|
||||
# 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(
|
||||
image_url,
|
||||
use_auth=False, # Example images don't need auth
|
||||
return_headers=True
|
||||
)
|
||||
|
||||
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}")
|
||||
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}")
|
||||
continue
|
||||
|
||||
# Save the file
|
||||
with open(save_path, 'wb') as f:
|
||||
f.write(content)
|
||||
|
||||
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
|
||||
# Return early to trigger metadata refresh attempt
|
||||
return False, True # (success, is_metadata_stale)
|
||||
else:
|
||||
error_msg = f"Failed to download file: {image_url}, error: {content}"
|
||||
logger.warning(error_msg)
|
||||
model_success = False # Mark the model as failed
|
||||
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
|
||||
|
||||
return model_success, False # (success, is_metadata_stale)
|
||||
|
||||
@staticmethod
|
||||
def _extract_retry_after(error_message: str) -> int:
|
||||
def _extract_retry_after(error_message: str) -> float:
|
||||
if not error_message:
|
||||
return 60
|
||||
match = re.search(r"retry after (\d+)s", str(error_message))
|
||||
return 30.0
|
||||
match = re.search(r"retry after ([\d.]+)s", str(error_message))
|
||||
if match:
|
||||
return max(1, int(match.group(1)))
|
||||
return 60
|
||||
return max(0.0, float(match.group(1)))
|
||||
return 30.0
|
||||
|
||||
@staticmethod
|
||||
async def download_model_images_with_tracking(model_hash, model_name, model_images, model_dir, optimize, downloader):
|
||||
model_success = True
|
||||
failed_images = []
|
||||
rate_limited_images = []
|
||||
any_successful_download = False
|
||||
def _is_transient_failure(message: str) -> bool:
|
||||
"""Return True for transport/offline/5xx errors that deserve a retry.
|
||||
|
||||
for i, image in enumerate(model_images):
|
||||
image_url = image.get('url')
|
||||
if not image_url:
|
||||
A 404 is a statement about the asset and is handled separately; an
|
||||
offline/shared-cooldown message describes the network path, so
|
||||
blacklisting the image for it (the old behavior) is wrong.
|
||||
"""
|
||||
|
||||
if not message:
|
||||
return False
|
||||
text = str(message)
|
||||
if text in {OFFLINE_FRIENDLY_MESSAGE, OFFLINE_COOLDOWN_ERROR}:
|
||||
return True
|
||||
lower = text.lower()
|
||||
if any(marker in lower for marker in _TRANSIENT_ERROR_MARKERS):
|
||||
return True
|
||||
match = re.search(r"status (\d{3})", text, re.IGNORECASE)
|
||||
if match:
|
||||
status = int(match.group(1))
|
||||
return status in {408, 425, 429} or 500 <= status < 600
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
async def download_model_images_with_tracking(
|
||||
model_hash,
|
||||
model_name,
|
||||
model_images,
|
||||
model_dir,
|
||||
optimize,
|
||||
downloader,
|
||||
limiter: MediaDownloadLimiter | None = None,
|
||||
retry_failed: bool = False,
|
||||
) -> ImageDownloadOutcome:
|
||||
"""Download all pending media for one model concurrently.
|
||||
|
||||
Each image is still one CDN request (Civitai has no batch-binary
|
||||
endpoint), but the requests are pipelined through ``limiter`` instead
|
||||
of being serialized. A 429 arms the limiter's shared cooldown, so one
|
||||
rate-limited response pauses every worker rather than letting each
|
||||
URL sleep and retry on its own.
|
||||
|
||||
``retry_failed`` re-attempts image-level ``downloadFailed`` entries
|
||||
for Genur URLs. Genur images are the ones affected by the old
|
||||
network-error classification bug; retrying them lets the caller clear
|
||||
the flag only after a real download succeeds.
|
||||
|
||||
Returns:
|
||||
ImageDownloadOutcome: per-URL lists and a ``success`` property.
|
||||
URL lists contain the *original* metadata URLs so callers can
|
||||
match them against the model's ``civitai.images`` entries.
|
||||
"""
|
||||
|
||||
limiter = limiter or MediaDownloadLimiter()
|
||||
outcome = ImageDownloadOutcome()
|
||||
candidates: list[tuple[int, dict[str, Any], str, bool]] = []
|
||||
|
||||
for index, image in enumerate(model_images):
|
||||
original_url = image.get("url")
|
||||
if not original_url:
|
||||
continue
|
||||
|
||||
if image.get('downloadFailed'):
|
||||
# A file already on disk is authoritative; a stale failure flag
|
||||
# must not force a redundant request.
|
||||
if ExampleImagesProcessor._example_image_file_exists(
|
||||
model_dir, index, image.get("type")
|
||||
):
|
||||
logger.debug("File already exists, skipping download for %s", original_url)
|
||||
if (
|
||||
image.get("downloadFailed")
|
||||
and retry_failed
|
||||
and is_genur_media_url(original_url)
|
||||
):
|
||||
outcome.recovered.append(original_url)
|
||||
continue
|
||||
|
||||
if image.get("downloadFailed"):
|
||||
if retry_failed and is_genur_media_url(original_url):
|
||||
logger.debug(
|
||||
"Retrying previously failed Genur image %s for %s",
|
||||
original_url,
|
||||
model_name,
|
||||
)
|
||||
candidates.append((index, image, original_url, True))
|
||||
continue
|
||||
outcome.skipped_failed += 1
|
||||
logger.debug(
|
||||
"Skipping example image %s for %s because it previously failed to download",
|
||||
image_url,
|
||||
original_url,
|
||||
model_name,
|
||||
)
|
||||
continue
|
||||
|
||||
# Apply optimization for Civitai URLs if enabled
|
||||
original_url = image_url
|
||||
if optimize and 'civitai.com' in image_url:
|
||||
image_url = ExampleImagesProcessor.get_civitai_optimized_url(image_url)
|
||||
candidates.append((index, image, original_url, False))
|
||||
|
||||
# Skip the download when the file already exists on disk
|
||||
if ExampleImagesProcessor._example_image_file_exists(model_dir, i, image.get("type")):
|
||||
logger.debug("File already exists, skipping download for %s", image_url)
|
||||
continue
|
||||
async def download_one(
|
||||
index: int,
|
||||
image: dict[str, Any],
|
||||
original_url: str,
|
||||
was_flagged: bool,
|
||||
) -> None:
|
||||
target_url = ExampleImagesProcessor.build_example_media_url(
|
||||
original_url, image.get("type"), optimize
|
||||
)
|
||||
logger.debug(
|
||||
"Downloading media file %s for %s from %s",
|
||||
index,
|
||||
model_name,
|
||||
target_url,
|
||||
)
|
||||
|
||||
async def _attempt_download() -> tuple[bool, Any, Any]:
|
||||
logger.debug("Downloading media file %s for %s", i, model_name)
|
||||
return await downloader.download_to_memory(
|
||||
image_url,
|
||||
use_auth=False,
|
||||
return_headers=True,
|
||||
)
|
||||
|
||||
try:
|
||||
success, content, headers = await _attempt_download()
|
||||
for attempt in range(1, MAX_DOWNLOAD_ATTEMPTS + 1):
|
||||
await limiter.acquire()
|
||||
try:
|
||||
success, content, headers = await downloader.download_to_memory(
|
||||
target_url,
|
||||
use_auth=False,
|
||||
return_headers=True,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error("Error downloading file %s: %s", target_url, e)
|
||||
outcome.failed.append(original_url)
|
||||
return
|
||||
finally:
|
||||
await limiter.release()
|
||||
|
||||
if success:
|
||||
await limiter.report_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']
|
||||
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)
|
||||
continue
|
||||
|
||||
save_filename = f"image_{i}{media_ext}"
|
||||
save_path = os.path.join(model_dir, save_filename)
|
||||
return
|
||||
|
||||
save_path = os.path.join(model_dir, f"image_{index}{media_ext}")
|
||||
if os.path.exists(save_path):
|
||||
logger.debug("File already exists: %s", save_path)
|
||||
continue
|
||||
return
|
||||
|
||||
with open(save_path, 'wb') as f:
|
||||
f.write(content)
|
||||
any_successful_download = True
|
||||
try:
|
||||
await asyncio.to_thread(_write_media_file, save_path, content)
|
||||
except Exception as e:
|
||||
logger.error("Failed to save example image %s: %s", save_path, e)
|
||||
outcome.failed.append(original_url)
|
||||
return
|
||||
|
||||
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
|
||||
failed_images.append(image_url)
|
||||
return False, True, failed_images, rate_limited_images
|
||||
outcome.downloaded.append(original_url)
|
||||
if was_flagged:
|
||||
outcome.recovered.append(original_url)
|
||||
return
|
||||
|
||||
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)
|
||||
message = str(content)
|
||||
if ExampleImagesProcessor._is_not_found_error(message):
|
||||
logger.warning(
|
||||
"Failed to download file: %s, status code: 404 - Model metadata might be stale",
|
||||
target_url,
|
||||
)
|
||||
outcome.failed.append(original_url)
|
||||
outcome.is_stale = True
|
||||
return
|
||||
|
||||
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 "Rate limited (429)" in message:
|
||||
retry_after = ExampleImagesProcessor._extract_retry_after(message)
|
||||
cooldown = await limiter.report_rate_limit(retry_after)
|
||||
logger.warning(
|
||||
"Rate limited (429) for %s, shared cooldown %.1fs (attempt %d/%d)",
|
||||
target_url,
|
||||
cooldown,
|
||||
attempt,
|
||||
MAX_DOWNLOAD_ATTEMPTS,
|
||||
)
|
||||
if attempt >= MAX_DOWNLOAD_ATTEMPTS:
|
||||
outcome.rate_limited.append(original_url)
|
||||
return
|
||||
# Loop back into acquire(): it parks until the shared
|
||||
# cooldown expires, so the retry is not a herd.
|
||||
continue
|
||||
|
||||
if not (is_image or is_video):
|
||||
logger.debug("Skipping unsupported file type: %s", media_ext)
|
||||
break
|
||||
if ExampleImagesProcessor._is_transient_failure(message):
|
||||
logger.warning(
|
||||
"Transient failure for %s: %s (will retry next run)",
|
||||
target_url,
|
||||
content,
|
||||
)
|
||||
outcome.transient.append(original_url)
|
||||
return
|
||||
|
||||
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
|
||||
logger.warning("Failed to download file: %s, error: %s", target_url, content)
|
||||
outcome.failed.append(original_url)
|
||||
return
|
||||
|
||||
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
|
||||
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
|
||||
failed_images.append(image_url)
|
||||
if candidates:
|
||||
await asyncio.gather(
|
||||
*(
|
||||
download_one(index, image, url, was_flagged)
|
||||
for index, image, url, was_flagged in candidates
|
||||
)
|
||||
)
|
||||
|
||||
return outcome
|
||||
|
||||
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):
|
||||
"""Process local example images
|
||||
|
||||
@@ -10,6 +10,7 @@ import pytest
|
||||
|
||||
from py.services.settings_manager import SettingsManager, get_settings_manager
|
||||
from py.utils import example_images_download_manager as download_module
|
||||
from py.utils.example_images_processor import ImageDownloadOutcome
|
||||
|
||||
|
||||
class RecordingWebSocketManager:
|
||||
@@ -174,7 +175,7 @@ async def test_pause_resume_blocks_processing(
|
||||
await first_release.wait()
|
||||
else:
|
||||
second_call_started.set()
|
||||
return True, False, [], []
|
||||
return ImageDownloadOutcome()
|
||||
|
||||
async def fake_get_downloader():
|
||||
class _Downloader:
|
||||
@@ -298,7 +299,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 ImageDownloadOutcome()
|
||||
|
||||
async def fake_get_downloader():
|
||||
class _Downloader:
|
||||
@@ -1017,3 +1018,274 @@ def test_unblock_failed_example_image_models_without_progress_file(
|
||||
monkeypatch.setitem(settings_manager.settings, "active_library", "default")
|
||||
|
||||
assert download_module.unblock_failed_example_image_models(["a" * 64]) == 0
|
||||
|
||||
|
||||
async def _run_single_outcome(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
tmp_path,
|
||||
settings_manager,
|
||||
outcome: ImageDownloadOutcome,
|
||||
*,
|
||||
force: bool = False,
|
||||
):
|
||||
model_hash = "a" * 64
|
||||
model_path = tmp_path / "model.safetensors"
|
||||
model = {
|
||||
"sha256": model_hash,
|
||||
"model_name": "Outcome Model",
|
||||
"file_path": str(model_path),
|
||||
"file_name": "model.safetensors",
|
||||
"civitai": {"images": [{"url": "https://img.genur.art/sig/x/y"}]},
|
||||
}
|
||||
_patch_scanner(monkeypatch, StubScanner([model]))
|
||||
monkeypatch.setitem(settings_manager.settings, "example_images_path", str(tmp_path))
|
||||
|
||||
async def fake_process_local_examples(*_args, **_kwargs):
|
||||
return False
|
||||
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
async def fake_download(*_args, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return outcome
|
||||
|
||||
async def fake_get_downloader():
|
||||
return object()
|
||||
|
||||
monkeypatch.setattr(
|
||||
download_module.ExampleImagesProcessor,
|
||||
"process_local_examples",
|
||||
staticmethod(fake_process_local_examples),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
download_module.ExampleImagesProcessor,
|
||||
"download_model_images_with_tracking",
|
||||
staticmethod(fake_download),
|
||||
)
|
||||
monkeypatch.setattr(download_module, "get_downloader", fake_get_downloader)
|
||||
|
||||
manager = download_module.DownloadManager(ws_manager=RecordingWebSocketManager())
|
||||
await manager.start_download(
|
||||
{"model_types": ["lora"], "delay": 0, "force": force}
|
||||
)
|
||||
if manager._download_task is not None:
|
||||
await asyncio.wait_for(manager._download_task, timeout=2)
|
||||
return manager, model_hash, captured
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_transient_failure_leaves_model_retryable(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path, settings_manager
|
||||
) -> None:
|
||||
outcome = ImageDownloadOutcome(transient=["https://img.genur.art/sig/x/y"])
|
||||
manager, model_hash, _ = await _run_single_outcome(
|
||||
monkeypatch, tmp_path, settings_manager, outcome
|
||||
)
|
||||
|
||||
assert model_hash in manager._progress["retryable_models"]
|
||||
assert model_hash not in manager._progress["processed_models"]
|
||||
assert model_hash not in manager._progress["failed_models"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skipped_flagged_image_is_not_reported_as_recovered(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path, settings_manager
|
||||
) -> None:
|
||||
outcome = ImageDownloadOutcome(skipped_failed=1)
|
||||
manager, model_hash, _ = await _run_single_outcome(
|
||||
monkeypatch, tmp_path, settings_manager, outcome
|
||||
)
|
||||
|
||||
assert model_hash in manager._progress["failed_models"]
|
||||
assert model_hash not in manager._progress["processed_models"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_force_passes_retry_failed_to_processor(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path, settings_manager
|
||||
) -> None:
|
||||
manager, model_hash, captured = await _run_single_outcome(
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
settings_manager,
|
||||
ImageDownloadOutcome(),
|
||||
force=True,
|
||||
)
|
||||
|
||||
assert captured["retry_failed"] is True
|
||||
assert model_hash in manager._progress["processed_models"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clear_recovered_flags_only_touches_matching_images(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path
|
||||
) -> None:
|
||||
manager = download_module.DownloadManager(ws_manager=RecordingWebSocketManager())
|
||||
recovered_url = "https://img.genur.art/sig/x/y"
|
||||
other_url = "https://image.civitai.com/x/y/original=true/1.jpeg"
|
||||
model_data = {
|
||||
"file_path": str(tmp_path / "model.safetensors"),
|
||||
"civitai": {
|
||||
"images": [
|
||||
{
|
||||
"url": recovered_url,
|
||||
"type": "image",
|
||||
"downloadFailed": True,
|
||||
"downloadError": "not_found",
|
||||
},
|
||||
{
|
||||
"url": other_url,
|
||||
"type": "image",
|
||||
"downloadFailed": True,
|
||||
"downloadError": "not_found",
|
||||
},
|
||||
]
|
||||
},
|
||||
}
|
||||
|
||||
async def fake_get_updated_model(*_args, **_kwargs):
|
||||
return model_data
|
||||
|
||||
saved: list[dict[str, Any]] = []
|
||||
cached: list[dict[str, Any]] = []
|
||||
|
||||
async def fake_save(_path, payload):
|
||||
saved.append(payload)
|
||||
return True
|
||||
|
||||
async def fake_update_cache(_scanner, _path, payload):
|
||||
cached.append(payload)
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(
|
||||
download_module.MetadataUpdater,
|
||||
"get_updated_model",
|
||||
staticmethod(fake_get_updated_model),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
download_module.MetadataManager, "save_metadata", staticmethod(fake_save)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
download_module, "update_cache_from_metadata", fake_update_cache
|
||||
)
|
||||
|
||||
await manager._clear_failed_images_from_metadata(
|
||||
"hash", "Model", {recovered_url}, object()
|
||||
)
|
||||
|
||||
assert "downloadFailed" not in model_data["civitai"]["images"][0]
|
||||
assert "downloadError" not in model_data["civitai"]["images"][0]
|
||||
assert model_data["civitai"]["images"][1]["downloadFailed"] is True
|
||||
assert saved and "folder" not in saved[0]
|
||||
assert cached
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_force_bypasses_populated_folder_guard(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path, settings_manager
|
||||
) -> None:
|
||||
model_hash = "b" * 64
|
||||
model = {
|
||||
"sha256": model_hash,
|
||||
"model_name": "Partial Model",
|
||||
"file_path": str(tmp_path / "partial.safetensors"),
|
||||
"file_name": "partial.safetensors",
|
||||
"civitai": {"images": [{"url": "https://img.genur.art/sig/x/y"}]},
|
||||
}
|
||||
scanner = StubScanner([model])
|
||||
_patch_scanner(monkeypatch, scanner)
|
||||
monkeypatch.setitem(settings_manager.settings, "example_images_path", str(tmp_path))
|
||||
monkeypatch.setitem(settings_manager.settings, "active_library", "comfyui")
|
||||
|
||||
model_dir = download_module.ExampleImagePathResolver.get_model_folder(
|
||||
model_hash, "comfyui"
|
||||
)
|
||||
Path(model_dir).mkdir(parents=True, exist_ok=True)
|
||||
Path(model_dir, "image_0.webp").write_bytes(b"existing")
|
||||
|
||||
async def fake_process_local_examples(*_args, **_kwargs):
|
||||
return False
|
||||
|
||||
captured: dict[str, Any] = {}
|
||||
|
||||
async def fake_download(*_args, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return ImageDownloadOutcome()
|
||||
|
||||
monkeypatch.setattr(
|
||||
download_module.ExampleImagesProcessor,
|
||||
"process_local_examples",
|
||||
staticmethod(fake_process_local_examples),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
download_module.ExampleImagesProcessor,
|
||||
"download_model_images_with_tracking",
|
||||
staticmethod(fake_download),
|
||||
)
|
||||
|
||||
manager = download_module.DownloadManager(ws_manager=RecordingWebSocketManager())
|
||||
manager._progress["status"] = "running"
|
||||
manager._progress["processed_models"].add(model_hash)
|
||||
|
||||
await manager._process_model(
|
||||
"lora",
|
||||
model,
|
||||
scanner,
|
||||
str(tmp_path),
|
||||
False,
|
||||
object(),
|
||||
"comfyui",
|
||||
force=True,
|
||||
explicit_targets=False,
|
||||
limiter=None,
|
||||
)
|
||||
|
||||
assert captured["retry_failed"] is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clear_recovered_flags_roundtrips_through_sidecar(
|
||||
tmp_path, settings_manager
|
||||
) -> None:
|
||||
model_path = tmp_path / "roundtrip.safetensors"
|
||||
model_path.write_bytes(b"model")
|
||||
model_hash = "c" * 64
|
||||
recovered_url = "https://img.genur.art/sig/x/y"
|
||||
model_data = {
|
||||
"file_name": "roundtrip",
|
||||
"model_name": "Roundtrip Model",
|
||||
"file_path": str(model_path),
|
||||
"sha256": model_hash,
|
||||
"size": 1,
|
||||
"modified": 0,
|
||||
"base_model": "SDXL",
|
||||
"preview_url": "",
|
||||
"preview_nsfw_level": 0,
|
||||
"from_civitai": True,
|
||||
"civitai": {
|
||||
"images": [
|
||||
{
|
||||
"url": recovered_url,
|
||||
"type": "image",
|
||||
"downloadFailed": True,
|
||||
"downloadError": "not_found",
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
assert await download_module.MetadataManager.save_metadata(
|
||||
str(model_path), model_data
|
||||
)
|
||||
|
||||
scanner = StubScanner([{**model_data, "folder": ""}])
|
||||
manager = download_module.DownloadManager(ws_manager=RecordingWebSocketManager())
|
||||
await manager._clear_failed_images_from_metadata(
|
||||
model_hash, "Roundtrip Model", {recovered_url}, scanner
|
||||
)
|
||||
|
||||
loaded, _ = await download_module.MetadataManager.load_metadata(str(model_path))
|
||||
assert loaded is not None
|
||||
images = loaded.to_dict()["civitai"]["images"]
|
||||
assert images[0]["url"] == recovered_url
|
||||
assert "downloadFailed" not in images[0]
|
||||
assert "downloadError" not in images[0]
|
||||
|
||||
@@ -5,6 +5,7 @@ from pathlib import Path
|
||||
from typing import Any
|
||||
from py.services.settings_manager import get_settings_manager
|
||||
from py.utils import example_images_download_manager as download_module
|
||||
from py.utils.example_images_processor import ImageDownloadOutcome
|
||||
|
||||
class RecordingWebSocketManager:
|
||||
def __init__(self) -> None:
|
||||
@@ -73,12 +74,12 @@ async def test_reprocessing_triggered_when_folder_missing(monkeypatch, tmp_path)
|
||||
|
||||
monkeypatch.setattr(download_module.ExampleImagesProcessor, "process_local_examples", fake_process_local_examples)
|
||||
|
||||
async def fake_download_model_images(*args):
|
||||
async def fake_download_model_images(*args, **_kwargs):
|
||||
# Create the directory so it's "fixed"
|
||||
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 ImageDownloadOutcome()
|
||||
|
||||
monkeypatch.setattr(download_module.ExampleImagesProcessor, "download_model_images_with_tracking", fake_download_model_images)
|
||||
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
import pytest
|
||||
|
||||
from py.utils.civitai_utils import rewrite_preview_url
|
||||
from py.utils.civitai_utils import build_civitai_media_url, rewrite_preview_url
|
||||
|
||||
|
||||
class TestRewritePreviewUrl:
|
||||
@@ -148,3 +148,72 @@ class TestRewritePreviewUrl:
|
||||
for url in test_cases:
|
||||
result, was_rewritten = rewrite_preview_url(url, "image")
|
||||
assert was_rewritten is True, f"Failed for URL: {url}"
|
||||
|
||||
def test_width_parameter_overrides_default(self):
|
||||
url = "https://image.civitai.com/xG1nk/uuid/original=true/1.jpeg"
|
||||
result, was_rewritten = rewrite_preview_url(url, "image", width=800)
|
||||
assert was_rewritten is True
|
||||
assert result == "https://image.civitai.com/xG1nk/uuid/width=800,optimized=true/1.jpeg"
|
||||
|
||||
def test_width_parameter_applies_to_video_transcode(self):
|
||||
url = "https://image.civitai.com/xG1nk/uuid/original=true/1.mp4"
|
||||
result, was_rewritten = rewrite_preview_url(url, "video", width=450)
|
||||
assert was_rewritten is True
|
||||
assert (
|
||||
result
|
||||
== "https://image.civitai.com/xG1nk/uuid/transcode=true,width=450,optimized=true/1.mp4"
|
||||
)
|
||||
|
||||
|
||||
class TestBuildCivitaiMediaUrl:
|
||||
"""Tests for the example-media download URL builder."""
|
||||
|
||||
def test_rewrites_original_true_and_keeps_filename(self):
|
||||
url = "https://image.civitai.com/xG1nk/uuid/original=true/74821598.jpeg"
|
||||
assert (
|
||||
build_civitai_media_url(url, "image", width=800)
|
||||
== "https://image.civitai.com/xG1nk/uuid/width=800,optimized=true/74821598.jpeg"
|
||||
)
|
||||
|
||||
def test_rewrites_videos_with_transcode(self):
|
||||
url = "https://image.civitai.com/xG1nk/uuid/original=true/75044257.mp4"
|
||||
assert (
|
||||
build_civitai_media_url(url, "video", width=450)
|
||||
== "https://image.civitai.com/xG1nk/uuid/transcode=true,width=450,optimized=true/75044257.mp4"
|
||||
)
|
||||
|
||||
def test_replaces_legacy_params_segment(self):
|
||||
url = "https://image.civitai.com/xG1nk/uuid/optimized=true"
|
||||
assert (
|
||||
build_civitai_media_url(url, "image", width=800)
|
||||
== "https://image.civitai.com/xG1nk/uuid/width=800,optimized=true"
|
||||
)
|
||||
|
||||
def test_inserts_params_before_filename_without_params(self):
|
||||
url = "https://image.civitai.com/xG1nk/uuid/74821598.jpeg"
|
||||
assert (
|
||||
build_civitai_media_url(url, "image", width=800)
|
||||
== "https://image.civitai.com/xG1nk/uuid/width=800,optimized=true/74821598.jpeg"
|
||||
)
|
||||
|
||||
def test_appends_params_when_there_is_no_filename(self):
|
||||
url = "https://image.civitai.com/xG1nk/uuid"
|
||||
assert (
|
||||
build_civitai_media_url(url, "image", width=800)
|
||||
== "https://image.civitai.com/xG1nk/uuid/width=800,optimized=true"
|
||||
)
|
||||
|
||||
def test_handles_image_b2_subdomain(self):
|
||||
url = "https://image-b2.civitai.com/file/cache/original=true/sample.png"
|
||||
assert (
|
||||
build_civitai_media_url(url, "image", width=800)
|
||||
== "https://image-b2.civitai.com/file/cache/width=800,optimized=true/sample.png"
|
||||
)
|
||||
|
||||
def test_leaves_non_civitai_urls_alone(self):
|
||||
url = "https://example.com/original=true/sample.png"
|
||||
assert build_civitai_media_url(url, "image", width=800) == url
|
||||
|
||||
def test_optimize_false_returns_source_url(self):
|
||||
url = "https://image.civitai.com/xG1nk/uuid/original=true/1.jpeg"
|
||||
assert build_civitai_media_url(url, "image", optimize=False) == url
|
||||
|
||||
@@ -1,7 +1,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, Dict, Generator, Tuple
|
||||
@@ -134,14 +136,14 @@ async def test_download_model_images_with_tracking_skips_existing_files(tmp_path
|
||||
return True, b"\xff\xd8\xff" + b"data", {}
|
||||
|
||||
downloader = RecordingDownloader()
|
||||
success, is_stale, failed, rate_limited = await proc.download_model_images_with_tracking(
|
||||
outcome = await proc.download_model_images_with_tracking(
|
||||
"hash", "model", images, str(tmp_path), False, downloader
|
||||
)
|
||||
|
||||
assert success is True
|
||||
assert is_stale is False
|
||||
assert failed == []
|
||||
assert rate_limited == []
|
||||
assert outcome.success is True
|
||||
assert outcome.is_stale is False
|
||||
assert outcome.failed == []
|
||||
assert outcome.rate_limited == []
|
||||
# Only the missing image is requested; the existing one is skipped without a network call
|
||||
assert len(downloader.calls) == 1
|
||||
assert "c/d" in downloader.calls[0]
|
||||
@@ -464,3 +466,276 @@ async def test_resolve_hash_for_file_path_unknown_file_returns_empty(
|
||||
)
|
||||
== ""
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_model_images_uses_bounded_renditions(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path
|
||||
) -> None:
|
||||
"""Images use the configured width and videos are transcoded at their cap."""
|
||||
settings_manager = get_settings_manager()
|
||||
settings_manager.settings["example_images_max_width"] = 800
|
||||
settings_manager.settings["example_images_video_max_width"] = 450
|
||||
|
||||
images = [
|
||||
{
|
||||
"url": "https://image.civitai.com/xG1nk/uuid/original=true/1.jpeg",
|
||||
"type": "image",
|
||||
},
|
||||
{
|
||||
"url": "https://image.civitai.com/xG1nk/uuid/original=true/2.mp4",
|
||||
"type": "video",
|
||||
},
|
||||
]
|
||||
|
||||
class RecordingDownloader:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[str] = []
|
||||
|
||||
async def download_to_memory(self, url, use_auth=False, return_headers=False):
|
||||
self.calls.append(url)
|
||||
if url.endswith(".mp4"):
|
||||
return True, b"\x00\x00\x00\x18ftypmp42data", {"content-type": "video/mp4"}
|
||||
return True, b"\xff\xd8\xff" + b"data", {"content-type": "image/jpeg"}
|
||||
|
||||
downloader = RecordingDownloader()
|
||||
limiter = processor_module.MediaDownloadLimiter(max_concurrency=2)
|
||||
outcome = await (
|
||||
processor_module.ExampleImagesProcessor.download_model_images_with_tracking(
|
||||
"hash", "model", images, str(tmp_path), True, downloader, limiter=limiter
|
||||
)
|
||||
)
|
||||
|
||||
assert outcome.success is True
|
||||
assert outcome.is_stale is False
|
||||
assert outcome.failed == []
|
||||
assert outcome.rate_limited == []
|
||||
assert any("width=800,optimized=true/1.jpeg" in url for url in downloader.calls)
|
||||
assert any(
|
||||
"transcode=true,width=450,optimized=true/2.mp4" in url
|
||||
for url in downloader.calls
|
||||
)
|
||||
assert Path(tmp_path, "image_0.jpg").exists()
|
||||
assert Path(tmp_path, "image_1.mp4").exists()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_model_images_runs_requests_concurrently(tmp_path) -> None:
|
||||
"""The limiter's pool overlaps network waits instead of serializing them."""
|
||||
images = [
|
||||
{
|
||||
"url": f"https://image.civitai.com/xG1nk/uuid/original=true/{i}.jpeg",
|
||||
"type": "image",
|
||||
}
|
||||
for i in range(6)
|
||||
]
|
||||
|
||||
class TrackingDownloader:
|
||||
def __init__(self) -> None:
|
||||
self.active = 0
|
||||
self.max_active = 0
|
||||
|
||||
async def download_to_memory(self, url, use_auth=False, return_headers=False):
|
||||
self.active += 1
|
||||
self.max_active = max(self.max_active, self.active)
|
||||
try:
|
||||
await asyncio.sleep(0.02)
|
||||
return True, b"\xff\xd8\xff" + b"data", {"content-type": "image/jpeg"}
|
||||
finally:
|
||||
self.active -= 1
|
||||
|
||||
downloader = TrackingDownloader()
|
||||
limiter = processor_module.MediaDownloadLimiter(max_concurrency=3)
|
||||
outcome = await (
|
||||
processor_module.ExampleImagesProcessor.download_model_images_with_tracking(
|
||||
"hash", "model", images, str(tmp_path), False, downloader, limiter=limiter
|
||||
)
|
||||
)
|
||||
|
||||
assert outcome.success is True
|
||||
assert outcome.failed == []
|
||||
assert outcome.rate_limited == []
|
||||
assert 1 < downloader.max_active <= 3
|
||||
assert len(list(Path(tmp_path).glob("image_*.jpg"))) == 6
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_model_images_shares_429_cooldown(tmp_path) -> None:
|
||||
"""A 429 retries after the shared cooldown, not a per-task 60s sleep."""
|
||||
images = [
|
||||
{
|
||||
"url": f"https://image.civitai.com/xG1nk/uuid/original=true/{i}.jpeg",
|
||||
"type": "image",
|
||||
}
|
||||
for i in range(2)
|
||||
]
|
||||
|
||||
class RateLimitedThenOk:
|
||||
def __init__(self) -> None:
|
||||
self.attempts: dict[str, int] = {}
|
||||
|
||||
async def download_to_memory(self, url, use_auth=False, return_headers=False):
|
||||
attempt = self.attempts.get(url, 0) + 1
|
||||
self.attempts[url] = attempt
|
||||
if attempt == 1:
|
||||
return False, "Rate limited (429), retry after 0.01s", None
|
||||
return True, b"\xff\xd8\xff" + b"data", {"content-type": "image/jpeg"}
|
||||
|
||||
downloader = RateLimitedThenOk()
|
||||
limiter = processor_module.MediaDownloadLimiter(
|
||||
max_concurrency=2, min_cooldown=0.01, default_cooldown=0.01
|
||||
)
|
||||
started = time.monotonic()
|
||||
outcome = await (
|
||||
processor_module.ExampleImagesProcessor.download_model_images_with_tracking(
|
||||
"hash", "model", images, str(tmp_path), False, downloader, limiter=limiter
|
||||
)
|
||||
)
|
||||
elapsed = time.monotonic() - started
|
||||
|
||||
assert outcome.success is True
|
||||
assert outcome.failed == []
|
||||
assert outcome.rate_limited == []
|
||||
assert all(attempts == 2 for attempts in downloader.attempts.values())
|
||||
assert elapsed < 1.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_model_images_marks_404_stale_with_original_url(tmp_path) -> None:
|
||||
image = {
|
||||
"url": "https://image.civitai.com/xG1nk/uuid/original=true/1.jpeg",
|
||||
"type": "image",
|
||||
}
|
||||
|
||||
class NotFoundDownloader:
|
||||
async def download_to_memory(self, url, use_auth=False, return_headers=False):
|
||||
return False, "File not found", None
|
||||
|
||||
outcome = await (
|
||||
processor_module.ExampleImagesProcessor.download_model_images_with_tracking(
|
||||
"hash", "model", [image], str(tmp_path), True, NotFoundDownloader()
|
||||
)
|
||||
)
|
||||
|
||||
assert outcome.success is False
|
||||
assert outcome.is_stale is True
|
||||
assert outcome.failed == [image["url"]]
|
||||
assert outcome.rate_limited == []
|
||||
|
||||
|
||||
def test_resolve_genur_signed_url_uses_direct_small_rendition() -> None:
|
||||
signed = (
|
||||
"https://img.genur.art/sig/width:450/quality:85/"
|
||||
"aHR0cHM6Ly9jLmdlbnVyLmFydC83MjIwZGMwYy05N2YyLTQ0YzEtYWNkZi00ZDBhMzkxY2ZjNmQ="
|
||||
)
|
||||
assert (
|
||||
processor_module.resolve_genur_media_url(signed)
|
||||
== "https://c.genur.art/7220dc0c-97f2-44c1-acdf-4d0a391cfc6d_small.webp"
|
||||
)
|
||||
|
||||
|
||||
def test_resolve_genur_direct_url_adds_small_suffix() -> None:
|
||||
assert (
|
||||
processor_module.resolve_genur_media_url(
|
||||
"https://c.genur.art/7220dc0c-97f2-44c1-acdf-4d0a391cfc6d"
|
||||
)
|
||||
== "https://c.genur.art/7220dc0c-97f2-44c1-acdf-4d0a391cfc6d_small.webp"
|
||||
)
|
||||
existing = "https://c.genur.art/7220dc0c-97f2-44c1-acdf-4d0a391cfc6d_large.webp"
|
||||
assert processor_module.resolve_genur_media_url(existing) == existing
|
||||
|
||||
|
||||
def test_resolve_genur_leaves_other_urls_alone() -> None:
|
||||
url = "https://image.civitai.com/xG1nk/uuid/original=true/1.jpeg"
|
||||
assert processor_module.resolve_genur_media_url(url) == url
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_model_images_classifies_network_errors_as_transient(tmp_path) -> None:
|
||||
image = {
|
||||
"url": "https://img.genur.art/sig/width:450/quality:85/"
|
||||
"aHR0cHM6Ly9jLmdlbnVyLmFydC83MjIwZGMwYy05N2YyLTQ0YzEtYWNkZi00ZDBhMzkxY2ZjNmQ=",
|
||||
"type": "image",
|
||||
}
|
||||
|
||||
class OfflineDownloader:
|
||||
async def download_to_memory(self, url, use_auth=False, return_headers=False):
|
||||
return False, "Network offline, will retry automatically later", None
|
||||
|
||||
outcome = await (
|
||||
processor_module.ExampleImagesProcessor.download_model_images_with_tracking(
|
||||
"hash", "model", [image], str(tmp_path), True, OfflineDownloader()
|
||||
)
|
||||
)
|
||||
|
||||
assert outcome.success is False
|
||||
assert outcome.transient == [image["url"]]
|
||||
assert outcome.failed == []
|
||||
assert outcome.rate_limited == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_model_images_retries_flagged_genur_image(tmp_path) -> None:
|
||||
original = (
|
||||
"https://img.genur.art/sig/width:450/quality:85/"
|
||||
"aHR0cHM6Ly9jLmdlbnVyLmFydC83MjIwZGMwYy05N2YyLTQ0YzEtYWNkZi00ZDBhMzkxY2ZjNmQ="
|
||||
)
|
||||
image = {
|
||||
"url": original,
|
||||
"type": "image",
|
||||
"downloadFailed": True,
|
||||
"downloadError": "not_found",
|
||||
}
|
||||
|
||||
class GenurDownloader:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[str] = []
|
||||
|
||||
async def download_to_memory(self, url, use_auth=False, return_headers=False):
|
||||
self.calls.append(url)
|
||||
return True, b"RIFF\x00\x00\x00\x00WEBPdata", {"content-type": "image/webp"}
|
||||
|
||||
downloader = GenurDownloader()
|
||||
outcome = await (
|
||||
processor_module.ExampleImagesProcessor.download_model_images_with_tracking(
|
||||
"hash",
|
||||
"model",
|
||||
[image],
|
||||
str(tmp_path),
|
||||
True,
|
||||
downloader,
|
||||
retry_failed=True,
|
||||
)
|
||||
)
|
||||
|
||||
assert outcome.success is True
|
||||
assert outcome.recovered == [original]
|
||||
assert outcome.downloaded == [original]
|
||||
assert downloader.calls == [
|
||||
"https://c.genur.art/7220dc0c-97f2-44c1-acdf-4d0a391cfc6d_small.webp"
|
||||
]
|
||||
assert Path(tmp_path, "image_0.webp").exists()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_model_images_does_not_retry_flag_without_retry_option(tmp_path) -> None:
|
||||
image = {
|
||||
"url": "https://image.civitai.com/xG1nk/uuid/original=true/1.jpeg",
|
||||
"type": "image",
|
||||
"downloadFailed": True,
|
||||
"downloadError": "not_found",
|
||||
}
|
||||
|
||||
class ShouldNotRun:
|
||||
async def download_to_memory(self, url, use_auth=False, return_headers=False):
|
||||
raise AssertionError("flagged image must not be requested")
|
||||
|
||||
outcome = await (
|
||||
processor_module.ExampleImagesProcessor.download_model_images_with_tracking(
|
||||
"hash", "model", [image], str(tmp_path), True, ShouldNotRun()
|
||||
)
|
||||
)
|
||||
|
||||
assert outcome.success is False
|
||||
assert outcome.skipped_failed == 1
|
||||
assert outcome.downloaded == []
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
from py.utils.example_images_processor import MediaDownloadLimiter
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_limiter_bounds_in_flight_requests() -> None:
|
||||
limiter = MediaDownloadLimiter(max_concurrency=2)
|
||||
|
||||
await limiter.acquire()
|
||||
await limiter.acquire()
|
||||
|
||||
waiting = asyncio.create_task(limiter.acquire())
|
||||
await asyncio.sleep(0)
|
||||
assert waiting.done() is False
|
||||
|
||||
await limiter.release()
|
||||
await asyncio.wait_for(waiting, timeout=1)
|
||||
|
||||
await limiter.release()
|
||||
await limiter.release()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_limiter_parks_every_worker_during_shared_cooldown() -> None:
|
||||
limiter = MediaDownloadLimiter(
|
||||
max_concurrency=2, min_cooldown=0.0, default_cooldown=0.01
|
||||
)
|
||||
|
||||
await limiter.report_rate_limit(0.05)
|
||||
started = time.monotonic()
|
||||
await limiter.acquire()
|
||||
elapsed = time.monotonic() - started
|
||||
|
||||
assert elapsed >= 0.04
|
||||
await limiter.release()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_limiter_halves_on_429_and_recovers_one_slot_at_a_time() -> None:
|
||||
limiter = MediaDownloadLimiter(
|
||||
max_concurrency=4,
|
||||
min_cooldown=0.0,
|
||||
default_cooldown=0.01,
|
||||
recovery_successes=2,
|
||||
)
|
||||
|
||||
await limiter.report_rate_limit(0.01)
|
||||
assert limiter.limit == 2
|
||||
|
||||
await limiter.report_success()
|
||||
assert limiter.limit == 2
|
||||
await limiter.report_success()
|
||||
assert limiter.limit == 3
|
||||
|
||||
await limiter.report_success()
|
||||
await limiter.report_success()
|
||||
assert limiter.limit == 4
|
||||
|
||||
|
||||
def test_limiter_reset_restores_initial_limit() -> None:
|
||||
limiter = MediaDownloadLimiter(max_concurrency=3)
|
||||
limiter._limit = 1
|
||||
limiter._in_flight = 1
|
||||
limiter._cooldown_until = time.monotonic() + 30
|
||||
|
||||
limiter.reset()
|
||||
|
||||
assert limiter.limit == 3
|
||||
assert limiter._in_flight == 0
|
||||
assert limiter._cooldown_until == 0.0
|
||||
Reference in New Issue
Block a user