feat: download media in bounded domain batches

2026-09-13 16:58:45 +07:00
parent 787a5065bc
commit 8207d3ae02
7 changed files with 161 additions and 29 deletions
@@ -14,6 +14,16 @@ from .media_assets import audit_media_catalog, extract_media_candidates, media_c
 DEFAULT_ROOT = Path("data/media")
 MAX_ASSET_BYTES = 15 * 1024 * 1024
 MEDIA_PRIORITY = {"waterbody": 0, "fish": 1, "tackle": 2, "reference": 3}
+MAX_BATCH_ASSETS_PER_DOMAIN = 40
+MAX_CONSECUTIVE_FAILURES = 3
+
+
+def _queue_sort_key(item: dict) -> tuple:
+    return (
+        MEDIA_PRIORITY.get(str(item.get("entity_type")), 99),
+        not bool(item.get("label")), str(item.get("label") or "").casefold(),
+        str(item.get("asset_url")),
+    )


 def media_queue_plan(root: Path, state_file: Path, *, now: float | None = None) -> dict:
@@ -28,11 +38,7 @@ def media_queue_plan(root: Path, state_file: Path, *, now: float | None = None)

    domains: dict[str, dict] = {}
    for site, items in sorted(grouped.items()):
-        items.sort(key=lambda item: (
-            MEDIA_PRIORITY.get(str(item.get("entity_type")), 99),
-            not bool(item.get("label")), str(item.get("label") or "").casefold(),
-            str(item.get("asset_url")),
-        ))
+        items.sort(key=_queue_sort_key)
        last_attempt = state.get(site)
        retry_in = max(0, round(float(last_attempt) + MIN_FETCH_INTERVAL_SECONDS - current)) if isinstance(last_attempt, (int, float)) else 0
        next_item = items[0]
@@ -58,6 +64,7 @@ def media_queue_plan(root: Path, state_file: Path, *, now: float | None = None)
        "generated_at": datetime.fromtimestamp(current, timezone.utc).isoformat(),
        "network_requests": 0,
        "cooldown_seconds": MIN_FETCH_INTERVAL_SECONDS,
+        "batch_limit_per_domain": MAX_BATCH_ASSETS_PER_DOMAIN,
        "queued_total": sum(group["queued"] for group in domains.values()),
        "domains": domains,
        "coverage": coverage["entities"],
@@ -66,15 +73,12 @@ def media_queue_plan(root: Path, state_file: Path, *, now: float | None = None)
    }


-def _download_one(root: Path, state_file: Path, asset_url: str | None = None) -> str:
-    manifest_path = root / "manifest.json"
-    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
-    queued = next((item for item in manifest["assets"] if item.get("status") == "queued" and (asset_url is None or item["asset_url"] == asset_url)), None)
-    if queued is None:
-        return "queue is empty" if asset_url is None else "asset is not queued"
+def _save_manifest(path: Path, manifest: dict) -> None:
+    path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
+
+
+def _attempt_asset(root: Path, manifest_path: Path, manifest: dict, queued: dict) -> tuple[str, Exception | None]:
    url = queued["asset_url"]
-    _validate_url_before_io(url)
-    check_and_reserve(fetch_site_key(url), state_file=state_file)
    attempted_at = datetime.now(timezone.utc).isoformat()
    try:
        request = urllib.request.Request(url, headers={"User-Agent": USER_AGENT, "Accept": "image/*"})
@@ -89,19 +93,92 @@ def _download_one(root: Path, state_file: Path, asset_url: str | None = None) ->
        digest, relative, width, height, detected_mime = store_asset(root, body, content_type=content_type, source_url=url)
    except Exception as exc:
        code = exc.code if isinstance(exc, urllib.error.HTTPError) else None
-        status = "missing" if code in {404, 410} else "blocked" if code in {401, 403} else "invalid" if isinstance(exc, ValueError) else "queued"
+        status = "missing" if code in {404, 410} else "blocked" if code in {401, 403, 429} else "invalid" if isinstance(exc, ValueError) else "queued"
        queued.update({"status": status, "last_attempt_at": attempted_at, "last_error": str(exc)[:500]})
-        manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
-        raise
+        _save_manifest(manifest_path, manifest)
+        return f"{status} {url}: {str(exc)[:200]}", exc
    queued.update({"status": "stored", "sha256": digest, "local_path": relative, "content_type": detected_mime, "bytes": len(body), "width": width, "height": height, "fetched_at": attempted_at})
-    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
-    return f"stored {url} as {relative}"
+    _save_manifest(manifest_path, manifest)
+    return f"stored {url} as {relative}", None
+
+
+def _download_one(root: Path, state_file: Path, asset_url: str | None = None) -> str:
+    manifest_path = root / "manifest.json"
+    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+    queued = next((item for item in manifest["assets"] if item.get("status") == "queued" and (asset_url is None or item["asset_url"] == asset_url)), None)
+    if queued is None:
+        return "queue is empty" if asset_url is None else "asset is not queued"
+    url = queued["asset_url"]
+    _validate_url_before_io(url)
+    check_and_reserve(fetch_site_key(url), state_file=state_file)
+    message, error = _attempt_asset(root, manifest_path, manifest, queued)
+    if error is not None:
+        raise error
+    return message
+
+
+def _download_batch(root: Path, state_file: Path, *, limit: int = MAX_BATCH_ASSETS_PER_DOMAIN) -> dict:
+    """Download up to ``limit`` assets per domain under one reserved window."""
+    if not 1 <= limit <= MAX_BATCH_ASSETS_PER_DOMAIN:
+        raise ValueError(f"batch limit must be 1..{MAX_BATCH_ASSETS_PER_DOMAIN}")
+    manifest_path = root / "manifest.json"
+    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+    grouped: dict[str, list[dict]] = {}
+    for item in manifest.get("assets", []):
+        if item.get("status") != "queued":
+            continue
+        _validate_url_before_io(item["asset_url"])
+        grouped.setdefault(fetch_site_key(item["asset_url"]), []).append(item)
+
+    domains: dict[str, dict] = {}
+    for site, items in sorted(grouped.items()):
+        items.sort(key=_queue_sort_key)
+        try:
+            # One reservation governs the explicitly allowed asset batch.
+            check_and_reserve(site, state_file=state_file)
+        except RuntimeError as exc:
+            domains[site] = {"attempted": 0, "stored": 0, "failed": 0, "skipped": str(exc)}
+            continue
+        stored = failed = consecutive_failures = 0
+        messages: list[str] = []
+        stopped_reason: str | None = None
+        for item in items[:limit]:
+            message, error = _attempt_asset(root, manifest_path, manifest, item)
+            messages.append(message)
+            if error is None:
+                stored += 1
+                consecutive_failures = 0
+                continue
+            failed += 1
+            consecutive_failures += 1
+            status = item.get("status")
+            if status in {"blocked", "queued"}:
+                stopped_reason = f"domain stopped after {status} response"
+                break
+            if consecutive_failures >= MAX_CONSECUTIVE_FAILURES:
+                stopped_reason = f"domain stopped after {MAX_CONSECUTIVE_FAILURES} consecutive failures"
+                break
+        domains[site] = {
+            "attempted": stored + failed, "stored": stored, "failed": failed,
+            "stopped_reason": stopped_reason, "messages": messages,
+        }
+    return {
+        "batch_limit_per_domain": limit,
+        "cooldown_seconds": MIN_FETCH_INTERVAL_SECONDS,
+        "domains": domains,
+        "attempted_total": sum(item["attempted"] for item in domains.values()),
+        "stored_total": sum(item["stored"] for item in domains.values()),
+        "failed_total": sum(item["failed"] for item in domains.values()),
+    }


 def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Index authorized RF4 media without hotlinking")
    parser.add_argument("url", nargs="?")
-    parser.add_argument("--download-one", action="store_true", help="Store one queued asset while respecting site cooldown")
+    download_mode = parser.add_mutually_exclusive_group()
+    download_mode.add_argument("--download-one", action="store_true", help="Store one queued asset while respecting site cooldown")
+    download_mode.add_argument("--download-batch", action="store_true", help="Store up to 40 queued assets per domain in one reserved window")
+    parser.add_argument("--batch-limit", type=int, default=MAX_BATCH_ASSETS_PER_DOMAIN, choices=range(1, MAX_BATCH_ASSETS_PER_DOMAIN + 1), metavar="1..40")
    parser.add_argument("--asset-url", help="Download this exact queued asset instead of the first queue item")
    parser.add_argument("--reclassify", action="store_true", help="Reapply current conservative classifier without network access")
    parser.add_argument("--review-url", help="Review an asset already present in the manifest")
@@ -139,8 +216,14 @@ def main(argv: list[str] | None = None) -> int:
        if args.download_one:
            print(_download_one(args.root, args.state_file, args.asset_url))
            return 0
+        if args.download_batch:
+            if args.asset_url:
+                parser.error("--asset-url is only valid with --download-one")
+            report = _download_batch(args.root, args.state_file, limit=args.batch_limit)
+            print(json.dumps(report, ensure_ascii=False, indent=2))
+            return 1 if any(item.get("stopped_reason") for item in report["domains"].values()) else 0
        if not args.url:
-            parser.error("url is required unless --download-one is used")
+            parser.error("url is required unless a download or offline mode is used")
        check_and_reserve(fetch_site_key(args.url), state_file=args.state_file)
        html = fetch_html(args.url)
        candidates = extract_media_candidates(html, source_page=args.url)