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

"""Upload materialized sample from Modal volume to R2 with parallel fan-out.

Fans out one worker per worker_NNNN directory. Each worker uploads its
shard files to R2 in parallel. Metadata is uploaded by a single worker.

Usage:
    SOC134_SAMPLE_NAME=sample_500_docs \
      uv run modal run --detach scripts/modal/upload_sample_to_r2.py \
        --sample-name sample_500_docs \
        --run-id soc134_materialize_20260320_150348
"""

from __future__ import annotations

import json
import logging
import os
import time
from pathlib import Path

import modal

logger = logging.getLogger("upload_sample_to_r2")

R2_BUCKET = "soc127-dedup"
R2_ENDPOINT_URL = "https://0934ab8e84ac8f4e81decaf3eb121337.r2.cloudflarestorage.com"
R2_SECRET_NAME = "r2-credentials"
R2_SAMPLES_PREFIX = "soc134-samples"

_upload_image = modal.Image.debian_slim(python_version="3.12").pip_install(
    "boto3>=1.37.0"
)

_SAMPLE_NAME_ENV = "SOC134_SAMPLE_NAME"
_sample_name = os.environ.get(_SAMPLE_NAME_ENV, "sample_500_docs")
_volume_name = f"soc134-output-{_sample_name.replace('_', '-')}"

app = modal.App("soc134-upload-sample-to-r2")
r2_secret = modal.Secret.from_name(R2_SECRET_NAME)
output_volume = modal.Volume.from_name(_volume_name)

VOLUME_MOUNT = "/data"


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)


def _get_r2_client():
    import boto3

    _ensure_r2_env()
    return 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",
    )


def _upload_file(client, local_path: Path, r2_key: str, content_type: str) -> bool:
    try:
        client.head_object(Bucket=R2_BUCKET, Key=r2_key)
        return False
    except Exception:
        pass

    max_retries = 3
    for attempt in range(max_retries):
        try:
            data = local_path.read_bytes()
            client.put_object(
                Bucket=R2_BUCKET, Key=r2_key, Body=data, ContentType=content_type
            )
            return True
        except Exception:
            if attempt == max_retries - 1:
                raise
            time.sleep((attempt + 1) * 5)
    return False


@app.function(
    image=_upload_image,
    volumes={VOLUME_MOUNT: output_volume},
    secrets=[r2_secret],
    timeout=7200,
    cpu=2,
    memory=4096,
)
def upload_metadata(
    sample_name: str,
    run_id: str,
) -> dict[str, object]:
    output_volume.reload()
    client = _get_r2_client()
    run_root = Path(VOLUME_MOUNT) / run_id
    r2_prefix = f"{R2_SAMPLES_PREFIX}/{sample_name}"

    uploaded = 0
    for fname in ["manifest.parquet", "sample_contract.json", "bin_summary.csv"]:
        local_path = run_root / fname
        if not local_path.exists():
            continue
        r2_key = f"{r2_prefix}/{fname}"
        content_type = "application/octet-stream"
        if fname.endswith(".json"):
            content_type = "application/json"
        elif fname.endswith(".csv"):
            content_type = "text/csv"
        if _upload_file(client, local_path, r2_key, content_type):
            uploaded += 1

    return {"metadata_uploaded": uploaded}


@app.function(
    image=_upload_image,
    volumes={VOLUME_MOUNT: output_volume},
    secrets=[r2_secret],
    timeout=7200,
    cpu=2,
    memory=4096,
)
def upload_worker_dir(
    sample_name: str,
    run_id: str,
    worker_index: int,
) -> dict[str, object]:
    t0 = time.monotonic()
    output_volume.reload()
    client = _get_r2_client()

    worker_name = f"worker_{worker_index:04d}"
    worker_dir = Path(VOLUME_MOUNT) / run_id / worker_name
    r2_prefix = f"{R2_SAMPLES_PREFIX}/{sample_name}"

    if not worker_dir.exists():
        return {
            "worker_index": worker_index,
            "uploaded": 0,
            "skipped": 0,
            "failed": 0,
            "total_bytes": 0,
        }

    shard_files = sorted(worker_dir.glob("*.jsonl.zst"))
    uploaded = 0
    skipped = 0
    failed = 0
    total_bytes = 0
    errors: list[str] = []

    for shard_file in shard_files:
        r2_key = f"{r2_prefix}/{worker_name}/{shard_file.name}"
        try:
            if _upload_file(client, shard_file, r2_key, "application/zstd"):
                uploaded += 1
                total_bytes += shard_file.stat().st_size
            else:
                skipped += 1
        except Exception as exc:
            failed += 1
            errors.append(f"{r2_key}: {type(exc).__name__}: {exc}")

    elapsed = time.monotonic() - t0
    return {
        "worker_index": worker_index,
        "files": len(shard_files),
        "uploaded": uploaded,
        "skipped": skipped,
        "failed": failed,
        "total_bytes": total_bytes,
        "elapsed_seconds": round(elapsed, 1),
        "errors": errors[:5],
    }


@app.local_entrypoint()
def main(
    sample_name: str = "",
    run_id: str = "",
    chunk_count: int = 128,
) -> None:
    if not sample_name:
        print("ERROR: --sample-name is required")
        raise SystemExit(1)
    if not run_id:
        print("ERROR: --run-id is required")
        raise SystemExit(1)

    print(f"Sample: {sample_name}")
    print(f"Run ID: {run_id}")
    print(f"Volume: {_volume_name}")
    print(f"R2 prefix: {R2_SAMPLES_PREFIX}/{sample_name}/")
    print(f"Workers: {chunk_count}")
    print("")

    print("Uploading metadata...")
    meta_result = upload_metadata.remote(sample_name, run_id)
    print(f"  Metadata files uploaded: {meta_result['metadata_uploaded']}")

    print(f"Launching {chunk_count} upload workers...")
    results = list(
        upload_worker_dir.starmap(
            [(sample_name, run_id, i) for i in range(chunk_count)]
        )
    )

    total_uploaded = sum(r.get("uploaded", 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)
    errors = [r for r in results if r.get("failed", 0) > 0]

    print("\nR2 upload results:")
    print(f"  Uploaded:  {total_uploaded:,}")
    print(f"  Skipped:   {total_skipped:,}")
    print(f"  Failed:    {total_failed:,}")
    print(f"  Bytes:     {total_bytes / (1024**3):.1f} GB")

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

    print(
        json.dumps(
            {
                "sample_name": sample_name,
                "r2_prefix": f"{R2_SAMPLES_PREFIX}/{sample_name}",
                "uploaded": total_uploaded,
                "skipped": total_skipped,
                "failed": total_failed,
                "total_bytes": total_bytes,
            },
            indent=2,
        )
    )