TeX 53.6%
Python 46.2%
1"""Resumable, checksummed, node-shardable downloader for the HF dataset.23Design (CLAUDE.md §4.1): dry-run, month-range filters, retry with exponential4backoff, resumability (hf_hub native), sha256 verification against LFS metadata,5per-file logs, disk-space checks, idempotence, node-aware sharding6(--shard k/n downloads files where index % n == k, so cluster nodes split work7without coordination).8"""910from __future__ import annotations1112import hashlib13import os14import shutil15import time16from datetime import datetime, timezone17from pathlib import Path1819from huggingface_hub import HfApi, hf_hub_download2021from . import paths22from .logging_utils import get_logger23from .manifest import FileRecord, Manifest2425CHUNK = 1 << 2026MIN_FREE_BYTES = 20 * (1 << 30) # refuse to start a file with <20 GB free272829def _sha256(path: Path) -> str:30 h = hashlib.sha256()31 with open(path, "rb") as f:32 for chunk in iter(lambda: f.read(CHUNK), b""):33 h.update(chunk)34 return h.hexdigest()353637def month_range(start: str, end: str) -> list[tuple[int, int]]:38 y0, m0 = map(int, start.split("-"))39 y1, m1 = map(int, end.split("-"))40 out = []41 y, m = y0, m042 while (y, m) <= (y1, m1):43 out.append((y, m))44 m += 145 if m == 13:46 y, m = y + 1, 147 return out484950def build_manifest(start: str, end: str) -> Manifest:51 """Populate the manifest from pinned repo metadata (idempotent)."""52 cfg = paths.base_config()["dataset"]53 api = HfApi()54 info = api.dataset_info(cfg["id"], revision=cfg["revision"], files_metadata=True)55 meta = {s.rfilename: s for s in info.siblings}56 man = Manifest(paths.get("manifests") / "raw_files.json")57 for y, m in month_range(start, end):58 remote = cfg["file_pattern"].format(year=y, month=m)59 if remote not in meta:60 raise FileNotFoundError(f"{remote} not in repo at pinned revision")61 s = meta[remote]62 existing = man.records.get(remote)63 if existing and existing.status == "complete":64 continue65 man.upsert(FileRecord(66 dataset_id=cfg["id"], dataset_revision=cfg["revision"],67 remote_path=remote,68 local_path=str(Path("data/raw") / remote),69 year=y, month=m,70 expected_bytes=s.size,71 lfs_sha256=(s.lfs.sha256 if s.lfs else None),72 ))73 man.save()74 return man757677def download(start: str, end: str, *, dry_run: bool = False,78 shard: str | None = None, max_retries: int = 5) -> None:79 cfg = paths.base_config()["dataset"]80 log = get_logger("download", paths.get("logs") / "download" / "download.jsonl")81 man = build_manifest(start, end)82 todo = sorted(man.pending(), key=lambda r: (r.year, r.month))83 if shard:84 k, n = map(int, shard.split("/"))85 todo = [r for i, r in enumerate(todo) if i % n == k]86 total = sum(r.expected_bytes or 0 for r in todo)87 log.info(f"{len(todo)} files to download, {total/1e9:.1f} GB"88 + (f" (shard {shard})" if shard else ""))89 if dry_run:90 for r in todo:91 log.info(f"DRY RUN would fetch {r.remote_path} ({(r.expected_bytes or 0)/1e6:.0f} MB)")92 return9394 raw_root = paths.get("raw")95 for rec in todo:96 if shutil.disk_usage(raw_root).free < MIN_FREE_BYTES:97 log.error("disk free below safety threshold; stopping")98 break99 rec.status = "downloading"100 rec.download_started_at = datetime.now(timezone.utc).isoformat()101 man.upsert(rec); man.save()102 for attempt in range(max_retries):103 try:104 local = hf_hub_download(105 cfg["id"], rec.remote_path, repo_type="dataset",106 revision=cfg["revision"], local_dir=raw_root,107 token=os.environ.get("HF_TOKEN"),108 )109 digest = _sha256(Path(local))110 if rec.lfs_sha256 and digest != rec.lfs_sha256:111 raise IOError(f"checksum mismatch for {rec.remote_path}")112 rec.sha256 = digest113 rec.downloaded_bytes = Path(local).stat().st_size114 rec.status = "complete"115 rec.download_completed_at = datetime.now(timezone.utc).isoformat()116 log.info(f"complete {rec.remote_path} ({rec.downloaded_bytes/1e6:.0f} MB)")117 break118 except Exception as e: # noqa: BLE001 — retry any transient failure119 wait = 2 ** attempt120 log.warning(f"attempt {attempt+1}/{max_retries} failed for "121 f"{rec.remote_path}: {e}; retrying in {wait}s")122 time.sleep(wait)123 else:124 rec.status = "failed"125 log.error(f"FAILED {rec.remote_path} after {max_retries} attempts")126 man.upsert(rec); man.save()127 man.export_tabular()128129130def verify(start: str, end: str) -> bool:131 """Re-hash every complete file against LFS sha256. Returns True if all pass."""132 log = get_logger("verify", paths.get("logs") / "download" / "verify.jsonl")133 man = Manifest(paths.get("manifests") / "raw_files.json")134 ok = True135 for rec in sorted(man.records.values(), key=lambda r: (r.year, r.month)):136 p = paths.repo_root() / rec.local_path137 if rec.status != "complete":138 log.warning(f"not complete: {rec.remote_path} ({rec.status})")139 ok = False140 continue141 if not p.exists():142 log.error(f"missing local file: {p}")143 rec.status = "failed"; ok = False; continue144 digest = _sha256(p)145 if rec.lfs_sha256 and digest != rec.lfs_sha256:146 log.error(f"checksum mismatch: {rec.remote_path}")147 rec.status = "failed"; ok = False148 man.save()149 log.info("verification " + ("PASSED" if ok else "FAILED"))150 return ok151