#!/usr/bin/env python3
"""Fast 24-worker HTTP Range fetcher + packer for Qwen3.8-Flash-Next MTP (mtp-q2_0.gguf)."""
from __future__ import annotations
import concurrent.futures, hashlib, json, os, shutil, subprocess, sys, time
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "vendor" / "strata" / "tools"))
import mtp_fetch as mf
OUT_DIR = ROOT / "models" / "K512_iq2_xs" / "mtp_src"
MTP_GGUF = ROOT / "models" / "K512_iq2_xs" / "mtp-q2_0.gguf"
RT_DIR = ROOT / "models" / "K512_iq2_xs" / "mtp_rt"
def main():
t0 = time.monotonic()
mf.resolve_repo()
OUT_DIR.mkdir(parents=True, exist_ok=True)
inv_path = OUT_DIR / "mtp-inventory.json"
if inv_path.exists():
rows = json.loads(inv_path.read_text(encoding="utf-8"))["tensors"]
else:
print("Fetching MTP inventory headers...", flush=True)
rows = mf.inventory(str(OUT_DIR))
tdir = OUT_DIR / "tensors"
tdir.mkdir(parents=True, exist_ok=True)
# Build 16 MiB chunk tasks for all missing/incomplete tensors
chunk_sz = 16 << 20
to_fetch = []
for r in rows:
p = tdir / (r["name"] + ".bin")
want = mf.SHA256.get(r["name"]) if mf.pinned() else None
if p.exists() and p.stat().st_size == r["bytes"]:
if want is None or mf.sha256_of(str(p)) == want:
continue
# Pre-allocate file
with open(p, "wb") as f:
f.truncate(r["bytes"])
pos = r["start"]
while pos <= r["end"]:
end = min(pos + chunk_sz - 1, r["end"])
to_fetch.append((r["name"], r["shard"], p, pos - r["start"], pos, end))
pos = end + 1
total_dl = sum(e - s + 1 for _, _, _, _, s, e in to_fetch)
print(f"Downloading {len(to_fetch)} chunks ({total_dl/1e9:.2f} GB) with 24 workers...", flush=True)
def dl_chunk(task):
name, shard, p, file_off, start, end = task
data = mf.get(mf.REPO + shard, start, end)
fd = os.open(str(p), os.O_WRONLY)
try:
os.pwrite(fd, data, file_off)
finally:
os.close(fd)
return len(data)
done_bytes = 0
with concurrent.futures.ThreadPoolExecutor(max_workers=24) as ex:
futs = [ex.submit(dl_chunk, t) for t in to_fetch]
for idx, fut in enumerate(concurrent.futures.as_completed(futs), 1):
done_bytes += fut.result()
if idx % 20 == 0 or idx == len(futs):
dt = time.monotonic() - t0
print(f" [{idx}/{len(futs)}] {done_bytes/1e9:.2f}/{total_dl/1e9:.2f} GB ({done_bytes/1e6/max(dt,1e-3):.1f} MB/s)", flush=True)
# Verify SHA256 and write mtp-manifest.json
manifest = []
for r in rows:
p = tdir / (r["name"] + ".bin")
digest = mf.sha256_of(str(p))
want = mf.SHA256.get(r["name"]) if mf.pinned() else None
if want and digest != want:
raise RuntimeError(f"SHA256 mismatch on {r['name']}: {digest} != {want}")
manifest.append(dict(r, file=os.path.relpath(str(p), str(OUT_DIR)), sha256=digest))
(OUT_DIR / "mtp-manifest.json").write_text(json.dumps(manifest, indent=1), encoding="utf-8")
print(f"All 31 MTP tensors verified in {time.monotonic()-t0:.1f}s! Packing mtp-q2_0.gguf...", flush=True)
# Run mtp_pack.py
env = dict(os.environ, STRATA_GGUF_PY=str((ROOT / "vendor" / "llama.cpp" / "gguf-py").resolve()))
subprocess.run(
[sys.executable, str(ROOT / "vendor" / "strata" / "tools" / "mtp_pack.py"),
"--src", str(OUT_DIR), "--experts", "q2_0", "--out", str(MTP_GGUF)],
env=env,
check=True,
)
RT_DIR.mkdir(parents=True, exist_ok=True)
shutil.copy2(ROOT / "vendor" / "strata" / "data" / "draft_vocab_en.bin", RT_DIR / "draft_vocab.bin")
# Clean up raw 4.9 GB BF16 tensor dir after packing
shutil.rmtree(OUT_DIR, ignore_errors=True)
print(f"COMPLETE in {time.monotonic()-t0:.1f}s: {MTP_GGUF} ({MTP_GGUF.stat().st_size/1e6:.1f} MB)", flush=True)
if __name__ == "__main__":
main()