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

"""Aggregate stats and validate pipeline outputs.

Uses direct S3 API for all R2 operations. Reads per-shard stats
files to compute totals without scanning every data record.
"""

from __future__ import annotations

import json
import logging
import random
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path


from .config import R2_PREFIX
from .soc127_app import (
    app,
    get_s3_client,
    image_no_bloom,
    r2_s3_download,
    r2_s3_list_keys,
    r2_s3_put_json,
    r2_secret,
)

logger = logging.getLogger("aggregate_stats")

POOL_PHASE = "phase1_pool_shared"
NONPOOL_FINAL_PHASE = "phase2_nonpool_final"
SAMPLE_CHECK_COUNT = 50


def _read_s3_json(key: str) -> dict:
    from .config import R2_BUCKET_NAME

    resp = get_s3_client().get_object(Bucket=R2_BUCKET_NAME, Key=key)
    return json.loads(resp["Body"].read().decode("utf-8"))


def _read_stats_parallel(keys: list[str]) -> list[dict]:
    with ThreadPoolExecutor(max_workers=64) as pool:
        return list(pool.map(_read_s3_json, keys))


@app.function(
    image=image_no_bloom,
    secrets=[r2_secret],
    cpu=2,
    memory=8192,
    ephemeral_disk=524288,
    timeout=7200,
    retries=1,
)
def aggregate() -> dict[str, object]:
    from dolma.dedup.materialize import REQUIRED_PREFLIGHT_FAMILIES

    logger.info("Collecting pool stats...")
    pool_stats_keys = r2_s3_list_keys(
        f"{R2_PREFIX}/{POOL_PHASE}/", suffix=".stats.json"
    )
    logger.info("Found %d pool stats files", len(pool_stats_keys))
    pool_stats = _read_stats_parallel(pool_stats_keys)

    pool_kept = sum(int(s.get("records_kept", 0)) for s in pool_stats)
    pool_errors = sum(1 for s in pool_stats if s.get("status") == "error")
    pool_families: Counter[str] = Counter()
    for s in pool_stats:
        fam = s.get("source_family") or s.get("source_category", "unknown")
        pool_families[str(fam)] += int(s.get("records_kept", 0))

    logger.info(
        "Pool: %d stats, %d kept, %d errors", len(pool_stats), pool_kept, pool_errors
    )

    logger.info("Collecting nonpool dedup stats...")
    nonpool_stats_keys = r2_s3_list_keys(
        f"{R2_PREFIX}/{NONPOOL_FINAL_PHASE}/stats/", suffix=".stats.json"
    )
    logger.info("Found %d nonpool stats files", len(nonpool_stats_keys))
    nonpool_stats = _read_stats_parallel(nonpool_stats_keys)

    nonpool_kept = sum(int(s.get("docs_out", 0)) for s in nonpool_stats)
    nonpool_errors = sum(1 for s in nonpool_stats if s.get("status") == "error")
    nonpool_families: Counter[str] = Counter()
    for s in nonpool_stats:
        for fam, count in (s.get("source_family_counts") or {}).items():
            nonpool_families[str(fam)] += int(count)

    logger.info(
        "Nonpool: %d stats, %d kept, %d errors",
        len(nonpool_stats),
        nonpool_kept,
        nonpool_errors,
    )

    all_families = dict(sorted((pool_families + nonpool_families).items()))
    total_kept = pool_kept + nonpool_kept

    missing_families = [
        f for f in REQUIRED_PREFLIGHT_FAMILIES if all_families.get(f, 0) == 0
    ]

    pool_data_keys = r2_s3_list_keys(f"{R2_PREFIX}/{POOL_PHASE}/", suffix=".jsonl.zst")
    nonpool_data_keys = r2_s3_list_keys(
        f"{R2_PREFIX}/{NONPOOL_FINAL_PHASE}/", suffix=".jsonl.zst"
    )
    nonpool_data_keys = [
        k for k in nonpool_data_keys if "/stats/" not in k and "/done/" not in k
    ]

    logger.info(
        "Data files: %d pool + %d nonpool", len(pool_data_keys), len(nonpool_data_keys)
    )

    sample_ok = True
    sample_notes: list[str] = []
    all_data_keys = pool_data_keys + nonpool_data_keys
    if all_data_keys:
        from dolma.dedup.materialize import iter_shard_records

        sample_keys = random.sample(
            all_data_keys, min(SAMPLE_CHECK_COUNT, len(all_data_keys))
        )
        for sk in sample_keys:
            local = Path(f"/tmp/sample/{sk.rsplit('/', 1)[-1]}")
            local.parent.mkdir(parents=True, exist_ok=True)
            try:
                r2_s3_download(sk, local)
                checked = 0
                for _, rec in iter_shard_records(local):
                    if rec is None or not isinstance(rec, dict):
                        sample_notes.append(f"Invalid record in {sk}")
                        sample_ok = False
                        break
                    text = str(rec.get("text", ""))
                    if not text.strip():
                        sample_notes.append(f"Empty text in {sk}")
                        sample_ok = False
                        break
                    checked += 1
                    if checked >= 3:
                        break
            except Exception as exc:
                sample_notes.append(f"Error reading {sk}: {exc}")
                sample_ok = False
            finally:
                local.unlink(missing_ok=True)

    result: dict[str, object] = {
        "pool_stats_count": len(pool_stats),
        "pool_records_kept": pool_kept,
        "pool_errors": pool_errors,
        "pool_data_files": len(pool_data_keys),
        "pool_family_counts": dict(sorted(pool_families.items())),
        "nonpool_stats_count": len(nonpool_stats),
        "nonpool_records_kept": nonpool_kept,
        "nonpool_errors": nonpool_errors,
        "nonpool_data_files": len(nonpool_data_keys),
        "nonpool_family_counts": dict(sorted(nonpool_families.items())),
        "total_records_kept": total_kept,
        "source_family_counts": all_families,
        "missing_families": missing_families,
        "all_families_present": not missing_families,
        "sample_check_passed": sample_ok,
        "sample_notes": sample_notes,
        "validation_passed": not missing_families and sample_ok,
    }

    r2_s3_put_json(f"{R2_PREFIX}/stats.json", result)
    r2_s3_put_json(f"{R2_PREFIX}/validation_report.json", result)

    if not result["validation_passed"]:
        logger.warning("Validation issues: %s", json.dumps(result, indent=2))
    else:
        logger.info(
            "Validation passed: %d total records (%d pool + %d nonpool)",
            total_kept,
            pool_kept,
            nonpool_kept,
        )

    return result


@app.local_entrypoint()
def main():
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
    )
    result = aggregate.remote()
    print(json.dumps(result, indent=2))