File size: 7,332 Bytes
68f9568 | 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 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 | """多连接 Range 分片下载 rnacentral.1024（单连接被 HF 限速到 ~200-470 KB/s）。
为什么要它
----------
`_redownload_rnacentral.py` 是**单连接**顺序追加，实测 ~470 KB/s，3.7 GB 要 2 小时。
实测 HF 的限速是**按连接**的：24 并发能到 ~2.9 MB/s（约 6 倍）。
所以这里把剩余区间切成固定大小的分片，每片一个 `.part` 文件、独立可续传，
全部下完后按 offset 顺序追加到主文件，最后校验 sha256。
设计要点
--------
* 分片网格固定在首次启动时的 `base`（= 当时主文件已有的字节数），存进 `meta.json`；
重跑时前缀可能变长，已完全落进主文件的分片会跳过。
* 每片单独可续传：`.part` 文件已有多少就从哪里接着要。
* 只有在**全部**分片完整后才会碰主文件 —— 不会把半成品拼进去。
"""
from __future__ import annotations
import argparse
import concurrent.futures as cf
import hashlib
import json
import os
import shutil
import sys
import threading
import time
from pathlib import Path
import requests
P = Path(r"C:\qwen-api\RNAQwen\data\external\rnacentral_1024.parquet")
PARTS = Path(r"C:\qwen-api\RNAQwen\data\external\_rnacentral_parts")
URL = "https://huggingface.co/datasets/multimolecule/rnacentral.1024/resolve/main/data.parquet"
TARGET = 3_695_330_389
WANT = "cf1f389a2accda43ee3d118c5e9f485be584c7511df12bca1d7b9fe5816ba7b4"
_lock = threading.Lock()
_done_bytes = 0
_t0 = 0.0
def human(n: float) -> str:
return f"{n/1024/1024:.1f} MB"
def fetch_chunk(start: int, end: int, retries: int = 40) -> tuple[int, int]:
"""下载 [start, end] 闭区间到独立 .part，返回 (bytes, attempts)。"""
path = PARTS / f"{start:012d}.part"
want = end - start + 1
if path.exists() and path.stat().st_size == want:
return want, 0
have = path.stat().st_size if path.exists() else 0
last = ""
for attempt in range(retries):
try:
r = requests.get(URL, headers={"Range": f"bytes={start + have}-{end}"},
stream=True, timeout=(20, 30))
if r.status_code != 206:
raise RuntimeError(f"status={r.status_code}（期望 206）")
with path.open("ab") as fh:
for c in r.iter_content(512 * 1024):
fh.write(c)
have += len(c)
if have >= want:
return have, attempt + 1
last = f"短读 {have}/{want}"
except Exception as exc: # noqa: BLE001
last = f"{type(exc).__name__}: {str(exc)[:90]}"
time.sleep(min(1.0 + attempt * 0.4, 8.0))
have = path.stat().st_size if path.exists() else 0
raise RuntimeError(f"分片 {start} 重试 {retries} 次仍失败：{last}")
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--chunk-mb", type=int, default=16)
ap.add_argument("--workers", type=int, default=20)
ap.add_argument("--assemble", action="store_true", default=True)
ap.add_argument("--no-assemble", dest="assemble", action="store_false")
args = ap.parse_args()
PARTS.mkdir(parents=True, exist_ok=True)
meta_path = PARTS / "meta.json"
chunk = args.chunk_mb * 1024 * 1024
cur = P.stat().st_size if P.exists() else 0
if cur >= TARGET:
print(f"[dl] 主文件已是 {cur} 字节，跳过下载")
else:
if meta_path.exists():
meta = json.loads(meta_path.read_text())
base = int(meta["base"])
else:
base = cur
meta_path.write_text(json.dumps({"base": base, "target": TARGET, "chunk": chunk}))
print(f"[dl] 主文件前缀 {cur} 字节；分片网格 base={base}；目标 {TARGET}；"
f"待下 {human(TARGET-base)}；并发 {args.workers}；片大小 {args.chunk_mb} MB", flush=True)
if cur != base:
print(f"[dl] 警告：当前前缀 {cur} != 网格 base {base}，只拼接 end>=cur 的分片")
chunks = [(s, min(s + chunk - 1, TARGET - 1)) for s in range(base, TARGET, chunk)]
todo = [(s, e) for s, e in chunks if e >= cur]
print(f"[dl] 共 {len(chunks)} 片，本次处理 {len(todo)} 片", flush=True)
global _done_bytes, _t0
_done_bytes = sum((PARTS / f"{s:012d}.part").stat().st_size
for s, _ in chunks if (PARTS / f"{s:012d}.part").exists())
_t0 = time.time()
n_ok = 0
with cf.ThreadPoolExecutor(max_workers=args.workers) as ex:
futs = {ex.submit(fetch_chunk, s, e): (s, e) for s, e in todo}
for fut in cf.as_completed(futs):
s, e = futs[fut]
try:
got, att = fut.result()
except Exception as exc: # noqa: BLE001
print(f"[dl] 分片 {s} 失败：{exc}", flush=True)
continue
n_ok += 1
with _lock:
_done_bytes += e - s + 1
el = max(time.time() - _t0, 1e-6)
rate = _done_bytes / el
if n_ok % 10 == 0 or n_ok == len(todo):
eta_min = (TARGET - cur - _done_bytes) / max(rate, 1) / 60
print(f"[dl] {n_ok}/{len(todo)} 片 {human(_done_bytes)} "
f"{rate/1024/1024:.2f} MB/s ETA {eta_min:.1f} min", flush=True)
got_all = sum((PARTS / f"{s:012d}.part").stat().st_size
for s, _ in chunks if (PARTS / f"{s:012d}.part").exists())
want_all = TARGET - base
print(f"[dl] 分片合计 {got_all} / {want_all}", flush=True)
if got_all != want_all:
print("[dl] 分片不完整，停止（不拼主文件）")
return 1
# ---- 拼接 ----
if P.stat().st_size < TARGET and args.assemble:
chunk = json.loads(meta_path.read_text())["chunk"] if meta_path.exists() else chunk
base = json.loads(meta_path.read_text())["base"] if meta_path.exists() else 0
cur = P.stat().st_size
print(f"[asm] 追加 [{cur}, {TARGET}) 到主文件 ...", flush=True)
with P.open("ab") as out:
for s in range(base, TARGET, chunk):
e = min(s + chunk - 1, TARGET - 1)
if e < cur:
continue
part = PARTS / f"{s:012d}.part"
if not part.exists():
print(f"[asm] 缺分片 {part} → 中止")
return 1
with part.open("rb") as fh:
shutil.copyfileobj(fh, out, 4 * 1024 * 1024)
size = P.stat().st_size
print(f"[chk] 大小 {size} / {TARGET} {'OK' if size == TARGET else '不一致'}", flush=True)
h = hashlib.sha256()
t0 = time.time()
with P.open("rb") as fh:
while True:
b = fh.read(16 * 1024 * 1024)
if not b:
break
h.update(b)
got = h.hexdigest()
print(f"[chk] sha256 {got} ({time.time()-t0:.1f}s)", flush=True)
print(f"[chk] 期望 {WANT}", flush=True)
print(f"[chk] 大小匹配={size == TARGET} sha256匹配={got == WANT}", flush=True)
return 0 if (size == TARGET and got == WANT) else 1
if __name__ == "__main__":
sys.exit(main())
|