53 lines
2.0 KiB
Python
53 lines
2.0 KiB
Python
|
"""Download only pinned diagnostic assets into owned ignored data."""
|
||
|
|
import argparse
|
||
|
|
import hashlib
|
||
|
|
import json
|
||
|
|
from pathlib import Path
|
||
|
|
import time
|
||
|
|
import urllib.request
|
||
|
|
import uuid
|
||
|
|
|
||
|
|
ROOT = Path(__file__).resolve().parents[2]
|
||
|
|
DESTINATION = ROOT / ".searchi" / "rerank-eval" / "model"
|
||
|
|
|
||
|
|
|
||
|
|
def digest(path):
|
||
|
|
checksum = hashlib.sha256()
|
||
|
|
with path.open("rb") as handle:
|
||
|
|
for chunk in iter(lambda: handle.read(65536), b""):
|
||
|
|
checksum.update(chunk)
|
||
|
|
return checksum.hexdigest()
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser=argparse.ArgumentParser()
|
||
|
|
parser.add_argument("--reference",action="store_true")
|
||
|
|
args=parser.parse_args()
|
||
|
|
name="reference-manifest.json" if args.reference else "model-manifest.json"
|
||
|
|
manifest = json.loads(Path(__file__).with_name(name).read_text(encoding="utf-8"))
|
||
|
|
for item in manifest["files"]:
|
||
|
|
destination = DESTINATION / item["path"]
|
||
|
|
if destination.is_file() and destination.stat().st_size == item["bytes"] and digest(destination) == item["sha256"]:
|
||
|
|
print("verified cache", item["path"], flush=True)
|
||
|
|
continue
|
||
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
temporary = destination.with_name(destination.name + "." + uuid.uuid4().hex + ".partial")
|
||
|
|
total = 0
|
||
|
|
start = time.monotonic()
|
||
|
|
try:
|
||
|
|
with urllib.request.urlopen(item["url"], timeout=20) as response, temporary.open("xb") as output:
|
||
|
|
while chunk := response.read(65536):
|
||
|
|
total += len(chunk)
|
||
|
|
if total > item["bytes"] or time.monotonic() - start > 180:
|
||
|
|
raise ValueError("Download exceeded size/deadline")
|
||
|
|
output.write(chunk)
|
||
|
|
if total != item["bytes"] or digest(temporary) != item["sha256"]:
|
||
|
|
raise ValueError("Pinned asset hash/size mismatch")
|
||
|
|
temporary.replace(destination)
|
||
|
|
print("verified download", item["path"], total, flush=True)
|
||
|
|
finally:
|
||
|
|
temporary.unlink(missing_ok=True)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|