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

"""Modal worker for SOC-91 WebOrganizer enrichment."""

from __future__ import annotations

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

import modal

from config import (
    DEFAULT_BATCH_SIZE,
    DEFAULT_MAX_LENGTH,
    DEFAULT_TIMEOUT,
    FORMAT_NOURL_MODEL,
    FORMAT_URL_MODEL,
    GPU_CONFIGS,
    R2_BUCKET,
    R2_ENDPOINT_URL,
    R2_OUTPUT_PREFIX,
    R2_SECRET_NAME,
    TOPIC_NOURL_MODEL,
    TOPIC_URL_MODEL,
)

logger = logging.getLogger(__name__)

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

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


def _load_image():
    from image import image_with_models

    return image_with_models


image_with_models = _load_image()


def _resolve_dtype(gpu_key: str):
    import torch

    dtype_str = GPU_CONFIGS[gpu_key]["dtype"]
    if dtype_str == "bf16":
        return torch.bfloat16
    return torch.float16


def _batched(iterable, n):
    it = iter(iterable)
    while True:
        batch = list(islice(it, n))
        if not batch:
            break
        yield batch


@app.cls(
    image=image_with_models,
    gpu=GPU_CONFIGS["L4"]["gpu"],
    volumes={"/r2": r2_mount},
    timeout=DEFAULT_TIMEOUT,
    retries=2,
)
class EnrichWorker:
    @modal.enter()
    def load_models(self):
        from dolma.enrich import extract_url
        from dolma.format_model import FormatClassifier

        self.extract_url = extract_url
        dtype = _resolve_dtype("L4")
        common = dict(
            device="cuda",
            max_length=DEFAULT_MAX_LENGTH,
            torch_dtype=dtype,
            use_memory_efficient_attention=False,
            unpad_inputs=False,
            compile_model=False,
        )
        models = [
            TOPIC_URL_MODEL,
            TOPIC_NOURL_MODEL,
            FORMAT_URL_MODEL,
            FORMAT_NOURL_MODEL,
        ]
        logger.info("Loading %d classifiers...", len(models))
        loaded = []
        for m in models:
            loaded.append(FormatClassifier(model_name=m, model_name_nourl=m, **common))
        self.topic_url, self.topic_nourl, self.format_url, self.format_nourl = loaded
        logger.info("All classifiers loaded.")

    @modal.method()
    def process_shard(self, shard_path: str) -> dict:
        from dolma.local_io import iter_jsonlzst
        from dolma.sidecar import extract_sidecar_rows
        from dolma.sidecar_writer import (
            SidecarWriter,
            copy_to_final,
            is_sidecar_complete,
        )

        filename = Path(shard_path).name
        output_name = filename.replace(".jsonl.zst", ".parquet")
        final_path = Path(f"/r2/{R2_OUTPUT_PREFIX}/{output_name}")
        if is_sidecar_complete(final_path):
            return {"status": "skipped", "shard": shard_path}

        input_path = Path(f"/r2/{shard_path}")
        tmp_path = Path(f"/tmp/soc91_{output_name}")

        start = time.monotonic()
        processed = 0
        failed = 0
        source_family = "unknown"

        with SidecarWriter(tmp_path) as writer:
            for batch in _batched(iter_jsonlzst(input_path), DEFAULT_BATCH_SIZE):
                doc_ids = []
                texts = []
                urls = []
                for record in batch:
                    soc127 = record.get("_soc_127", {})
                    doc_id = soc127.get("doc_id", record.get("id"))
                    text = record.get("text", "")
                    if not isinstance(text, str):
                        text = ""
                    doc_ids.append(doc_id)
                    texts.append(text)
                    urls.append(self.extract_url(record))

                    if processed == 0 and failed == 0:
                        source_family = soc127.get("source_family", "unknown")

                no_urls = [None] * len(batch)

                try:
                    t_url_p, t_url_m = self.topic_url.predict_batch(urls, texts)
                    t_no_p, t_no_m = self.topic_nourl.predict_batch(no_urls, texts)
                    f_url_p, f_url_m = self.format_url.predict_batch(urls, texts)
                    f_no_p, f_no_m = self.format_nourl.predict_batch(no_urls, texts)
                except Exception:
                    logger.exception(
                        "Batch failed in shard %s at doc %d",
                        shard_path,
                        processed,
                    )
                    failed += len(batch)
                    continue

                rows = extract_sidecar_rows(
                    doc_ids,
                    t_url_p,
                    t_url_m,
                    t_no_p,
                    t_no_m,
                    f_url_p,
                    f_url_m,
                    f_no_p,
                    f_no_m,
                )
                writer.write_batch(rows)
                processed += len(batch)

        copy_to_final(tmp_path, final_path)
        tmp_path.unlink(missing_ok=True)
        elapsed = time.monotonic() - start

        final_writer = SidecarWriter(final_path)
        final_writer._total_rows = processed
        final_writer._start_time = start
        final_writer.write_stats(
            extra={
                "shard": shard_path,
                "source_family": source_family,
                "docs_failed": failed,
            }
        )
        final_writer.write_done()

        return {
            "status": "ok",
            "shard": shard_path,
            "source_family": source_family,
            "docs_classified": processed,
            "docs_failed": failed,
            "elapsed_seconds": round(elapsed, 2),
        }