fix(ai): retry bundle downloads on a different transport after checksum mismatch (#715)

The published archives are intact (manifest hash, HF LFS hash, and CDN
bytes all agree), but the accelerated hf_xet client assembles files with
parallel offset writes, which some storage backends (network mounts,
FUSE bind mounts) corrupt silently. The old retry re-ran that same
client with force_download, so an install on such storage failed the
checksum forever with no way out.

The mismatch retry now goes through the plain sequential downloader,
and discards stale .partial/.meta resume sidecars first so a previously
killed run cannot weld old bytes onto the fresh attempt. Checksum errors
report the actual digest next to the expected one plus the Offline
Import workaround, IO errors during verification surface as parseable
error frames instead of tracebacks, and a retry-download failure keeps
the mismatch that started it in the message.

Fixes #714
This commit is contained in:
SnapOtter
2026-08-02 17:25:53 +08:00
committed by GitHub
parent a9bb76fbc7
commit 72bc88a9f8
2 changed files with 432 additions and 52 deletions
+101 -52
View File
@@ -141,8 +141,8 @@ def get_site_packages_dir(venv_path: str) -> str:
# -- SHA256 verification --
def verify_sha256(filepath: str, expected: str) -> bool:
"""Stream-hash a file and compare to expected hex digest."""
def file_sha256(filepath: str) -> str:
"""Stream-hash a file and return the hex digest."""
h = hashlib.sha256()
with open(filepath, "rb") as f:
while True:
@@ -150,7 +150,16 @@ def verify_sha256(filepath: str, expected: str) -> bool:
if not chunk:
break
h.update(chunk)
return h.hexdigest() == expected
return h.hexdigest()
def _unlink_quietly(path: str) -> None:
"""Best-effort delete for cleanup on error paths; a failed unlink must
never mask the error being reported."""
try:
os.unlink(path)
except OSError:
pass
# -- Download with resume --
@@ -225,7 +234,6 @@ def download_with_hf_hub(
expected_size: int,
progress_start: int,
progress_end: int,
force_download: bool = False,
) -> bool:
"""Download through huggingface_hub when available.
@@ -264,7 +272,6 @@ def download_with_hf_hub(
filename=archive_file,
repo_type="model",
local_dir=local_dir,
force_download=force_download,
)
except Exception as e:
emit_progress(
@@ -470,6 +477,89 @@ def _retry_or_raise(err, attempt, max_retries, bytes_downloaded, meta_path,
raise RuntimeError(f"Failed to download after {max_retries} attempts: {err}")
def download_and_verify(
bundle_repo: str,
archive_file: str,
tar_path: str,
expected_sha256: str,
compressed_size: int,
) -> None:
"""Download the bundle archive to tar_path and verify its SHA256.
The first attempt prefers the accelerated Hugging Face client, which
assembles the file with parallel offset writes. On some storage backends
(network mounts, FUSE bind mounts) that write pattern can corrupt the
assembled file even though every transfer step reports success, and
re-running the same client just reproduces the corruption (issue #714).
So a checksum mismatch is retried over the plain sequential downloader:
single-stream append-only writes are the most widely compatible pattern.
(When the accelerated client was unavailable, the first attempt was
already sequential and the retry is simply a fresh download.)
"""
url = f"https://huggingface.co/{bundle_repo}/resolve/main/{archive_file}"
manual_hint = (
f"You can download the bundle manually from:\n"
f" {url}\n"
f"Then upload it via Settings > AI Features > Offline Import."
)
def read_back_sha256() -> str:
# An IO error while hashing a multi-GB file is a live possibility on
# the flaky mounts this retry exists for; it must surface as a JSON
# error frame the Node bridge can parse, not a raw traceback.
try:
return file_sha256(tar_path)
except OSError as e:
fail(
f"Could not read back the downloaded archive to verify it "
f"({e}). This points at the storage backing the data "
f"volume.\n\n{manual_hint}"
)
try:
if not download_with_hf_hub(
bundle_repo, archive_file, tar_path, compressed_size, 2, 85
):
download_with_resume(url, tar_path, compressed_size, 2, 85)
except RuntimeError as e:
fail(f"{e}\n\n{manual_hint}")
emit_progress(86, "Verifying integrity...")
actual = read_back_sha256()
if actual == expected_sha256:
return
# A mismatch is positive evidence of corruption, so discard every byte of
# the first attempt: the assembled file AND the .partial/.meta resume
# sidecars a previously killed run may have left, which the sequential
# downloader would otherwise weld stale bytes from. Cleanup is best-effort
# (the downloader renames over tar_path anyway); the retry matters more.
for stale in (tar_path, tar_path + ".partial", tar_path + ".meta"):
_unlink_quietly(stale)
emit_progress(86, "Checksum mismatch, retrying with the sequential downloader...")
try:
download_with_resume(url, tar_path, compressed_size, 2, 85)
except RuntimeError as e:
fail(
f"Checksum mismatch on the first download (expected "
f"{expected_sha256}, got {actual}), and the sequential "
f"re-download then failed: {e}\n\n{manual_hint}"
)
actual = read_back_sha256()
if actual != expected_sha256:
_unlink_quietly(tar_path)
fail(
f"SHA256 checksum mismatch after re-download.\n"
f"Expected: {expected_sha256}\n"
f"Actual: {actual}\n"
f"The re-download used the plain sequential downloader, so a repeat "
f"mismatch usually points at the storage backing the data volume "
f"(network mounts and FUSE filesystems can alter large files) or a "
f"proxy rewriting the transfer.\n\n{manual_hint}"
)
# -- Safe tar extraction --
def safe_extract(tar_path: str, staging_dir: str) -> None:
@@ -1050,16 +1140,17 @@ def _install() -> None:
# Verify checksum
emit_progress(10, "Verifying checksum...")
if not verify_sha256(tar_path, expected_sha256):
actual_sha256 = file_sha256(tar_path)
if actual_sha256 != expected_sha256:
fail(
f"SHA256 checksum mismatch for local file.\n"
f"Expected: {expected_sha256}\n"
f"Actual: {actual_sha256}\n"
f"This usually means the manifest and archive are out of sync."
)
else:
# Remote mode: download from HuggingFace
bundle_repo = manifest.get("bundleRepo", "deepsafe/feature-bundles")
url = f"https://huggingface.co/{bundle_repo}/resolve/main/{archive_file}"
# Disk space check (early sanity bail before a multi-GB download).
# estimate_extracted covers the extractedSize:0 case so the budget can't
@@ -1080,51 +1171,9 @@ def _install() -> None:
# drifted/upgraded venv doesn't silently fall back to slow urllib.
ensure_hf_hub(venv_path)
try:
if not download_with_hf_hub(
bundle_repo,
archive_file,
tar_path,
compressed_size,
2,
85,
):
download_with_resume(url, tar_path, compressed_size, 2, 85)
except RuntimeError as e:
fail(
f"{e}\n\n"
f"You can download the bundle manually from:\n"
f" {url}\n"
f"Then upload it via Settings > AI Features > Offline Import."
)
# Verify checksum
emit_progress(86, "Verifying integrity...")
if not verify_sha256(tar_path, expected_sha256):
# Delete and retry once from scratch
os.unlink(tar_path)
emit_progress(86, "Checksum mismatch, retrying download...")
try:
if not download_with_hf_hub(
bundle_repo,
archive_file,
tar_path,
compressed_size,
2,
85,
force_download=True,
):
download_with_resume(url, tar_path, compressed_size, 2, 85)
except RuntimeError as e:
fail(str(e))
if not verify_sha256(tar_path, expected_sha256):
os.unlink(tar_path)
fail(
f"SHA256 checksum mismatch after re-download.\n"
f"Expected: {expected_sha256}\n"
f"The archive may be corrupted. Try again later."
)
download_and_verify(
bundle_repo, archive_file, tar_path, expected_sha256, compressed_size
)
# -- Extract to staging --
staging_dir = os.path.join(ai_dir, f"staging-{bundle_id}")
@@ -1,8 +1,12 @@
import hashlib
import importlib.util
import json
import os
import sys
import types
import pytest
def load_installer():
script_path = os.path.join(os.path.dirname(__file__), "..", "install_feature.py")
@@ -169,3 +173,330 @@ def test_ensure_hf_hub_self_heals_missing_client(monkeypatch, tmp_path):
spec = next(part for part in cmd if part.startswith("huggingface-hub["))
assert "hf_xet" in spec
assert "hf_transfer" in spec
# -- download_and_verify: checksum-mismatch retry contract (issue #714) --
GOOD_BYTES = b"good archive bytes"
BAD_BYTES = b"corrupt archive bytes"
def _sha256(data: bytes) -> str:
return hashlib.sha256(data).hexdigest()
def _last_error(capsys) -> str:
"""Parse the JSON error line fail() writes to stderr."""
err_lines = [line for line in capsys.readouterr().err.strip().splitlines() if line]
payload = json.loads(err_lines[-1])
return payload["error"]
def test_download_and_verify_accepts_good_first_download(monkeypatch, tmp_path):
installer = load_installer()
dest = tmp_path / "background-removal-amd64-gpu.tar.gz"
calls = {"hf": 0, "resume": 0}
def fake_hf(*args, **kwargs):
calls["hf"] += 1
with open(args[2], "wb") as f:
f.write(GOOD_BYTES)
return True
def fake_resume(url, tar_path, *rest):
calls["resume"] += 1
monkeypatch.setattr(installer, "download_with_hf_hub", fake_hf)
monkeypatch.setattr(installer, "download_with_resume", fake_resume)
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
installer.download_and_verify(
"deepsafe/feature-bundles",
"v2.0.0/background-removal-amd64-gpu.tar.gz",
str(dest),
_sha256(GOOD_BYTES),
len(GOOD_BYTES),
)
assert calls == {"hf": 1, "resume": 0}
assert dest.read_bytes() == GOOD_BYTES
def test_checksum_mismatch_retries_with_sequential_downloader(monkeypatch, tmp_path):
"""A corrupt accelerated download must be retried over the plain sequential
downloader, not by re-running the transport that just produced bad bytes
(issue #714: hf_xet parallel writes corrupt on some bind mounts, so a
force_download re-run fails the same way forever)."""
installer = load_installer()
dest = tmp_path / "background-removal-amd64-gpu.tar.gz"
calls = {"hf": 0, "resume": 0}
def fake_hf(*args, **kwargs):
calls["hf"] += 1
with open(args[2], "wb") as f:
f.write(BAD_BYTES)
return True
def fake_resume(url, tar_path, *rest):
calls["resume"] += 1
assert not os.path.exists(tar_path), "corrupt file must be gone before retry"
with open(tar_path, "wb") as f:
f.write(GOOD_BYTES)
monkeypatch.setattr(installer, "download_with_hf_hub", fake_hf)
monkeypatch.setattr(installer, "download_with_resume", fake_resume)
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
installer.download_and_verify(
"deepsafe/feature-bundles",
"v2.0.0/background-removal-amd64-gpu.tar.gz",
str(dest),
_sha256(GOOD_BYTES),
len(GOOD_BYTES),
)
assert calls["hf"] == 1, "accelerated client must not be re-run after a mismatch"
assert calls["resume"] == 1
assert dest.read_bytes() == GOOD_BYTES
def test_mismatch_retry_discards_stale_resume_sidecars(monkeypatch, tmp_path):
"""A mismatch is positive evidence of corruption, so the retry must not
resume from .partial/.meta sidecars a previously killed run left behind
(welding stale bytes onto the fresh download would fail the checksum again
and misdiagnose the cause as the user's storage)."""
installer = load_installer()
dest = tmp_path / "background-removal-amd64-gpu.tar.gz"
(tmp_path / "background-removal-amd64-gpu.tar.gz.partial").write_bytes(b"stale")
(tmp_path / "background-removal-amd64-gpu.tar.gz.meta").write_text(
'{"bytesDownloaded": 5}'
)
def fake_hf(*args, **kwargs):
with open(args[2], "wb") as f:
f.write(BAD_BYTES)
return True
def fake_resume(url, tar_path, *rest):
assert not os.path.exists(tar_path + ".partial"), "stale .partial must be wiped"
assert not os.path.exists(tar_path + ".meta"), "stale .meta must be wiped"
with open(tar_path, "wb") as f:
f.write(GOOD_BYTES)
monkeypatch.setattr(installer, "download_with_hf_hub", fake_hf)
monkeypatch.setattr(installer, "download_with_resume", fake_resume)
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
installer.download_and_verify(
"deepsafe/feature-bundles",
"v2.0.0/background-removal-amd64-gpu.tar.gz",
str(dest),
_sha256(GOOD_BYTES),
len(GOOD_BYTES),
)
assert dest.read_bytes() == GOOD_BYTES
def test_hf_unavailable_falls_back_to_sequential_download(monkeypatch, tmp_path):
installer = load_installer()
dest = tmp_path / "background-removal-amd64-gpu.tar.gz"
seen = {}
def fake_resume(url, tar_path, *rest):
seen["url"] = url
with open(tar_path, "wb") as f:
f.write(GOOD_BYTES)
monkeypatch.setattr(installer, "download_with_hf_hub", lambda *a, **k: False)
monkeypatch.setattr(installer, "download_with_resume", fake_resume)
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
installer.download_and_verify(
"deepsafe/feature-bundles",
"v2.0.0/background-removal-amd64-gpu.tar.gz",
str(dest),
_sha256(GOOD_BYTES),
len(GOOD_BYTES),
)
assert seen["url"] == (
"https://huggingface.co/deepsafe/feature-bundles"
"/resolve/main/v2.0.0/background-removal-amd64-gpu.tar.gz"
)
assert dest.read_bytes() == GOOD_BYTES
def test_first_download_failure_reports_manual_hint(monkeypatch, tmp_path, capsys):
installer = load_installer()
dest = tmp_path / "background-removal-amd64-gpu.tar.gz"
def fake_resume(url, tar_path, *rest):
raise RuntimeError("Failed to download after 5 attempts: boom")
monkeypatch.setattr(installer, "download_with_hf_hub", lambda *a, **k: False)
monkeypatch.setattr(installer, "download_with_resume", fake_resume)
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
with pytest.raises(SystemExit):
installer.download_and_verify(
"deepsafe/feature-bundles",
"v2.0.0/background-removal-amd64-gpu.tar.gz",
str(dest),
_sha256(GOOD_BYTES),
len(GOOD_BYTES),
)
error = _last_error(capsys)
assert "Failed to download after 5 attempts: boom" in error
assert "resolve/main/v2.0.0/background-removal-amd64-gpu.tar.gz" in error
assert "Offline Import" in error
def test_retry_download_failure_keeps_mismatch_context(monkeypatch, tmp_path, capsys):
"""When the sequential re-download after a mismatch itself fails, the error
must still say a checksum mismatch started it (with the digests), so a
#714-class corruption is distinguishable from a plain flaky download."""
installer = load_installer()
dest = tmp_path / "background-removal-amd64-gpu.tar.gz"
def fake_hf(*args, **kwargs):
with open(args[2], "wb") as f:
f.write(BAD_BYTES)
return True
def fake_resume(url, tar_path, *rest):
raise RuntimeError("Failed to download after 5 attempts: boom")
monkeypatch.setattr(installer, "download_with_hf_hub", fake_hf)
monkeypatch.setattr(installer, "download_with_resume", fake_resume)
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
with pytest.raises(SystemExit):
installer.download_and_verify(
"deepsafe/feature-bundles",
"v2.0.0/background-removal-amd64-gpu.tar.gz",
str(dest),
_sha256(GOOD_BYTES),
len(GOOD_BYTES),
)
error = _last_error(capsys)
assert "Checksum mismatch" in error
assert _sha256(GOOD_BYTES) in error
assert _sha256(BAD_BYTES) in error
assert "Failed to download after 5 attempts: boom" in error
assert "Offline Import" in error
def test_verify_read_error_fails_with_json_error(monkeypatch, tmp_path, capsys):
"""EIO while hashing the multi-GB archive (a live possibility on the flaky
mounts this retry exists for) must surface as a JSON error frame the Node
bridge can parse, not a raw traceback."""
installer = load_installer()
dest = tmp_path / "background-removal-amd64-gpu.tar.gz"
def fake_hf(*args, **kwargs):
with open(args[2], "wb") as f:
f.write(GOOD_BYTES)
return True
def broken_hash(path):
raise OSError(5, "Input/output error")
monkeypatch.setattr(installer, "download_with_hf_hub", fake_hf)
monkeypatch.setattr(installer, "download_with_resume", lambda *a: None)
monkeypatch.setattr(installer, "file_sha256", broken_hash)
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
with pytest.raises(SystemExit):
installer.download_and_verify(
"deepsafe/feature-bundles",
"v2.0.0/background-removal-amd64-gpu.tar.gz",
str(dest),
_sha256(GOOD_BYTES),
len(GOOD_BYTES),
)
error = _last_error(capsys)
assert "Input/output error" in error
assert "Offline Import" in error
def test_checksum_mismatch_failure_reports_both_hashes(monkeypatch, tmp_path, capsys):
"""When both transports produce wrong bytes, the error must state the
expected AND actual digests and point at the offline-import workaround."""
installer = load_installer()
dest = tmp_path / "background-removal-amd64-gpu.tar.gz"
def fake_hf(*args, **kwargs):
with open(args[2], "wb") as f:
f.write(BAD_BYTES)
return True
def fake_resume(url, tar_path, *rest):
with open(tar_path, "wb") as f:
f.write(BAD_BYTES)
monkeypatch.setattr(installer, "download_with_hf_hub", fake_hf)
monkeypatch.setattr(installer, "download_with_resume", fake_resume)
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
with pytest.raises(SystemExit):
installer.download_and_verify(
"deepsafe/feature-bundles",
"v2.0.0/background-removal-amd64-gpu.tar.gz",
str(dest),
_sha256(GOOD_BYTES),
len(GOOD_BYTES),
)
error = _last_error(capsys)
assert _sha256(GOOD_BYTES) in error
assert _sha256(BAD_BYTES) in error
assert "Offline Import" in error
assert not dest.exists()
def test_local_bundle_checksum_error_reports_actual_hash(monkeypatch, tmp_path, capsys):
installer = load_installer()
archive_entry = {
"file": "v2.0.0/background-removal.tar.gz",
"sha256": _sha256(GOOD_BYTES),
"compressedSize": len(GOOD_BYTES),
"extractedSize": 0,
}
manifest_path = tmp_path / "manifest.json"
manifest_path.write_text(
json.dumps(
{
"bundles": {
"background-removal": {
"archives": {
"amd64-gpu": archive_entry,
"arm64-cpu": archive_entry,
}
}
}
}
)
)
models_dir = tmp_path / "ai" / "models"
models_dir.mkdir(parents=True)
local = tmp_path / "local.tar.gz"
local.write_bytes(BAD_BYTES)
monkeypatch.setenv("SNAPOTTER_BUNDLE_LOCAL_PATH", str(local))
monkeypatch.setattr(
sys,
"argv",
["install_feature.py", "background-removal", str(manifest_path), str(models_dir)],
)
monkeypatch.setattr(installer, "emit_progress", lambda p, s: None)
with pytest.raises(SystemExit):
installer._install()
error = _last_error(capsys)
assert _sha256(GOOD_BYTES) in error
assert _sha256(BAD_BYTES) in error