Spaces:
Running on Zero
Running on Zero
| """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) | |