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

"""One-time download of SOC-95 manifest parquets from R2 to a persistent volume.

Usage:
    uv run modal run scripts/modal/cache_manifest.py
"""

from __future__ import annotations

import logging
import os
import time
from pathlib import Path

import modal

logger = logging.getLogger("cache_manifest")

R2_BUCKET_NAME = "soc127-dedup"
R2_MANIFEST_PREFIX = "soc95-manifest/data"
R2_ENDPOINT_URL = "https://0934ab8e84ac8f4e81decaf3eb121337.r2.cloudflarestorage.com"
R2_SECRET_NAME = "r2-credentials"

_local_path = Path(__file__).resolve()
if len(_local_path.parents) > 2:
    _cache_image = modal.Image.debian_slim(python_version="3.12").pip_install(
        "boto3>=1.37.0"
    )
else:
    _cache_image = modal.Image.debian_slim(python_version="3.12")

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

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

MANIFEST_MOUNT = "/manifests"


def _get_s3_client():
    import boto3

    env_aliases = {
        "AWS_ACCESS_KEY_ID": ("AWS_ACCESS_KEY_ID", "R2_ACCESS_KEY_ID", "access_key_id"),
        "AWS_SECRET_ACCESS_KEY": (
            "AWS_SECRET_ACCESS_KEY",
            "R2_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

    endpoint = os.environ.get("AWS_ENDPOINT_URL", R2_ENDPOINT_URL)

    return boto3.client(
        "s3",
        endpoint_url=endpoint,
        aws_access_key_id=os.environ["AWS_ACCESS_KEY_ID"],
        aws_secret_access_key=os.environ["AWS_SECRET_ACCESS_KEY"],
        region_name="auto",
    )


@app.function(
    image=_cache_image,
    volumes={MANIFEST_MOUNT: manifest_volume},
    secrets=[r2_secret],
    timeout=300,
)
def list_manifest_keys() -> list[str]:
    s3 = _get_s3_client()
    keys: list[str] = []
    prefix = R2_MANIFEST_PREFIX.rstrip("/") + "/"
    paginator = s3.get_paginator("list_objects_v2")
    for page in paginator.paginate(Bucket=R2_BUCKET_NAME, Prefix=prefix):
        for obj in page.get("Contents", []):
            k = obj["Key"]
            if k.endswith(".parquet"):
                keys.append(k)
    return sorted(keys)


@app.function(
    image=_cache_image,
    volumes={MANIFEST_MOUNT: manifest_volume},
    secrets=[r2_secret],
    timeout=3600,
    cpu=2,
    memory=8192,
)
def cache_manifest_chunk(
    keys: list[str],
    chunk_index: int,
) -> dict[str, object]:
    t0 = time.monotonic()
    s3 = _get_s3_client()

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

    for key in keys:
        dest = Path(MANIFEST_MOUNT) / key
        if dest.exists() and dest.stat().st_size > 0:
            skipped += 1
            continue

        max_retries = 3
        for attempt in range(max_retries):
            try:
                resp = s3.get_object(Bucket=R2_BUCKET_NAME, Key=key)
                data = resp["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
                else:
                    time.sleep((attempt + 1) * 3)

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

    return {
        "chunk_index": chunk_index,
        "total_keys": len(keys),
        "downloaded": downloaded,
        "skipped": skipped,
        "failed": failed,
        "total_bytes": total_bytes,
        "elapsed_seconds": round(elapsed, 1),
        "errors": errors[:5],
    }


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.local_entrypoint()
def main(
    chunk_count: int = 64,
    dry_run: bool = False,
) -> None:
    print("Listing manifest parquets from R2...")
    all_keys = list_manifest_keys.remote()
    print(f"Found {len(all_keys):,} manifest parquet files")

    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...")

    results = list(cache_manifest_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("\nManifest cache results:")
    print(f"  Total files:   {len(all_keys):,}")
    print(f"  Downloaded:    {total_downloaded:,}")
    print(f"  Skipped:       {total_skipped:,}")
    print(f"  Failed:        {total_failed:,}")
    print(f"  Bytes written: {total_bytes / (1024**2):.1f} MB")

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