social-data-attribution-w2/scripts/manifests/pre_split_pool_shards.py

#!/usr/bin/env python3
"""Pre-split pool shards into per-job manifest files for parallel SLURM execution.

Lists all .jsonl.zst paths from allenai/dolma3_pool, selects a fraction
with a seeded shuffle, then partitions evenly into N non-overlapping groups.
Each group is written to a text file (one shard path per line).
"""

from __future__ import annotations

import argparse
import os
import random
from pathlib import Path


def list_pool_shards(token: str | None = None) -> list[str]:
    from huggingface_hub import HfApi

    api = HfApi(token=token)
    paths = []
    for entry in api.list_repo_tree(
        "allenai/dolma3_pool", repo_type="dataset", recursive=True
    ):
        if hasattr(entry, "rfilename") and entry.rfilename.endswith(".jsonl.zst"):
            paths.append(entry.rfilename)
    return sorted(paths)


def main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(description="Pre-split pool shards for SLURM jobs")
    parser.add_argument(
        "--shard-fraction",
        type=float,
        default=0.05,
        help="Fraction of shards to select",
    )
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--num-jobs", type=int, default=5)
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("scripts/slurm"),
        help="Output directory",
    )
    args = parser.parse_args(argv)

    token = os.environ.get("HF_TOKEN")
    print("Listing shards from allenai/dolma3_pool ...")
    all_paths = list_pool_shards(token=token)
    print(f"Found {len(all_paths)} total shards")

    rng = random.Random(args.seed)
    shuffled = list(all_paths)
    rng.shuffle(shuffled)
    count = max(1, int(len(shuffled) * args.shard_fraction))
    selected = shuffled[:count]
    print(f"Selected {count} shards (fraction={args.shard_fraction}, seed={args.seed})")

    args.output_dir.mkdir(parents=True, exist_ok=True)
    partitions: list[list[str]] = [[] for _ in range(args.num_jobs)]
    for idx, path in enumerate(selected):
        partitions[idx % args.num_jobs].append(path)

    for job_idx, partition in enumerate(partitions):
        out_path = args.output_dir / f"pool_shards_job_{job_idx}.txt"
        out_path.write_text("\n".join(partition) + "\n")
        print(f"  Job {job_idx}: {len(partition)} shards -> {out_path}")

    print(f"Done. {args.num_jobs} manifest files written to {args.output_dir}")


if __name__ == "__main__":
    main()