social-data-attribution-w2/scripts/modal/cache_corpus_shards.py

"""One-time download of all source shards from R2 to a persistent Modal volume.

Fans out workers to download source shards in parallel.  Each worker handles
a chunk of shards and writes them to a shared volume preserving the original
R2 key structure.  Workers skip files that already exist (idempotent).

Usage:
    uv run modal run scripts/modal/cache_corpus_shards.py --chunk-count 256
"""

from __future__ import annotations

import logging
import os
import time
from pathlib import Path

import modal

from config import (
    R2_BUCKET,
    R2_ENDPOINT_URL,
    R2_INPUT_PREFIXES,
    R2_SECRET_NAME,
)

logger = logging.getLogger("cache_corpus_shards")

_local_path = Path(__file__).resolve()
if len(_local_path.parents) > 2:
    _REPO_ROOT = _local_path.parents[2]
    _SRC_ROOT = str(_REPO_ROOT / "src")
    _CONFIG_PY = str(_REPO_ROOT / "scripts" / "modal" / "config.py")
    _cache_image = (
        modal.Image.debian_slim(python_version="3.12")
        .pip_install("boto3>=1.37.0")
        .env({"PYTHONPATH": "/root/src:/root"})
        .add_local_file(_CONFIG_PY, remote_path="/root/config.py", copy=True)
        .add_local_dir(_SRC_ROOT, remote_path="/root/src", copy=True)
    )
else:
    _cache_image = modal.Image.debian_slim(python_version="3.12")

app = modal.App("soc134-cache-corpus")

r2_secret = modal.Secret.from_name(R2_SECRET_NAME)
corpus_volume = modal.Volume.from_name("soc134-corpus-cache", create_if_missing=True)

CORPUS_MOUNT = "/corpus"
WORKER_TIMEOUT = 7200
WORKER_CPU = 2
WORKER_MEMORY = 8192


def _ensure_r2_env() -> None:
    env_aliases = {
        "R2_ACCESS_KEY_ID": ("R2_ACCESS_KEY_ID", "AWS_ACCESS_KEY_ID", "access_key_id"),
        "R2_SECRET_ACCESS_KEY": (
            "R2_SECRET_ACCESS_KEY",
            "AWS_SECRET_ACCESS_KEY",
            "secret_access_key",
        ),
    }
    for target, aliases in env_aliases.items():
        if target in os.environ and os.environ[target]:
            continue
        for alias in aliases:
            value = os.environ.get(alias)
            if value:
                os.environ[target] = value
                break
        if target not in os.environ:
            raise KeyError(target)

    if "R2_ENDPOINT_URL" not in os.environ:
        os.environ["R2_ENDPOINT_URL"] = R2_ENDPOINT_URL
    if "R2_BUCKET" not in os.environ:
        os.environ["R2_BUCKET"] = R2_BUCKET


def _list_all_source_keys(client, bucket: str) -> list[str]:
    keys: list[str] = []
    for prefix in R2_INPUT_PREFIXES:
        normalized = prefix.rstrip("/") + "/"
        paginator = client.get_paginator("list_objects_v2")
        for page in paginator.paginate(Bucket=bucket, Prefix=normalized):
            for item in page.get("Contents", []):
                key = item["Key"]
                if key.endswith(".jsonl.zst"):
                    keys.append(key)
    return sorted(set(keys))


def _slice_for_chunk(items: list[str], chunk_index: int, chunk_count: int) -> list[str]:
    chunk_size = len(items) // chunk_count
    remainder = len(items) % chunk_count
    start = chunk_index * chunk_size + min(chunk_index, remainder)
    end = start + chunk_size + (1 if chunk_index < remainder else 0)
    return items[start:end]


@app.function(
    image=_cache_image,
    volumes={CORPUS_MOUNT: corpus_volume},
    secrets=[r2_secret],
    timeout=300,
    cpu=1,
)
def list_source_shards() -> list[str]:
    _ensure_r2_env()
    import boto3

    client = boto3.client(
        "s3",
        endpoint_url=R2_ENDPOINT_URL,
        aws_access_key_id=os.environ["R2_ACCESS_KEY_ID"],
        aws_secret_access_key=os.environ["R2_SECRET_ACCESS_KEY"],
        region_name="auto",
    )
    return _list_all_source_keys(client, R2_BUCKET)


@app.function(
    image=_cache_image,
    volumes={CORPUS_MOUNT: corpus_volume},
    secrets=[r2_secret],
    timeout=WORKER_TIMEOUT,
    cpu=WORKER_CPU,
    memory=WORKER_MEMORY,
)
def cache_chunk(
    shard_keys: list[str],
    chunk_index: int,
) -> dict[str, object]:
    t0 = time.monotonic()
    _ensure_r2_env()
    import boto3

    client = boto3.client(
        "s3",
        endpoint_url=R2_ENDPOINT_URL,
        aws_access_key_id=os.environ["R2_ACCESS_KEY_ID"],
        aws_secret_access_key=os.environ["R2_SECRET_ACCESS_KEY"],
        region_name="auto",
    )

    downloaded = 0
    skipped = 0
    failed = 0
    total_bytes = 0
    errors: list[str] = []

    for i, key in enumerate(shard_keys):
        dest = Path(CORPUS_MOUNT) / key
        if dest.exists() and dest.stat().st_size > 0:
            skipped += 1
            continue

        max_retries = 3
        for attempt in range(max_retries):
            try:
                response = client.get_object(Bucket=R2_BUCKET, Key=key)
                data = response["Body"].read()

                dest.parent.mkdir(parents=True, exist_ok=True)
                tmp = dest.with_suffix(".tmp")
                tmp.write_bytes(data)
                tmp.rename(dest)

                downloaded += 1
                total_bytes += len(data)
                break
            except Exception as exc:
                if attempt == max_retries - 1:
                    error_msg = f"{key}: {type(exc).__name__}: {exc}"
                    errors.append(error_msg)
                    failed += 1
                    logger.error("Failed shard %s: %s", key, error_msg)
                else:
                    wait = (attempt + 1) * 5
                    time.sleep(wait)

        if (i + 1) % 50 == 0:
            corpus_volume.commit()
            logger.info(
                "Chunk %d progress: %d/%d (downloaded=%d, skipped=%d, failed=%d)",
                chunk_index,
                i + 1,
                len(shard_keys),
                downloaded,
                skipped,
                failed,
            )

    corpus_volume.commit()
    elapsed = time.monotonic() - t0

    return {
        "chunk_index": chunk_index,
        "total_shards": len(shard_keys),
        "downloaded": downloaded,
        "skipped": skipped,
        "failed": failed,
        "total_bytes": total_bytes,
        "elapsed_seconds": round(elapsed, 1),
        "errors": errors[:10],
    }


@app.local_entrypoint()
def main(
    chunk_count: int = 256,
    dry_run: bool = False,
) -> None:
    print("Listing all source shards from R2...")
    all_keys = list_source_shards.remote()
    print(
        f"Found {len(all_keys):,} source shards across {len(R2_INPUT_PREFIXES)} prefixes"
    )

    if dry_run:
        print("Dry run - not downloading")
        return

    chunks = [_slice_for_chunk(all_keys, i, chunk_count) for i in range(chunk_count)]
    non_empty = [(keys, i) for i, keys in enumerate(chunks) if keys]
    print(f"Launching {len(non_empty)} workers (chunk_count={chunk_count})...")

    results = list(cache_chunk.starmap(non_empty))

    total_downloaded = sum(r.get("downloaded", 0) for r in results)
    total_skipped = sum(r.get("skipped", 0) for r in results)
    total_failed = sum(r.get("failed", 0) for r in results)
    total_bytes = sum(r.get("total_bytes", 0) for r in results)

    print("\nCorpus cache results:")
    print(f"  Total shards:  {len(all_keys):,}")
    print(f"  Downloaded:    {total_downloaded:,}")
    print(f"  Skipped (exist): {total_skipped:,}")
    print(f"  Failed:        {total_failed:,}")
    print(f"  Bytes written: {total_bytes / (1024**3):.1f} GB")

    if total_failed > 0:
        failed_results = [r for r in results if r.get("failed", 0) > 0]
        for r in failed_results[:5]:
            for err in r.get("errors", [])[:3]:
                print(f"  Error: {err}")