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

"""Draw stratified working samples from the full SOC-95 manifest on R2.

The pipeline is preemption-resilient: each chunk worker writes its results
to a Modal volume. If the orchestrator gets preempted and restarts, chunks
detect existing results and return immediately. The merge step reads from
the volume and completes in seconds.

Usage:
    uv run modal run scripts/modal/draw_samples.py \
        --configs '500,1000,5000,10000'

    # With exclusion manifests (staggered sampling)
    uv run modal run scripts/modal/draw_samples.py \
        --configs '10000' --sample-suffix batch_1 \
        --exclude-manifest /samples/sample_10000_docs/working_sample_manifest.parquet

    # Write to local filesystem instead of Modal volume
    uv run modal run scripts/modal/draw_samples.py \
        --configs '500' --local-output
"""

from __future__ import annotations

import hashlib
import io
import json
import logging
import os
import shutil
from pathlib import Path
from typing import TypedDict

import modal

logger = logging.getLogger("draw_samples")

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("soc149-draw-samples")

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
)
samples_volume = modal.Volume.from_name("soc149-drawn-samples", create_if_missing=True)

MANIFEST_MOUNT = "/manifests"
SAMPLES_MOUNT = "/samples"
CHUNK_RESULTS_DIR = "_chunk_results"

_s3_client_cache = None


class ExclusionInfo(TypedDict):
    exclude_paths_list: list[str]
    exclude_set_path: str
    excluded_doc_count: int


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)


def _parse_exclude_manifest_keys(exclude_manifest_keys: str) -> list[str]:
    parsed = [k.strip() for k in exclude_manifest_keys.split(",") if k.strip()]
    return list(dict.fromkeys(parsed))


def _exclude_manifest_digest(exclude_paths_list: list[str]) -> str:
    payload = "\n".join(sorted(exclude_paths_list))
    return hashlib.blake2b(payload.encode(), digest_size=16).hexdigest()


def _compute_run_id(
    configs: str,
    seed: int,
    chunk_count: int,
    min_token_count: int,
    max_token_count: int,
    exclude_manifest_keys: str,
    sample_suffix: str,
) -> str:
    params = (
        f"{configs}:{seed}:{chunk_count}:{min_token_count}:"
        f"{max_token_count}:{exclude_manifest_keys}:{sample_suffix}"
    )
    digest = hashlib.blake2b(params.encode(), digest_size=8).hexdigest()
    return f"draw_{digest}"


def _prepare_exclusion_set(exclude_manifest_keys: str) -> ExclusionInfo:
    import pandas as pd
    import pyarrow.parquet as pq

    exclude_paths_list = _parse_exclude_manifest_keys(exclude_manifest_keys)
    if not exclude_paths_list:
        return {
            "exclude_paths_list": [],
            "exclude_set_path": "",
            "excluded_doc_count": 0,
        }

    all_exclude_ids: set[str] = set()
    s3 = None
    for path_or_key in exclude_paths_list:
        if path_or_key.startswith("/"):
            local_p = Path(path_or_key)
            if not local_p.exists():
                raise FileNotFoundError(
                    f"Exclusion manifest not found on volume: {path_or_key}"
                )
            logger.info("Loading exclusion manifest from volume: %s", path_or_key)
            tbl = pq.read_table(str(local_p), columns=["doc_id"])
        else:
            if s3 is None:
                s3 = get_s3_client()
            logger.info("Loading exclusion manifest from R2: %s", path_or_key)
            resp = s3.get_object(Bucket=R2_BUCKET_NAME, Key=path_or_key)
            body = resp["Body"].read()
            tbl = pq.read_table(io.BytesIO(body), columns=["doc_id"])
        ids = set(tbl.column("doc_id").to_pylist())
        logger.info(
            "  %d doc_ids from %s (total: %d)",
            len(ids),
            path_or_key,
            len(all_exclude_ids | ids),
        )
        all_exclude_ids |= ids

    excluded_doc_count = len(all_exclude_ids)
    logger.info("Total exclusion set: %d doc_ids", excluded_doc_count)

    # Exclusion manifests are materialized once under /manifests/exclusions so
    # workers filter against the same doc_id set across retries and resumed runs.
    exclude_dir = Path(MANIFEST_MOUNT) / "exclusions"
    exclude_dir.mkdir(parents=True, exist_ok=True)
    digest = _exclude_manifest_digest(exclude_paths_list)
    exclude_parquet_path = exclude_dir / f"{digest}.parquet"

    if not exclude_parquet_path.exists() or exclude_parquet_path.stat().st_size == 0:
        tmp_path = exclude_dir / f"{digest}.{os.getpid()}.{excluded_doc_count}.tmp"
        pd.DataFrame({"doc_id": list(all_exclude_ids)}).to_parquet(
            str(tmp_path), index=False
        )
        tmp_path.rename(exclude_parquet_path)
        manifest_volume.commit()
        logger.info("Wrote exclusion set to %s", exclude_parquet_path)
    else:
        logger.info("Using existing exclusion set at %s", exclude_parquet_path)

    return {
        "exclude_paths_list": exclude_paths_list,
        "exclude_set_path": str(exclude_parquet_path),
        "excluded_doc_count": excluded_doc_count,
    }


@app.function(
    image=image,
    volumes={MANIFEST_MOUNT: manifest_volume, SAMPLES_MOUNT: samples_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 pandas as pd
    import pyarrow.parquet as pq

    chunk = json.loads(chunk_json)
    keys = chunk["keys"]
    chunk_id = chunk["chunk_id"]
    max_dpb = chunk["max_docs_per_bin"]
    seed = chunk["seed"]
    min_token_count = chunk.get("min_token_count")
    max_token_count = chunk.get("max_token_count")
    exclude_set_path = chunk.get("exclude_set_path")
    run_id = chunk.get("run_id", "unknown")

    output_path = (
        Path(SAMPLES_MOUNT)
        / CHUNK_RESULTS_DIR
        / run_id
        / f"chunk_{chunk_id:04d}.parquet"
    )

    # Workers checkpoint parquet chunks under /samples/_chunk_results/<run_id>.
    # Delete that run directory before retrying a stale or failed Modal run.
    samples_volume.reload()
    if output_path.exists() and output_path.stat().st_size > 0:
        row_count = pq.read_metadata(str(output_path)).num_rows
        logger.info("Chunk %d: cached (%d rows)", chunk_id, row_count)
        return json.dumps({"chunk_id": chunk_id, "status": "cached", "rows": row_count})

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

    exclude_ids: set[str] = set()
    if exclude_set_path:
        ep = Path(exclude_set_path)
        if not ep.exists():
            raise FileNotFoundError(
                f"Exclusion set path is not available in the container: {exclude_set_path}"
            )
        import pyarrow.parquet as _pq

        _tbl = _pq.read_table(str(ep), columns=["doc_id"])
        exclude_ids = set(_tbl.column("doc_id").to_pylist())

    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()
        df = df.dropna(subset=["weborganizer_topic"])
        if min_token_count is not None:
            df = df[df["token_count"] >= min_token_count]
        if max_token_count is not None:
            df = df[df["token_count"] <= max_token_count]
        if exclude_ids:
            df = df[~df["doc_id"].isin(exclude_ids)]
        if len(df) > 0:
            frames.append(df)

    if not frames:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        pd.DataFrame(
            columns=["priority", "doc_id", "token_count", "shard_path", "bin_key"]
        ).to_parquet(str(output_path), index=False)
        samples_volume.commit()
        return json.dumps({"chunk_id": chunk_id, "status": "ok", "rows": 0})

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

    def strip_label(s):
        if isinstance(s, str) and s.startswith("__label__"):
            return s[9:]
        return s

    all_df["_topic"] = all_df["weborganizer_topic"].map(strip_label)
    all_df["_format"] = all_df["weborganizer_format"].map(strip_label)
    all_df["bin_key"] = all_df["_topic"] + "|" + all_df["_format"]

    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",
        )
    )

    result_rows = []
    for bin_key, group in all_df.groupby("bin_key"):
        top = group.nlargest(max_dpb, "priority")
        for _, row in top.iterrows():
            result_rows.append(
                {
                    "priority": int(row["priority"]),
                    "doc_id": row["doc_id"],
                    "token_count": int(row["token_count"]),
                    "shard_path": row["shard_path"],
                    "bin_key": bin_key,
                }
            )

    result_df = pd.DataFrame(result_rows)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    result_df.to_parquet(str(output_path), index=False)
    samples_volume.commit()

    logger.info("Chunk %d: wrote %d rows", chunk_id, len(result_df))
    return json.dumps({"chunk_id": chunk_id, "status": "ok", "rows": len(result_df)})


@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")


@app.function(
    image=image,
    volumes={MANIFEST_MOUNT: manifest_volume, SAMPLES_MOUNT: samples_volume},
    secrets=[r2_secret],
    timeout=900,
    cpu=1,
    memory=8192,
)
def prepare_exclusion_set_remote(exclude_manifest_keys: str) -> ExclusionInfo:
    return _prepare_exclusion_set(exclude_manifest_keys)


def _merge_chunk_results(
    run_id: str,
    max_dpb: int,
    expected_chunk_count: int,
) -> dict[str, list]:
    import pandas as pd

    chunk_dir = Path(SAMPLES_MOUNT) / CHUNK_RESULTS_DIR / run_id
    samples_volume.reload()
    chunk_files = sorted(chunk_dir.glob("chunk_*.parquet"))
    logger.info("Reading %d chunk result files from volume", len(chunk_files))
    if len(chunk_files) < expected_chunk_count:
        raise RuntimeError(
            "Missing chunk result files for "
            f"{run_id}: expected {expected_chunk_count}, found {len(chunk_files)}. "
            f"Delete {chunk_dir} before retrying a stale or failed Modal run."
        )

    merged: dict[str, list] = {}
    for cf in chunk_files:
        df = pd.read_parquet(cf)
        if len(df) == 0:
            continue
        for bin_key, group in df.groupby("bin_key"):
            entries = []
            for _, row in group.iterrows():
                entries.append(
                    [
                        int(row["priority"]),
                        row["doc_id"],
                        int(row["token_count"]),
                        row["shard_path"],
                    ]
                )
            merged.setdefault(str(bin_key), []).extend(entries)

    logger.info("Merging %d bins, selecting top-%d per bin", len(merged), max_dpb)
    for bin_key in merged:
        merged[bin_key].sort(key=lambda x: x[0], reverse=True)
        merged[bin_key] = merged[bin_key][:max_dpb]

    return merged


def _build_sample_outputs(
    merged: dict[str, list],
    config_list: list[int],
    seed: int,
    effective_min: int | None,
    effective_max: int | None,
    exclude_paths_list: list[str],
    excluded_doc_count: int,
    sample_suffix: str,
    output_root: Path,
    commit_volume: bool = False,
) -> list[dict]:
    import pandas as pd

    from dolma.constants import FORMATS, TOPICS

    sample_results = []
    for dpb in config_list:
        sample_name = (
            f"sample_{dpb}_docs_{sample_suffix}"
            if sample_suffix
            else f"sample_{dpb}_docs"
        )
        logger.info("--- Building %s (%d docs/bin) ---", sample_name, dpb)

        rows = []
        bin_summaries = []
        for topic in TOPICS:
            for fmt in FORMATS:
                bin_key = f"{topic}|{fmt}"
                t_idx = TOPICS.index(topic)
                f_idx = FORMATS.index(fmt)
                bin_id = t_idx * len(FORMATS) + f_idx + 1

                candidates = merged.get(bin_key, [])
                selected = candidates[:dpb]

                for entry in selected:
                    rows.append(
                        {
                            "doc_id": entry[1],
                            "token_count": entry[2],
                            "shard_path": entry[3],
                            "bin_id": bin_id,
                            "bin_topic": topic,
                            "bin_format": fmt,
                        }
                    )

                bin_summaries.append(
                    {
                        "bin_id": bin_id,
                        "topic": topic,
                        "format": fmt,
                        "requested_docs_per_bin": dpb,
                        "realized_docs": len(selected),
                        "realized_tokens": sum(e[2] for e in selected),
                        "underfilled": len(selected) < dpb,
                    }
                )

        sample_df = pd.DataFrame(rows)
        summary_df = pd.DataFrame(bin_summaries)

        total_docs = len(sample_df)
        total_tokens = int(sample_df["token_count"].sum()) if total_docs > 0 else 0
        underfilled = int(summary_df["underfilled"].sum())
        covered = int((summary_df["realized_docs"] > 0).sum())

        contract = {
            "WORKING_SAMPLE_TOKEN_FLOOR_PER_BIN": 0,
            "WORKING_SAMPLE_DOCS_PER_BIN": dpb,
            "WORKING_SAMPLE_GLOBAL_TOKEN_BUDGET": None,
            "WORKING_SAMPLE_MIN_TOKEN_COUNT": effective_min,
            "WORKING_SAMPLE_MAX_TOKEN_COUNT": effective_max,
            "WORKING_SAMPLE_REALIZED_TOKEN_TOTAL": total_tokens,
            "WORKING_SAMPLE_REALIZED_DOC_COUNT": total_docs,
            "WORKING_SAMPLE_UNDERFILLED_BIN_COUNT": underfilled,
            "WORKING_SAMPLE_COVERED_BIN_COUNT": covered,
            "WORKING_SAMPLE_TOTAL_BIN_COUNT": len(TOPICS) * len(FORMATS),
            "WORKING_SAMPLE_SAMPLING_SEED": seed,
        }
        if exclude_paths_list:
            contract["WORKING_SAMPLE_EXCLUDE_MANIFEST_PATHS"] = exclude_paths_list
            contract["WORKING_SAMPLE_EXCLUDED_DOC_COUNT"] = excluded_doc_count

        sample_dir = output_root / sample_name
        sample_dir.mkdir(parents=True, exist_ok=True)
        sample_df.to_parquet(
            sample_dir / "working_sample_manifest.parquet", index=False
        )
        summary_df.to_csv(sample_dir / "bin_summary.csv", index=False)
        (sample_dir / "sample_contract.json").write_text(
            json.dumps(contract, indent=2) + "\n"
        )
        if commit_volume:
            samples_volume.commit()

        logger.info(
            "%s: %s docs, %s tokens, %d/%d bins, %d underfilled",
            sample_name,
            f"{total_docs:,}",
            f"{total_tokens:,}",
            covered,
            len(TOPICS) * len(FORMATS),
            underfilled,
        )

        sample_results.append(
            {
                "sample_name": sample_name,
                "docs": total_docs,
                "tokens": total_tokens,
                "underfilled": underfilled,
                "covered": covered,
            }
        )

    return sample_results


@app.function(
    image=image,
    volumes={MANIFEST_MOUNT: manifest_volume, SAMPLES_MOUNT: samples_volume},
    secrets=[r2_secret],
    timeout=7200,
    cpu=1,
    memory=16384,
)
def run_draw_pipeline(
    configs: str,
    seed: int,
    chunk_count: int,
    min_token_count: int,
    max_token_count: int,
    exclude_manifest_keys: str = "",
    sample_suffix: str = "",
) -> str:
    effective_min = min_token_count if min_token_count > 0 else None
    effective_max = max_token_count if max_token_count > 0 else None

    config_list = [int(v.strip()) for v in configs.split(",")]
    max_dpb = max(config_list)

    run_id = _compute_run_id(
        configs,
        seed,
        chunk_count,
        min_token_count,
        max_token_count,
        exclude_manifest_keys,
        sample_suffix,
    )

    logger.info(
        "Drawing %d samples (max %d docs/bin): %s [run_id=%s]",
        len(config_list),
        max_dpb,
        config_list,
        run_id,
    )
    logger.info("Token filter: min=%s, max=%s", effective_min, effective_max)

    exclusion_info = _prepare_exclusion_set(exclude_manifest_keys)
    exclude_set_path = exclusion_info["exclude_set_path"]
    exclude_paths_list = exclusion_info["exclude_paths_list"]
    excluded_doc_count = exclusion_info["excluded_doc_count"]

    keys = r2_s3_list_keys(R2_MANIFEST_PREFIX, suffix=".parquet")
    logger.info("Found %d manifest parquets", len(keys))

    if not keys:
        logger.error("No manifest parquets found at prefix %s", R2_MANIFEST_PREFIX)
        return json.dumps({"status": "error", "error": "No manifest parquets found"})

    chunk_size = (len(keys) + chunk_count - 1) // chunk_count
    chunks = []
    for i in range(0, len(keys), chunk_size):
        chunk_payload = {
            "chunk_id": i // chunk_size,
            "keys": keys[i : i + chunk_size],
            "max_docs_per_bin": max_dpb,
            "seed": seed,
            "min_token_count": effective_min,
            "max_token_count": effective_max,
            "run_id": run_id,
        }
        if exclude_set_path:
            chunk_payload["exclude_set_path"] = exclude_set_path
        chunks.append(json.dumps(chunk_payload))
    logger.info("Dispatching %d chunks (%d files each)", len(chunks), chunk_size)

    completed = 0
    cached = 0
    total_rows = 0
    for status_json in sample_chunk.map(chunks):
        status = json.loads(status_json)
        completed += 1
        total_rows += status.get("rows", 0)
        if status.get("status") == "cached":
            cached += 1
        if completed % 50 == 0:
            logger.info(
                "Progress: %d/%d chunks (%d cached, %d rows so far)",
                completed,
                len(chunks),
                cached,
                total_rows,
            )

    logger.info(
        "All %d chunks done (%d cached, %d new). Total rows: %d",
        completed,
        cached,
        completed - cached,
        total_rows,
    )

    merged = _merge_chunk_results(run_id, max_dpb, expected_chunk_count=len(chunks))

    sample_results = _build_sample_outputs(
        merged=merged,
        config_list=config_list,
        seed=seed,
        effective_min=effective_min,
        effective_max=effective_max,
        exclude_paths_list=exclude_paths_list,
        excluded_doc_count=excluded_doc_count,
        sample_suffix=sample_suffix,
        output_root=Path(SAMPLES_MOUNT),
        commit_volume=True,
    )

    chunk_dir = Path(SAMPLES_MOUNT) / CHUNK_RESULTS_DIR / run_id
    shutil.rmtree(str(chunk_dir), ignore_errors=True)
    samples_volume.commit()
    logger.info("Cleaned up chunk results for run %s", run_id)

    logger.info("All samples written to Modal volume: soc149-drawn-samples")
    return json.dumps({"status": "done", "samples": sample_results})


@app.local_entrypoint()
def main(
    configs: str = "500,1000,5000,10000",
    seed: int = 42,
    chunk_count: int = 256,
    min_token_count: int = 512,
    max_token_count: int = 0,
    exclude_manifest: str = "",
    sample_suffix: str = "",
    local_output: bool = False,
):
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s %(message)s",
    )

    if local_output:
        effective_min = min_token_count if min_token_count > 0 else None
        effective_max = max_token_count if max_token_count > 0 else None
        config_list = [int(v.strip()) for v in configs.split(",")]
        max_dpb = max(config_list)

        exclude_paths_list = _parse_exclude_manifest_keys(exclude_manifest)
        exclude_set_path = ""
        excluded_doc_count = 0

        run_id = _compute_run_id(
            configs,
            seed,
            chunk_count,
            min_token_count,
            max_token_count,
            exclude_manifest,
            sample_suffix,
        )

        logger.info(
            "Drawing %d samples locally (max %d docs/bin) [run_id=%s]",
            len(config_list),
            max_dpb,
            run_id,
        )
        if exclude_paths_list:
            logger.info("Exclusion manifests: %s", exclude_paths_list)

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

        if not keys:
            logger.error("No manifest parquets found")
            return

        if exclude_paths_list:
            exclusion_info = prepare_exclusion_set_remote.remote(exclude_manifest)
            exclude_paths_list = exclusion_info["exclude_paths_list"]
            exclude_set_path = exclusion_info["exclude_set_path"]
            excluded_doc_count = exclusion_info["excluded_doc_count"]

        chunk_size = (len(keys) + chunk_count - 1) // chunk_count
        chunks = []
        for i in range(0, len(keys), chunk_size):
            chunk_payload = {
                "chunk_id": i // chunk_size,
                "keys": keys[i : i + chunk_size],
                "max_docs_per_bin": max_dpb,
                "seed": seed,
                "min_token_count": effective_min,
                "max_token_count": effective_max,
                "run_id": run_id,
            }
            if exclude_set_path:
                chunk_payload["exclude_set_path"] = exclude_set_path
            chunks.append(json.dumps(chunk_payload))
        logger.info("Dispatching %d chunks", len(chunks))

        completed = 0
        cached = 0
        for status_json in sample_chunk.map(chunks):
            status = json.loads(status_json)
            completed += 1
            if status.get("status") == "cached":
                cached += 1
            if completed % 50 == 0:
                logger.info(
                    "Progress: %d/%d chunks (%d cached)", completed, len(chunks), cached
                )

        logger.info("All %d chunks done (%d cached)", completed, cached)

        merged = _merge_chunk_results(run_id, max_dpb, expected_chunk_count=len(chunks))

        sample_results = _build_sample_outputs(
            merged=merged,
            config_list=config_list,
            seed=seed,
            effective_min=effective_min,
            effective_max=effective_max,
            exclude_paths_list=exclude_paths_list,
            excluded_doc_count=excluded_doc_count,
            sample_suffix=sample_suffix,
            output_root=Path("data/samples"),
            commit_volume=False,
        )

        chunk_dir = Path(SAMPLES_MOUNT) / CHUNK_RESULTS_DIR / run_id
        shutil.rmtree(str(chunk_dir), ignore_errors=True)
        samples_volume.commit()

        for s in sample_results:
            logger.info(
                "%s: %s docs, %s tokens",
                s["sample_name"],
                f"{s['docs']:,}",
                f"{s['tokens']:,}",
            )
        logger.info("All samples written to data/samples/")
    else:
        logger.info("Dispatching full draw pipeline to Modal...")
        result_json = run_draw_pipeline.remote(
            configs,
            seed,
            chunk_count,
            min_token_count,
            max_token_count,
            exclude_manifest_keys=exclude_manifest,
            sample_suffix=sample_suffix,
        )
        result = json.loads(result_json)
        for s in result["samples"]:
            print(
                f"{s['sample_name']}: {s['docs']:,} docs, "
                f"{s['tokens']:,} tokens, "
                f"{s['covered']}/{576} bins, "
                f"{s['underfilled']} underfilled"
            )
        print("Samples written to Modal volume: soc149-drawn-samples")