sora 3f2ca8bfd5 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>
2026-09-17 08:03:34 +00:00

570 lines
23 KiB
Python

"""Raw record loading and the FieldSpec-based default conversion.
Supported sources:
- local file: .jsonl / .json / .csv / .tsv / .parquet
- local dir: looks for ``{subset}_{split}.jsonl`` etc. (and csv/tsv/parquet)
- hub id: HuggingFace ``datasets`` (optional dependency, imported on demand)
- modelscope: ``params={'hub': 'modelscope'}`` — native ModelScope raw-file
HTTP download (no ``modelscope`` package needed, immune to its datasets
version pinning). File selection mirrors the local-dir convention and
understands HF-style shards (``test-00000-of-00001.parquet``).
"""
import csv
import json
import os
import subprocess
import sys
import time
import re
import shutil
import urllib.request
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Union
from .sample import Sample
from .spec import DatasetSpec, FieldSpec
_SUPPORTED_EXTS = ['.jsonl', '.json', '.csv', '.tsv', '.parquet']
# params keys consumed by the loader itself; never forwarded to load_dataset()
_RESERVED_PARAMS = {'hub', 'ms_files', 'hf_files', 'filter_column'}
_MS_API = 'https://www.modelscope.cn/api/v1/datasets'
def get_cache_root() -> Path:
"""Canonical evalharness cache root ($EVALHARNESS_CACHE or ~/.cache/evalharness)."""
return Path(os.environ.get('EVALHARNESS_CACHE', '~/.cache/evalharness')).expanduser()
def set_cache_root(path) -> None:
"""Override the cache root at runtime (also moves the .raw blob store)."""
os.environ['EVALHARNESS_CACHE'] = str(path)
def load_raw_records(spec: DatasetSpec, raw_dir: Optional[Path] = None) -> List[Dict[str, Any]]:
"""Load raw records (list of dicts) from the spec's source.
When ``raw_dir`` is given, a copy of the *native* source data is placed
there (downloaded blob / copy of the local file / record dump for hub
sources) so every cache entry is self-contained: raw in, samples out.
"""
source = spec.source
if os.path.exists(source):
records, native_files = _load_local(source, spec)
if raw_dir is not None and native_files:
_link_or_copy_all(native_files, raw_dir)
elif spec.params.get('hub') == 'modelscope':
records = _load_from_modelscope(spec, raw_dir)
elif spec.params.get('hub') == 'hf_raw':
records = _load_from_hf_raw(spec, raw_dir)
else:
records = _load_from_hub(spec)
if raw_dir is not None:
# hub-native packaging lives in the HF cache; keep an exact,
# pre-conversion record dump so the entry is self-contained
raw_dir.mkdir(parents=True, exist_ok=True)
with open(raw_dir / 'records.jsonl', 'w', encoding='utf-8') as f:
for r in records:
f.write(json.dumps(r, ensure_ascii=False, default=str) + '\n')
records = _apply_subset_filter(spec, records)
return records
def _load_local(path: str, spec: DatasetSpec) -> tuple:
"""Load from a local path; also return the native file(s) to preserve."""
if os.path.isfile(path):
return _read_file(path), [Path(path)]
# directory: follow the <subset>_<split>.<ext> / <subset>.<ext> convention
for ext in _SUPPORTED_EXTS:
for candidate in (f'{spec.subset}_{spec.split}{ext}', f'{spec.subset}{ext}', f'{spec.split}{ext}'):
full = os.path.join(path, candidate)
if os.path.exists(full):
return _read_file(full), [Path(full)]
expected = [os.path.join(path, f'{spec.subset}_{spec.split}{e}') for e in _SUPPORTED_EXTS]
available = sorted(f for f in os.listdir(path) if os.path.splitext(f)[1] in _SUPPORTED_EXTS)
raise FileNotFoundError(
f'no dataset file found for subset={spec.subset!r} split={spec.split!r} in {path!r}.\n'
f'Expected one of:\n - ' + '\n - '.join(expected) + '\n'
f'Available: {available or "(none)"}'
)
def _link_or_copy_all(files: List[Path], raw_dir: Path) -> None:
"""Place native files in raw_dir (hardlink when possible, else copy)."""
raw_dir.mkdir(parents=True, exist_ok=True)
for src in files:
dst = raw_dir / src.name
if dst.exists() or src.parent.resolve() == raw_dir.resolve():
continue
try:
os.link(src, dst)
except OSError:
shutil.copy2(src, dst)
def _read_file(path: str) -> List[Dict[str, Any]]:
ext = os.path.splitext(path)[1]
if ext == '.parquet':
try:
import pyarrow.parquet as pq
except ImportError:
raise ImportError(
f'{path} is parquet; install the reader first: pip install "evalharness[all]"'
)
return pq.read_table(path).to_pylist()
with open(path, encoding='utf-8') as f:
if ext == '.jsonl':
return [json.loads(line) for line in f if line.strip()]
if ext == '.json':
data = json.load(f)
return data if isinstance(data, list) else [data]
if ext in ('.csv', '.tsv'):
return list(csv.DictReader(f, delimiter='\t' if ext == '.tsv' else ','))
raise ValueError(f'unsupported file format: {path}')
def _apply_subset_filter(spec: DatasetSpec, records: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
"""When ``params['filter_column']`` is set, subset selects that column's value.
Used when a mirror packs all subsets into one file with a category column
(e.g. evalscope/cmmlu: one parquet, ``category`` = subject).
"""
col = spec.params.get('filter_column')
if not col or spec.subset in ('default', 'all'):
return records
return [r for r in records if r.get(col) == spec.subset]
# ---------------- ModelScope raw-file loading (no modelscope package) ----------------
def _ms_api_json(url: str) -> Dict[str, Any]:
req = urllib.request.Request(url, headers={'User-Agent': 'evalharness/0.1'})
with urllib.request.urlopen(req, timeout=60) as resp:
return json.loads(resp.read().decode('utf-8'))
def _ms_list_files(repo: str, root: str = '', depth: int = 0) -> List[str]:
"""Recursively list blob paths under a ModelScope dataset repo directory."""
url = f'{_MS_API}/{repo}/repo/tree?Revision=master' + (f'&Root={root}' if root else '')
try:
data = _ms_api_json(url).get('Data') or {}
except Exception:
return []
files: List[str] = []
for entry in data.get('Files') or []:
if entry.get('Type') == 'blob':
files.append(entry['Path'])
elif entry.get('Type') == 'tree' and depth < 4:
files.extend(_ms_list_files(repo, entry['Path'], depth + 1))
return files
def _ms_match_files(spec: DatasetSpec, files: List[str]) -> List[str]:
"""Pick repo files for this subset/split, mirroring the local-dir convention.
Priority: explicit ``params['ms_files']`` > exact name match > HF-style
shard match (``{split}-00000-of-00001`` / ``{subset}_{split}-...`` /
``{subset}_0``) > single-data-file repos. When the subset names a parent
directory, matches inside it win (e.g. livecodebench ``release_latest/``).
"""
explicit = spec.params.get('ms_files')
if explicit:
# paths support {subset}/{split} templating (e.g. tau2 domains)
return [p.format(subset=spec.subset, split=spec.split) for p in explicit]
def match_one(path: str) -> bool:
base = os.path.basename(path)
name = os.path.splitext(base)[0]
for ext in _SUPPORTED_EXTS:
if not base.endswith(ext):
continue
stem = base[: -len(ext)]
if stem in (f'{spec.subset}_{spec.split}', spec.subset, spec.split):
return True
if stem == f'{spec.subset}{spec.split}':
return True
shard = (
# tolerate a content-hash suffix after the shard pattern
# (e.g. test-00000-of-00001-6282153cc50a2626.parquet)
rf'{re.escape(spec.split)}-\d+-of-\d+(-[0-9a-f]+)?$'
rf'|{re.escape(spec.subset)}_{re.escape(spec.split)}[-_].*'
rf'|{re.escape(spec.subset)}[-_]\d+$'
)
if re.fullmatch(shard, stem, flags=re.ASCII):
return True
return False
candidates = [f for f in files if match_one(f)]
if not candidates:
data_files = [f for f in files if os.path.splitext(f)[1] in _SUPPORTED_EXTS]
if len(data_files) == 1:
return data_files
return []
in_subset_dir = [f for f in candidates if os.path.dirname(f) == spec.subset]
return sorted(in_subset_dir or candidates)
def _parquet_ok(path: Path) -> bool:
"""Cheap integrity check: parquet files end with magic 'PAR1'."""
try:
with open(path, 'rb') as f:
f.seek(-4, 2)
return f.read(4) == b'PAR1'
except OSError:
return False
def _download_with_progress(resp, out, filename: str) -> None:
"""Stream a download with a rich byte-level progress bar.
Used for the minute-scale dataset fetches (ModelScope blobs, HF raw
files); silent (plain streaming) when rich is unavailable or output
is redirected."""
total = 0
try:
total = int(resp.headers.get('Content-Length') or 0)
except (TypeError, ValueError):
total = 0
if not sys.stdout.isatty():
while True:
chunk = resp.read(1 << 20)
if not chunk:
break
out.write(chunk)
return
try:
from rich.progress import (BarColumn, DownloadColumn, Progress,
SpinnerColumn, TextColumn,
TimeElapsedColumn, TransferSpeedColumn)
pg = Progress(SpinnerColumn(),
TextColumn('[cyan]{task.fields[name]}[/cyan]'),
BarColumn(), DownloadColumn(), TransferSpeedColumn(),
TextColumn(''), TimeElapsedColumn(),
transient=True)
with pg:
t = pg.add_task('download', total=total or None, name=filename)
while True:
chunk = resp.read(1 << 20)
if not chunk:
break
out.write(chunk)
pg.advance(t, len(chunk))
return
except ImportError:
pass
while True:
chunk = resp.read(1 << 20)
if not chunk:
break
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)
if dest.exists() and dest.stat().st_size > 0 \
and (os.path.splitext(dest)[1] != '.parquet' or _parquet_ok(dest)):
return dest
if dest.exists(): # truncated/corrupt (e.g. an interrupted download)
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:
_download_with_progress(resp, out, dest.name)
os.replace(tmp, dest)
return dest
def _load_from_modelscope(spec: DatasetSpec, raw_dir: Optional[Path] = None) -> List[Dict[str, Any]]:
import hashlib
files = _ms_list_files(spec.source)
if not files:
raise FileNotFoundError(f'no files found on ModelScope dataset {spec.source!r}')
selected = _ms_match_files(spec, files)
if not selected:
raise FileNotFoundError(
f'no data file for subset={spec.subset!r} split={spec.split!r} in '
f'modelscope {spec.source!r}. Available (first 10): {files[:10]}'
)
# shared blob store: download once per repo, hardlink into each cache entry
blob_dir = get_cache_root() / '.raw' / hashlib.md5(spec.source.encode()).hexdigest()[:10]
records: List[Dict[str, Any]] = []
blobs: List[Path] = []
for path in selected:
blobs.append(_ms_download(spec.source, path, blob_dir))
records.extend(_read_file(str(blobs[-1])))
if raw_dir is not None:
_link_or_copy_all(blobs, raw_dir)
return records
# ---------------- HuggingFace raw-file loading (no auth, exact bytes) ----------------
# For repos whose layout datasets.load_dataset cannot express (raw json trees,
# shard dirs without configs) or that only exist as plain files. Honors
# $HF_ENDPOINT (e.g. https://hf-mirror.com).
def _hf_base() -> str:
return os.environ.get('HF_ENDPOINT', 'https://huggingface.co').rstrip('/')
def _hf_list_files(repo: str, root: str = '', depth: int = 0) -> List[str]:
"""Recursively list file paths under an HF dataset repo directory."""
url = f'{_hf_base()}/api/datasets/{repo}/tree/main' + (f'/{root}' if root else '') + '?limit=1000'
req = urllib.request.Request(url, headers={'User-Agent': 'evalharness/0.1'})
try:
with urllib.request.urlopen(req, timeout=60) as resp:
entries = json.loads(resp.read().decode('utf-8'))
except Exception:
return []
files: List[str] = []
for entry in entries:
if entry.get('type') == 'file':
files.append(entry['path'])
elif entry.get('type') == 'directory' and depth < 4:
files.extend(_hf_list_files(repo, entry['path'], depth + 1))
return files
def _hf_match_files(spec: DatasetSpec, files: List[str]) -> List[str]:
"""Pick repo files for this subset/split (see module docstring conventions)."""
explicit = spec.params.get('hf_files')
if explicit:
return [p.format(subset=spec.subset, split=spec.split) for p in explicit]
def ext_ok(p: str) -> bool:
return os.path.splitext(p)[1] in _SUPPORTED_EXTS
data = [f for f in files if ext_ok(f)]
# multi-window subset, e.g. 'release_v2_v4' = union of release_v2..v4
# shards (LCB release-window convention: per-release files are disjoint)
m = re.fullmatch(r'(release_v\d+)((?:_v\d+)+)', spec.subset)
if m:
members = [m.group(1)] + [f'release{x}' for x in m.group(2).split('_') if x]
picked = []
for mem in members:
picked += [
f for f in data
if re.fullmatch(rf'{re.escape(mem)}[-_]\d+(-of-\d+)?'
rf'|{re.escape(mem)}_{re.escape(spec.split)}[-_].*'
rf'|{re.escape(mem)}',
os.path.splitext(os.path.basename(f))[0])
]
if picked:
# deterministic: member order, shard order within each member
return [f for mem in members for f in sorted(
p for p in picked
if os.path.splitext(os.path.basename(p))[0].startswith(mem))]
if spec.subset != 'default':
# subset wins exclusively: never ALSO match bare-split shards in the
# same dir (repos like sam-paech LCB mix test-*.parquet and
# release_v*-*.parquet in one data/ folder)
in_dir = [f for f in data if os.path.dirname(f) == spec.subset]
if in_dir:
return sorted(in_dir)
exact = []
for f in data:
stem = os.path.splitext(os.path.basename(f))[0]
if stem in (f'{spec.subset}_{spec.split}', spec.subset):
exact.append(f)
if exact:
return sorted(exact)
shards = [
f for f in data
if re.fullmatch(rf'{re.escape(spec.subset)}[-_]\d+(-of-\d+)?'
rf'|{re.escape(spec.subset)}_{re.escape(spec.split)}[-_].*',
os.path.splitext(os.path.basename(f))[0])
]
if shards:
return sorted(shards)
else:
exact = [
f for f in data
if os.path.splitext(os.path.basename(f))[0] == spec.split
]
if exact:
return sorted(exact)
shards = [
f for f in data
if re.fullmatch(rf'{re.escape(spec.split)}[-_]\d+(-of-\d+)?',
os.path.splitext(os.path.basename(f))[0])
]
if shards:
return sorted(shards)
if len(data) == 1:
return data
return []
def _hf_download(repo: str, path: str, dest_dir: Path) -> Path:
"""Download one repo file (follows the CDN redirect) into the blob store."""
dest = dest_dir / os.path.basename(path)
if dest.exists() and dest.stat().st_size > 0 \
and (os.path.splitext(dest)[1] != '.parquet' or _parquet_ok(dest)):
return dest
if dest.exists(): # truncated/corrupt (e.g. an interrupted download)
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:
_download_with_progress(resp, out, dest.name)
os.replace(tmp, dest)
return dest
def _load_from_hf_raw(spec: DatasetSpec, raw_dir: Optional[Path] = None) -> List[Dict[str, Any]]:
import hashlib
# the mirror's tree API degrades at night and returns TRUNCATED listings
# (a repo with 9 parquets listed as [.gitattributes, README.md]); the
# blobs from a previous run are already in the shared store -- use them
blob_dir = get_cache_root() / '.raw' / hashlib.md5(spec.source.encode()).hexdigest()[:10]
def _cached_blobs() -> List[str]:
if not blob_dir.exists():
return []
have = [f.name for f in blob_dir.iterdir()
if os.path.splitext(f)[1] in _SUPPORTED_EXTS]
pref = [f for f in have if f.startswith(spec.subset)] or have
return pref
files = _hf_list_files(spec.source)
selected = _hf_match_files(spec, files) if files else []
if not selected:
cached = _cached_blobs()
if cached:
print(f'· listing unavailable/degraded -- using {len(cached)} '
f'cached file(s) from {blob_dir}', flush=True)
records: List[Dict[str, Any]] = []
for f in sorted(cached):
records.extend(_read_file(str(blob_dir / f)))
if raw_dir is not None:
_link_or_copy_all([blob_dir / f for f in sorted(cached)], raw_dir)
return records
raise FileNotFoundError(
f'no data file for subset={spec.subset!r} split={spec.split!r} in '
f'HF {spec.source!r}. Available (first 10): {files[:10]}'
)
blob_dir = get_cache_root() / '.raw' / hashlib.md5(spec.source.encode()).hexdigest()[:10]
records: List[Dict[str, Any]] = []
blobs: List[Path] = []
for path in selected:
blobs.append(_hf_download(spec.source, path, blob_dir))
records.extend(_read_file(str(blobs[-1])))
if raw_dir is not None:
_link_or_copy_all(blobs, raw_dir)
return records
def _load_from_hub(spec: DatasetSpec) -> List[Dict[str, Any]]:
try:
import datasets
except ImportError:
raise ImportError(
f'dataset {spec.name!r} lives on a hub ({spec.source!r}); '
'install it first: pip install evalharness (light deps are default)'
)
_probe_hub_reachable(spec)
kwargs = {k: v for k, v in spec.params.items() if k not in _RESERVED_PARAMS}
# filter_column subsets select rows by column value post-load, so the hub
# load itself must always use the default config.
subset = None if (spec.subset == 'default' or spec.params.get('filter_column')) else spec.subset
print(f'· downloading {spec.source} from HF hub ({_hf_base()}) ...', flush=True)
t0 = time.monotonic()
ds = datasets.load_dataset(spec.source, subset, split=spec.split, revision=spec.version, **kwargs)
print(f'· downloaded + parsed {len(ds)} records in {time.monotonic() - t0:.0f}s', flush=True)
return [dict(r) for r in ds]
_HF_MIRROR_FALLBACKS = ['https://hf-mirror.com']
def _probe_hub_reachable(spec: DatasetSpec) -> None:
"""Fail fast when the HF endpoint is unreachable -- otherwise the hub
client retries silently for minutes and looks like a hang.
When the DEFAULT endpoint is down but a mirror answers, switch to it
automatically (one-line notice, same behavior as --hf-endpoint)."""
base = _hf_base()
err = None
for candidate in [base] + (_HF_MIRROR_FALLBACKS if base == 'https://huggingface.co' else []):
try:
req = urllib.request.Request(candidate, method='HEAD',
headers={'User-Agent': 'evalharness/0.1'})
urllib.request.urlopen(req, timeout=8)
if candidate != base:
os.environ['HF_ENDPOINT'] = candidate
print(f'· {base} unreachable -- falling back to {candidate}',
flush=True)
return
except Exception as e:
err = e
raise RuntimeError(
f'HuggingFace endpoint unreachable: {base}'
+ (f' (mirrors tried: {", ".join(_HF_MIRROR_FALLBACKS)})'
if base == 'https://huggingface.co' else '')
+ f'\n reason: {type(err).__name__}: {str(err)[:80]}\n'
f' dataset {spec.source!r} cannot download. Fix:\n'
f' evalharness eval run ... --hf-endpoint <a reachable endpoint>') from err
def field_spec_to_record_fn(fields: FieldSpec) -> Callable[[Dict[str, Any]], Sample]:
"""Build the default record->Sample converter from a FieldSpec."""
def convert(record: Dict[str, Any]) -> Sample:
target = record.get(fields.target, '')
if isinstance(target, (int, float)):
target = str(target)
metadata = {k: record.get(k) for k in fields.metadata}
return Sample(
input=record.get(fields.input, ''),
choices=record.get(fields.choices) if fields.choices in record else None,
target=target,
id=record.get(fields.id) if fields.id else None,
metadata=metadata,
)
return convert