Spaces:
Running on Zero
Running on Zero
File size: 5,115 Bytes
6ef0a95 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """Ad-hoc check: the shared transfer engine used by fetch_image_model.py.
Proves the four guarantees the image downloader relies on, against a real
(public, ungated) Hub repo so nothing is mocked:
1. first run fetches every file, byte-exact against the published size,
2. a second run skips everything (download only what is missing),
3. an interrupted run resumes from its ledger instead of starting over,
4. a corrupt file of the CORRECT length is rejected by the sha256 gate and
deleted, so no invalid file is ever promoted into place.
python scripts/_selftest_transfer.py
"""
import hashlib
import json
import shutil
import sys
import tempfile
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent))
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import fetch_i2v_model as engine # noqa: E402
# Ungated and small, but a real ~90 MB safetensors weight file, so an
# interruption genuinely leaves a partial transfer behind.
REPO = "sentence-transformers/all-MiniLM-L6-v2"
WEIGHT = "model.safetensors"
ok = True
def check(label, condition):
global ok
ok = ok and bool(condition)
print(f" [{'PASS' if condition else 'FAIL'}] {label}")
work = Path(tempfile.mkdtemp())
try:
manifest = engine.fetch_manifest(REPO, allow=[WEIGHT], ignore=[])
size = manifest["files"][WEIGHT]
expected_sha = manifest["oids"][WEIGHT]
dest = work / WEIGHT
print(f"target: {REPO}/{WEIGHT} ({size / 1e6:.1f} MB)")
print("1. first run fetches the file, byte-exact")
check("no errors", engine.download_snapshot(manifest, work) == [])
check("size matches the published size", dest.stat().st_size == size)
good = dest.read_bytes()
check("sha256 matches the Hub's published hash",
hashlib.sha256(good).hexdigest() == expected_sha)
print("2. second run skips it (only missing files are downloaded)")
before = engine._PROGRESS[0]
check("no errors", engine.download_snapshot(manifest, work) == [])
check("zero bytes re-transferred", engine._PROGRESS[0] == before)
print("3. an interrupted run resumes instead of restarting")
half = size // 2
dest.unlink()
part = dest.with_name(dest.name + ".part")
with part.open("wb") as handle:
handle.truncate(size) # pre-allocated to the full size
with part.open("r+b") as handle: # first half is genuinely good
handle.write(good[:half])
part.with_name(dest.name + ".parts.json").write_text(json.dumps({
"repo": manifest["repo"], "revision": manifest["revision"],
"size": size, "done": [[0, half]],
}))
baseline = engine._PROGRESS[0]
engine._download_file(
manifest["repo"], manifest["revision"], WEIGHT, dest, size,
expect_sha256=expected_sha,
)
re_fetched = engine._PROGRESS[0] - baseline
check(f"only the missing tail was fetched ({re_fetched / 1e6:.1f} MB of {size / 1e6:.1f} MB)",
0 < re_fetched < size * 0.75)
check("resumed file is byte-exact",
hashlib.sha256(dest.read_bytes()).hexdigest() == expected_sha)
check("scratch file and ledger cleaned up",
not part.exists() and not part.with_name(dest.name + ".parts.json").exists())
print("4. a file the ledger calls complete but whose bytes are wrong is rejected")
dest.unlink()
with part.open("wb") as handle:
handle.truncate(size) # right length...
with part.open("r+b") as handle: # ...wrong bytes
handle.seek(size // 2)
handle.write(b"CORRUPTION")
# The ledger claims every byte already landed, so the downloader has
# nothing left to fetch and the sha256 gate is the ONLY thing that can
# catch this. This is the case that once discarded ~2.6 GB of perfectly
# good weights when the hash was read from the CDN `ETag` instead.
part.with_name(dest.name + ".parts.json").write_text(json.dumps({
"repo": manifest["repo"], "revision": manifest["revision"],
"size": size, "done": [[0, size]],
}))
check("size check alone cannot catch it", part.stat().st_size == size)
baseline = engine._PROGRESS[0]
try:
engine._download_file(
manifest["repo"], manifest["revision"], WEIGHT, dest, size,
expect_sha256=expected_sha,
)
check("corrupt file rejected by the sha256 gate", False)
except RuntimeError as exc:
check("corrupt file rejected by the sha256 gate",
"sha256 mismatch" in str(exc))
check("nothing was re-downloaded to paper over it",
engine._PROGRESS[0] == baseline)
check("corrupt scratch and ledger deleted",
not part.exists()
and not part.with_name(dest.name + ".parts.json").exists())
check("no invalid file promoted into place", not dest.exists())
check("verify_snapshot still reports it missing",
engine.verify_snapshot(work, manifest, [WEIGHT], []) != [])
finally:
shutil.rmtree(work, ignore_errors=True)
print("\nRESULT:", "ALL CHECKS PASSED" if ok else "FAILURES PRESENT")
raise SystemExit(0 if ok else 1)
|