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

#!/usr/bin/env python3
"""Extract doc IDs from 6T shards into shard-local gzip files."""

from __future__ import annotations

import argparse
import gzip
import io
import json
import logging
import os
import random
import time
from pathlib import Path

import zstandard as zstd
from huggingface_hub import hf_hub_download

from dolma.provenance import extract_doc_id

DEFAULT_DATASET = "allenai/dolma3_mix-6T-1025-7B"
logger = logging.getLogger("extract_ids")


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Extract doc IDs from shard manifest")
    parser.add_argument("--dataset", default=DEFAULT_DATASET)
    parser.add_argument("--manifest", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument(
        "--download-dir",
        type=Path,
        default=None,
        help="Local dir for shard downloads (default: $TMPDIR or /tmp)",
    )
    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=50)
    parser.add_argument("--local-root", type=Path)
    parser.add_argument(
        "--skip-existing", action=argparse.BooleanOptionalAction, default=True
    )
    parser.add_argument(
        "--startup-delay",
        type=int,
        default=0,
        help="Max random startup delay in seconds to stagger HF API calls",
    )
    parser.add_argument("--verbose", action="store_true", default=False)
    return parser.parse_args(argv)


def read_manifest(path: Path) -> list[str]:
    return [
        line.strip()
        for line in path.read_text(encoding="utf-8").splitlines()
        if line.strip()
    ]


def shard_slice(
    shards: list[str], task_id: int, task_count: int, shards_per_task: int
) -> list[str]:
    if task_id < 0 or task_count <= 0:
        raise ValueError("task_id must be >= 0 and task_count must be > 0")
    if shards_per_task <= 0:
        raise ValueError("shards_per_task must be > 0")
    start = task_id * shards_per_task
    stop = min(len(shards), start + shards_per_task)
    if start >= len(shards):
        return []
    return shards[start:stop]


def resolve_task(args: argparse.Namespace) -> tuple[int, int]:
    env_task_id = os.environ.get("SLURM_ARRAY_TASK_ID")
    env_task_count = os.environ.get("SLURM_ARRAY_TASK_COUNT")
    task_id = args.task_id if args.task_id is not None else int(env_task_id or "0")
    if args.task_count is not None:
        return task_id, args.task_count
    if env_task_count:
        return task_id, int(env_task_count)
    return task_id, 1


def shard_output_path(output_dir: Path, shard_path: str) -> Path:
    safe_name = shard_path.replace("/", "__")
    return output_dir / f"{safe_name}.ids.gz"


def download_shard(
    dataset: str,
    shard_path: str,
    download_dir: Path | None = None,
) -> Path:
    cache_dir = str(download_dir) if download_dir else None
    return Path(
        hf_hub_download(
            repo_id=dataset,
            filename=shard_path,
            repo_type="dataset",
            cache_dir=cache_dir,
        )
    )


def extract_shard_ids(
    dataset: str,
    shard_path: str,
    output_path: Path,
    download_dir: Path | None = None,
    local_root: Path | None = None,
) -> tuple[int, int]:
    local_path = local_root / shard_path if local_root else None
    if local_path and local_path.exists():
        source = local_path
    else:
        source = download_shard(dataset, shard_path, download_dir)
    seen_lines = 0
    extracted = 0
    output_path.parent.mkdir(parents=True, exist_ok=True)
    with gzip.open(output_path, "wt", encoding="utf-8", compresslevel=6) as out_f:
        with source.open("rb") as in_f:
            dctx = zstd.ZstdDecompressor()
            with dctx.stream_reader(in_f) as reader:
                with io.TextIOWrapper(reader, encoding="utf-8") as text_reader:
                    for line in text_reader:
                        seen_lines += 1
                        doc_id = extract_doc_id(line)
                        if doc_id is None:
                            continue
                        out_f.write(f"{doc_id}\n")
                        extracted += 1
    return seen_lines, extracted


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",
    )
    task_id, task_count = resolve_task(args)
    shards = read_manifest(args.manifest)
    selected = shard_slice(shards, task_id, task_count, args.shards_per_task)
    if not selected:
        logger.info("No shards assigned for task %d", task_id)
        return 0
    pending = []
    skipped = 0
    if args.skip_existing:
        for shard in selected:
            out_path = shard_output_path(args.output_dir, shard)
            if out_path.exists() and out_path.stat().st_size > 0:
                skipped += 1
            else:
                pending.append(shard)
    else:
        pending = list(selected)
    if not pending:
        logger.info("Task %d: all %d shards already exist, skipping", task_id, skipped)
        stats = {
            "dataset": args.dataset,
            "task_id": task_id,
            "task_count": task_count,
            "assigned_shards": len(selected),
            "processed_shards": 0,
            "skipped_existing_shards": skipped,
            "lines_seen": 0,
            "ids_extracted": 0,
        }
        stats_path = args.output_dir / f"extract_ids_task_{task_id:04d}.json"
        stats_path.write_text(
            json.dumps(stats, indent=2, sort_keys=True) + "\n", encoding="utf-8"
        )
        return 0
    if args.startup_delay > 0:
        random.seed(task_id)
        delay = random.uniform(0, args.startup_delay)
        logger.info(
            "Task %d startup delay: %.1fs (%d pending shards)",
            task_id,
            delay,
            len(pending),
        )
        time.sleep(delay)
    download_dir = args.download_dir or Path(os.environ.get("TMPDIR", "/tmp"))
    total_lines = 0
    total_ids = 0
    processed = 0
    failed_shards = 0
    for shard in pending:
        out_path = shard_output_path(args.output_dir, shard)
        try:
            lines, ids = extract_shard_ids(
                dataset=args.dataset,
                shard_path=shard,
                output_path=out_path,
                download_dir=download_dir,
                local_root=args.local_root,
            )
            total_lines += lines
            total_ids += ids
            processed += 1
        except Exception:
            logger.exception("Failed shard %s, continuing", shard)
            failed_shards += 1
            if out_path.exists():
                out_path.unlink()
    stats = {
        "dataset": args.dataset,
        "task_id": task_id,
        "task_count": task_count,
        "assigned_shards": len(selected),
        "processed_shards": processed,
        "skipped_existing_shards": skipped,
        "failed_shards": failed_shards,
        "lines_seen": total_lines,
        "ids_extracted": total_ids,
    }
    stats_path = args.output_dir / f"extract_ids_task_{task_id:04d}.json"
    stats_path.write_text(
        json.dumps(stats, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    logger.info(
        "Task %d processed %d shards and extracted %d IDs",
        task_id,
        processed,
        total_ids,
    )
    logger.info("Task stats: %s", stats_path)
    return 0


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