social-data-attribution-w2/scripts/attribution/sync_sample.py

"""Sync a sample dataset from R2 to a local directory.

Usage:
    bash scripts/bootstrap/with_r2_credentials.sh \
        python scripts/attribution/sync_sample.py \
        --prefix soc134-samples/sample_500_docs \
        --output ~/scratch/trackstar/sample_500_docs \
        --workers 64
"""

import argparse
import logging
import os
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

import boto3

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
)
log = logging.getLogger(__name__)

BUCKET = "soc127-dedup"
ENDPOINT = "https://0934ab8e84ac8f4e81decaf3eb121337.r2.cloudflarestorage.com"

_counter_lock = threading.Lock()
_downloaded = 0
_skipped = 0
_total = 0


def _make_client():
    return boto3.client(
        "s3",
        endpoint_url=ENDPOINT,
        aws_access_key_id=os.environ["R2_ACCESS_KEY_ID"],
        aws_secret_access_key=os.environ["R2_SECRET_ACCESS_KEY"],
        region_name="auto",
    )


def _download_one(key: str, size: int, dest: Path) -> bool:
    global _downloaded, _skipped
    if dest.exists() and dest.stat().st_size == size:
        with _counter_lock:
            _skipped += 1
        return False

    dest.parent.mkdir(parents=True, exist_ok=True)
    client = _make_client()
    client.download_file(BUCKET, key, str(dest))

    with _counter_lock:
        _downloaded += 1
        done = _downloaded + _skipped
        if _downloaded % 500 == 0:
            log.info(
                "Progress: %d/%d done (%d downloaded, %d skipped)",
                done,
                _total,
                _downloaded,
                _skipped,
            )
    return True


def sync_prefix(prefix: str, output_dir: Path, workers: int) -> int:
    global _total
    client = _make_client()
    output_dir.mkdir(parents=True, exist_ok=True)
    paginator = client.get_paginator("list_objects_v2")

    tasks = []
    for page in paginator.paginate(Bucket=BUCKET, Prefix=prefix):
        for obj in page.get("Contents", []):
            key = obj["Key"]
            rel = key[len(prefix) :].lstrip("/")
            if not rel:
                continue
            tasks.append((key, obj["Size"], output_dir / rel))

    _total = len(tasks)
    log.info("Found %d objects to sync with %d workers", _total, workers)

    errors = 0
    with ThreadPoolExecutor(max_workers=workers) as pool:
        futures = {
            pool.submit(_download_one, key, size, dest): rel
            for key, size, dest in tasks
        }
        for future in as_completed(futures):
            try:
                future.result()
            except Exception as exc:
                errors += 1
                log.error("Failed %s: %s", futures[future], exc)

    log.info(
        "Done: %d downloaded, %d skipped, %d errors", _downloaded, _skipped, errors
    )
    return _downloaded


def main() -> None:
    parser = argparse.ArgumentParser(description="Sync R2 prefix to local directory")
    parser.add_argument(
        "--prefix",
        required=True,
        help="R2 key prefix (e.g. soc134-samples/sample_500_docs)",
    )
    parser.add_argument(
        "--output",
        type=Path,
        required=True,
        help="Local output directory",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=64,
        help="Number of parallel download threads (default: 64)",
    )
    args = parser.parse_args()
    sync_prefix(args.prefix, args.output, args.workers)
    if _downloaded + _skipped < _total:
        log.error(
            "Incomplete sync: expected %d, got %d", _total, _downloaded + _skipped
        )
        raise SystemExit(1)


if __name__ == "__main__":
    main()