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

"""Count how many times each document appears across the 6T mix shards.

Phase K1: Extract doc_ids from mix shards into bucketed files on R2.
Phase K2: Aggregate per-bucket frequencies.

The result is a set of bucket files where each line is a JSON object
with {"doc_id": "...", "mix_frequency": N} for every doc_id that
appeared in the mix.
"""

from __future__ import annotations

import hashlib
import json
import logging
import traceback
from collections import defaultdict
from pathlib import Path


from .config import (
    DOLMA_6T_MIX_DATASET_ID,
    R2_PREFIX,
    WORKER_CPU,
    WORKER_MEMORY,
    WORKER_RETRIES,
    WORKER_TIMEOUT,
)
from .soc127_app import (
    app,
    copy_to_r2,
    hf_secret,
    image_no_bloom,
    r2_base_path,
    r2_mount,
    r2_secret,
    write_r2_json,
    write_r2_text,
)

logger = logging.getLogger("count_mix_frequency")

FREQ_PREFIX = f"{R2_PREFIX}/mix_frequency"
EXTRACT_PHASE = f"{FREQ_PREFIX}/extract"
AGG_PHASE = f"{FREQ_PREFIX}/freq"
FREQ_BUCKET_COUNT = 64
EXTRACT_BATCH_SIZE = 64


def _bucket_for_id(doc_id: str) -> int:
    h = hashlib.blake2b(doc_id.encode("utf-8"), digest_size=4)
    return int.from_bytes(h.digest(), "big") % FREQ_BUCKET_COUNT


@app.function(
    image=image_no_bloom,
    secrets=[hf_secret, r2_secret],
    volumes={"/r2": r2_mount},
    cpu=WORKER_CPU,
    memory=WORKER_MEMORY,
    timeout=WORKER_TIMEOUT,
    retries=WORKER_RETRIES,
)
def extract_batch(shard_paths: list[str], batch_index: int) -> dict:
    r2 = r2_base_path()
    batch_id = f"batch_{batch_index:04d}"
    done_key = f"{EXTRACT_PHASE}/done/{batch_id}.done"

    if (r2 / done_key).exists():
        logger.info("Skipping completed batch %s", batch_id)
        return {"batch_index": batch_index, "skipped": True}

    try:
        return _extract_batch_inner(shard_paths, batch_index, batch_id, done_key, r2)
    except Exception as exc:
        error_msg = f"{type(exc).__name__}: {exc}"
        tb = traceback.format_exc()
        logger.error("FATAL batch %s: %s\n%s", batch_id, error_msg, tb)
        return {
            "batch_index": batch_index,
            "batch_id": batch_id,
            "status": "error",
            "error": error_msg,
        }


def _extract_batch_inner(
    shard_paths: list[str],
    batch_index: int,
    batch_id: str,
    done_key: str,
    r2: Path,
) -> dict:
    from huggingface_hub import hf_hub_download

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

    buckets: dict[int, list[str]] = defaultdict(list)
    stats = {
        "batch_index": batch_index,
        "batch_id": batch_id,
        "shards_processed": 0,
        "shards_failed": 0,
        "total_records": 0,
        "total_ids_extracted": 0,
        "records_invalid": 0,
        "records_missing_id": 0,
    }

    for shard_path in shard_paths:
        try:
            local_path = Path(
                hf_hub_download(
                    repo_id=DOLMA_6T_MIX_DATASET_ID,
                    filename=shard_path,
                    repo_type="dataset",
                    cache_dir="/tmp/hf_cache",
                )
            )
            for _, record in iter_shard_records(local_path):
                stats["total_records"] += 1
                if record is None or not isinstance(record, dict):
                    stats["records_invalid"] += 1
                    continue
                doc_id, _ = resolve_record_doc_id(record, shard_path)
                if doc_id is None:
                    stats["records_missing_id"] += 1
                    continue
                bucket_id = _bucket_for_id(doc_id)
                buckets[bucket_id].append(doc_id)
                stats["total_ids_extracted"] += 1
            stats["shards_processed"] += 1
        except Exception as exc:
            logger.warning("Shard %s failed in batch %s: %s", shard_path, batch_id, exc)
            stats["shards_failed"] += 1

    for bucket_id, doc_ids in buckets.items():
        if not doc_ids:
            continue
        bucket_dir = f"bucket_{bucket_id:02d}"
        filename = f"{batch_id}.ids.zst"
        tmp_path = Path(f"/tmp/freq_extract/{bucket_dir}/{filename}")
        tmp_path.parent.mkdir(parents=True, exist_ok=True)

        with open_zstd_writer(tmp_path) as writer:
            for did in doc_ids:
                writer.write(did + "\n")

        relative = f"{EXTRACT_PHASE}/{bucket_dir}/{filename}"
        copy_to_r2(tmp_path, r2, relative)

    stats_key = f"{EXTRACT_PHASE}/stats/{batch_id}.stats.json"
    write_r2_json(r2, stats_key, stats)
    write_r2_text(r2, done_key, "ok\n")

    logger.info(
        "Batch %s: %d shards, %d ids extracted, %d failed",
        batch_id,
        stats["shards_processed"],
        stats["total_ids_extracted"],
        stats["shards_failed"],
    )
    return stats


@app.function(
    image=image_no_bloom,
    secrets=[r2_secret],
    volumes={"/r2": r2_mount},
    cpu=WORKER_CPU,
    memory=WORKER_MEMORY,
    timeout=WORKER_TIMEOUT,
    retries=WORKER_RETRIES,
)
def aggregate_bucket(bucket_id: int) -> dict:
    r2 = r2_base_path()
    bucket_dir = f"bucket_{bucket_id:02d}"
    done_key = f"{AGG_PHASE}/done/{bucket_dir}.done"

    if (r2 / done_key).exists():
        logger.info("Skipping completed bucket %s", bucket_dir)
        return {"bucket_id": bucket_id, "skipped": True}

    try:
        return _aggregate_bucket_inner(bucket_id, bucket_dir, done_key, r2)
    except Exception as exc:
        error_msg = f"{type(exc).__name__}: {exc}"
        tb = traceback.format_exc()
        logger.error("FATAL bucket %d: %s\n%s", bucket_id, error_msg, tb)
        return {
            "bucket_id": bucket_id,
            "status": "error",
            "error": error_msg,
        }


def _aggregate_bucket_inner(
    bucket_id: int, bucket_dir: str, done_key: str, r2: Path
) -> dict:
    from dolma.dedup.materialize import open_zstd_writer

    extract_dir = r2 / EXTRACT_PHASE / bucket_dir
    if not extract_dir.exists():
        logger.warning("No extraction files for bucket %d", bucket_id)
        write_r2_text(r2, done_key, "empty\n")
        return {"bucket_id": bucket_id, "unique_ids": 0, "total_occurrences": 0}

    freq: dict[str, int] = defaultdict(int)
    files_read = 0

    for ids_file in sorted(extract_dir.iterdir()):
        if not ids_file.name.endswith(".ids.zst"):
            continue
        import zstandard

        dctx = zstandard.ZstdDecompressor()
        with open(ids_file, "rb") as fh:
            with dctx.stream_reader(fh) as reader:
                import io

                for line in io.TextIOWrapper(reader, encoding="utf-8"):
                    doc_id = line.strip()
                    if doc_id:
                        freq[doc_id] += 1
        files_read += 1

    tmp_path = Path(f"/tmp/freq_agg/{bucket_dir}.jsonl.zst")
    tmp_path.parent.mkdir(parents=True, exist_ok=True)

    with open_zstd_writer(tmp_path) as writer:
        for doc_id, count in sorted(freq.items()):
            writer.write(json.dumps({"doc_id": doc_id, "mix_frequency": count}) + "\n")

    relative = f"{AGG_PHASE}/{bucket_dir}.jsonl.zst"
    copy_to_r2(tmp_path, r2, relative)

    stats = {
        "bucket_id": bucket_id,
        "files_read": files_read,
        "unique_ids": len(freq),
        "total_occurrences": sum(freq.values()),
        "max_frequency": max(freq.values()) if freq else 0,
    }
    stats_key = f"{AGG_PHASE}/stats/{bucket_dir}.stats.json"
    write_r2_json(r2, stats_key, stats)
    write_r2_text(r2, done_key, "ok\n")

    logger.info(
        "Bucket %d: %d unique ids, %d total occurrences, max freq %d",
        bucket_id,
        stats["unique_ids"],
        stats["total_occurrences"],
        stats["max_frequency"],
    )
    return stats


@app.function(
    image=image_no_bloom,
    secrets=[hf_secret, r2_secret],
    volumes={"/r2": r2_mount},
    timeout=86400,
)
def run_frequency_count(skip_extract: bool = False) -> None:
    from tqdm import tqdm

    from .soc127_app import read_manifest_from_r2

    r2 = r2_base_path()

    if not skip_extract:
        mix_shards = read_manifest_from_r2(r2, "mix")
        total_shards = len(mix_shards)
        batches = []
        for i in range(0, total_shards, EXTRACT_BATCH_SIZE):
            batches.append(mix_shards[i : i + EXTRACT_BATCH_SIZE])
        batch_indices = list(range(len(batches)))

        logger.info(
            "Extracting doc_ids from %d shards in %d batches",
            total_shards,
            len(batches),
        )
        errors = 0
        pbar = tqdm(total=len(batches), desc="extract", unit="batch")
        for result in extract_batch.map(batches, batch_indices):
            pbar.update(1)
            if result.get("status") == "error":
                errors += 1
        pbar.close()
        logger.info(
            "Extraction done. %d errors out of %d batches", errors, len(batches)
        )

    logger.info("Aggregating %d frequency buckets", FREQ_BUCKET_COUNT)
    bucket_ids = list(range(FREQ_BUCKET_COUNT))
    agg_errors = 0
    total_unique = 0
    total_occurrences = 0
    pbar = tqdm(total=FREQ_BUCKET_COUNT, desc="aggregate", unit="bucket")
    for result in aggregate_bucket.map(bucket_ids):
        pbar.update(1)
        if result.get("status") == "error":
            agg_errors += 1
        else:
            total_unique += result.get("unique_ids", 0)
            total_occurrences += result.get("total_occurrences", 0)
    pbar.close()

    summary = {
        "total_unique_ids": total_unique,
        "total_occurrences": total_occurrences,
        "avg_frequency": total_occurrences / total_unique if total_unique else 0,
        "bucket_count": FREQ_BUCKET_COUNT,
        "agg_errors": agg_errors,
    }
    write_r2_json(r2, f"{FREQ_PREFIX}/frequency_summary.json", summary)

    logger.info(
        "Frequency count done: %d unique ids, %d total occurrences, avg %.2f",
        total_unique,
        total_occurrences,
        summary["avg_frequency"],
    )


@app.local_entrypoint()
def main(skip_extract: bool = False):
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
    )
    run_frequency_count.remote(skip_extract=skip_extract)