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

"""Phase 2a: Process non-pool mix shards through the Bloom filter."""

from __future__ import annotations

import json
import logging
import time
from pathlib import Path

import modal

from .config import (
    BLOOM_IMAGE_PATH,
    DOLMA_6T_MIX_DATASET_ID,
    R2_PREFIX,
    WORKER_CPU,
    WORKER_EPHEMERAL_DISK,
    WORKER_MEMORY,
    WORKER_RETRIES,
    WORKER_TIMEOUT,
    r2_done_path,
    r2_stats_path,
    shard_hash,
)
from .soc127_app import (
    app,
    copy_to_r2,
    hf_secret,
    image,
    r2_base_path,
    r2_done_exists,
    r2_mount,
    r2_secret,
    write_r2_json,
    write_r2_text,
)

logger = logging.getLogger("process_nonpool_raw")

PHASE = "phase2_nonpool_raw"


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

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

    @modal.method()
    def process_shard(self, shard_path: str) -> dict[str, object]:
        import traceback

        r2 = r2_base_path()

        if r2_done_exists(r2, PHASE, shard_path):
            logger.info("Skipping completed shard %s", shard_path)
            return {"shard_path": shard_path, "skipped": True}

        try:
            return self._process_shard_inner(shard_path, r2)
        except Exception as exc:
            error_msg = f"{type(exc).__name__}: {exc}"
            tb = traceback.format_exc()
            logger.error("FATAL shard %s: %s\n%s", shard_path, error_msg, tb)
            error_stats = {
                "shard_path": shard_path,
                "status": "error",
                "error": error_msg,
                "error_type": type(exc).__name__,
            }
            try:
                write_r2_json(r2, r2_stats_path(PHASE, shard_path), error_stats)
            except Exception:
                logger.error("Failed to write error stats for %s", shard_path)
            return error_stats

    def _process_shard_inner(self, shard_path: str, r2: Path) -> dict[str, object]:
        from huggingface_hub import hf_hub_download

        from dolma.dedup.materialize import (
            iter_shard_records,
            open_zstd_writer,
            resolve_record_doc_id,
        )
        from dolma.provenance import shard_folder_name, source_family

        preflight = json.loads(
            (r2 / R2_PREFIX / "preflight_report.json").read_text(encoding="utf-8")
        )
        if not preflight.get("all_required_approved"):
            raise ValueError("Preflight not approved")

        local_path = Path(
            hf_hub_download(
                repo_id=DOLMA_6T_MIX_DATASET_ID,
                filename=shard_path,
                repo_type="dataset",
                cache_dir="/tmp/hf_cache",
            )
        )

        family = source_family(shard_path)
        folder = shard_folder_name(shard_path)
        filename = f"{shard_hash(shard_path)}.jsonl.zst"
        output_relative = f"{R2_PREFIX}/{PHASE}/{family}/{filename}"
        tmp_output = Path(f"/tmp/output/{filename}")
        tmp_output.parent.mkdir(parents=True, exist_ok=True)

        stats = {
            "dataset": DOLMA_6T_MIX_DATASET_ID,
            "input_shard": shard_path,
            "source_family": family,
            "source_folder": folder,
            "records_seen": 0,
            "records_kept": 0,
            "records_invalid_json": 0,
            "records_removed_text": 0,
            "records_missing_doc_id": 0,
            "records_not_in_bloom": 0,
            "records_duplicate_ids": 0,
        }

        seen_ids: set[str] = set()

        with open_zstd_writer(tmp_output) as writer:
            for _, record in iter_shard_records(local_path):
                stats["records_seen"] += 1
                if record is None or not isinstance(record, dict):
                    stats["records_invalid_json"] += 1
                    continue
                doc_id, doc_id_field = resolve_record_doc_id(record, shard_path)
                if doc_id is None:
                    stats["records_missing_doc_id"] += 1
                    raise ValueError(f"Shard {shard_path} missing resolved doc id")
                text = str(record.get("text", "")).strip()
                if text == "[REMOVED]":
                    stats["records_removed_text"] += 1
                    continue
                if doc_id not in self.bloom:
                    stats["records_not_in_bloom"] += 1
                    continue
                if doc_id in seen_ids:
                    stats["records_duplicate_ids"] += 1
                    continue
                seen_ids.add(doc_id)
                enriched = dict(record)
                enriched["_soc_127"] = {
                    "doc_id": doc_id,
                    "doc_id_field": doc_id_field,
                    "input_shard": shard_path,
                    "phase": "mix_nonpool",
                    "source_family": family,
                    "source_folder": folder,
                }
                writer.write(json.dumps(enriched, sort_keys=True) + "\n")
                stats["records_kept"] += 1

        copy_to_r2(tmp_output, r2, output_relative)
        write_r2_json(r2, r2_stats_path(PHASE, shard_path), stats)
        write_r2_text(r2, r2_done_path(PHASE, shard_path), "ok\n")

        logger.info(
            "Shard %s: kept %d / %d (dupes: %d, removed: %d)",
            shard_path,
            stats["records_kept"],
            stats["records_seen"],
            stats["records_duplicate_ids"],
            stats["records_removed_text"],
        )
        return stats


@app.function(
    image=image,
    secrets=[hf_secret, r2_secret],
    volumes={"/r2": r2_mount},
    timeout=14400,
)
def run_nonpool_raw(limit: int = 0) -> None:
    from tqdm import tqdm

    from .soc127_app import read_manifest_from_r2

    r2 = r2_base_path()
    shards = read_manifest_from_r2(r2, "mix_nonpool")
    if limit > 0:
        shards = shards[:limit]

    total = len(shards)
    processor = NonpoolRawProcessor()
    completed = 0
    skipped = 0
    errors = 0
    error_shards: list[str] = []
    pbar = tqdm(total=total, desc="nonpool_raw", unit="shard")
    for result in processor.process_shard.map(shards):
        pbar.update(1)
        if result.get("status") == "error":
            errors += 1
            error_shards.append(str(result.get("shard_path", "?")))
        elif result.get("skipped"):
            skipped += 1
        else:
            completed += 1
    pbar.close()
    logger.info(
        "nonpool_raw done: %d completed, %d skipped, %d errors out of %d",
        completed,
        skipped,
        errors,
        total,
    )
    if error_shards:
        logger.warning("Failed shards (%d): %s", errors, error_shards[:20])


@app.local_entrypoint()
def main(limit: int = 0):
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
    )
    run_nonpool_raw.remote(limit=limit)