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

"""Fail-closed preflight checks for SOC-127 on Modal."""

from __future__ import annotations

import logging
import time
from collections import Counter
from pathlib import Path

import modal

from .config import (
    BLOOM_IMAGE_PATH,
    DOLMA_6T_MIX_DATASET_ID,
    DOLMA_POOL_DATASET_ID,
    R2_PREFIX,
    WORKER_MEMORY,
    WORKER_TIMEOUT,
)
from .soc127_app import (
    app,
    hf_secret,
    image,
    image_no_bloom,
    r2_base_path,
    r2_mount,
    r2_secret,
    read_manifest_from_r2,
    write_r2_json,
)

logger = logging.getLogger("preflight")

SAMPLE_SIZE = 10_000
MIN_ID_RATE = 0.99
MIN_HIT_RATE = 0.01
SHARDS_PER_FAMILY = 8


def _select_probe_shards(
    shards: list[str], families: tuple[str, ...], per_family: int
) -> dict[str, list[str]]:
    from dolma.provenance import shard_folder_name, source_family

    selected: dict[str, list[str]] = {f: [] for f in families}
    seen_folders: dict[str, set[str]] = {f: set() for f in families}
    stack_python: str | None = None
    stack_regular: 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 is None:
                stack_python = shard_path
            elif folder != "stack_edu-Python" and stack_regular is None:
                stack_regular = shard_path
            continue
        if len(selected[family]) >= per_family or folder in seen_folders[family]:
            continue
        selected[family].append(shard_path)
        seen_folders[family].add(folder)
    if "stack_edu" in selected:
        if stack_python is not None:
            selected["stack_edu"].append(stack_python)
        if stack_regular is not None:
            selected["stack_edu"].append(stack_regular)
    missing = [f for f in families if not selected[f]]
    if missing:
        raise ValueError(f"Missing probe shard(s) for: {', '.join(missing)}")
    return selected


@app.cls(
    image=image,
    secrets=[hf_secret, r2_secret],
    volumes={"/r2": r2_mount},
    timeout=WORKER_TIMEOUT,
    memory=WORKER_MEMORY,
)
class PreflightProber:
    @modal.enter()
    def load_bloom(self):
        from dolma.provenance import BloomIndex

        t0 = time.monotonic()
        self.bloom = BloomIndex.load(Path(BLOOM_IMAGE_PATH))
        self.bloom_load_seconds = time.monotonic() - t0
        logger.info("Bloom loaded in %.1fs", self.bloom_load_seconds)

    @modal.method()
    def probe_shard(
        self, dataset_id: str, shard_paths: list[str], sample_size: int
    ) -> dict[str, object]:
        from huggingface_hub import hf_hub_download
        from dolma.dedup.materialize import iter_shard_records, resolve_record_doc_id

        valid = invalid_json = with_doc_id = hits = 0
        fields: Counter[str] = Counter()
        remaining = sample_size
        for idx, shard_path in enumerate(shard_paths):
            local = Path(
                hf_hub_download(
                    repo_id=dataset_id,
                    filename=shard_path,
                    repo_type="dataset",
                    cache_dir="/tmp/hf_cache",
                )
            )
            target = max(1, -(-remaining // (len(shard_paths) - idx)))
            shard_n = 0
            for _, record in iter_shard_records(local):
                if record is None or not isinstance(record, dict):
                    invalid_json += 1
                    continue
                valid += 1
                shard_n += 1
                doc_id, field = resolve_record_doc_id(record, shard_path)
                if doc_id is not None:
                    with_doc_id += 1
                    fields[str(field)] += 1
                    if doc_id in self.bloom:
                        hits += 1
                if valid >= sample_size or shard_n >= target:
                    break
            remaining = max(0, sample_size - valid)
            if remaining == 0:
                break

        doc_id_rate = with_doc_id / valid if valid else 0.0
        hit_rate = hits / with_doc_id if with_doc_id else 0.0
        return {
            "approved": (
                valid >= sample_size
                and doc_id_rate >= MIN_ID_RATE
                and hit_rate >= MIN_HIT_RATE
            ),
            "dataset": dataset_id,
            "sample_shards": shard_paths,
            "valid_records": valid,
            "invalid_json_records": invalid_json,
            "records_with_doc_id": with_doc_id,
            "bloom_hits": hits,
            "doc_id_rate": doc_id_rate,
            "bloom_hit_rate": hit_rate,
            "doc_id_fields_seen": dict(sorted(fields.items())),
            "bloom_load_seconds": self.bloom_load_seconds,
        }


@app.function(
    image=image_no_bloom,
    secrets=[r2_secret],
    volumes={"/r2": r2_mount},
    timeout=3600,
    memory=4096,
)
def run_preflight() -> dict[str, object]:
    from dolma.dedup.materialize import REQUIRED_PREFLIGHT_FAMILIES

    r2 = r2_base_path()
    pool_probes = _select_probe_shards(
        read_manifest_from_r2(r2, "pool_shared"),
        REQUIRED_PREFLIGHT_FAMILIES[:2],
        SHARDS_PER_FAMILY,
    )
    mix_probes = _select_probe_shards(
        read_manifest_from_r2(r2, "mix_nonpool"),
        REQUIRED_PREFLIGHT_FAMILIES[2:],
        SHARDS_PER_FAMILY,
    )
    prober = PreflightProber()
    results: dict[str, object] = {}
    for family, shard_paths in {**pool_probes, **mix_probes}.items():
        dataset_id = (
            DOLMA_POOL_DATASET_ID if family in pool_probes else DOLMA_6T_MIX_DATASET_ID
        )
        results[family] = prober.probe_shard.remote(
            dataset_id, shard_paths, SAMPLE_SIZE
        )
    blocked = [f for f in REQUIRED_PREFLIGHT_FAMILIES if not results[f]["approved"]]
    report = {
        "required_families": list(REQUIRED_PREFLIGHT_FAMILIES),
        "approved_families": [
            f for f in REQUIRED_PREFLIGHT_FAMILIES if results[f]["approved"]
        ],
        "blocked_families": blocked,
        "all_required_approved": not blocked,
        "results": results,
    }
    write_r2_json(r2, f"{R2_PREFIX}/preflight_report.json", report)
    if blocked:
        logger.warning("Blocked families: %s", ", ".join(blocked))
    else:
        logger.info("All required families approved")
    return report


@app.local_entrypoint()
def main():
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
    )
    report = run_preflight.remote()
    for family, result in report["results"].items():
        status = "APPROVED" if result["approved"] else "BLOCKED"
        logger.info(
            "%s: %s (doc_id=%.3f, bloom=%.3f)",
            family,
            status,
            result["doc_id_rate"],
            result["bloom_hit_rate"],
        )