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

"""Phase 5: Publish sidecar dataset to Hugging Face.

Partitions sidecar Parquet files, generates taxonomy files,
and uploads via upload_large_folder().
"""

from __future__ import annotations

import json
import logging
import shutil
from pathlib import Path

import modal

from config import (
    R2_BUCKET,
    R2_ENDPOINT_URL,
    R2_OUTPUT_PREFIX,
    R2_SECRET_NAME,
    R2_STATS_PREFIX,
)

logger = logging.getLogger(__name__)

HF_REPO = "HCAI-Lab/dolma3-6t-unique_weborganizer_labels"
FILES_PER_DIR = 5000

_local_path = Path(__file__).resolve()
if len(_local_path.parents) > 2:
    _REPO_ROOT = _local_path.parents[2]
    _SRC_DOLMA = str(_REPO_ROOT / "src" / "dolma")
    _CONFIG_PY = str(_REPO_ROOT / "scripts" / "modal" / "config.py")

    _publish_image = (
        modal.Image.debian_slim(python_version="3.12")
        .pip_install("pyarrow>=18.0.0", "huggingface_hub>=0.25")
        .env({"PYTHONPATH": "/root/src:/root"})
        .add_local_file(_CONFIG_PY, remote_path="/root/config.py", copy=True)
        .add_local_dir(_SRC_DOLMA, remote_path="/root/src/dolma")
    )
else:
    _publish_image = modal.Image.debian_slim(python_version="3.12")

app = modal.App("soc91-publish")

r2_mount = modal.CloudBucketMount(
    R2_BUCKET,
    bucket_endpoint_url=R2_ENDPOINT_URL,
    secret=modal.Secret.from_name(R2_SECRET_NAME),
)


@app.function(
    image=_publish_image,
    volumes={"/r2": r2_mount},
    timeout=7200,
    secrets=[modal.Secret.from_name("huggingface-secret")],
)
def prepare_and_upload(
    repo_id: str = HF_REPO,
    dry_run: bool = True,
) -> dict:
    from dolma.taxonomy import write_taxonomy_files

    output_dir = Path(f"/r2/{R2_OUTPUT_PREFIX}")
    staging = Path("/tmp/soc91_publish")
    if staging.exists():
        shutil.rmtree(staging)
    staging.mkdir(parents=True)

    parquet_files = sorted(output_dir.rglob("*.parquet"))
    logger.info("Found %d parquet files to publish", len(parquet_files))

    for i, src in enumerate(parquet_files):
        partition = f"{i // FILES_PER_DIR:03d}"
        dst_dir = staging / "data" / partition
        dst_dir.mkdir(parents=True, exist_ok=True)
        dst = dst_dir / src.name
        shutil.copy2(str(src), str(dst))

    taxonomy_dir = staging / "taxonomy"
    write_taxonomy_files(taxonomy_dir)

    stats_path = Path(f"/r2/{R2_STATS_PREFIX}/aggregate_stats.json")
    if stats_path.exists():
        shutil.copy2(str(stats_path), str(staging / "stats.json"))

    file_count = sum(1 for _ in staging.rglob("*") if _.is_file())
    dir_count = sum(1 for _ in staging.rglob("*") if _.is_dir())
    max_per_dir = max(
        (
            sum(1 for f in d.iterdir() if f.is_file())
            for d in staging.rglob("*")
            if d.is_dir()
        ),
        default=0,
    )

    summary = {
        "repo_id": repo_id,
        "total_files": file_count,
        "total_dirs": dir_count,
        "max_files_per_dir": max_per_dir,
        "parquet_count": len(parquet_files),
        "dry_run": dry_run,
    }

    if not dry_run:
        from huggingface_hub import HfApi

        api = HfApi()
        api.create_repo(repo_id, repo_type="dataset", exist_ok=True)
        api.upload_large_folder(
            repo_id=repo_id,
            repo_type="dataset",
            folder_path=str(staging),
        )
        summary["uploaded"] = True
        logger.info("Upload complete to %s", repo_id)
    else:
        logger.info("Dry run: skipping upload to %s", repo_id)
        summary["uploaded"] = False

    return summary


@app.local_entrypoint()
def main(
    repo: str = HF_REPO,
    dry_run: bool = True,
):
    result = prepare_and_upload.remote(repo_id=repo, dry_run=dry_run)
    print(json.dumps(result, indent=2))
    if result["dry_run"]:
        print("\nThis was a dry run. Use --no-dry-run to upload.")