Merge branch 'f-020-dl' into f-020-int

setup.py

@@ -39,9 +39,11 @@ stop at; --yes alone is not.
 from __future__ import annotations
 
 import argparse
+import calendar
 import ctypes
 import glob
 import hashlib
+import http.client
 import json
 import math
 import os
@@ -52,8 +54,10 @@ import struct
 import subprocess
 import sys
 import textwrap
+import threading
 import time
 import urllib.error
+import urllib.parse
 import urllib.request
 import zipfile
 from pathlib import Path
@@ -1904,6 +1908,389 @@ def drop_archive(z: Path) -> None:
     z.with_name(z.name + ".done").unlink(missing_ok=True)
 
 
+# Parallel ranged downloads (0.2.0).  One TCP stream often cannot fill a fast line (a long or lossy route, a CDN that limits
+# each connection): download() takes the file in DL_CHUNK pieces over several connections, writes each piece at its own offset,
+# and keeps the finished pieces in <name>.part.chunks so a stopped download continues where it was.  Anything that is not a
+# ranged file of known length (no Range support, no length, a small file) is one stream, as before.
+DL_CONNECTIONS = 8              # STRATA_DL_CONNECTIONS=1 restores one stream (for testers)
+DL_CHUNK = 32 << 20             # bytes per ranged request; a failed request loses at most this much
+DL_MIN_PARALLEL = 256 << 20     # smaller files are one stream
+DL_CHUNK_TRIES = 5              # tries per piece (1, 2, 4, 8 s apart) before the round is given up and run again
+DL_ROUNDS = 30                  # rounds (10 s apart) before the download is given up, like the 30 tries of one stream
+DL_BACKOFF = 1.0                # seconds before the second try of a piece, doubled each time (tests set 0)
+DL_SAVE_EVERY = 5.0             # seconds between writes of the progress file
+
+
+def dl_connections() -> int:
+    """How many connections a big download uses: STRATA_DL_CONNECTIONS, else DL_CONNECTIONS (1 = one stream)."""
+    raw = (os.environ.get("STRATA_DL_CONNECTIONS") or "").strip()
+    try:
+        n = int(raw) if raw else DL_CONNECTIONS
+    except ValueError:
+        n = DL_CONNECTIONS
+    return max(1, min(n, 32))
+
+
+class _NoRange(Exception):
+    """The server does not answer a byte range with that range (no 206, or another range): one stream instead."""
+
+
+class _DiskError(Exception):
+    """The .part file cannot be written (disk full, no permission): retrying would not help."""
+
+
+class _PieceFailed(Exception):
+    """A piece failed DL_CHUNK_TRIES times."""
+
+
+def _url_ttl(url: str):
+    """Seconds a signed URL (the CDN's, after the redirect) is still valid by its own query, else None."""
+    q = urllib.parse.parse_qs(urllib.parse.urlsplit(url).query)
+    try:
+        if "Expires" in q:
+            ttl = float(q["Expires"][0]) - time.time()
+        elif "X-Amz-Expires" in q and "X-Amz-Date" in q:
+            t0 = calendar.timegm(time.strptime(q["X-Amz-Date"][0], "%Y%m%dT%H%M%SZ"))
+            ttl = t0 + float(q["X-Amz-Expires"][0]) - time.time()
+        else:
+            return None
+    except (ValueError, OverflowError):
+        return None
+    return ttl if 0 < ttl < 30 * 86400 else None       # a PC clock that is far off says nothing: a 403 re-resolves
+
+
+class _Source:
+    """The URL the pieces are read from: the redirect (Hugging Face answers with a signed, time-limited CDN address) is
+    resolved once and reused; it is resolved again a little before it expires and after a 401/403/410."""
+
+    def __init__(self, url: str, total: int):
+        self.url, self.total = url, total
+        self.final, self.deadline = None, None
+        self.lock = threading.Lock()
+
+    def resolve(self, stale=None) -> str:
+        with self.lock:
+            if stale is not None and self.final is not None and self.final != stale:
+                return self.final                        # another connection already got a fresh address
+            req = urllib.request.Request(self.url, headers={"User-Agent": "strata-setup", "Range": "bytes=0-0"})
+            with _urlopen(req, timeout=60) as r:
+                cr = (r.headers.get("Content-Range") or "").strip()
+                if r.status != 206 or cr not in (f"bytes 0-0/{self.total}", "bytes 0-0/*"):
+                    raise _NoRange(f"HTTP {r.status}, Content-Range {cr!r}")
+                r.read(1)
+                final = r.geturl() or self.url
+            ttl = _url_ttl(final)
+            self.final = final
+            self.deadline = time.monotonic() + max(ttl - 120, 0) if ttl else None
+            return final
+
+    def get(self) -> str:
+        with self.lock:
+            cur, expired = self.final, self.deadline is not None and time.monotonic() >= self.deadline
+        return self.resolve(stale=cur) if cur is None or expired else cur
+
+
+class _Pieces:
+    """The shared state of one parallel download: finished pieces, bytes of the pieces in flight, the stop flags."""
+
+    def __init__(self, total: int, done: set):
+        self.total, self.n = total, -(-total // DL_CHUNK)
+        self.done, self.flight = set(done), {}
+        self.lock = threading.Lock()
+        self.stop = threading.Event()                    # the whole download stops (Ctrl-C, a fatal error)
+        self.abort_round = threading.Event()             # a piece failed for good: end this round, run the next
+        self.fatal, self.errors = None, []
+
+    def size(self, i: int) -> int:
+        return min((i + 1) * DL_CHUNK, self.total) - i * DL_CHUNK
+
+    def add(self, i: int, n: int):
+        with self.lock:
+            self.flight[i] = self.flight.get(i, 0) + n
+
+    def finish(self, i: int):
+        with self.lock:
+            self.flight.pop(i, None)
+            self.done.add(i)
+
+    def have(self) -> int:
+        with self.lock:
+            return sum(self.size(i) for i in self.done) + sum(self.flight.values())
+
+    def snapshot(self) -> set:
+        with self.lock:
+            return set(self.done)
+
+
+def _ranges(done) -> list:
+    out = []
+    for i in sorted(done):
+        if out and out[-1][1] == i:
+            out[-1][1] = i + 1
+        else:
+            out.append([i, i + 1])
+    return out
+
+
+def _prefix_bytes(done: set, total: int) -> int:
+    """Bytes of the unbroken run of finished pieces from the start of the file."""
+    i = 0
+    while i in done:
+        i += 1
+    return min(i * DL_CHUNK, total)
+
+
+def _load_pieces(part: Path, total: int) -> set:
+    """The finished pieces of an earlier run.  <name>.part.chunks lists them; a .part without it comes from one stream
+    (an older setup, or STRATA_DL_CONNECTIONS=1): its bytes are a run from the start, so those pieces are finished."""
+    side = part.with_name(part.name + ".chunks")
+    n = -(-total // DL_CHUNK)
+    try:
+        size = part.stat().st_size if part.is_file() else None
+    except OSError:
+        size = None
+    if side.is_file():
+        try:
+            data = json.loads(side.read_text(encoding="utf-8"))
+            if data.get("v") == 1 and data.get("total") == total and data.get("chunk") == DL_CHUNK and size is not None:
+                done = {i for a, b in data["done"] for i in range(int(a), min(int(b), n))}
+                return {i for i in done if min((i + 1) * DL_CHUNK, total) <= size}
+        except (OSError, ValueError, TypeError, KeyError):
+            pass
+        side.unlink(missing_ok=True)                    # not this file's list: start the file over
+        part.unlink(missing_ok=True)
+        return set()
+    if size is None:
+        return set()
+    if size > total:
+        part.unlink(missing_ok=True)
+        return set()
+    return set(range(n)) if size >= total else set(range(size // DL_CHUNK))
+
+
+def _save_pieces(part: Path, total: int, done: set):
+    """Flush the data first, then record the pieces: a list never names bytes the disk does not have."""
+    try:
+        with open(part, "r+b") as g:
+            os.fsync(g.fileno())
+        side = part.with_name(part.name + ".chunks")
+        tmp = side.with_name(side.name + ".tmp")
+        tmp.write_text(json.dumps({"v": 1, "total": total, "chunk": DL_CHUNK, "done": _ranges(done)}), encoding="utf-8")
+        os.replace(tmp, side)
+    except OSError:
+        pass                                            # only costs re-downloading a few pieces after a stop
+
+
+def _single_ready(part: Path):
+    """Before one stream continues a .part: keep only the unbroken run from the start of the file (what a parallel
+    download left may have gaps behind it) and drop the list."""
+    side = part.with_name(part.name + ".chunks")
+    if not side.is_file():
+        return
+    try:
+        data = json.loads(side.read_text(encoding="utf-8"))
+        keep = _prefix_bytes({i for a, b in data["done"] for i in range(int(a), int(b))}, int(data["total"]))
+        if data.get("chunk") != DL_CHUNK:
+            keep = 0
+    except (OSError, ValueError, TypeError, KeyError):
+        keep = 0
+    try:
+        if part.is_file():
+            with open(part, "r+b") as g:
+                g.truncate(min(keep, part.stat().st_size))
+    except OSError:
+        part.unlink(missing_ok=True)
+    side.unlink(missing_ok=True)
+
+
+def _fetch_piece(src: _Source, st: _Pieces, f, i: int):
+    start = i * DL_CHUNK
+    end = start + st.size(i) - 1
+    pos, tries = start, 0
+    with st.lock:
+        st.flight[i] = 0
+    while True:
+        if st.stop.is_set():
+            return
+        url, err, wait = None, None, DL_BACKOFF * 2 ** tries
+        try:
+            url = src.get()
+            req = urllib.request.Request(url, headers={"User-Agent": "strata-setup", "Range": f"bytes={pos}-{end}"})
+            with _urlopen(req, timeout=60) as r:
+                if r.status != 206 or not (r.headers.get("Content-Range") or "").strip().startswith(f"bytes {pos}-{end}/"):
+                    raise _NoRange(f"HTTP {r.status}")
+                while pos <= end:
+                    if st.stop.is_set():
+                        return
+                    b = r.read(min(1 << 20, end + 1 - pos))
+                    if not b:
+                        raise IOError("the connection closed early")
+                    try:
+                        f.seek(pos)
+                        f.write(b)
+                    except OSError as e:
+                        raise _DiskError(str(e)) from e
+                    pos += len(b)
+                    st.add(i, len(b))
+            st.finish(i)
+            return
+        except (_NoRange, _DiskError):
+            raise
+        except urllib.error.HTTPError as e:
+            err = e
+            if e.code in (401, 403, 410) and url is not None:       # the signed address ran out: ask for a new one
+                try:
+                    src.resolve(stale=url)
+                    wait = 0
+                except _NoRange:
+                    raise
+                except (OSError, http.client.HTTPException):
+                    pass
+            elif e.code in (429, 503):
+                try:
+                    wait = max(wait, min(float(e.headers.get("Retry-After") or 0), 60))
+                except (TypeError, ValueError):
+                    pass
+        except (OSError, http.client.HTTPException) as e:
+            err = e
+        tries += 1
+        if tries >= DL_CHUNK_TRIES:
+            raise _PieceFailed(str(err))
+        st.stop.wait(wait)
+
+
+def _piece_worker(src: _Source, st: _Pieces, todo: list, part: Path):
+    try:
+        with open(part, "r+b", buffering=0) as f:
+            while not (st.stop.is_set() or st.abort_round.is_set()):
+                with st.lock:
+                    if not todo:
+                        return
+                    i = todo.pop(0)
+                _fetch_piece(src, st, f, i)
+    except _PieceFailed as e:
+        st.errors.append(e)
+        st.abort_round.set()
+    except (_NoRange, _DiskError) as e:
+        st.fatal = st.fatal or e
+        st.stop.set()
+    except OSError as e:                                # the .part itself could not be opened
+        st.fatal = st.fatal or _DiskError(str(e))
+        st.stop.set()
+
+
+def parallel_download(url: str, part: Path, total: int, label: str, conns: int, name: str) -> bool:
+    """Fill `part` (total bytes) over `conns` ranged connections.  True when every byte is there; False when the server
+    does not do ranges (the caller reads the file in one stream)."""
+    src = _Source(url, total)
+    try:
+        src.resolve()
+    except _NoRange:
+        return False
+    except (OSError, http.client.HTTPException):
+        return False                                    # the one-stream path has its own retries and messages
+    done = _load_pieces(part, total)
+    st = _Pieces(total, done)
+    part.parent.mkdir(parents=True, exist_ok=True)
+    if not part.exists():
+        part.touch()
+    last_print = last_save = 0.0
+
+    def report(force=False):
+        nonlocal last_print, last_save
+        now = time.time()
+        if now - last_print > 2 or force:
+            last_print = now
+            have = st.have()
+            print(f"\r  {label}: {have / 1e9:6.2f} / {total / 1e9:.2f} GB ({100 * have / total:.0f}%)   ", end="", flush=True)
+        if now - last_save > DL_SAVE_EVERY:
+            last_save = now
+            _save_pieces(part, total, st.snapshot())
+
+    try:
+        for rnd in range(DL_ROUNDS):
+            todo = [i for i in range(st.n) if i not in st.done]
+            if not todo:
+                break
+            st.abort_round.clear()
+            st.errors.clear()
+            threads = [threading.Thread(target=_piece_worker, args=(src, st, todo, part), daemon=True)
+                       for _ in range(min(conns, len(todo)))]
+            for t in threads:
+                t.start()
+            while any(t.is_alive() for t in threads):
+                for t in threads:
+                    t.join(0.25)
+                report()
+            if st.fatal is not None:
+                break
+            if st.errors:
+                print()
+                warn(f"download interrupted ({st.errors[0]}); retrying in 10 s ...")
+                _save_pieces(part, total, st.snapshot())
+                time.sleep(10)
+    except BaseException:
+        st.stop.set()                                   # Ctrl-C: let the connections stop, keep what is finished
+        _save_pieces(part, total, st.snapshot())
+        raise
+    report(force=True)
+    print()
+    if isinstance(st.fatal, _DiskError):
+        _save_pieces(part, total, st.snapshot())
+        fail(f"cannot write {part.name}: {st.fatal}", "free some disk space (or pick another --models-dir) and run it again")
+    if isinstance(st.fatal, _NoRange):
+        _save_pieces(part, total, st.snapshot())
+        return False                                    # the server stopped doing ranges: one stream from the unbroken start
+    if len(st.done) < st.n:
+        _save_pieces(part, total, st.snapshot())
+        have = st.have()
+        fail(f"could not finish downloading {name}: {have:,} of {total:,} bytes",
+             "check your internet connection and run it again (the download resumes where it stopped)")
+    try:
+        with open(part, "r+b") as g:
+            g.truncate(total)
+            os.fsync(g.fileno())
+    except OSError:
+        pass
+    part.with_name(part.name + ".chunks").unlink(missing_ok=True)
+    return True
+
+
+def download_single(url: str, part: Path, total: int, label: str):
+    """One stream with a Range resume from the size of `part` (the way every download worked before parallel pieces)."""
+    _single_ready(part)
+    have = part.stat().st_size if part.exists() else 0
+    for attempt in range(30):
+        if total and have >= total:                    # stopped after the last byte, before the rename: nothing to
+            break                                      # ask for (a range past the end is a 416, retried 30 times)
+        try:
+            req = urllib.request.Request(url, headers={"User-Agent": "strata-setup", "Range": f"bytes={have}-"})
+            with _urlopen(req, timeout=60) as r, open(part, "ab" if have else "wb") as f:
+                if have and r.status != 206:                     # the server ignored the range: start over
+                    f.seek(0)
+                    f.truncate()
+                    have = 0
+                last = 0.0
+                while True:
+                    b = r.read(8 << 20)
+                    if not b:
+                        break
+                    f.write(b)
+                    have += len(b)
+                    if time.time() - last > 2:
+                        last = time.time()
+                        size = f"{have / 1e9:6.2f} / {total / 1e9:.2f} GB ({100 * have / total:.0f}%)" if total \
+                            else f"{have / 1e6:7.1f} MB"
+                        print(f"\r  {label}: {size}   ", end="", flush=True)
+            print()
+            if not total or have >= total:
+                break
+        except (OSError, http.client.HTTPException) as e:
+            print()
+            warn(f"download interrupted ({e}); retrying in 10 s ...")
+            time.sleep(10)
+
+
 def download(url, dst: Path, what=None):
     """Resumable HTTP(S) download with a progress line; `file://` and plain paths are copied (tests, mirrors).
     A finished file gets a <name>.done mark, so a later run skips it without asking the server."""
@@ -1962,36 +2349,10 @@ def download(url, dst: Path, what=None):
         mark(dst)
         ok(f"{what or dst.name} already downloaded")
         return
-    have = part.stat().st_size if part.exists() else 0
-    for attempt in range(30):
-        if total and have >= total:                    # stopped after the last byte, before the rename: nothing to
-            break                                      # ask for (a range past the end is a 416, retried 30 times)
-        try:
-            req = urllib.request.Request(url, headers={"User-Agent": "strata-setup", "Range": f"bytes={have}-"})
-            with _urlopen(req, timeout=60) as r, open(part, "ab" if have else "wb") as f:
-                if have and r.status != 206:                     # the server ignored the range: start over
-                    f.seek(0)
-                    f.truncate()
-                    have = 0
-                last = 0.0
-                while True:
-                    b = r.read(8 << 20)
-                    if not b:
-                        break
-                    f.write(b)
-                    have += len(b)
-                    if time.time() - last > 2:
-                        last = time.time()
-                        size = f"{have / 1e9:6.2f} / {total / 1e9:.2f} GB ({100 * have / total:.0f}%)" if total \
-                            else f"{have / 1e6:7.1f} MB"
-                        print(f"\r  {what or dst.name}: {size}   ", end="", flush=True)
-            print()
-            if not total or have >= total:
-                break
-        except OSError as e:
-            print()
-            warn(f"download interrupted ({e}); retrying in 10 s ...")
-            time.sleep(10)
+    label = what or dst.name
+    conns = dl_connections()
+    if not (total >= DL_MIN_PARALLEL and conns > 1 and parallel_download(url, part, total, label, conns, dst.name)):
+        download_single(url, part, total, label)
     if total and part.stat().st_size != total:
         fail(f"could not finish downloading {dst.name}: {part.stat().st_size:,} bytes on disk, the server says {total:,}",
              "check your internet connection and run it again (the download resumes where it stopped)")

tools/mtp_fetch.py

@@ -17,6 +17,8 @@ downloaded catches it (the drafter then accepts nothing, silently).  A range rea
 Content-Range is the range asked for, and every tensor of the pinned revision is checked against SHA256 below.
 """
 import argparse
+import collections
+import concurrent.futures
 import hashlib
 import json
 import os
@@ -133,6 +135,42 @@ def get(url, start=None, end=None, retries=4):
             print("retry %s: %s" % (url, e), file=sys.stderr)
 
 
+def connections():
+    """Ranges read at once: STRATA_DL_CONNECTIONS (the setting of setup's downloads), else 8; 1 reads them one by one."""
+    try:
+        n = int((os.environ.get("STRATA_DL_CONNECTIONS") or "").strip() or 8)
+    except ValueError:
+        n = 8
+    return max(1, min(n, 32))
+
+
+def ranged(url, start, last, chunk, workers):
+    """The bytes start..last of `url` in order, read in `chunk`-sized ranges with `workers` of them in flight.  Each range is
+    checked by get() (a 206 with the range asked for), so a server that ignores Range still fails here."""
+    spans = []
+    pos = start
+    while pos <= last:
+        spans.append((pos, min(pos + chunk - 1, last)))
+        pos = spans[-1][1] + 1
+    if workers <= 1 or len(spans) <= 1:
+        for a, b in spans:
+            yield get(url, a, b)
+        return
+    todo = iter(spans)
+    window = collections.deque()
+    with concurrent.futures.ThreadPoolExecutor(min(workers, len(spans))) as ex:
+        for a, b in todo:
+            window.append(ex.submit(get, url, a, b))
+            if len(window) >= workers:
+                break
+        while window:
+            data = window.popleft().result()
+            for a, b in todo:
+                window.append(ex.submit(get, url, a, b))
+                break
+            yield data
+
+
 def check_range(status, content_range, start, end):
     """#327: a range read is only what was asked for when the server says so: 206 and `bytes START-END/...`."""
     if status != 206:
@@ -229,7 +267,8 @@ def fetch(out, only):
     rows = inv["tensors"] if same else inventory(out)
     os.makedirs(tdir, exist_ok=True)
     manifest = []
-    chunk = 64 << 20
+    workers = connections()
+    chunk = (32 if workers > 1 else 64) << 20
     for r in rows:
         if only and only not in r["name"]:
             continue
@@ -245,10 +284,9 @@ def fetch(out, only):
                 have = 0
             with open(path, "ab") as f:
                 pos = r["start"] + have
-                while pos <= r["end"]:
-                    end = min(pos + chunk - 1, r["end"])
-                    f.write(get(REPO + r["shard"], pos, end))
-                    pos = end + 1
+                for data in ranged(REPO + r["shard"], pos, r["end"], chunk, workers):
+                    f.write(data)
+                    pos += len(data)
                     print("%s %.0f%%" % (r["name"], 100 * (pos - r["start"]) / r["bytes"]), file=sys.stderr)
             if os.path.getsize(path) != r["bytes"]:
                 sys.exit("%s: size %d != %d" % (path, os.path.getsize(path), r["bytes"]))

tools/strata_mcp.py

@@ -724,7 +724,8 @@ class Strata:
             else:
                 remove.append({"path": str(p), "kind": "folder" if is_dir else "file", "bytes": self.tree_bytes(p)})
                 if not is_dir:
-                    for extra in (p.with_name(p.name + ".done"), p.with_name(p.name + ".part")):
+                    for extra in (p.with_name(p.name + ".done"), p.with_name(p.name + ".part"),
+                                  p.with_name(p.name + ".part.chunks")):
                         if extra.exists():
                             remove.append({"path": str(extra), "kind": "file", "bytes": self.tree_bytes(extra)})
         tag = config.stem[len("strata-"):]

tools/test_mtp_fetch.py

@@ -128,6 +128,34 @@ class Ranges(FetchCase):
         self.assertEqual({r["name"]: r["sha256"] for r in manifest}, HASHES)
         self.assertEqual(M.verify(str(self.out)), [])
 
+    def test_ranges_are_read_in_parallel_and_come_back_in_order(self):
+        data = bytes(range(256)) * 40
+        seen = []
+
+        def get(url, a, b):
+            seen.append((a, b))
+            return data[a:b + 1]
+
+        for workers in (1, 4, 32):
+            seen.clear()
+            with self.subTest(workers=workers), mock.patch.object(M, "get", get):
+                got = b"".join(M.ranged("u", 100, len(data) - 1, 1000, workers))
+            self.assertEqual(got, data[100:])
+            self.assertEqual(sorted(seen)[0], (100, 1099))
+        with mock.patch.dict(M.os.environ, {"STRATA_DL_CONNECTIONS": "1"}):
+            self.assertEqual(M.connections(), 1)
+        with mock.patch.dict(M.os.environ, {"STRATA_DL_CONNECTIONS": ""}):
+            self.assertEqual(M.connections(), 8)
+
+    def test_a_failing_range_stops_the_parallel_read(self):
+        def get(url, a, b):
+            if a >= 2000:
+                raise IOError("boom")
+            return bytes(b - a + 1)
+
+        with mock.patch.object(M, "get", get), self.assertRaises(IOError):
+            list(M.ranged("u", 0, 9999, 1000, 4))
+
     def test_wrong_bytes_are_fetched_again_once_then_refused(self):
         self.fetch(Mirror())
         bad = bytes(len(A))                       # an install the #327 mirror corrupted: right size, wrong bytes

tools/test_setup_download.py

@@ -0,0 +1,473 @@
+"""Tests for setup.py's parallel ranged downloads (0.2.0): a local threaded HTTP server that supports Range stands in
+for Hugging Face's CDN.  Covers the parallel path, a resume after a kill and after a failed round, a server without Range,
+a piece that fails once, a signed address that runs out, the SHA-256 pin, STRATA_DL_CONNECTIONS=1 and the progress line.
+
+    python -m unittest tools.test_setup_download
+"""
+from __future__ import annotations
+
+import contextlib
+import hashlib
+import http.server
+import io
+import json
+import os
+import re
+import subprocess
+import sys
+import tempfile
+import threading
+import time
+import unittest
+from pathlib import Path
+from unittest import mock
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+sys.path.insert(0, str(ROOT / "tools"))
+import setup  # noqa: E402
+import strata_mcp  # noqa: E402
+
+CHUNK = 64 << 10                                       # the tests' piece size (setup uses 32 MiB)
+DATA = bytes((i * 7 + (i >> 8)) & 0xFF for i in range(CHUNK * 40 + 12345))       # 40 whole pieces and a short last one
+NAME = "t.gguf"
+
+
+class Server:
+    """A threaded HTTP server for one file.  `faults(path, start, end, nth)` may answer a request with an error."""
+
+    def __init__(self, data=DATA, ranges=True, length=True, redirect=False, delay=0.002, faults=None):
+        self.data, self.ranges, self.length, self.redirect, self.delay = data, ranges, length, redirect, delay
+        self.faults = faults
+        self.lock = threading.Lock()
+        self.gets, self.heads, self.redirects = [], 0, 0           # gets: (start, end) of every ranged GET
+        self.active = self.max_active = 0
+        self.nth = {}
+        outer = self
+
+        class Handler(http.server.BaseHTTPRequestHandler):
+            def log_message(self, *a):
+                pass
+
+            def do_HEAD(self):
+                outer.handle(self, head=True)
+
+            def do_GET(self):
+                outer.handle(self, head=False)
+
+        self.httpd = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+        self.httpd.daemon_threads = True
+        self.port = self.httpd.server_address[1]
+        self.thread = threading.Thread(target=self.httpd.serve_forever, daemon=True)
+
+    @property
+    def url(self):
+        return f"http://127.0.0.1:{self.port}/{NAME}"
+
+    def __enter__(self):
+        self.thread.start()
+        return self
+
+    def __exit__(self, *exc):
+        self.httpd.shutdown()
+        self.httpd.server_close()
+
+    def handle(self, h, head):
+        path = h.path.split("?")[0]
+        if self.redirect and path == "/" + NAME:
+            with self.lock:
+                self.redirects += 1
+            h.send_response(302)
+            h.send_header("Location", f"http://127.0.0.1:{self.port}/cdn/{NAME}?Expires={int(time.time()) + 3600}&Signature=x")
+            h.send_header("Content-Length", "0")
+            h.end_headers()
+            return
+        size = len(self.data)
+        rng = h.headers.get("Range")
+        start, end = 0, size - 1
+        if rng and self.ranges:
+            m = re.match(r"bytes=(\d+)-(\d*)$", rng)
+            start = int(m.group(1))
+            end = min(int(m.group(2)), size - 1) if m.group(2) else size - 1
+        if head:
+            with self.lock:
+                self.heads += 1
+            h.send_response(200)
+            if self.length:
+                h.send_header("Content-Length", str(size))
+            h.send_header("Content-Type", "application/octet-stream")
+            h.end_headers()
+            return
+        with self.lock:
+            self.gets.append((start, end) if rng and self.ranges else (0, size - 1))
+            key = (start, end)
+            self.nth[key] = n = self.nth.get(key, 0) + 1
+            self.active += 1
+            self.max_active = max(self.max_active, self.active)
+        try:
+            fault = self.faults(path, start, end, n) if self.faults else None
+            if fault and fault[0] == "status":
+                h.send_response(fault[1])
+                h.send_header("Content-Length", "0")
+                h.end_headers()
+                return
+            body = self.data[start:end + 1]
+            if rng and self.ranges:
+                h.send_response(206)
+                h.send_header("Content-Range", f"bytes {start}-{end}/{size}")
+            else:
+                h.send_response(200)
+            h.send_header("Content-Length", str(len(body)))
+            h.end_headers()
+            cut = fault[1] if fault and fault[0] == "cut" else None
+            try:
+                for i in range(0, len(body), 16384):
+                    if cut is not None and i >= cut:
+                        h.connection.shutdown(2)           # the connection drops in the middle of the body
+                        return
+                    h.wfile.write(body[i:i + 16384])
+                    if self.delay:
+                        time.sleep(self.delay)
+            except OSError:
+                pass
+        finally:
+            with self.lock:
+                self.active -= 1
+
+
+def pin_of(data):
+    return len(data), hashlib.sha256(data).hexdigest()
+
+
+@contextlib.contextmanager
+def small_pieces(**extra):
+    """Setup's piece size, size floor and waits scaled down to the test file."""
+    consts = dict(DL_CHUNK=CHUNK, DL_MIN_PARALLEL=0, DL_BACKOFF=0.0, DL_SAVE_EVERY=0.05, DL_CONNECTIONS=4)
+    consts.update(extra)
+    with contextlib.ExitStack() as st:
+        for k, v in consts.items():
+            st.enter_context(mock.patch.object(setup, k, v))
+        st.enter_context(mock.patch.object(setup.time, "sleep", lambda s: None))
+        env = {k: v for k, v in os.environ.items() if k not in ("STRATA_DL_CONNECTIONS", "STRATA_ALLOW_UNPINNED")}
+        env["STRATA_SOURCE"] = "huggingface"
+        st.enter_context(mock.patch.dict(os.environ, env, clear=True))
+        yield
+
+
+def run(server, d, pin="auto", name=NAME):
+    """setup.download() of the server's file into folder d; returns (exit code or None, stdout)."""
+    dst = Path(d) / name
+    want = pin_of(DATA) if pin == "auto" else pin
+    out = io.StringIO()
+    code = None
+    with mock.patch.object(setup, "file_pin", lambda u: want), contextlib.redirect_stdout(out):
+        try:
+            setup.download(server.url, dst, "the model")
+        except SystemExit as e:
+            code = e.code if e.code is not None else 0
+    return code, out.getvalue()
+
+
+def asked(srv):
+    """The pieces (indexes) the server was asked for, one entry per request; the one-byte probe is left out."""
+    return [s // CHUNK for s, e in srv.gets if (s, e) != (0, 0)]
+
+
+def files(d):
+    return sorted(p.name for p in Path(d).iterdir())
+
+
+class ParallelDownload(unittest.TestCase):
+    def test_a_big_file_comes_over_several_connections(self):
+        with small_pieces(), Server(redirect=True) as srv, tempfile.TemporaryDirectory() as d:
+            code, out = run(srv, d)
+            self.assertIsNone(code, out)
+            self.assertEqual((Path(d) / NAME).read_bytes(), DATA)
+            self.assertEqual(files(d), [NAME, NAME + ".done"])             # no .part, no piece list left behind
+            self.assertIn("sha256 " + hashlib.sha256(DATA).hexdigest(), (Path(d) / (NAME + ".done")).read_text())
+            self.assertGreater(srv.max_active, 1)
+            self.assertLessEqual(srv.max_active, 4)
+            self.assertEqual(sorted(asked(srv)), list(range(41)))           # every piece once, none twice
+            self.assertEqual(srv.redirects, 2)                              # HEAD and the one probe: the pieces use the CDN address
+
+    def test_the_redirect_is_resolved_once(self):
+        with small_pieces(), Server(redirect=True) as srv, tempfile.TemporaryDirectory() as d:
+            code, out = run(srv, d)
+            self.assertIsNone(code, out)
+            self.assertLessEqual(srv.redirects, 2)
+
+    def test_a_small_file_is_one_stream(self):
+        with small_pieces(DL_MIN_PARALLEL=len(DATA) + 1), Server() as srv, tempfile.TemporaryDirectory() as d:
+            code, out = run(srv, d)
+            self.assertIsNone(code, out)
+            self.assertEqual((Path(d) / NAME).read_bytes(), DATA)
+            self.assertEqual(srv.max_active, 1)
+            self.assertEqual(len(srv.gets), 1)
+
+    def test_one_connection_variable_restores_one_stream(self):
+        with small_pieces(), Server() as srv, tempfile.TemporaryDirectory() as d:
+            with mock.patch.dict(os.environ, {"STRATA_DL_CONNECTIONS": "1"}):
+                code, out = run(srv, d)
+            self.assertIsNone(code, out)
+            self.assertEqual((Path(d) / NAME).read_bytes(), DATA)
+            self.assertEqual(srv.max_active, 1)
+            self.assertEqual(srv.gets, [(0, len(DATA) - 1)])                # one request for the whole file
+
+    def test_connection_count_variable(self):
+        with small_pieces():
+            self.assertEqual(setup.dl_connections(), 4)
+            for raw, want in (("16", 16), ("0", 1), ("-3", 1), ("99", 32), ("x", 4), ("", 4)):
+                with mock.patch.dict(os.environ, {"STRATA_DL_CONNECTIONS": raw}):
+                    self.assertEqual(setup.dl_connections(), want, raw)
+
+
+class NoRangeSupport(unittest.TestCase):
+    def test_a_server_that_ignores_range_is_read_in_one_stream(self):
+        with small_pieces(), Server(ranges=False) as srv, tempfile.TemporaryDirectory() as d:
+            code, out = run(srv, d)
+            self.assertIsNone(code, out)
+            self.assertEqual((Path(d) / NAME).read_bytes(), DATA)
+            self.assertEqual(srv.max_active, 1)
+            self.assertEqual(files(d), [NAME, NAME + ".done"])
+
+    def test_a_server_without_a_length_is_read_in_one_stream(self):
+        with small_pieces(), Server(length=False) as srv, tempfile.TemporaryDirectory() as d:
+            code, out = run(srv, d, pin=None)                                # no length, so nothing to check a pin against first
+            self.assertIsNone(code, out)
+            self.assertEqual((Path(d) / NAME).read_bytes(), DATA)
+            self.assertEqual(srv.max_active, 1)
+
+    def test_range_stopping_in_the_middle_falls_back_to_the_unbroken_start(self):
+        def faults(path, start, end, nth):
+            return None
+
+        with small_pieces(), Server(faults=faults) as srv, tempfile.TemporaryDirectory() as d:
+            real = setup._fetch_piece
+            calls = []
+
+            def fetch(src, st, f, i):
+                calls.append(i)
+                if len(calls) == 6:
+                    raise setup._NoRange("HTTP 200")
+                return real(src, st, f, i)
+
+            with mock.patch.object(setup, "_fetch_piece", fetch):
+                code, out = run(srv, d)
+            self.assertIsNone(code, out)
+            self.assertEqual((Path(d) / NAME).read_bytes(), DATA)           # the pin check passed: no gaps were kept
+            self.assertEqual(files(d), [NAME, NAME + ".done"])
+
+
+class Resume(unittest.TestCase):
+    def test_a_failed_round_keeps_the_finished_pieces(self):
+        broken = {"on": True}
+
+        def faults(path, start, end, nth):
+            if broken["on"] and start >= 20 * CHUNK:
+                return ("status", 500)
+            return None
+
+        with small_pieces(DL_ROUNDS=1, DL_CHUNK_TRIES=1), Server(faults=faults) as srv, tempfile.TemporaryDirectory() as d:
+            code, out = run(srv, d)
+            self.assertEqual(code, 1)                                       # could not finish: setup stops with a message
+            self.assertIn("could not finish downloading", out)
+            part = Path(d) / (NAME + ".part")
+            side = Path(d) / (NAME + ".part.chunks")
+            self.assertTrue(part.exists() and side.exists())
+            saved = json.loads(side.read_text())
+            done = {i for a, b in saved["done"] for i in range(a, b)}
+            self.assertTrue(done and all(i < 20 for i in done), done)
+            srv.gets.clear()
+            broken["on"] = False
+            code, out = run(srv, d)
+            self.assertIsNone(code, out)
+            self.assertEqual((Path(d) / NAME).read_bytes(), DATA)
+            again = set(asked(srv))
+            self.assertFalse(again & done, "pieces that were finished are fetched again")
+            self.assertEqual(again | done, set(range(41)))
+            self.assertEqual(files(d), [NAME, NAME + ".done"])
+
+    def test_a_killed_download_continues(self):
+        runner = (
+            "import sys; sys.path[:0] = [%r, %r]\n"
+            "import setup\n"
+            "setup.DL_CHUNK = %d; setup.DL_MIN_PARALLEL = 0; setup.DL_SAVE_EVERY = 0.05; setup.DL_CONNECTIONS = 4\n"
+            "pin = (%d, %r)\n"
+            "setup.file_pin = lambda u: pin\n"
+            "setup.download(sys.argv[1], __import__('pathlib').Path(sys.argv[2]) / %r, 'the model')\n"
+            % (str(ROOT), str(ROOT / "tools"), CHUNK, len(DATA), hashlib.sha256(DATA).hexdigest(), NAME))
+        with small_pieces(), Server(delay=0.05) as srv, tempfile.TemporaryDirectory() as d:
+            env = dict(os.environ, STRATA_SOURCE="huggingface")
+            proc = subprocess.Popen([sys.executable, "-c", runner, srv.url, d], env=env,
+                                    stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+            try:
+                deadline = time.time() + 60
+                side = Path(d) / (NAME + ".part.chunks")
+                while time.time() < deadline and proc.poll() is None:
+                    try:
+                        if sum(b - a for a, b in json.loads(side.read_text())["done"]) >= 12:
+                            break
+                    except (OSError, ValueError, KeyError):
+                        pass
+                    time.sleep(0.05)
+            finally:
+                proc.kill()
+                proc.wait()
+            self.assertTrue((Path(d) / (NAME + ".part")).exists(), "the download finished before it was killed")
+            saved = json.loads(side.read_text())
+            done = {i for a, b in saved["done"] for i in range(a, b)}
+            self.assertGreaterEqual(len(done), 12)
+            srv.gets.clear()
+            code, out = run(srv, d)
+            self.assertIsNone(code, out)
+            self.assertEqual((Path(d) / NAME).read_bytes(), DATA)
+            self.assertFalse(set(asked(srv)) & done, "pieces that were finished before the kill are fetched again")
+
+    def test_a_part_from_one_stream_continues_after_its_last_whole_piece(self):
+        with small_pieces(), Server() as srv, tempfile.TemporaryDirectory() as d:
+            have = CHUNK * 10 + 777                                         # an older setup's .part: bytes from the start
+            (Path(d) / (NAME + ".part")).write_bytes(DATA[:have])
+            code, out = run(srv, d)
+            self.assertIsNone(code, out)
+            self.assertEqual((Path(d) / NAME).read_bytes(), DATA)
+            self.assertEqual(sorted(asked(srv)), list(range(10, 41)))       # pieces 0-9 were not asked for again
+
+    def test_one_stream_continues_a_parallel_part_from_its_unbroken_start(self):
+        with small_pieces(), Server() as srv, tempfile.TemporaryDirectory() as d:
+            part = Path(d) / (NAME + ".part")
+            junk = bytearray(DATA)
+            junk[8 * CHUNK:9 * CHUNK] = b"\0" * CHUNK                       # a gap behind the later pieces
+            part.write_bytes(bytes(junk[:20 * CHUNK]))
+            (Path(d) / (NAME + ".part.chunks")).write_text(json.dumps(
+                {"v": 1, "total": len(DATA), "chunk": CHUNK, "done": [[0, 8], [9, 20]]}))
+            with mock.patch.dict(os.environ, {"STRATA_DL_CONNECTIONS": "1"}):
+                code, out = run(srv, d)
+            self.assertIsNone(code, out)
+            self.assertEqual((Path(d) / NAME).read_bytes(), DATA)
+            self.assertEqual(srv.gets, [(8 * CHUNK, len(DATA) - 1)])        # resumed after piece 7, the gap refetched
+            self.assertEqual(files(d), [NAME, NAME + ".done"])
+
+    def test_a_piece_list_of_another_file_starts_over(self):
+        with small_pieces(), Server() as srv, tempfile.TemporaryDirectory() as d:
+            (Path(d) / (NAME + ".part")).write_bytes(b"x" * (CHUNK * 3))
+            (Path(d) / (NAME + ".part.chunks")).write_text(json.dumps(
+                {"v": 1, "total": 999, "chunk": CHUNK, "done": [[0, 3]]}))
+            code, out = run(srv, d)
+            self.assertIsNone(code, out)
+            self.assertEqual((Path(d) / NAME).read_bytes(), DATA)
+
+
+class Failures(unittest.TestCase):
+    def test_a_piece_that_fails_once_is_tried_again(self):
+        def faults(path, start, end, nth):
+            if start == 5 * CHUNK and nth == 1:
+                return ("status", 500)
+            if start == 9 * CHUNK and nth == 1:
+                return ("cut", CHUNK // 2)                                  # drops half way: continues from where it was
+            return None
+
+        with small_pieces(), Server(faults=faults) as srv, tempfile.TemporaryDirectory() as d:
+            code, out = run(srv, d)
+            self.assertIsNone(code, out)
+            self.assertEqual((Path(d) / NAME).read_bytes(), DATA)
+            self.assertEqual(sum(1 for s, _ in srv.gets if s == 5 * CHUNK), 2)
+            self.assertEqual(sum(1 for s, _ in srv.gets if s == 9 * CHUNK), 1)
+            resumed = [s for s, _ in srv.gets if 9 * CHUNK < s < 10 * CHUNK]
+            self.assertEqual(len(resumed), 1)                               # the rest of the piece, not all of it again
+            self.assertGreater(resumed[0], 9 * CHUNK)
+
+    def test_a_signed_address_that_ran_out_is_asked_for_again(self):
+        def faults(path, start, end, nth):
+            if path.startswith("/cdn/") and nth == 1 and start in (4 * CHUNK, 11 * CHUNK):
+                return ("status", 403)
+            return None
+
+        with small_pieces(), Server(redirect=True, faults=faults) as srv, tempfile.TemporaryDirectory() as d:
+            code, out = run(srv, d)
+            self.assertIsNone(code, out)
+            self.assertEqual((Path(d) / NAME).read_bytes(), DATA)
+            self.assertGreater(srv.redirects, 2)                            # HEAD + first probe + a new address after the 403
+
+    def test_a_pinned_hash_mismatch_deletes_the_file_and_stops(self):
+        wrong = bytes(reversed(DATA))
+        with small_pieces(), Server() as srv, tempfile.TemporaryDirectory() as d:
+            code, out = run(srv, d, pin=pin_of(wrong))
+            self.assertEqual(code, 1)
+            self.assertEqual(files(d), [])                                  # no file, no .done, no .part, no piece list
+
+    def test_a_pinned_hash_mismatch_with_the_same_size(self):
+        other = bytes((b ^ 0x55) for b in DATA)
+        with small_pieces(), Server(data=other) as srv, tempfile.TemporaryDirectory() as d:
+            code, out = run(srv, d)
+            self.assertEqual(code, 1)
+            self.assertIn("wrong SHA-256", out)
+            self.assertEqual(files(d), [])
+
+    def test_a_full_disk_stops_with_a_message(self):
+        with small_pieces(), Server() as srv, tempfile.TemporaryDirectory() as d:
+            real = setup._piece_worker
+
+            def worker(src, st, todo, part):
+                st.fatal = setup._DiskError("No space left on device")
+                st.stop.set()
+
+            with mock.patch.object(setup, "_piece_worker", worker):
+                code, out = run(srv, d)
+            self.assertEqual(code, 1)
+            self.assertIn("cannot write", out)
+
+    def test_the_file_is_not_preallocated(self):
+        """The app counts a .part's size as the bytes done, so it must not jump to the full size at the start."""
+        seen = []
+
+        def faults(path, start, end, nth):
+            if start == 30 * CHUNK:
+                p = Path(self.d) / (NAME + ".part")
+                seen.append(p.stat().st_size if p.exists() else None)
+            return None
+
+        with small_pieces(DL_CONNECTIONS=2), Server(faults=faults, delay=0.01) as srv, tempfile.TemporaryDirectory() as d:
+            self.d = d
+            code, out = run(srv, d)
+            self.assertIsNone(code, out)
+            self.assertTrue(seen and seen[0] is not None and seen[0] <= 32 * CHUNK, seen)
+
+
+class Progress(unittest.TestCase):
+    def test_the_progress_line_is_the_one_the_app_reads(self):
+        with small_pieces(), Server() as srv, tempfile.TemporaryDirectory() as d:
+            code, out = run(srv, d)
+            self.assertIsNone(code, out)
+            lines = [x for x in out.replace("\r", "\n").split("\n") if x.strip()]
+            bar = [x for x in lines if "GB (" in x]
+            self.assertTrue(bar)
+            for x in bar:
+                self.assertRegex(x, r"^  the model: +[\d.]+ / [\d.]+ GB \(\d+%\) +$")
+            parsed = strata_mcp.install_progress(bar)["download"]
+            self.assertEqual(parsed["file"], "the model")
+            self.assertEqual(parsed["percent"], 100)
+            self.assertEqual(parsed["total_gb"], round(len(DATA) / 1e9, 2))
+
+    def test_the_format_matches_the_one_stream_path(self):
+        with small_pieces(), Server() as srv, tempfile.TemporaryDirectory() as d1, tempfile.TemporaryDirectory() as d2:
+            _, par = run(srv, d1)
+            with mock.patch.dict(os.environ, {"STRATA_DL_CONNECTIONS": "1"}):
+                _, one = run(srv, d2)
+            def shape(out):
+                return {re.sub(r"[\d.]+", "N", x) for x in out.replace("\r", "\n").split("\n") if "GB (" in x}
+            self.assertEqual(shape(par), shape(one))
+
+
+class SignedAddress(unittest.TestCase):
+    def test_expiry_of_a_cdn_address(self):
+        now = time.time()
+        self.assertAlmostEqual(setup._url_ttl(f"https://cdn.example/f?Expires={int(now) + 600}&Signature=x"), 600, delta=3)
+        self.assertIsNone(setup._url_ttl("https://cdn.example/f"))
+        self.assertIsNone(setup._url_ttl(f"https://cdn.example/f?Expires={int(now) - 5}"))      # already out, or a wrong clock
+        self.assertIsNone(setup._url_ttl("https://cdn.example/f?Expires=soon"))
+        stamp = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime(now))
+        self.assertAlmostEqual(setup._url_ttl(f"https://cdn.example/f?X-Amz-Date={stamp}&X-Amz-Expires=900"), 900, delta=3)
+
+
+if __name__ == "__main__":
+    unittest.main()