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

#!/usr/bin/env python3
"""Materialize non-pool 6T documents behind the SOC-127 preflight gate."""

from __future__ import annotations

import argparse
import json
import logging
from pathlib import Path

from dolma.constants import DOLMA_6T_MIX_DATASET_ID
from dolma.dedup.materialize import (
    NONPOOL_PREFLIGHT_FAMILIES,
    done_path_for,
    download_shard,
    ensure_preflight_approved,
    iter_shard_records,
    load_preflight,
    maybe_sleep_for_startup,
    open_zstd_writer,
    read_manifest,
    resolve_record_doc_id,
    resolve_task,
    shard_slice,
    stats_path_for,
    write_done_marker,
    write_json,
)
from dolma.provenance import BloomIndex, shard_folder_name, source_family

logger = logging.getLogger("materialize_mix_docs")

DEFAULT_OUTPUT = Path(
    "/storage/ice-shared/cs7634/staff/TDA/soc-90/unique_docs_work/mix_nonpool_raw"
)
DEFAULT_PREFLIGHT = 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="Materialize non-pool mix documents")
    parser.add_argument("--dataset", default=DOLMA_6T_MIX_DATASET_ID)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--bloom-file", type=Path, required=True)
    parser.add_argument("--preflight-file", type=Path, default=DEFAULT_PREFLIGHT)
    parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT)
    parser.add_argument("--download-dir", type=Path, default=None)
    parser.add_argument("--local-root", type=Path, default=None)
    parser.add_argument("--task-id", type=int, default=None)
    parser.add_argument("--task-count", type=int, default=None)
    parser.add_argument("--shards-per-task", type=int, default=5)
    parser.add_argument("--startup-delay", type=int, default=60)
    parser.add_argument(
        "--skip-existing", action=argparse.BooleanOptionalAction, default=True
    )
    parser.add_argument(
        "--drop-removed-text", action=argparse.BooleanOptionalAction, default=True
    )
    parser.add_argument("--verbose", action="store_true", default=False)
    return parser.parse_args(argv)


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",
    )
    preflight = load_preflight(args.preflight_file)
    ensure_preflight_approved(preflight, NONPOOL_PREFLIGHT_FAMILIES)
    task_id, task_count = resolve_task(args.task_id, args.task_count)
    output_path = args.output_dir / f"task_{task_id:05d}.jsonl.zst"
    done_path = done_path_for(output_path)
    if args.skip_existing and done_path.exists():
        logger.info("Skipping existing task output %s", output_path.name)
        return 0
    maybe_sleep_for_startup(task_id, args.startup_delay)
    shards = read_manifest(args.manifest)
    selected = shard_slice(shards, task_id, args.shards_per_task)
    stats = {
        "dataset": args.dataset,
        "task_id": task_id,
        "task_count": task_count,
        "assigned_shards": len(selected),
        "input_shards": selected,
        "source_families": sorted({source_family(shard) for shard in selected}),
        "records_seen": 0,
        "records_kept": 0,
        "records_invalid_json": 0,
        "records_removed_text": 0,
        "records_duplicate_ids": 0,
        "records_not_in_bloom": 0,
        "records_missing_doc_id": 0,
    }
    if not selected:
        write_json(stats_path_for(output_path), stats)
        write_done_marker(done_path)
        return 0
    bloom = BloomIndex.load(args.bloom_file)
    seen_ids: set[str] = set()
    try:
        with open_zstd_writer(output_path) as writer:
            for shard_path in selected:
                local_path = download_shard(
                    args.dataset, shard_path, args.download_dir, args.local_root
                )
                family = source_family(shard_path)
                folder = shard_folder_name(shard_path)
                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"Approved shard {shard_path} is missing a resolved doc id"
                        )
                    if (
                        args.drop_removed_text
                        and str(record.get("text", "")).strip() == "[REMOVED]"
                    ):
                        stats["records_removed_text"] += 1
                        continue
                    if doc_id not in 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
        write_json(stats_path_for(output_path), stats)
        write_done_marker(done_path)
        return 0
    except Exception:
        if output_path.exists():
            output_path.unlink()
        raise


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