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

"""Draw a uniform random document sample for preconditioner building.

Fan-out: workers read SOC-95 manifest parquet chunks from the cached
Modal volume (R2 fallback), apply a token-count floor, compute a
deterministic hash priority per doc, and return the top candidates.
The local entrypoint merges all chunk results and takes the global
top-N to produce a 100K-doc manifest.

Usage:
    uv run modal run scripts/modal/draw_preconditioner_sample.py
    uv run modal run scripts/modal/draw_preconditioner_sample.py \
        --target-docs 100000 --min-tokens 512 --seed 42
"""

from __future__ import annotations

import io
import json
import logging
import os
from pathlib import Path

import modal

logger = logging.getLogger("draw_preconditioner_sample")

R2_BUCKET_NAME = "soc127-dedup"
R2_MANIFEST_PREFIX = "soc95-manifest/data"

_local_path = Path(__file__).resolve()
_can_resolve_repo = len(_local_path.parents) > 2
_SRC_ROOT = str(_local_path.parents[2] / "src") if _can_resolve_repo else "/root/src"

app = modal.App("soc151-draw-preconditioner-sample")

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install("boto3", "pyarrow", "pandas")
    .env({"PYTHONPATH": "/root/src"})
    .add_local_dir(_SRC_ROOT, remote_path="/root/src", copy=True)
)

r2_secret = modal.Secret.from_name("r2-credentials")
manifest_volume = modal.Volume.from_name(
    "soc134-manifest-cache", create_if_missing=True
)

MANIFEST_MOUNT = "/manifests"

_s3_client_cache = None


def get_s3_client():
    global _s3_client_cache
    if _s3_client_cache is None:
        import boto3

        _s3_client_cache = boto3.client(
            "s3",
            endpoint_url=os.environ["AWS_ENDPOINT_URL"],
            aws_access_key_id=os.environ["AWS_ACCESS_KEY_ID"],
            aws_secret_access_key=os.environ["AWS_SECRET_ACCESS_KEY"],
        )
    return _s3_client_cache


def r2_s3_list_keys(prefix: str, suffix: str = "") -> list[str]:
    s3 = get_s3_client()
    keys: list[str] = []
    paginator = s3.get_paginator("list_objects_v2")
    for page in paginator.paginate(Bucket=R2_BUCKET_NAME, Prefix=prefix):
        for obj in page.get("Contents", []):
            k = obj["Key"]
            if suffix and not k.endswith(suffix):
                continue
            keys.append(k)
    return sorted(keys)


@app.function(
    image=image,
    volumes={MANIFEST_MOUNT: manifest_volume},
    secrets=[r2_secret],
    cpu=2,
    memory=16384,
    timeout=3600,
    retries=2,
    max_containers=128,
)
def sample_chunk(chunk_json: str) -> str:
    import hashlib

    import pyarrow.parquet as pq

    from dolma.token_filter import apply_token_filter

    chunk = json.loads(chunk_json)
    keys = chunk["keys"]
    chunk_id = chunk["chunk_id"]
    per_chunk_cap = chunk["per_chunk_cap"]
    min_tokens = chunk["min_tokens"]
    seed = chunk["seed"]

    manifest_volume.reload()
    columns = [
        "doc_id",
        "token_count",
        "shard_path",
    ]

    import pandas as pd

    frames = []
    for key in keys:
        local_path = Path(MANIFEST_MOUNT) / key
        if local_path.exists() and local_path.stat().st_size > 0:
            body = local_path.read_bytes()
        else:
            s3 = get_s3_client()
            resp = s3.get_object(Bucket=R2_BUCKET_NAME, Key=key)
            body = resp["Body"].read()
        table = pq.read_table(io.BytesIO(body), columns=columns)
        df = table.to_pandas()
        if len(df) > 0:
            frames.append(df)

    if not frames:
        return json.dumps(
            {
                "chunk_id": chunk_id,
                "total_docs_read": 0,
                "docs_below_min_tokens": 0,
                "candidates_returned": 0,
                "candidates": [],
            }
        )

    all_df = pd.concat(frames, ignore_index=True)
    del frames
    total_read = len(all_df)

    all_df, dropped = apply_token_filter(all_df, min_tokens)

    if len(all_df) == 0:
        return json.dumps(
            {
                "chunk_id": chunk_id,
                "total_docs_read": total_read,
                "docs_below_min_tokens": dropped,
                "candidates_returned": 0,
                "candidates": [],
            }
        )

    seed_str = str(seed)
    all_df["_priority"] = all_df["doc_id"].apply(
        lambda d: int.from_bytes(
            hashlib.blake2b(f"{d}:{seed_str}".encode(), digest_size=8).digest(),
            "big",
        )
    )

    top = all_df.nlargest(per_chunk_cap, "_priority")

    candidates = []
    for _, row in top.iterrows():
        candidates.append(
            [
                int(row["_priority"]),
                row["doc_id"],
                int(row["token_count"]),
                row["shard_path"],
            ]
        )

    return json.dumps(
        {
            "chunk_id": chunk_id,
            "total_docs_read": total_read,
            "docs_below_min_tokens": dropped,
            "candidates_returned": len(candidates),
            "candidates": candidates,
        }
    )


@app.function(image=image, secrets=[r2_secret], timeout=300)
def list_keys_remote(r2_prefix: str) -> list[str]:
    return r2_s3_list_keys(r2_prefix, suffix=".parquet")


def _derive_source_family(shard_path: str) -> str:
    if "phase1_pool_shared" in shard_path:
        return "pool_shared"
    if "phase2_nonpool_final" in shard_path:
        return "nonpool_final"
    return "unknown"


@app.local_entrypoint()
def main(
    target_docs: int = 100_000,
    min_tokens: int = 512,
    seed: int = 42,
    chunk_count: int = 256,
    per_chunk_cap: int = 100_000,
):
    import numpy as np
    import pandas as pd

    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
    )

    logger.info(
        "Drawing %d uniform random docs (min_tokens=%d, seed=%d)",
        target_docs,
        min_tokens,
        seed,
    )

    logger.info("Listing R2 manifest keys ...")
    keys = list_keys_remote.remote(r2_prefix=R2_MANIFEST_PREFIX)
    logger.info("Found %d manifest parquets", len(keys))

    chunk_size = (len(keys) + chunk_count - 1) // chunk_count
    chunks = []
    for i in range(0, len(keys), chunk_size):
        chunks.append(
            json.dumps(
                {
                    "chunk_id": i // chunk_size,
                    "keys": keys[i : i + chunk_size],
                    "per_chunk_cap": per_chunk_cap,
                    "min_tokens": min_tokens,
                    "seed": seed,
                }
            )
        )
    logger.info("Dispatching %d chunks (%d files each)", len(chunks), chunk_size)

    all_candidates: list[list] = []
    total_read = 0
    total_dropped = 0
    completed = 0

    for result_json in sample_chunk.map(chunks):
        result = json.loads(result_json)
        all_candidates.extend(result["candidates"])
        total_read += result["total_docs_read"]
        total_dropped += result["docs_below_min_tokens"]
        completed += 1
        if completed % 50 == 0:
            logger.info(
                "Collected %d / %d chunks (%d candidates so far)",
                completed,
                len(chunks),
                len(all_candidates),
            )

    logger.info(
        "Total: %s docs read, %s below min_tokens, %s candidates",
        f"{total_read:,}",
        f"{total_dropped:,}",
        f"{len(all_candidates):,}",
    )

    all_candidates.sort(key=lambda x: x[0], reverse=True)
    selected = all_candidates[:target_docs]
    del all_candidates

    rows = []
    for entry in selected:
        rows.append(
            {
                "doc_id": entry[1],
                "token_count": entry[2],
                "shard_path": entry[3],
            }
        )

    sample_df = pd.DataFrame(rows)
    total_docs = len(sample_df)
    token_counts = sample_df["token_count"]
    total_tokens = int(token_counts.sum())
    unique_shards = sample_df["shard_path"].nunique()

    contract = {
        "SAMPLE_TYPE": "uniform_random",
        "SAMPLE_PURPOSE": "preconditioner_building",
        "SAMPLE_TARGET_DOCS": target_docs,
        "SAMPLE_MIN_TOKENS": min_tokens,
        "SAMPLE_REALIZED_DOC_COUNT": total_docs,
        "SAMPLE_REALIZED_TOKEN_TOTAL": total_tokens,
        "SAMPLE_TOKEN_STATS": {
            "min": int(token_counts.min()),
            "max": int(token_counts.max()),
            "median": int(token_counts.median()),
            "mean": round(float(token_counts.mean()), 1),
        },
        "SAMPLE_UNIQUE_SHARDS": unique_shards,
        "SAMPLE_TOTAL_DOCS_READ": total_read,
        "SAMPLE_TOTAL_DOCS_BELOW_MIN_TOKENS": total_dropped,
        "SAMPLE_SAMPLING_SEED": seed,
    }

    source_counts: dict[str, int] = {}
    for sp in sample_df["shard_path"]:
        family = _derive_source_family(sp)
        source_counts[family] = source_counts.get(family, 0) + 1

    percentiles = [5, 25, 50, 75, 95]
    pct_values = np.percentile(token_counts.values, percentiles)
    token_percentiles = {f"p{p}": int(v) for p, v in zip(percentiles, pct_values)}

    validation = {
        "source_family_distribution": dict(sorted(source_counts.items())),
        "token_count_percentiles": token_percentiles,
    }

    out_dir = Path("data/samples/preconditioner_100k")
    out_dir.mkdir(parents=True, exist_ok=True)
    sample_df.to_parquet(out_dir / "working_sample_manifest.parquet", index=False)
    (out_dir / "sample_contract.json").write_text(json.dumps(contract, indent=2) + "\n")
    (out_dir / "validation_stats.json").write_text(
        json.dumps(validation, indent=2) + "\n"
    )

    logger.info(
        "Sample drawn: %s docs, %s tokens, %d unique shards",
        f"{total_docs:,}",
        f"{total_tokens:,}",
        unique_shards,
    )
    logger.info("Source distribution: %s", source_counts)
    logger.info("Token percentiles: %s", token_percentiles)
    logger.info("Output: %s", out_dir)