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

"""Phase 4: Post-production validation for SOC-91.

Checks coverage, row counts, and label distributions.
"""

from __future__ import annotations

import json
import logging
from collections import Counter, defaultdict
from pathlib import Path

import modal

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

logger = logging.getLogger(__name__)

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

    _validate_image = (
        modal.Image.debian_slim(python_version="3.12")
        .pip_install("pyarrow>=18.0.0")
        .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:
    _validate_image = modal.Image.debian_slim(python_version="3.12")

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

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


@app.function(
    image=_validate_image,
    volumes={"/r2": r2_mount},
    timeout=3600,
)
def run_validation(sample_rows: int = 10) -> dict:
    import pyarrow.parquet as pq

    from dolma.taxonomy import format_taxonomy, topic_taxonomy

    output_dir = Path(f"/r2/{R2_OUTPUT_PREFIX}")

    input_shards = set()
    for prefix in R2_INPUT_PREFIXES:
        input_dir = Path(f"/r2/{prefix}")
        if not input_dir.exists():
            continue
        for path in input_dir.rglob("*.jsonl.zst"):
            input_shards.add(path.name)

    missing: list[str] = []
    incomplete: list[str] = []
    total_rows = 0
    topic_hist: Counter = Counter()
    format_hist: Counter = Counter()
    by_source: dict[str, dict] = defaultdict(lambda: {"shards": 0, "rows": 0})

    topic_names = topic_taxonomy()
    format_names = format_taxonomy()

    for shard in sorted(input_shards):
        parquet_name = shard.replace(".jsonl.zst", ".parquet")
        parquet_path = output_dir / parquet_name
        done_path = Path(f"{parquet_path}.done")

        if not parquet_path.exists():
            missing.append(shard)
            continue
        if not done_path.exists():
            incomplete.append(shard)
            continue

        stats_path = Path(f"{parquet_path}.stats.json")
        source = "unknown"
        if stats_path.exists():
            try:
                stats_data = json.loads(stats_path.read_text())
                source = stats_data.get("source_family", "unknown")
            except (json.JSONDecodeError, OSError):
                pass

        table = pq.read_table(str(parquet_path))
        n = table.num_rows
        total_rows += n
        by_source[source]["shards"] += 1
        by_source[source]["rows"] += n

        topic_ids = table.column("topic_nourl_label_id").to_pylist()
        format_ids = table.column("format_nourl_label_id").to_pylist()
        for tid in topic_ids:
            if tid is not None:
                label = topic_names.get(tid, f"unknown_{tid}")
                topic_hist[label] += 1
        for fid in format_ids:
            if fid is not None:
                label = format_names.get(fid, f"unknown_{fid}")
                format_hist[label] += 1

    report = {
        "total_input_shards": len(input_shards),
        "total_output_rows": total_rows,
        "missing_outputs": len(missing),
        "incomplete_outputs": len(incomplete),
        "missing_shards": missing[:20],
        "incomplete_shards": incomplete[:20],
        "by_source": dict(by_source),
        "topic_histogram": dict(topic_hist.most_common()),
        "format_histogram": dict(format_hist.most_common()),
    }

    stats_dir = Path(f"/r2/{R2_STATS_PREFIX}")
    stats_dir.mkdir(parents=True, exist_ok=True)
    out_path = stats_dir / "validation_report.json"
    with out_path.open("w", encoding="utf-8") as f:
        json.dump(report, f, indent=2)
        f.write("\n")

    return report


@app.local_entrypoint()
def main():
    report = run_validation.remote()
    print(f"Input shards: {report['total_input_shards']}")
    print(f"Output rows:  {report['total_output_rows']}")
    print(f"Missing:      {report['missing_outputs']}")
    print(f"Incomplete:   {report['incomplete_outputs']}")
    if report["missing_shards"]:
        print(f"Missing (first 20): {report['missing_shards']}")
    print("\nTopic distribution (NoURL):")
    for label, count in list(report["topic_histogram"].items())[:10]:
        print(f"  {label}: {count}")
    print("\nFormat distribution (NoURL):")
    for label, count in list(report["format_histogram"].items())[:10]:
        print(f"  {label}: {count}")