social-data-attribution-w2/scripts/slurm/dedup/preflight_materialize_docs.py

#!/usr/bin/env python3
"""Fail-closed preflight checks for SOC-127 materialization."""

from __future__ import annotations

import argparse
import logging
from collections import Counter
from pathlib import Path

from dolma.constants import DOLMA_6T_MIX_DATASET_ID, DOLMA_POOL_DATASET_ID
from dolma.dedup.materialize import (
    DEFAULT_MIN_HIT_RATE,
    DEFAULT_MIN_ID_RATE,
    DEFAULT_SAMPLE_SIZE,
    REQUIRED_PREFLIGHT_FAMILIES,
    download_shard,
    read_manifest,
    resolve_record_doc_id,
    write_json,
)
from dolma.provenance import BloomIndex, shard_folder_name, source_family

logger = logging.getLogger("preflight_materialize_docs")

DEFAULT_OUTPUT = Path(
    "/storage/ice-shared/cs7634/staff/TDA/soc-90/unique_docs_work/preflight.json"
)


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run SOC-127 preflight checks")
    parser.add_argument("--pool-dataset", default=DOLMA_POOL_DATASET_ID)
    parser.add_argument("--mix-dataset", default=DOLMA_6T_MIX_DATASET_ID)
    parser.add_argument("--pool-manifest", type=Path, required=True)
    parser.add_argument("--mix-manifest", type=Path, required=True)
    parser.add_argument("--bloom-file", type=Path, required=True)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--download-dir", type=Path, default=None)
    parser.add_argument("--pool-local-root", type=Path, default=None)
    parser.add_argument("--mix-local-root", type=Path, default=None)
    parser.add_argument("--sample-size", type=int, default=DEFAULT_SAMPLE_SIZE)
    parser.add_argument("--sample-shards-per-family", type=int, default=4)
    parser.add_argument("--min-id-rate", type=float, default=DEFAULT_MIN_ID_RATE)
    parser.add_argument("--min-hit-rate", type=float, default=DEFAULT_MIN_HIT_RATE)
    parser.add_argument("--verbose", action="store_true", default=False)
    return parser.parse_args(argv)


def select_probe_shards(
    shards: list[str], families: tuple[str, ...], sample_shards_per_family: int
) -> dict[str, list[str]]:
    selected = {family: [] for family in families}
    seen_folders = {family: set() for family in families}
    stack_python_shard: str | None = None
    stack_regular_shard: str | None = None
    for shard_path in shards:
        family = source_family(shard_path)
        if family not in selected:
            continue
        folder = shard_folder_name(shard_path)
        if family == "stack_edu":
            if folder == "stack_edu-Python" and stack_python_shard is None:
                stack_python_shard = shard_path
            elif folder != "stack_edu-Python" and stack_regular_shard is None:
                stack_regular_shard = shard_path
            continue
        if len(selected[family]) >= sample_shards_per_family:
            continue
        if folder in seen_folders[family]:
            continue
        selected[family].append(shard_path)
        seen_folders[family].add(folder)
    if "stack_edu" in selected:
        if stack_python_shard is not None:
            selected["stack_edu"].append(stack_python_shard)
        if stack_regular_shard is not None:
            selected["stack_edu"].append(stack_regular_shard)
    missing = [family for family in families if not selected[family]]
    if missing:
        raise ValueError(f"Missing sample shard(s) for: {', '.join(missing)}")
    return selected


def probe_shards(
    dataset_id: str,
    shard_paths: list[str],
    bloom: BloomIndex,
    sample_size: int,
    min_id_rate: float,
    min_hit_rate: float,
    download_dir: Path | None,
    local_root: Path | None,
) -> dict[str, object]:
    valid_records = 0
    invalid_json_records = 0
    records_with_top_level_id = 0
    records_with_doc_id = 0
    bloom_hits = 0
    doc_id_fields_seen: Counter[str] = Counter()
    from dolma.dedup.materialize import iter_shard_records

    remaining = sample_size
    for index, shard_path in enumerate(shard_paths):
        local_path = download_shard(dataset_id, shard_path, download_dir, local_root)
        shards_left = len(shard_paths) - index
        target_for_shard = max(1, -(-remaining // shards_left))
        shard_records = 0
        for _, record in iter_shard_records(local_path):
            if record is None or not isinstance(record, dict):
                invalid_json_records += 1
                continue
            valid_records += 1
            shard_records += 1
            top_level_id = record.get("id")
            if isinstance(top_level_id, str) and top_level_id:
                records_with_top_level_id += 1
            doc_id, doc_id_field = resolve_record_doc_id(record, shard_path)
            if doc_id is not None:
                records_with_doc_id += 1
                doc_id_fields_seen[str(doc_id_field)] += 1
                if doc_id in bloom:
                    bloom_hits += 1
            if valid_records >= sample_size or shard_records >= target_for_shard:
                break
        remaining = max(0, sample_size - valid_records)
        if remaining == 0:
            break
    id_rate = records_with_top_level_id / valid_records if valid_records else 0.0
    doc_id_rate = records_with_doc_id / valid_records if valid_records else 0.0
    hit_rate = bloom_hits / records_with_doc_id if records_with_doc_id else 0.0
    approved = (
        valid_records >= sample_size
        and doc_id_rate >= min_id_rate
        and hit_rate >= min_hit_rate
    )
    return {
        "approved": approved,
        "dataset": dataset_id,
        "sample_shard": shard_paths[0],
        "sample_shards": shard_paths,
        "valid_records": valid_records,
        "invalid_json_records": invalid_json_records,
        "records_with_top_level_id": records_with_top_level_id,
        "records_with_doc_id": records_with_doc_id,
        "bloom_hits": bloom_hits,
        "top_level_id_rate": id_rate,
        "doc_id_rate": doc_id_rate,
        "doc_id_fields_seen": dict(sorted(doc_id_fields_seen.items())),
        "bloom_hit_rate": hit_rate,
        "sample_complete": valid_records >= sample_size,
    }


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
    )
    bloom = BloomIndex.load(args.bloom_file)
    pool_samples = select_probe_shards(
        read_manifest(args.pool_manifest),
        REQUIRED_PREFLIGHT_FAMILIES[:2],
        args.sample_shards_per_family,
    )
    mix_samples = select_probe_shards(
        read_manifest(args.mix_manifest),
        REQUIRED_PREFLIGHT_FAMILIES[2:],
        args.sample_shards_per_family,
    )
    results: dict[str, object] = {}
    for family, shard_paths in {**pool_samples, **mix_samples}.items():
        dataset_id = args.pool_dataset if family in pool_samples else args.mix_dataset
        local_root = (
            args.pool_local_root if family in pool_samples else args.mix_local_root
        )
        results[family] = probe_shards(
            dataset_id=dataset_id,
            shard_paths=shard_paths,
            bloom=bloom,
            sample_size=args.sample_size,
            min_id_rate=args.min_id_rate,
            min_hit_rate=args.min_hit_rate,
            download_dir=args.download_dir,
            local_root=local_root,
        )
    blocked = [
        family
        for family in REQUIRED_PREFLIGHT_FAMILIES
        if not results[family]["approved"]
    ]
    write_json(
        args.output,
        {
            "required_families": list(REQUIRED_PREFLIGHT_FAMILIES),
            "approved_families": [
                family
                for family in REQUIRED_PREFLIGHT_FAMILIES
                if results[family]["approved"]
            ],
            "blocked_families": blocked,
            "all_required_approved": not blocked,
            "results": results,
        },
    )
    if blocked:
        logger.warning("Blocked families: %s", ", ".join(blocked))
        return 1
    logger.info("Preflight approved all required families")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())