#!/usr/bin/env python3
from __future__ import annotations
"""
download_tbv_shards.py
======================
多线程下载 Argoverse 2 Map Change (TbV) 数据集的分片压缩包，
以及 sensor 测试集/验证集。
数据源（公开，无需凭证）:
TbV: https://s3.amazonaws.com/argoverse/datasets/av2/tars/tbv/TbV_v1.0_shard{i}.tar.gz
i = 0..15 共 16 个分片。
Test: https://s3.amazonaws.com/argoverse/datasets/av2/tars/sensor/test-{i:03d}.tar
i = 0..2
Val: https://s3.amazonaws.com/argoverse/datasets/av2/tars/sensor/val-{i:03d}.tar
i = 0..2
特性:
- 并发线程数默认 = 系统可用 CPU 线程数
- 实时仪表盘：原地刷新，不滚屏、不刷长
- HTTP Range 断点续传：Ctrl+C 后重跑自动接着下
* 用 Content-Range 校验服务端偏移，确保拼接正确
* 本地文件大于远端 / 不支持 Range / 越界，均会安全回退到全量下载
- 失败自动重试（指数退避）；若最终仍失败，则等待后重新入队，不会退出
- 实时网速：每线程独立滑动窗口样本，切文件即重置，反映真实瞬时吞吐
用法:
python download_tbv_shards.py # 下载全部（16 TbV + 3 test + 3 val）
python download_tbv_shards.py -o /data/tbv # 指定输出目录
python download_tbv_shards.py -w 8 # 手动指定 8 个线程
python download_tbv_shards.py --start 0 --end 15
python download_tbv_shards.py --no-tui # 关闭实时仪表盘
python download_tbv_shards.py --only tbv # 只下载 TbV
python download_tbv_shards.py --only sensor # 只下载 test/val
"""
import argparse
import os
import re
import sys
import threading
import time
from collections import deque
from concurrent.futures import ThreadPoolExecutor, wait, FIRST_COMPLETED
from dataclasses import dataclass
try:
import requests
except ImportError:
sys.exit("请先安装 requests: pip install requests")
TBV_URL = "https://s3.amazonaws.com/argoverse/datasets/av2/tars/tbv/TbV_v1.0_shard{i}.tar.gz"
TEST_URL = "https://s3.amazonaws.com/argoverse/datasets/av2/tars/sensor/test-{i:03d}.tar"
VAL_URL = "https://s3.amazonaws.com/argoverse/datasets/av2/tars/sensor/val-{i:03d}.tar"
CHUNK = 1 << 20 # 1MB 读写块
FAIL_RETRY_WAIT = 60 # 彻底失败后等待再入队的秒数
SPEED_WINDOW = 8.0 # 速度滑动窗口长度（秒）
IDLE_THRESHOLD = 10.0 # 超过这么久没有新样本 → 速度视为 0
MAX_SAMPLES = 500 # 每个 worker 保留的最大样本数（防内存无界）
# bytes start-end/total 例: "bytes 1048576-2097151/12345678"
CONTENT_RANGE_RE = re.compile(r"bytes\s+(\d+)-(\d+)/(\d+|\*)", re.IGNORECASE)
# ----------------------------- 状态管理 -----------------------------
@dataclass
class WorkerState:
name: str = "idle"
done_bytes: int = 0
total_bytes: int = 0
speed: float = 0.0
status: str = "等待任务"
class DownloadStats:
"""
线程安全的总/分线程进度与速度统计。
速度采用"每 worker 独立样本队列 + 滑动窗口差分":
- 队列元素 = (时间戳, 该文件当前累计字节数)
- 速度 = (末尾样本 - 窗口起始样本) / (时间差)
- 切换文件 / 每次重试开始时调用 begin_file()，清空队列，避免跨文件样本污染
- 超过 IDLE_THRESHOLD 无新样本 → 速度归零，避免"卡在旧速度"
"""
def __init__(self, total_bytes: int, total_files: int):
self.lock = threading.Lock()
self.total_bytes = total_bytes
self.total_files = total_files
self.completed_bytes = 0
self.done_files = 0
self.failed_files = 0
self.skipped_files = 0
self._worker_samples: dict[int, deque] = {}
self.workers: list[WorkerState] = []
self.start_time = time.time()
# ---- worker 生命周期 ----
def register_worker(self) -> WorkerState:
ws = WorkerState()
with self.lock:
self.workers.append(ws)
self._worker_samples[id(ws)] = deque()
return ws
def begin_file(self, ws: WorkerState, name: str, total: int,
initial_bytes: int, status: str):
"""开始/重试某个文件：重置样本队列，把已有字节数作为起点。"""
now = time.time()
with self.lock:
ws.name = name
ws.total_bytes = total
ws.done_bytes = initial_bytes
ws.status = status
ws.speed = 0.0
samples = self._worker_samples[id(ws)]
samples.clear()
samples.append((now, initial_bytes))
def record(self, ws: WorkerState, n: int):
now = time.time()
with self.lock:
ws.done_bytes += n
samples = self._worker_samples[id(ws)]
samples.append((now, ws.done_bytes))
while len(samples) > MAX_SAMPLES:
samples.popleft()
def finish_file(self, ws: WorkerState, size: int, status: str):
with self.lock:
self.completed_bytes += size
self.done_files += 1
if status == "FAIL":
self.failed_files += 1
elif status == "SKIP":
self.skipped_files += 1
ws.status = {"OK": "完成", "FAIL": "失败", "SKIP": "跳过"}.get(status, status)
ws.name = "idle"
ws.done_bytes = 0
ws.total_bytes = 0
ws.speed = 0.0
self._worker_samples[id(ws)].clear()
# ---- 速度计算 ----
def _speed_locked(self, ws: WorkerState, now: float) -> float:
samples = self._worker_samples.get(id(ws))
if not samples:
return 0.0
# 长时间无新数据 → 视为空闲
if now - samples[-1][0] > IDLE_THRESHOLD:
return 0.0
cutoff = now - SPEED_WINDOW
base_idx = 0
for i, (t, _) in enumerate(samples):
if t >= cutoff:
base_idx = i
break
# 窗口内样本不足 2 个 → 用最近两个
if base_idx >= len(samples) - 1:
base_idx = max(0, len(samples) - 2)
t0, b0 = samples[base_idx]
t1, b1 = samples[-1]
if t1 <= t0:
return 0.0
return (b1 - b0) / (t1 - t0)
def snapshot(self):
"""一次性返回 (workers 副本, 聚合速度, 已下载字节数)，全部在锁内完成。"""
now = time.time()
with self.lock:
workers = []
total_sp = 0.0
for w in self.workers:
sp = self._speed_locked(w, now)
total_sp += sp
workers.append(WorkerState(
name=w.name, done_bytes=w.done_bytes,
total_bytes=w.total_bytes, speed=sp, status=w.status,
))
done = self.completed_bytes + sum(w.done_bytes for w in self.workers)
return workers, total_sp, done
# ----------------------------- 工具函数 -----------------------------
def human_size(n: float) -> str:
for unit in ("B", "KB", "MB", "GB", "TB"):
if n < 1024 or unit == "TB":
return f"{n:.2f}{unit}"
n /= 1024
return f"{n:.2f}TB"
def human_time(sec: float) -> str:
if sec != sec or sec == float("inf"): # NaN / inf
return "--:--:--"
sec = int(sec)
h, sec = divmod(sec, 3600)
m, s = divmod(sec, 60)
return f"{h:02d}:{m:02d}:{s:02d}"
def progress_bar(ratio: float, width: int = 30) -> str:
ratio = max(0.0, min(1.0, ratio))
filled = int(width * ratio)
return "█" * filled + "░" * (width - filled)
# ----------------------------- 下载逻辑 -----------------------------
def head_size(url: str) -> int:
r = requests.head(url, timeout=30, allow_redirects=True)
r.raise_for_status()
return int(r.headers["Content-Length"])
def _remove_quietly(path: str):
try:
os.remove(path)
except OSError:
pass
def download_one(url: str, out_dir: str, idle_q, stats: DownloadStats, retries: int = 5):
ws = idle_q.popleft()
try:
return _download_one(url, out_dir, ws, stats, retries)
finally:
with stats.lock:
ws.name = "idle"
ws.done_bytes = 0
ws.total_bytes = 0
ws.speed = 0.0
ws.status = "等待任务"
stats._worker_samples[id(ws)].clear()
idle_q.append(ws)
def _download_one(url: str, out_dir: str, ws: WorkerState,
stats: DownloadStats, retries: int):
fname = url.rsplit("/", 1)[-1]
dest = os.path.join(out_dir, fname)
# --- HEAD 拿到远端大小；失败不退出，持续重试 ---
total = 0
while total == 0:
try:
total = head_size(url)
except Exception:
with stats.lock:
ws.name = fname
ws.status = "HEAD 失败，等待重试"
time.sleep(10)
# --- 检查磁盘已有部分 ---
have = os.path.getsize(dest) if os.path.exists(dest) else 0
if have > total:
# 本地比远端还大：明显是被污染/损坏的文件，删掉重下
_remove_quietly(dest)
have = 0
stats.begin_file(ws, fname, total, have, "检查中")
if have == total:
stats.finish_file(ws, total, "SKIP")
return "SKIP"
for attempt in range(1, retries + 1):
stats.begin_file(ws, fname, total, have, f"下载中(第{attempt}次)")
try:
headers = {"Range": f"bytes={have}-"} if have > 0 else {}
resume_from = have
with requests.get(url, stream=True, timeout=(30, 300), headers=headers) as r:
if resume_from > 0:
if r.status_code == 206:
# 严格校验服务端返回的 Content-Range 起点
cr = r.headers.get("Content-Range", "")
m = CONTENT_RANGE_RE.match(cr)
if not m or int(m.group(1)) != resume_from:
# 偏移不一致 → 不能拼，丢掉重来
have = 0
continue
server_total = m.group(3)
if server_total != "*" and int(server_total) != total:
total = int(server_total) # 远端变了，跟随更新
elif r.status_code == 200:
# 服务端忽略 Range → 从头覆盖写
have = 0
else:
r.raise_for_status()
else:
r.raise_for_status()
mode = "ab" if have > 0 else "wb"
with open(dest, mode) as f:
for chunk in r.iter_content(CHUNK):
if not chunk:
continue
f.write(chunk)
stats.record(ws, len(chunk))
actual = os.path.getsize(dest)
if actual == total:
stats.finish_file(ws, total, "OK")
return "OK"
if actual > total:
# 超出了预期大小 → 损坏，清掉重来
_remove_quietly(dest)
have = 0
else:
have = actual # 部分完成，下次从这儿续
except Exception:
have = os.path.getsize(dest) if os.path.exists(dest) else 0
if have > total:
_remove_quietly(dest)
have = 0
if attempt < retries:
time.sleep(min(2 ** attempt, 30))
# 所有重试均失败：不退出，等待后由主循环重新入队
with stats.lock:
ws.status = f"失败，{FAIL_RETRY_WAIT}s 后重试"
ws.name = fname
time.sleep(FAIL_RETRY_WAIT)
with stats.lock:
stats.failed_files += 1
return "RETRY"
# ----------------------------- 实时仪表盘 -----------------------------
class Dashboard:
"""原地刷新固定行数的仪表盘，不滚屏、不刷长。"""
def __init__(self, stats: DownloadStats, refresh: float = 0.5):
self.stats = stats
self.refresh = refresh
self.lines_drawn = 0
def _render(self) -> list[str]:
workers, total_speed, done_bytes = self.stats.snapshot()
with self.stats.lock:
total_bytes = self.stats.total_bytes
done_files = self.stats.done_files
failed = self.stats.failed_files
skipped = self.stats.skipped_files
total_files = self.stats.total_files
elapsed = time.time() - self.stats.start_time
ratio = done_bytes / total_bytes if total_bytes else 0
eta = (total_bytes - done_bytes) / total_speed if total_speed > 0.01 else float("inf")
lines = [
f" TbV/sensor 下载 文件 {done_files}/{total_files} "
f"失败 {failed} 跳过 {skipped}",
f" 总进度 [{progress_bar(ratio)}] {ratio*100:6.2f}% "
f"{human_size(done_bytes)} / {human_size(total_bytes)}",
f" 总速度: {human_size(total_speed)}/s 已运行: "
f"{human_time(elapsed)} 预计剩余: {human_time(eta)}",
f" {'线程':>4} {'状态':<14} {'速度':>11} {'单文件进度':<26} 文件",
]
for i, ws in enumerate(workers):
if ws.total_bytes > 0:
p = ws.done_bytes / ws.total_bytes
fbar = f"[{progress_bar(p, 14)}] {p*100:5.1f}%"
else:
fbar = "-" * 26
name = ws.name if len(ws.name) <= 40 else "..." + ws.name[-37:]
lines.append(
f" {i:>4} {ws.status:<14} {human_size(ws.speed)+'/s':>11} {fbar} {name}"
)
lines.append(" Ctrl+C 退出(断点续传，重跑即可)")
return lines
def draw(self):
lines = self._render()
out = []
if self.lines_drawn == 0:
out.append("\n".join(lines))
else:
out.append(f"\033[{self.lines_drawn}A")
for ln in lines:
out.append("\033[2K" + ln + "\n")
if len(lines) < self.lines_drawn:
for _ in range(self.lines_drawn - len(lines)):
out.append("\033[2K\n")
out.append(f"\033[{self.lines_drawn - len(lines)}A")
sys.stdout.write("".join(out))
sys.stdout.flush()
self.lines_drawn = len(lines)
def loop(self, stop: threading.Event):
while not stop.is_set():
self.draw()
stop.wait(self.refresh)
# ----------------------------- 主流程 -----------------------------
def build_urls(args) -> list[str]:
urls = []
if args.only in ("all", "tbv"):
urls += [TBV_URL.format(i=i) for i in range(args.start, args.end + 1)]
if args.only in ("all", "sensor"):
urls += [TEST_URL.format(i=i) for i in range(0, 3)]
urls += [VAL_URL.format(i=i) for i in range(0, 3)]
return urls
def main():
_cpu = getattr(os, "process_cpu_count", None) or os.cpu_count() or 1
ap = argparse.ArgumentParser(description="多线程下载 TbV + sensor test/val")
ap.add_argument("-o", "--out", default="tbv")
ap.add_argument("-w", "--workers", type=int, default=max(1, _cpu))
ap.add_argument("--start", type=int, default=0, help="TbV 起始 shard 编号")
ap.add_argument("--end", type=int, default=15, help="TbV 结束 shard 编号(含)")
ap.add_argument("--only", choices=["all", "tbv", "sensor"], default="all",
help="只下载 TbV 或只下载 sensor(test/val)")
ap.add_argument("--no-tui", action="store_true", help="关闭实时仪表盘")
args = ap.parse_args()
urls = build_urls(args)
print(f"[INFO] 系统可用线程数: {_cpu}，本次并发: {args.workers}")
print(f"[INFO] 共 {len(urls)} 个文件")
print("[INFO] 正在获取各文件大小 (HEAD) ...")
sizes = {}
for u in urls:
try:
sizes[u] = head_size(u)
print(f" {u.rsplit('/',1)[-1]:<28} {human_size(sizes[u])}")
except Exception as e:
print(f" {u.rsplit('/',1)[-1]:<28} 获取失败: {e}")
if len(sizes) < len(urls):
print("[WARN] 部分文件大小未知，总进度将不准确（仍可下载）")
stats = DownloadStats(sum(sizes.values()), len(urls))
os.makedirs(args.out, exist_ok=True)
stop = threading.Event()
if not args.no_tui:
dash = Dashboard(stats)
threading.Thread(target=dash.loop, args=(stop,), daemon=True).start()
try:
idle_q = deque(stats.register_worker() for _ in range(args.workers))
with ThreadPoolExecutor(max_workers=args.workers) as pool:
pending = deque(urls)
futures = {}
# 初次填满并发窗口
while pending and len(futures) < args.workers:
u = pending.popleft()
futures[pool.submit(download_one, u, args.out, idle_q, stats)] = u
while futures:
done, _ = wait(futures.keys(), timeout=1.0,
return_when=FIRST_COMPLETED)
if not done:
continue
for fut in done:
u = futures.pop(fut)
try:
res = fut.result()
except Exception:
res = "RETRY"
if res == "RETRY":
# 失败不退出：重新入队
pending.append(u)
# 补充新任务
while pending and len(futures) < args.workers:
nu = pending.popleft()
futures[pool.submit(
download_one, nu, args.out, idle_q, stats
)] = nu
except KeyboardInterrupt:
print("\n[EXIT] 收到 Ctrl+C，已退出（已下载部分保留，重跑自动续传）")
finally:
stop.set()
time.sleep(0.3) # 让仪表盘至少刷一次最新状态
print(f"[DONE] 文件 {stats.done_files}/{stats.total_files}，失败 {stats.failed_files}，"
f"耗时 {human_time(time.time()-stats.start_time)}")
if __name__ == "__main__":
main()