christopher-kapic's picture
Upload folder using huggingface_hub
fdc6474 verified
Raw
History Blame Contribute Delete
2.52 kB
#!/usr/bin/env python3
"""Parallel HTTP range downloader with resume. Reads manifest.json entries
{url, start, length, out}; each entry resumes from existing partial file size.
"""
import json
import os
import sys
import time
import urllib.request
from concurrent.futures import ThreadPoolExecutor
CHUNK = 8 << 20
def download(e):
out, url, start, length = e["out"], e["url"], e["start"], e["length"]
os.makedirs(os.path.dirname(out), exist_ok=True)
have = os.path.getsize(out) if os.path.exists(out) else 0
if have >= length:
return 0
for attempt in range(8):
try:
req = urllib.request.Request(url)
req.add_header("Range", f"bytes={start+have}-{start+length-1}")
with urllib.request.urlopen(req, timeout=120) as r:
with open(out, "ab") as f:
while True:
buf = r.read(CHUNK)
if not buf:
break
f.write(buf)
have += len(buf)
if have >= length:
return length
except Exception as ex:
time.sleep(min(2 ** attempt, 60))
have = os.path.getsize(out) if os.path.exists(out) else 0
raise RuntimeError(f"failed: {out} ({have}/{length})")
def main():
manifest = json.load(open(sys.argv[1]))
workers = int(sys.argv[2]) if len(sys.argv) > 2 else 24
total = sum(e["length"] for e in manifest)
done_bytes = 0
t0 = time.time()
failed = []
def wrapped(e):
try:
download(e)
return e["length"]
except Exception as ex:
failed.append((e["out"], str(ex)))
return 0
with ThreadPoolExecutor(workers) as ex:
for i, n in enumerate(ex.map(wrapped, manifest)):
done_bytes += n
if i % 200 == 0 or i == len(manifest) - 1:
el = time.time() - t0
print(f"[{i+1}/{len(manifest)}] {done_bytes/1e9:.1f} GB "
f"({done_bytes/max(el,1)/1e6:.0f} MB/s avg)", flush=True)
if failed:
print("FAILED entries:")
for o, m in failed[:20]:
print(" ", o, m)
sys.exit(1)
# verify
bad = [e for e in manifest
if not os.path.exists(e["out"]) or os.path.getsize(e["out"]) < e["length"]]
print(f"verify: {len(manifest)-len(bad)}/{len(manifest)} complete")
sys.exit(1 if bad else 0)
if __name__ == "__main__":
main()