fix(docker): embedding 模型下载加镜像源 fallback，根治 HF 429 (#1640)

wehos · Hongzhi Wen · claude · web-flow · commit 94cbf47fb1c5 · 2026-06-05T02:41:01.000+08:00
#1630 的 30s 退避扛不住 huggingface.co 按小时的 per-IP 限流，真正挡住 429 的 是 GHA layer cache；cache 被驱逐后退回直连 HF，于是间歇性 429（full 比 standard 更频繁——多下 849MB fp32，下载层更大更易被驱逐）。 prepare_embedding_model.py 加镜像源 fallback： - 抽出 _download_one（单 URL 退避重试，逻辑不变；urlopen 改带稳定 UA 的 Request，防 CDN 拒默认 Python-urllib UA）。 - _download 外包一层，对每个文件按 huggingface.co -> hf-mirror.com 顺序尝试， HF 持续 429 / 404 / 不可达就切下一镜像，每源各跑完整退避重试，全源失败才抛。 顺序可由 HF_ENDPOINTS / HF_ENDPOINT 覆盖，默认内置，两个 Dockerfile + 两个 desktop workflow 共享脚本，无需改动。 新增 tests/unit/test_prepare_embedding_model.py 覆盖 fallback / env 覆盖 / 404 穿透 / 全源失败 / 首源命中不打镜像 / UA。 Co-authored-by: Hongzhi Wen <cartabio.coder1@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
diff --git a/scripts/prepare_embedding_model.py b/scripts/prepare_embedding_model.py
@@ -25,15 +25,32 @@
 DEFAULT_OUTPUT_ROOT = Path("data") / "embedding_models"
 PREPARED_MARKER = ".prepared.json"
 
-# Download resilience. huggingface.co rate-limits by source IP, and the Docker
-# build runs several arch/variant jobs from a shared proxy egress, so a single
-# urlopen routinely hit HTTP 429 and killed the whole build with no recovery.
-# Retry transient failures (429 / 5xx / connection errors) with exponential
-# backoff, honoring a numeric Retry-After when the server sends one.
+# Download resilience. huggingface.co rate-limits anonymous requests per source
+# IP by the hour, and CI runs from shared runner / proxy egress IPs that are
+# chronically throttled, so a direct fetch routinely hit HTTP 429 and killed the
+# whole build with no recovery. Two defenses stack here:
+#   1. Mirror fallback: each file is tried against every endpoint in order
+#      (huggingface.co first, then the hf-mirror.com reverse proxy). A source
+#      that 429s or is unreachable falls through to the next instead of failing
+#      the build. Override the list/order via HF_ENDPOINTS (comma-separated) or a
+#      single HF_ENDPOINT (the huggingface_hub convention).
+#   2. Per-endpoint backoff: within one endpoint, retry transient failures
+#      (429 / 5xx / connection errors) with exponential backoff, honoring a
+#      numeric Retry-After when the server sends one.
+# A bounded ~30s backoff alone can't outwait an hourly per-IP limit; the mirror
+# fallback is what actually breaks the deadlock when the runner IP is throttled.
 _RETRYABLE_STATUS = frozenset({429, 500, 502, 503, 504})
 _MAX_ATTEMPTS = 5
 _BACKOFF_BASE_SECONDS = 2.0
 _BACKOFF_CAP_SECONDS = 60.0
+# Mirror endpoints tried in order. hf-mirror.com mirrors the full
+# /{repo}/resolve/{revision}/{file} layout and is not under the same per-IP
+# throttle, so it recovers builds when huggingface.co rate-limits the shared CI
+# egress. Kept as a plain default so neither Dockerfile nor the desktop
+# workflows need to pass anything.
+DEFAULT_ENDPOINTS = ("https://huggingface.co", "https://hf-mirror.com")
+# Some CDNs reject the default "Python-urllib/x.y" agent; send a stable one.
+_USER_AGENT = "neko-embedding-prepare/1.0 (+https://github.com/Project-N-E-K-O/N.E.K.O)"
 # 40-char lowercase hex git SHA. Tags / branch refs / short SHAs are rejected
 # so the profile id stays a strict compatibility contract — anything that can
 # move under our feet, even tags (which can be force-pushed), is excluded.
@@ -85,18 +102,36 @@ def _backoff_seconds(attempt: int) -> float:
     return min(_BACKOFF_BASE_SECONDS * (2 ** (attempt - 1)), _BACKOFF_CAP_SECONDS)
 
 
-def _download(url: str, dest: Path, *, force: bool) -> None:
-    if dest.exists() and dest.stat().st_size > 0 and not force:
-        print(f"[embedding-model] keep existing {dest}")
-        return
+def _endpoints() -> list[str]:
+    """Ordered HF-compatible base URLs to try for each file.
+
+    ``HF_ENDPOINTS`` (comma-separated) takes precedence and fully replaces the
+    default order; a single ``HF_ENDPOINT`` (the huggingface_hub convention) is
+    honored next and pins to that one mirror. Otherwise the built-in
+    huggingface.co -> hf-mirror.com fallback is used.
+    """
+    raw = os.environ.get("HF_ENDPOINTS") or os.environ.get("HF_ENDPOINT")
+    if raw:
+        eps = [item.strip().rstrip("/") for item in raw.split(",") if item.strip()]
+        if eps:
+            return eps
+    return list(DEFAULT_ENDPOINTS)
 
+
+def _download_one(url: str, dest: Path) -> None:
+    """Fetch one URL into ``dest`` with bounded exponential-backoff retry.
+
+    Raises ``RuntimeError`` when retries are exhausted or the server returns a
+    non-retryable status (e.g. 404), so the caller can fall back to the next
+    mirror.
+    """
     dest.parent.mkdir(parents=True, exist_ok=True)
     tmp = dest.with_suffix(dest.suffix + ".tmp")
-    print(f"[embedding-model] download {url}")
+    request = urllib.request.Request(url, headers={"User-Agent": _USER_AGENT})
 
     for attempt in range(1, _MAX_ATTEMPTS + 1):
         try:
-            with urllib.request.urlopen(url, timeout=120) as response:
+            with urllib.request.urlopen(request, timeout=120) as response:
                 with tmp.open("wb") as f:
                     while True:
                         chunk = response.read(1024 * 1024)
@@ -145,6 +180,44 @@ def _download(url: str, dest: Path, *, force: bool) -> None:
         time.sleep(delay)
 
 
+def _download(
+    rel: str,
+    dest: Path,
+    *,
+    repo: str,
+    revision: str,
+    endpoints: list[str],
+    force: bool,
+) -> None:
+    """Download one repo file into ``dest``, trying each mirror in order.
+
+    Each endpoint gets its own bounded backoff retry; a source that exhausts its
+    retries (e.g. a persistent 429) or returns a non-retryable status falls
+    through to the next mirror. Only when every endpoint fails does this raise.
+    """
+    if dest.exists() and dest.stat().st_size > 0 and not force:
+        print(f"[embedding-model] keep existing {dest}")
+        return
+
+    failures: list[str] = []
+    for index, base in enumerate(endpoints, 1):
+        url = f"{base}/{repo}/resolve/{revision}/{rel}"
+        suffix = f" (source {index}/{len(endpoints)})" if len(endpoints) > 1 else ""
+        print(f"[embedding-model] download {url}{suffix}")
+        try:
+            _download_one(url, dest)
+            return
+        except RuntimeError as exc:
+            failures.append(str(exc))
+            if index < len(endpoints):
+                print(f"[embedding-model] source {base} failed; falling back to next mirror")
+
+    raise RuntimeError(
+        f"failed to download {rel} from all {len(endpoints)} source(s): "
+        + " | ".join(failures)
+    )
+
+
 def _verify(profile_dir: Path, files: list[str]) -> None:
     missing = [
         str(profile_dir / rel)
@@ -234,9 +307,16 @@ def main(argv: list[str] | None = None) -> int:
             f"forcing re-download for {args.repo}@{args.revision}",
         )
 
+    endpoints = _endpoints()
     for rel in files:
-        url = f"https://huggingface.co/{args.repo}/resolve/{args.revision}/{rel}"
-        _download(url, profile_dir / rel, force=args.force or revision_changed)
+        _download(
+            rel,
+            profile_dir / rel,
+            repo=args.repo,
+            revision=args.revision,
+            endpoints=endpoints,
+            force=args.force or revision_changed,
+        )
     _verify(profile_dir, files)
     _write_marker(profile_dir, args.repo, args.revision)
     print(f"[embedding-model] profile ready: {profile_dir}")
diff --git a/tests/unit/test_prepare_embedding_model.py b/tests/unit/test_prepare_embedding_model.py
@@ -0,0 +1,166 @@
+# -*- coding: utf-8 -*-
+"""Unit tests for scripts/prepare_embedding_model.py — the build-time embedding
+model downloader.
+
+Focus is the mirror-fallback layer added to survive huggingface.co's per-IP
+HTTP 429 throttling of shared CI egress IPs: each file is tried against every
+endpoint in order (huggingface.co -> hf-mirror.com by default), each endpoint
+getting its own bounded backoff retry, and only an all-source failure raises.
+
+Pure stdlib + mock: no network, no numpy/onnxruntime, so this runs on every
+workstation regardless of the embedding bundle state.
+"""
+from __future__ import annotations
+
+import importlib.util
+import io
+import urllib.error
+import urllib.request
+from pathlib import Path
+
+import pytest
+
+
+_SCRIPT = Path(__file__).resolve().parents[2] / "scripts" / "prepare_embedding_model.py"
+
+
+@pytest.fixture
+def prep(monkeypatch):
+    """Load the script as a module and neutralize real backoff sleeps."""
+    spec = importlib.util.spec_from_file_location("prepare_embedding_model", _SCRIPT)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    monkeypatch.setattr(module.time, "sleep", lambda *_: None)
+    # Default env: tests opt into HF_ENDPOINT(S) explicitly where relevant.
+    monkeypatch.delenv("HF_ENDPOINTS", raising=False)
+    monkeypatch.delenv("HF_ENDPOINT", raising=False)
+    return module
+
+
+class _FakeResp(io.BytesIO):
+    status = 200
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        self.close()
+
+
+def _http_error(url, code):
+    return urllib.error.HTTPError(url, code, "err", {}, None)
+
+
+def _install_urlopen(monkeypatch, prep, behavior, *, record=None):
+    """Patch the module's urlopen with a callable mapping url -> bytes | Exception."""
+
+    def fake_urlopen(request, timeout=120):
+        url = request.full_url if hasattr(request, "full_url") else request
+        if record is not None:
+            record.append(url)
+        result = behavior(url)
+        if isinstance(result, Exception):
+            raise result
+        return _FakeResp(result)
+
+    monkeypatch.setattr(prep.urllib.request, "urlopen", fake_urlopen)
+
+
+# --- endpoint resolution ---------------------------------------------------
+
+def test_endpoints_default(prep):
+    assert prep._endpoints() == ["https://huggingface.co", "https://hf-mirror.com"]
+
+
+def test_hf_endpoint_pins_single_mirror(prep, monkeypatch):
+    monkeypatch.setenv("HF_ENDPOINT", "https://hf-mirror.com")
+    assert prep._endpoints() == ["https://hf-mirror.com"]
+
+
+def test_hf_endpoints_list_takes_precedence_and_is_cleaned(prep, monkeypatch):
+    monkeypatch.setenv("HF_ENDPOINT", "https://ignored.example")
+    monkeypatch.setenv("HF_ENDPOINTS", "https://a.example/ , https://b.example ,, ")
+    assert prep._endpoints() == ["https://a.example", "https://b.example"]
+
+
+# --- download fallback -----------------------------------------------------
+
+def test_falls_back_to_mirror_on_persistent_429(prep, monkeypatch, tmp_path):
+    def behavior(url):
+        if "huggingface.co" in url:
+            return _http_error(url, 429)
+        return b"FROM-MIRROR"
+
+    _install_urlopen(monkeypatch, prep, behavior)
+    dest = tmp_path / "tokenizer.json"
+    prep._download(
+        "tokenizer.json", dest, repo="r/m", revision="abc",
+        endpoints=["https://huggingface.co", "https://hf-mirror.com"], force=False,
+    )
+    assert dest.read_bytes() == b"FROM-MIRROR"
+
+
+def test_non_retryable_404_on_first_source_still_falls_through(prep, monkeypatch, tmp_path):
+    def behavior(url):
+        if "huggingface.co" in url:
+            return _http_error(url, 404)
+        return b"OK2"
+
+    _install_urlopen(monkeypatch, prep, behavior)
+    dest = tmp_path / "x"
+    prep._download(
+        "x", dest, repo="r/m", revision="abc",
+        endpoints=["https://huggingface.co", "https://hf-mirror.com"], force=False,
+    )
+    assert dest.read_bytes() == b"OK2"
+
+
+def test_all_sources_failing_raises_listing_every_source(prep, monkeypatch, tmp_path):
+    _install_urlopen(monkeypatch, prep, lambda url: _http_error(url, 429))
+    with pytest.raises(RuntimeError) as excinfo:
+        prep._download(
+            "f.bin", tmp_path / "f.bin", repo="r/m", revision="abc",
+            endpoints=["https://huggingface.co", "https://hf-mirror.com"], force=False,
+        )
+    message = str(excinfo.value)
+    assert "all 2 source(s)" in message
+    assert "huggingface.co" in message and "hf-mirror.com" in message
+
+
+def test_primary_success_does_not_contact_mirror(prep, monkeypatch, tmp_path):
+    seen: list[str] = []
+    _install_urlopen(monkeypatch, prep, lambda url: b"PRIMARY", record=seen)
+    dest = tmp_path / "x"
+    prep._download(
+        "x", dest, repo="r/m", revision="abc",
+        endpoints=["https://huggingface.co", "https://hf-mirror.com"], force=False,
+    )
+    assert dest.read_bytes() == b"PRIMARY"
+    assert len(seen) == 1 and "huggingface.co" in seen[0]
+
+
+def test_existing_nonempty_file_is_kept_without_any_request(prep, monkeypatch, tmp_path):
+    dest = tmp_path / "x"
+    dest.write_bytes(b"already-here")
+
+    def explode(url):
+        raise AssertionError("should not hit the network when file already exists")
+
+    _install_urlopen(monkeypatch, prep, explode)
+    prep._download(
+        "x", dest, repo="r/m", revision="abc",
+        endpoints=["https://huggingface.co", "https://hf-mirror.com"], force=False,
+    )
+    assert dest.read_bytes() == b"already-here"
+
+
+def test_sends_stable_user_agent(prep, monkeypatch, tmp_path):
+    captured: dict[str, str] = {}
+
+    def fake_urlopen(request, timeout=120):
+        captured["ua"] = request.get_header("User-agent")
+        return _FakeResp(b"data")
+
+    monkeypatch.setattr(prep.urllib.request, "urlopen", fake_urlopen)
+    prep._download_one("https://huggingface.co/r/m/resolve/abc/x", tmp_path / "x")
+    assert captured["ua"] == prep._USER_AGENT
