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

"""Phase 0: Preflight audit of SOC-127 input shards.

Samples shards per source family and reports metadata coverage,
URL availability, and text length distribution.

Source family grouping uses a hybrid approach:
- Pool shards: extracted from path (phase1_pool_shared/{source_family}/...)
- Nonpool shards: read from first record's _soc_127.source_family
"""

from __future__ import annotations

import json
import logging
from collections import defaultdict
from itertools import islice
from pathlib import Path

import modal

logger = logging.getLogger(__name__)

SAMPLE_SHARDS_PER_SOURCE = 5
DOCS_PER_SHARD = 1000

POOL_PREFIX = "soc127/phase1_pool_shared"
NONPOOL_PREFIX = "soc127/phase2_nonpool_final"

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

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

    _audit_image = (
        modal.Image.debian_slim(python_version="3.12")
        .pip_install("zstandard>=0.24.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:
    _audit_image = modal.Image.debian_slim(python_version="3.12")


def _load_config():
    from config import (
        R2_BUCKET,
        R2_ENDPOINT_URL,
        R2_SECRET_NAME,
        R2_STATS_PREFIX,
    )

    return R2_BUCKET, R2_ENDPOINT_URL, R2_SECRET_NAME, R2_STATS_PREFIX


R2_BUCKET, R2_ENDPOINT_URL, R2_SECRET_NAME, R2_STATS_PREFIX = _load_config()

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


def _source_from_pool_path(rel_path: str) -> str:
    parts = rel_path.split("/")
    if len(parts) >= 2 and parts[0] == "phase1_pool_shared":
        return parts[1]
    return "unknown"


def _read_source_family(shard_path: Path, max_records: int = 10) -> str:
    from dolma.local_io import iter_jsonlzst

    for record in islice(iter_jsonlzst(shard_path), max_records):
        soc127 = record.get("_soc_127", {})
        family = soc127.get("source_family")
        if isinstance(family, str) and family:
            return family
    return "unknown"


@app.function(
    image=_audit_image,
    volumes={"/r2": r2_mount},
    timeout=1800,
)
def run_audit(
    shards_per_source: int = SAMPLE_SHARDS_PER_SOURCE,
    docs_per_shard: int = DOCS_PER_SHARD,
) -> dict:
    from dolma.local_io import iter_jsonlzst

    by_source: dict[str, list[Path]] = defaultdict(list)
    pool_count = 0
    nonpool_count = 0

    pool_dir = Path(f"/r2/{POOL_PREFIX}")
    if pool_dir.exists():
        for path in pool_dir.rglob("*.jsonl.zst"):
            rel = str(path.relative_to(Path("/r2/soc127")))
            family = _source_from_pool_path(rel)
            by_source[family].append(path)
            pool_count += 1

    nonpool_dir = Path(f"/r2/{NONPOOL_PREFIX}")
    if nonpool_dir.exists():
        for path in sorted(nonpool_dir.glob("*.jsonl.zst")):
            family = _read_source_family(path)
            by_source[family].append(path)
            nonpool_count += 1

    report: dict[str, object] = {"sources": {}}

    for source, paths in sorted(by_source.items()):
        sample = paths[:shards_per_source]
        source_stats = {
            "shards_available": len(paths),
            "shards_sampled": len(sample),
            "total_docs_sampled": 0,
            "has_soc127_metadata": 0,
            "has_url": 0,
            "text_lengths": [],
        }

        for shard_path in sample:
            for record in islice(iter_jsonlzst(shard_path), docs_per_shard):
                source_stats["total_docs_sampled"] += 1

                soc127 = record.get("_soc_127")
                if isinstance(soc127, dict) and soc127.get("doc_id"):
                    source_stats["has_soc127_metadata"] += 1

                meta = record.get("metadata", {})
                if not isinstance(meta, dict):
                    meta = {}

                url = None
                for key in ("warc_url", "url", "source_url"):
                    val = meta.get(key)
                    if isinstance(val, str) and val:
                        url = val
                        break
                if url:
                    source_stats["has_url"] += 1

                text = record.get("text", "")
                if isinstance(text, str):
                    source_stats["text_lengths"].append(len(text))

        n = source_stats["total_docs_sampled"]
        if n > 0:
            lengths = source_stats["text_lengths"]
            sorted_len = sorted(lengths)
            source_stats["text_length_stats"] = {
                "min": sorted_len[0] if sorted_len else 0,
                "max": sorted_len[-1] if sorted_len else 0,
                "median": sorted_len[len(sorted_len) // 2] if sorted_len else 0,
                "mean": round(sum(sorted_len) / len(sorted_len), 1)
                if sorted_len
                else 0,
            }
            source_stats["soc127_coverage_pct"] = round(
                100 * source_stats["has_soc127_metadata"] / n, 1
            )
            source_stats["url_coverage_pct"] = round(
                100 * source_stats["has_url"] / n, 1
            )
        del source_stats["text_lengths"]
        report["sources"][source] = source_stats

    report["total_source_families"] = len(by_source)
    report["total_shards"] = pool_count + nonpool_count
    report["pool_shards"] = pool_count
    report["nonpool_shards"] = nonpool_count

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

    return report


@app.local_entrypoint()
def main():
    report = run_audit.remote()
    print(json.dumps(report, indent=2))