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

"""Phase 1: GPU and truncation benchmark for SOC-91.

Runs a small sample through each GPU config and truncation length,
measuring throughput, VRAM, and label agreement.
"""

from __future__ import annotations

import json
import logging
import time
from itertools import islice
from pathlib import Path

import modal

from config import (
    DEFAULT_BATCH_SIZE,
    FORMAT_NOURL_MODEL,
    FORMAT_URL_MODEL,
    GPU_CONFIGS,
    R2_BUCKET,
    R2_ENDPOINT_URL,
    R2_SECRET_NAME,
    TOPIC_NOURL_MODEL,
    TOPIC_URL_MODEL,
)
from image import image_with_models

r2_mount = modal.CloudBucketMount(
    R2_BUCKET,
    bucket_endpoint_url=R2_ENDPOINT_URL,
    secret=modal.Secret.from_name(R2_SECRET_NAME),
)

logger = logging.getLogger(__name__)

app = modal.App("soc91-benchmark")

BENCHMARK_DOCS = 500
TRUNCATION_LENGTHS = [512, 1024, 2048]


def _run_benchmark_on_gpu(
    gpu_key: str,
    shard_path: str,
    max_length: int,
    max_docs: int,
) -> dict:
    import torch

    from dolma.enrich import extract_url
    from dolma.format_model import FormatClassifier
    from dolma.local_io import iter_jsonlzst

    gpu_cfg = GPU_CONFIGS[gpu_key]
    dtype = torch.bfloat16 if gpu_cfg["dtype"] == "bf16" else torch.float16
    common = dict(
        device="cuda",
        max_length=max_length,
        torch_dtype=dtype,
        use_memory_efficient_attention=False,
        unpad_inputs=False,
        compile_model=False,
    )

    load_start = time.monotonic()
    topic_url = FormatClassifier(
        model_name=TOPIC_URL_MODEL,
        model_name_nourl=TOPIC_URL_MODEL,
        **common,
    )
    topic_nourl = FormatClassifier(
        model_name=TOPIC_NOURL_MODEL,
        model_name_nourl=TOPIC_NOURL_MODEL,
        **common,
    )
    format_url = FormatClassifier(
        model_name=FORMAT_URL_MODEL,
        model_name_nourl=FORMAT_URL_MODEL,
        **common,
    )
    format_nourl = FormatClassifier(
        model_name=FORMAT_NOURL_MODEL,
        model_name_nourl=FORMAT_NOURL_MODEL,
        **common,
    )
    load_elapsed = time.monotonic() - load_start

    input_path = Path(f"/r2/{shard_path}")
    records = list(islice(iter_jsonlzst(input_path), max_docs))

    batch_size = 16 if gpu_cfg["memory"] <= 16 else DEFAULT_BATCH_SIZE

    infer_start = time.monotonic()
    processed = 0
    for i in range(0, len(records), batch_size):
        batch = records[i : i + batch_size]
        texts = [r.get("text", "") or "" for r in batch]
        urls = [extract_url(r) for r in batch]
        no_urls = [None] * len(batch)

        topic_url.predict_batch(urls, texts)
        topic_nourl.predict_batch(no_urls, texts)
        format_url.predict_batch(urls, texts)
        format_nourl.predict_batch(no_urls, texts)
        processed += len(batch)

    infer_elapsed = time.monotonic() - infer_start
    peak_vram = torch.cuda.max_memory_allocated() / (1024**3)

    return {
        "gpu": gpu_key,
        "max_length": max_length,
        "batch_size": batch_size,
        "docs_processed": processed,
        "load_seconds": round(load_elapsed, 2),
        "infer_seconds": round(infer_elapsed, 2),
        "docs_per_second": round(processed / max(infer_elapsed, 1e-6), 2),
        "peak_vram_gb": round(peak_vram, 2),
        "shard": shard_path,
    }


_benchmark_fns: dict[str, object] = {}

for gpu_key, gpu_cfg in GPU_CONFIGS.items():

    @app.function(
        image=image_with_models,
        gpu=gpu_cfg["gpu"],
        volumes={"/r2": r2_mount},
        timeout=1800,
        name=f"benchmark_{gpu_key.lower()}",
    )
    def _benchmark_fn(
        shard_path: str,
        max_length: int = 1024,
        max_docs: int = BENCHMARK_DOCS,
        _gpu_key: str = gpu_key,
    ) -> dict:
        return _run_benchmark_on_gpu(_gpu_key, shard_path, max_length, max_docs)

    _benchmark_fns[gpu_key] = _benchmark_fn


@app.local_entrypoint()
def main(shard: str = "", max_docs: int = BENCHMARK_DOCS):
    if not shard:
        print("Usage: modal run benchmark.py --shard <path>")
        print(
            "Example: modal run benchmark.py --shard "
            "soc127/phase1_pool_shared/common_crawl/000/shard.jsonl.zst"
        )
        return

    handles = []
    for gpu_key, fn in _benchmark_fns.items():
        for trunc in TRUNCATION_LENGTHS:
            print(f"Spawning: GPU={gpu_key}, max_length={trunc}")
            handle = fn.spawn(
                shard_path=shard,
                max_length=trunc,
                max_docs=max_docs,
                _gpu_key=gpu_key,
            )
            handles.append((gpu_key, trunc, handle))

    print(f"\n{len(handles)} jobs spawned, collecting results...\n")

    results = []
    for gpu_key, trunc, handle in handles:
        try:
            result = handle.get()
            results.append(result)
            print(
                f"  GPU={gpu_key}, len={trunc}: "
                f"{result['docs_per_second']} docs/sec, "
                f"VRAM={result['peak_vram_gb']}GB"
            )
        except Exception as e:
            err_msg = str(e)[:200]
            print(f"  GPU={gpu_key}, len={trunc}: FAILED: {err_msg}")
            results.append(
                {
                    "gpu": gpu_key,
                    "max_length": trunc,
                    "status": "failed",
                    "error": err_msg,
                }
            )

    print("\n--- Results ---")
    print(json.dumps(results, indent=2))