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

"""Phase 1: Process pool-shared 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_POOL_DATASET_ID,
    HF_SUBFOLDER_MAX,
    R2_PREFIX,
    WORKER_CPU,
    WORKER_EPHEMERAL_DISK,
    WORKER_MEMORY,
    WORKER_RETRIES,
    WORKER_TIMEOUT,
    r2_done_path,
    r2_stats_path,
    subfolder_index,
)
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_pool_shared")

PHASE = "phase1_pool_shared"


@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 PoolShardProcessor:
    @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, shard_index: int) -> 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, shard_index, 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,
                "shard_index": shard_index,
                "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, shard_index: int, 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

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

        family = source_family(shard_path)
        folder = shard_folder_name(shard_path)
        part = subfolder_index(shard_index, HF_SUBFOLDER_MAX)
        filename = shard_path.replace("/", "__")
        output_relative = f"{R2_PREFIX}/{PHASE}/{family}/{part}/{filename}"
        tmp_output = Path(f"/tmp/output/{filename}")
        tmp_output.parent.mkdir(parents=True, exist_ok=True)

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

        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
                    continue
                if doc_id not in self.bloom:
                    stats["records_not_in_bloom"] += 1
                    continue
                enriched = dict(record)
                enriched["_soc_127"] = {
                    "doc_id": doc_id,
                    "doc_id_field": doc_id_field,
                    "input_shard": shard_path,
                    "phase": "pool",
                    "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",
            shard_path,
            stats["records_kept"],
            stats["records_seen"],
        )
        return stats


MAP_BATCH_SIZE = 20000


@app.function(
    image=image,
    secrets=[hf_secret, r2_secret],
    volumes={"/r2": r2_mount},
    timeout=86400,
)
def run_pool_shared(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, "pool_shared")
    if limit > 0:
        shards = shards[:limit]

    processor = PoolShardProcessor()
    total = len(shards)
    completed = 0
    skipped = 0
    errors = 0
    error_shards: list[str] = []
    pbar = tqdm(total=total, desc="pool_shared", unit="shard")
    for batch_start in range(0, total, MAP_BATCH_SIZE):
        batch_end = min(batch_start + MAP_BATCH_SIZE, total)
        batch_shards = shards[batch_start:batch_end]
        batch_indices = list(range(batch_start, batch_end))
        logger.info("Batch %d-%d of %d", batch_start, batch_end, total)
        for result in processor.process_shard.map(batch_shards, batch_indices):
            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(
        "pool_shared 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_pool_shared.remote(limit=limit)