aria2c multi-connection dataset downloads (CDN edge roulette fix)

The mirror's CDN assigns 6KB/s or 4.8MB/s to the SAME file depending
on which edge a connection lands on -- a single-connection download
is one dice roll that can stall for the whole file. aria2c (-x8 -s8)
splits the file so each segment rolls independently, and
--lowest-speed-limit=50K re-opens stalled segments. Falls back to the
urllib path when aria2c is absent. Live test: the mrcr file that sat
at 2.1MB/190MB for 6 minutes came down in 43s (4.4MB/s).

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
sora 2026-09-17 08:03:34 +00:00
parent 9f765fc5ce
commit 3f2ca8bfd5

View File

@ -13,6 +13,7 @@ Supported sources:
import csv
import json
import os
import subprocess
import sys
import time
import re
@ -261,6 +262,36 @@ def _download_with_progress(resp, out, filename: str) -> None:
out.write(chunk)
def _aria2_fetch(url: str, dest_dir: Path, name: str) -> bool:
"""Multi-connection fetch via aria2c when available.
The mirror's CDN assigns wildly different edges per connection (a
190MB file ran at 6KB/s on one connection and 4.8MB/s on a fresh
one). aria2c splits the file into segments -- each segment rolls its
own edge dice -- and --lowest-speed-limit self-heals a stalled piece
by re-opening it. Returns False (caller falls back to urllib) when
aria2c is missing or fails.
"""
import shutil as _sh
if not _sh.which('aria2c'):
return False
tmp = dest_dir / (name + '.aria2.part')
cmd = ['aria2c', '-x', '8', '-s', '8', '-k', '4M', '--continue=true',
'--file-allocation=none', '--console-log-level=warn',
'--summary-interval=0', '--retry-wait=3', '--max-tries=5',
'--lowest-speed-limit=50K', '--timeout=20',
'-d', str(dest_dir), '-o', tmp.name, url]
r = subprocess.run(cmd, capture_output=True, text=True)
if r.returncode == 0 and tmp.exists():
os.replace(tmp, dest_dir / name)
return True
tmp.unlink(missing_ok=True)
tmp2 = dest_dir / (tmp.name + '.aria2')
tmp2.unlink(missing_ok=True)
return False
def _ms_download(repo: str, path: str, dest_dir: Path) -> Path:
"""Download one repo file into the raw cache (content-addressed, reused)."""
dest = dest_dir / os.path.basename(path)
@ -271,6 +302,8 @@ def _ms_download(repo: str, path: str, dest_dir: Path) -> Path:
dest.unlink()
dest_dir.mkdir(parents=True, exist_ok=True)
url = f'{_MS_API}/{repo}/repo?Revision=master&FilePath={path}'
if _aria2_fetch(url, dest_dir, dest.name):
return dest_dir / dest.name
tmp = dest.with_name(dest.name + f'.part-{os.getpid()}')
req = urllib.request.Request(url, headers={'User-Agent': 'evalharness/0.1'})
with urllib.request.urlopen(req, timeout=600) as resp, open(tmp, 'wb') as out:
@ -411,6 +444,8 @@ def _hf_download(repo: str, path: str, dest_dir: Path) -> Path:
dest.unlink()
dest_dir.mkdir(parents=True, exist_ok=True)
url = f'{_hf_base()}/datasets/{repo}/resolve/main/{path}'
if _aria2_fetch(url, dest_dir, dest.name):
return dest_dir / dest.name
tmp = dest.with_name(dest.name + f'.part-{os.getpid()}')
req = urllib.request.Request(url, headers={'User-Agent': 'evalharness/0.1'})
with urllib.request.urlopen(req, timeout=1800) as resp, open(tmp, 'wb') as out: