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

"""Multi-sample materialization from the persistent corpus cache volume.

DEPRECATED: This script writes all samples to one shared volume and hits
Modal's 500K inode limit with 3+ samples.  Use the per-sample approach
instead:

    bash scripts/modal/run_materialize_sample.sh data/samples/sample_500_docs

See materialize_working_sample.py with SOC134_OUTPUT_VOLUME env var.

Original single-pass approach: for each source shard, extract matching
doc_ids across all samples and write per-sample output files.  Reads from
the corpus volume (local NVMe) with R2 fallback.
"""

from __future__ import annotations

import json
import logging
import os
import time
import traceback
from datetime import datetime, timezone
from pathlib import Path

import modal

from config import (
    R2_BUCKET,
    R2_ENDPOINT_URL,
    R2_SECRET_NAME,
)

logger = logging.getLogger("materialize_all_samples")

_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")
    _materialize_image = (
        modal.Image.debian_slim(python_version="3.12")
        .pip_install(
            "boto3>=1.37.0",
            "pyarrow>=18.0.0",
            "zstandard>=0.24.0",
            "pandas>=2.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_ROOT, remote_path="/root/src", copy=True)
    )
    _hf_upload_image = (
        modal.Image.debian_slim(python_version="3.12")
        .pip_install(
            "huggingface-hub>=0.32",
            "hf-xet",
            "zstandard>=0.24.0",
        )
        .env(
            {
                "PYTHONPATH": "/root/src:/root",
                "HF_XET_HIGH_PERFORMANCE": "1",
            }
        )
        .add_local_file(_CONFIG_PY, remote_path="/root/config.py", copy=True)
        .add_local_dir(_SRC_ROOT, remote_path="/root/src", copy=True)
    )
else:
    _materialize_image = modal.Image.debian_slim(python_version="3.12")
    _hf_upload_image = modal.Image.debian_slim(python_version="3.12")

app = modal.App("soc134-materialize-all-samples")

r2_secret = modal.Secret.from_name(R2_SECRET_NAME)
hf_secret = modal.Secret.from_name("huggingface-secret")
materialize_volume = modal.Volume.from_name(
    "soc134-materialize-cache", create_if_missing=True
)
corpus_volume = modal.Volume.from_name("soc134-corpus-cache", create_if_missing=True)

VOLUME_MOUNT = "/cache"
CORPUS_MOUNT = "/corpus"
WORKER_CPU = 2
WORKER_MEMORY = 16384
WORKER_TIMEOUT = 7200


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 _create_r2_client():
    from dolma.quality.r2 import R2Config, create_r2_client

    config = R2Config.from_env(output_prefix="soc134-materialized")
    return create_r2_client(config), config


@app.function(
    image=_materialize_image,
    volumes={VOLUME_MOUNT: materialize_volume},
    timeout=300,
    cpu=1,
)
def stage_all_manifests(
    run_id: str,
    samples: list[dict[str, bytes | str]],
) -> dict[str, object]:
    staged = []
    for sample in samples:
        name = sample["name"]
        sample_dir = Path(VOLUME_MOUNT) / run_id / name
        sample_dir.mkdir(parents=True, exist_ok=True)

        manifest_path = sample_dir / "manifest.parquet"
        manifest_path.write_bytes(sample["manifest_bytes"])

        if sample.get("contract_json"):
            (sample_dir / "sample_contract.json").write_text(sample["contract_json"])
        if sample.get("bin_summary_csv"):
            (sample_dir / "bin_summary.csv").write_text(sample["bin_summary_csv"])

        staged.append({"name": name, "manifest_size": len(sample["manifest_bytes"])})

    materialize_volume.commit()
    return {"run_id": run_id, "staged": staged}


@app.function(
    image=_materialize_image,
    volumes={
        VOLUME_MOUNT: materialize_volume,
        CORPUS_MOUNT: corpus_volume,
    },
    secrets=[r2_secret],
    timeout=WORKER_TIMEOUT,
    cpu=WORKER_CPU,
    memory=WORKER_MEMORY,
)
def run_chunk_multi(
    run_id: str,
    sample_names: list[str],
    chunk_index: int,
    chunk_count: int,
) -> dict[str, object]:
    t0 = time.monotonic()

    try:
        _ensure_r2_env()
        client, config = _create_r2_client()

        from dolma.materialize_sample import (
            materialize_all,
            write_materialize_stats,
        )

        materialize_volume.reload()
        corpus_volume.reload()
        corpus_dir = Path(CORPUS_MOUNT)

        sample_results = {}
        for name in sample_names:
            manifest_path = Path(VOLUME_MOUNT) / run_id / name / "manifest.parquet"
            worker_dir = (
                Path(VOLUME_MOUNT) / run_id / name / f"worker_{chunk_index:04d}"
            )
            worker_dir.mkdir(parents=True, exist_ok=True)

            result = materialize_all(
                manifest_path=manifest_path,
                output_dir=worker_dir,
                r2_client=client,
                bucket=config.bucket,
                chunk_index=chunk_index,
                chunk_count=chunk_count,
                skip_existing=True,
                corpus_dir=corpus_dir,
            )

            write_materialize_stats(result, worker_dir)

            sample_results[name] = {
                "found_docs": result.total_found,
                "expected_docs": result.total_expected,
                "missing_docs": result.total_missing,
                "shards_processed": len(result.shard_stats),
                "bytes_written": result.total_bytes,
            }

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

        return {
            "run_id": run_id,
            "chunk_index": chunk_index,
            "elapsed_seconds": round(elapsed, 1),
            "samples": sample_results,
        }
    except Exception as exc:
        error_msg = f"{type(exc).__name__}: {exc}"
        logger.error(
            "FATAL chunk %d: %s\n%s",
            chunk_index,
            error_msg,
            traceback.format_exc(),
        )
        return {
            "run_id": run_id,
            "chunk_index": chunk_index,
            "status": "error",
            "error": error_msg,
        }


@app.function(
    image=_hf_upload_image,
    volumes={VOLUME_MOUNT: materialize_volume},
    secrets=[hf_secret],
    timeout=7200,
    cpu=2,
    memory=16384,
    ephemeral_disk=524288,
)
def upload_sample_to_hf(
    run_id: str,
    sample_name: str,
    hf_repo: str,
    chunk_count: int,
) -> dict[str, object]:
    try:
        from huggingface_hub import CommitOperationAdd, HfApi

        materialize_volume.reload()

        api = HfApi(token=os.environ.get("HF_TOKEN"))
        api.create_repo(repo_id=hf_repo, repo_type="dataset", exist_ok=True)

        sample_root = Path(VOLUME_MOUNT) / run_id / sample_name

        metadata_ops = []
        contract = {}
        contract_path = sample_root / "sample_contract.json"
        if contract_path.exists():
            contract_text = contract_path.read_text()
            contract = json.loads(contract_text)
            metadata_ops.append(
                CommitOperationAdd(
                    path_in_repo="sample_contract.json",
                    path_or_fileobj=contract_text.encode(),
                )
            )

        bin_summary_path = sample_root / "bin_summary.csv"
        if bin_summary_path.exists():
            metadata_ops.append(
                CommitOperationAdd(
                    path_in_repo="bin_summary.csv",
                    path_or_fileobj=bin_summary_path.read_bytes(),
                )
            )

        manifest_path = sample_root / "manifest.parquet"
        if manifest_path.exists():
            metadata_ops.append(
                CommitOperationAdd(
                    path_in_repo="working_sample_manifest.parquet",
                    path_or_fileobj=str(manifest_path),
                )
            )

        readme_text = _build_readme(hf_repo, contract)
        metadata_ops.append(
            CommitOperationAdd(
                path_in_repo="README.md",
                path_or_fileobj=readme_text.encode(),
            )
        )

        if metadata_ops:
            api.preupload_lfs_files(
                repo_id=hf_repo,
                additions=[
                    op for op in metadata_ops if op.path_in_repo.endswith(".parquet")
                ],
                repo_type="dataset",
            )
            api.create_commit(
                repo_id=hf_repo,
                operations=metadata_ops,
                commit_message="Add sampling metadata and data card",
                repo_type="dataset",
            )

        all_files: list[tuple[Path, str]] = []
        for chunk_idx in range(chunk_count):
            worker_dir = sample_root / f"worker_{chunk_idx:04d}"
            if not worker_dir.exists():
                continue
            for jsonl_file in sorted(worker_dir.glob("*.jsonl.zst")):
                hf_path = f"data/{jsonl_file.name}"
                all_files.append((jsonl_file, hf_path))

        if not all_files:
            return {"status": "error", "error": "No JSONL.zst files found"}

        batch_size = 100
        total_uploaded = 0
        total_bytes = 0

        for batch_start in range(0, len(all_files), batch_size):
            batch = all_files[batch_start : batch_start + batch_size]
            operations = [
                CommitOperationAdd(
                    path_in_repo=hf_path,
                    path_or_fileobj=str(local_path),
                )
                for local_path, hf_path in batch
            ]

            api.preupload_lfs_files(
                repo_id=hf_repo,
                additions=operations,
                repo_type="dataset",
            )

            batch_num = batch_start // batch_size
            api.create_commit(
                repo_id=hf_repo,
                operations=operations,
                commit_message=f"Add materialized shards batch {batch_num}",
                repo_type="dataset",
            )

            batch_bytes = sum(p.stat().st_size for p, _ in batch)
            total_uploaded += len(batch)
            total_bytes += batch_bytes

        return {
            "sample": sample_name,
            "hf_repo": hf_repo,
            "files_uploaded": total_uploaded,
            "bytes_uploaded": total_bytes,
        }
    except Exception as exc:
        error_msg = f"{type(exc).__name__}: {exc}"
        logger.error(
            "FATAL HF upload %s: %s\n%s", sample_name, error_msg, traceback.format_exc()
        )
        return {"sample": sample_name, "status": "error", "error": error_msg}


def _build_readme(hf_repo: str, contract: dict) -> str:
    repo_name = hf_repo.split("/")[-1] if "/" in hf_repo else hf_repo
    dpb = contract.get("WORKING_SAMPLE_DOCS_PER_BIN", 0)
    tfpb = contract.get("WORKING_SAMPLE_TOKEN_FLOOR_PER_BIN", 0)
    mode = f"{dpb:,} docs/bin" if dpb else f"{tfpb:,} token floor/bin"
    return (
        "\n".join(
            [
                "---",
                "license: odc-by",
                "task_categories:",
                "  - text-generation",
                "language:",
                "  - en",
                "---",
                "",
                f"# {repo_name}",
                "",
                "Materialized working sample from the Dolma3 6T deduplicated corpus.",
                "",
                "## Parameters",
                "",
                "| Parameter | Value |",
                "|-----------|-------|",
                f"| Sampling mode | {mode} |",
                f"| Seed | {contract.get('WORKING_SAMPLE_SAMPLING_SEED', 'N/A')} |",
                "",
                "## Counts",
                "",
                "| Metric | Value |",
                "|--------|-------|",
                f"| Total documents | {contract.get('WORKING_SAMPLE_REALIZED_DOC_COUNT', 0):,} |",
                f"| Total tokens | {contract.get('WORKING_SAMPLE_REALIZED_TOKEN_TOTAL', 0):,} |",
                f"| Bins covered | {contract.get('WORKING_SAMPLE_COVERED_BIN_COUNT', 0)}/{contract.get('WORKING_SAMPLE_TOTAL_BIN_COUNT', 576)} |",
                f"| Bins underfilled | {contract.get('WORKING_SAMPLE_UNDERFILLED_BIN_COUNT', 0)} |",
                "",
                "## Format",
                "",
                "JSONL files compressed with zstandard under `data/`.",
                "Each record contains the original Dolma document fields (id, text, metadata).",
                "",
                "Sampling metadata files:",
                "- `working_sample_manifest.parquet` - sampled doc_ids with bin assignments",
                "- `sample_contract.json` - aggregate sampling contract",
                "- `bin_summary.csv` - per-bin fill rates (576 bins = 24 topics x 24 formats)",
            ]
        )
        + "\n"
    )


@app.local_entrypoint()
def main(
    sample_dir: str = "data/samples",
    hf_org: str = "HCAI-Lab",
    chunk_count: int = 128,
    run_id: str = "",
    skip_upload: bool = False,
) -> None:
    sample_base = Path(sample_dir)
    if not sample_base.exists():
        print(f"ERROR: sample directory not found: {sample_base}")
        raise SystemExit(1)

    sample_dirs = sorted(
        [
            d
            for d in sample_base.iterdir()
            if d.is_dir() and (d / "working_sample_manifest.parquet").exists()
        ]
    )

    if not sample_dirs:
        print(f"ERROR: no samples found in {sample_base}")
        raise SystemExit(1)

    current_run_id = run_id or datetime.now(timezone.utc).strftime(
        "soc134_multi_%Y%m%d_%H%M%S"
    )

    print(f"Run ID: {current_run_id}")
    print(f"Samples: {len(sample_dirs)}")
    for sd in sample_dirs:
        print(f"  - {sd.name}")
    print(f"Chunks: {chunk_count}")

    samples = []
    for sd in sample_dirs:
        manifest_path = sd / "working_sample_manifest.parquet"
        contract_path = sd / "sample_contract.json"
        bin_summary_path = sd / "bin_summary.csv"

        samples.append(
            {
                "name": sd.name,
                "manifest_bytes": manifest_path.read_bytes(),
                "contract_json": contract_path.read_text()
                if contract_path.exists()
                else "",
                "bin_summary_csv": bin_summary_path.read_text()
                if bin_summary_path.exists()
                else "",
            }
        )

    sample_names = [s["name"] for s in samples]

    print("Staging all manifests to volume...")
    stage_result = stage_all_manifests.remote(current_run_id, samples)
    for s in stage_result["staged"]:
        print(f"  {s['name']}: {s['manifest_size']:,} bytes")

    print(f"\nLaunching {chunk_count} workers across {len(sample_names)} samples...")
    worker_results = list(
        run_chunk_multi.starmap(
            [(current_run_id, sample_names, i, chunk_count) for i in range(chunk_count)]
        )
    )

    errors = [r for r in worker_results if r.get("status") == "error"]

    print(
        f"\nResults ({len(worker_results) - len(errors)}/{len(worker_results)} workers ok):"
    )
    for name in sample_names:
        total_found = sum(
            r.get("samples", {}).get(name, {}).get("found_docs", 0)
            for r in worker_results
            if r.get("status") != "error"
        )
        total_expected = sum(
            r.get("samples", {}).get(name, {}).get("expected_docs", 0)
            for r in worker_results
            if r.get("status") != "error"
        )
        total_missing = sum(
            r.get("samples", {}).get(name, {}).get("missing_docs", 0)
            for r in worker_results
            if r.get("status") != "error"
        )
        print(
            f"  {name}: {total_found:,}/{total_expected:,} docs (missing: {total_missing:,})"
        )

    if errors:
        print(f"\n{len(errors)} workers failed:")
        for e in errors[:5]:
            print(f"  Chunk {e['chunk_index']}: {e['error']}")

    if not skip_upload and not errors:
        hf_repo_map = {name: f"{hf_org}/dolma3-6t-{name}" for name in sample_names}

        print("\nUploading to HuggingFace...")
        for name in sample_names:
            hf_repo = hf_repo_map[name]
            print(f"  Uploading {name} -> {hf_repo}")
            result = upload_sample_to_hf.remote(
                current_run_id, name, hf_repo, chunk_count
            )
            if result.get("status") == "error":
                print(f"    FAILED: {result['error']}")
            else:
                print(f"    OK: {result.get('files_uploaded', 0)} files")
    elif errors:
        print("\nSkipping HF upload due to worker errors")

    print(
        json.dumps(
            {
                "run_id": current_run_id,
                "samples": sample_names,
                "workers_completed": len(worker_results) - len(errors),
                "workers_failed": len(errors),
            },
            indent=2,
        )
    )