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

"""SOC-91 Modal app: orchestrates shard enrichment via spawn_map."""

from __future__ import annotations

import json
import logging
from pathlib import Path

from config import R2_INPUT_PREFIXES, R2_OUTPUT_PREFIX, R2_STATS_PREFIX
from worker import EnrichWorker, app, r2_mount

logger = logging.getLogger(__name__)


@app.function(volumes={"/r2": r2_mount}, timeout=600)
def list_shards() -> dict[str, list[str]]:
    output_dir = Path(f"/r2/{R2_OUTPUT_PREFIX}")

    all_shards: list[str] = []
    for prefix in R2_INPUT_PREFIXES:
        input_dir = Path(f"/r2/{prefix}")
        if not input_dir.exists():
            continue
        for path in input_dir.rglob("*.jsonl.zst"):
            rel = str(path.relative_to(Path("/r2")))
            all_shards.append(rel)

    pending: list[str] = []
    completed: list[str] = []
    for shard in sorted(all_shards):
        filename = Path(shard).name
        done_name = filename.replace(".jsonl.zst", ".parquet.done")
        done_path = output_dir / done_name
        if done_path.exists():
            completed.append(shard)
        else:
            pending.append(shard)

    logger.info(
        "Shards: %d total, %d completed, %d pending",
        len(all_shards),
        len(completed),
        len(pending),
    )
    return {
        "total": len(all_shards),
        "completed": len(completed),
        "pending_count": len(pending),
        "pending": pending,
    }


@app.function(volumes={"/r2": r2_mount}, timeout=300)
def write_aggregate_stats(results: list[dict]) -> str:
    stats_dir = Path(f"/r2/{R2_STATS_PREFIX}")
    stats_dir.mkdir(parents=True, exist_ok=True)

    ok = [r for r in results if r.get("status") == "ok"]
    skipped = [r for r in results if r.get("status") == "skipped"]
    failed = [r for r in results if r.get("status") not in ("ok", "skipped")]

    aggregate = {
        "total_shards": len(results),
        "ok": len(ok),
        "skipped": len(skipped),
        "failed": len(failed),
        "total_docs_classified": sum(r.get("docs_classified", 0) for r in ok),
        "total_docs_failed": sum(r.get("docs_failed", 0) for r in ok),
        "failed_shards": [r.get("shard") for r in failed],
    }

    out_path = stats_dir / "aggregate_stats.json"
    with out_path.open("w", encoding="utf-8") as f:
        json.dump(aggregate, f, indent=2)
        f.write("\n")
    return str(out_path)


@app.local_entrypoint()
def main(
    mode: str = "status",
    limit: int = 0,
    detach: bool = False,
):
    if mode == "status":
        manifest = list_shards.remote()
        print(
            f"Total: {manifest['total']}, "
            f"Completed: {manifest['completed']}, "
            f"Pending: {manifest['pending_count']}"
        )

    elif mode == "pilot":
        manifest = list_shards.remote()
        pending = manifest["pending"]
        n = limit if limit > 0 else min(20, len(pending))
        subset = pending[:n]
        print(f"Pilot: processing {len(subset)} shards")
        worker = EnrichWorker()
        results = list(worker.process_shard.map(subset))
        for r in results:
            print(r)
        write_aggregate_stats.remote(results)

    elif mode == "production":
        manifest = list_shards.remote()
        pending = manifest["pending"]
        if limit > 0:
            pending = pending[:limit]
        print(f"Production: processing {len(pending)} shards")
        worker = EnrichWorker()
        if detach:
            worker.process_shard.spawn_map(pending)
            print("Detached. Jobs submitted.")
        else:
            results = list(worker.process_shard.map(pending))
            write_aggregate_stats.remote(results)
            ok = sum(1 for r in results if r.get("status") == "ok")
            print(f"Done: {ok}/{len(results)} shards processed")

    else:
        print(f"Unknown mode: {mode}. Use: status, pilot, production")