File size: 5,928 Bytes
87b8147 db6878f 87b8147 61ccd39 87b8147 61ccd39 cd0d564 87b8147 61ccd39 87b8147 | 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 | #!/usr/bin/env python3.13
"""Download a Stokastisk release bundle from Hugging Face (CDN-backed `resolve/` URLs) - standard library only.
python3.13 fetch_release.py stokastisk-0.1.0-kaggle-t4-sm75 [DEST] # prints the bundle directory
* ONE file (the release's `<release>-kaggle-input.zip`, ~37 MB) as a single stream by default (workers=1; `workers=N` opts into N parallel range requests, which stalled against the Hub CDN in testing); falls back
to a single stream when the server does not support ranges, and to `file://` for offline tests;
* SHA256 verified against the `.sha256` sidecar BEFORE extraction; path-traversal-safe extraction; the extracted directory must carry
the release name and `manifests/RELEASE-MANIFEST.json` must say the same name;
* nothing is installed here: run `<bundle>/scripts/kaggle/02_install.py <bundle>` afterwards (it verifies every file in SHA256SUMS first).
Default destination is /tmp (NOT /kaggle/working, which has a 20 GiB quota)."""
import hashlib, json, os, pathlib, shutil, sys, time, urllib.request, zipfile
from concurrent.futures import ThreadPoolExecutor
HF_REPO = "waqasm86/stokastisk"
HF_BASE = f"https://huggingface.co/{HF_REPO}/resolve"
def _open(url, headers=None):
return urllib.request.urlopen(urllib.request.Request(url, headers=headers or {"User-Agent": "stokastisk-fetch/1"}), timeout=60)
def _range_ok(url):
if not url.startswith("http"):
return None
with _open(url, {"Range": "bytes=0-0", "User-Agent": "stokastisk-fetch/1"}) as r:
cr = r.headers.get("Content-Range") # "bytes 0-0/37080766"
return int(cr.rsplit("/", 1)[1]) if r.status == 206 and cr else None
def _download(url, dest, workers=1, retries=3):
"""Download to `<dest>.part`, check free disk space and the advertised length, then replace `dest` atomically; the partial file is removed on any failure."""
dest = pathlib.Path(dest); part = dest.with_name(dest.name + ".part")
part.unlink(missing_ok=True)
try:
total = _range_ok(url)
if total and shutil.disk_usage(dest.parent).free < total + (64 << 20):
raise RuntimeError(f"not enough free disk space in {dest.parent} for {total / 1e6:.0f} MB")
_download_to(url, part, workers, retries)
if total and part.stat().st_size != total:
raise RuntimeError(f"incomplete download: {part.stat().st_size} of {total} bytes")
os.replace(part, dest)
except BaseException:
part.unlink(missing_ok=True)
raise
def _download_to(url, dest, workers=1, retries=3):
dest = pathlib.Path(dest); total = _range_ok(url)
t0 = time.time()
if total and total > (4 << 20):
chunk = -(-total // workers); parts = [(i * chunk, min(total, (i + 1) * chunk) - 1) for i in range(workers) if i * chunk < total]
with open(dest, "wb") as f:
f.truncate(total)
def part(rng):
a, b = rng
for attempt in range(retries):
try:
with _open(url, {"Range": f"bytes={a}-{b}", "User-Agent": "stokastisk-fetch/1"}) as r, open(dest, "r+b") as f:
f.seek(a); shutil.copyfileobj(r, f, 1 << 20)
return
except Exception:
if attempt == retries - 1: raise
time.sleep(2 ** attempt)
with ThreadPoolExecutor(len(parts)) as ex:
list(ex.map(part, parts))
mode = f"{len(parts)} parallel ranges"
else:
for attempt in range(retries):
try:
with _open(url) as r, open(dest, "wb") as f:
shutil.copyfileobj(r, f, 1 << 20)
break
except Exception:
if attempt == retries - 1: raise
time.sleep(2 ** attempt)
mode = "single stream"
n = dest.stat().st_size
print(f"downloaded {n / 1e6:.1f} MB in {time.time() - t0:.1f}s ({n / 1e6 / max(time.time() - t0, 1e-6):.1f} MB/s, {mode})")
def _sha256(path):
h = hashlib.sha256()
with open(path, "rb") as f:
for c in iter(lambda: f.read(1 << 20), b""):
h.update(c)
return h.hexdigest()
def fetch_release(release, dest="/tmp/stokastisk-release", base=None, revision="main", workers=1):
"""Return the verified, extracted bundle directory (…/<release>). `base` overrides the Hub URL prefix (tests use file://)."""
base = (base or HF_BASE).rstrip("/")
root = f"{base}/{revision}/releases/{release}" if base.startswith("http") else f"{base}/releases/{release}"
dest = pathlib.Path(dest); dest.mkdir(parents=True, exist_ok=True)
zname = f"{release}-kaggle-input.zip"; zpath = dest / zname
if zpath.exists(): zpath.unlink()
with _open(f"{root}/{zname}.sha256") as r:
expected = r.read().decode().split()[0]
_download(f"{root}/{zname}", zpath, workers=workers)
got = _sha256(zpath)
if got != expected:
zpath.unlink(); raise RuntimeError(f"SHA256 mismatch for {zname}: expected {expected}, got {got}")
out = dest / release
if out.exists(): shutil.rmtree(out)
with zipfile.ZipFile(zpath) as z:
assert z.testzip() is None
for n in z.namelist():
if n.startswith("/") or ".." in pathlib.PurePosixPath(n).parts or not n.startswith(release + "/"):
raise RuntimeError(f"unsafe or foreign path in archive: {n}")
z.extractall(dest)
m = json.loads((out / "manifests" / "RELEASE-MANIFEST.json").read_text())
if m.get("name") != release:
raise RuntimeError(f"manifest name {m.get('name')!r} != {release!r}")
print(f"release {release}: SHA256 {got[:16]}... verified, extracted to {out}")
return out
if __name__ == "__main__":
if len(sys.argv) < 2:
raise SystemExit(__doc__)
print(fetch_release(sys.argv[1], *(sys.argv[2:3] or ["/tmp/stokastisk-release"])))
|