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

"""Modal runner for SOC-101 sidecar EDA."""

from __future__ import annotations

import json
import os
import shutil
from datetime import datetime, timezone
from pathlib import Path

import modal

from config import (
    R2_BUCKET,
    R2_ENDPOINT_URL,
    R2_OUTPUT_PREFIX,
    R2_SECRET_NAME,
    R2_STATS_PREFIX,
)

_local_path = Path(__file__).resolve()
if len(_local_path.parents) > 2:
    _REPO_ROOT = _local_path.parents[2]
    _SRC_ROOT = str(_REPO_ROOT / "src")
    _CONFIG_PY = str(_REPO_ROOT / "scripts" / "modal" / "config.py")
    _eda_image = (
        modal.Image.debian_slim(python_version="3.12")
        .pip_install("boto3>=1.37.0", "pyarrow>=18.0.0")
        .env({"PYTHONPATH": "/root/src:/root"})
        .add_local_file(_CONFIG_PY, remote_path="/root/config.py", copy=True)
        .add_local_dir(_SRC_ROOT, remote_path="/root/src", copy=True)
    )
else:
    _eda_image = modal.Image.debian_slim(python_version="3.12")

app = modal.App("soc101-sidecar-eda")

r2_mount = modal.CloudBucketMount(
    R2_BUCKET,
    bucket_endpoint_url=R2_ENDPOINT_URL,
    secret=modal.Secret.from_name(R2_SECRET_NAME),
)
eda_volume = modal.Volume.from_name("soc101-sidecar-eda-cache", create_if_missing=True)


def _run_cli(argv: list[str]) -> None:
    from dolma.sidecar_eda.cli import main

    rc = main(argv)
    if rc != 0:
        raise RuntimeError(f"sidecar EDA CLI failed with exit code {rc}")


def _read_json(path: Path) -> dict[str, object]:
    return json.loads(path.read_text(encoding="utf-8"))


def _ensure_r2_env() -> None:
    env_aliases = {
        "R2_ACCESS_KEY_ID": ("R2_ACCESS_KEY_ID", "AWS_ACCESS_KEY_ID", "access_key_id"),
        "R2_SECRET_ACCESS_KEY": (
            "R2_SECRET_ACCESS_KEY",
            "AWS_SECRET_ACCESS_KEY",
            "secret_access_key",
        ),
    }
    for target, aliases in env_aliases.items():
        if target in os.environ and os.environ[target]:
            continue
        for alias in aliases:
            value = os.environ.get(alias)
            if value:
                os.environ[target] = value
                break
        if target not in os.environ:
            raise KeyError(target)


@app.function(
    image=_eda_image,
    volumes={"/r2": r2_mount, "/cache": eda_volume},
    secrets=[modal.Secret.from_name(R2_SECRET_NAME)],
    timeout=3600,
    cpu=2.0,
)
def build_manifest(run_id: str, max_shards: int | None = None) -> dict[str, object]:
    from dolma.sidecar_eda.modal import (
        build_manifest_argv,
        modal_manifest_path,
        modal_run_root,
    )

    run_root = modal_run_root(R2_STATS_PREFIX, run_id)
    manifest_path = modal_manifest_path(run_root)
    run_root.mkdir(parents=True, exist_ok=True)
    _ensure_r2_env()
    _run_cli(
        build_manifest_argv(
            output_dir=run_root,
            manifest_path=manifest_path,
            bucket=R2_BUCKET,
            prefix=R2_OUTPUT_PREFIX,
            endpoint_url=R2_ENDPOINT_URL,
            max_shards=max_shards,
        )
    )
    return {
        "run_id": run_id,
        "run_root": str(run_root),
        "manifest_path": str(manifest_path),
        "benchmark": _read_json(run_root / "benchmark.json"),
    }


@app.function(
    image=_eda_image,
    volumes={"/r2": r2_mount, "/cache": eda_volume},
    secrets=[modal.Secret.from_name(R2_SECRET_NAME)],
    timeout=7200,
    cpu=4.0,
)
def run_chunk(run_id: str, chunk_index: int, chunk_count: int) -> dict[str, object]:
    from dolma.sidecar_eda.modal import (
        build_worker_argv,
        modal_cache_run_root,
        modal_manifest_path,
        modal_run_root,
        modal_worker_output_dir,
    )

    run_root = modal_run_root(R2_STATS_PREFIX, run_id)
    manifest_path = modal_manifest_path(run_root)
    cache_run_root = modal_cache_run_root(run_id)
    worker_output_dir = modal_worker_output_dir(cache_run_root, chunk_index)
    r2_worker_output_dir = modal_worker_output_dir(run_root, chunk_index)
    _ensure_r2_env()
    _run_cli(
        build_worker_argv(
            output_dir=run_root,
            manifest_path=manifest_path,
            worker_output_dir=worker_output_dir,
            bucket=R2_BUCKET,
            prefix=R2_OUTPUT_PREFIX,
            endpoint_url=R2_ENDPOINT_URL,
            chunk_index=chunk_index,
            chunk_count=chunk_count,
        )
    )
    r2_worker_output_dir.mkdir(parents=True, exist_ok=True)
    for name in ("state.json", "benchmark.json"):
        shutil.copyfile(worker_output_dir / name, r2_worker_output_dir / name)
    eda_volume.commit()
    return {
        "run_id": run_id,
        "chunk_index": chunk_index,
        "worker_output_dir": str(worker_output_dir),
        "benchmark": _read_json(worker_output_dir / "benchmark.json"),
    }


@app.function(
    image=_eda_image,
    volumes={"/r2": r2_mount, "/cache": eda_volume},
    timeout=3600,
    cpu=2.0,
)
def merge_run(run_id: str) -> dict[str, object]:
    from dolma.sidecar_eda.modal import (
        build_merge_argv,
        modal_cache_run_root,
        modal_run_root,
        modal_workers_dir,
    )

    run_root = modal_run_root(R2_STATS_PREFIX, run_id)
    eda_volume.reload()
    _run_cli(
        build_merge_argv(
            output_dir=run_root,
            workers_dir=modal_workers_dir(modal_cache_run_root(run_id)),
        )
    )
    return {
        "run_id": run_id,
        "run_root": str(run_root),
        "benchmark": _read_json(run_root / "benchmark.json"),
        "summary": _read_json(run_root / "summary.json"),
    }


@app.local_entrypoint()
def main(run_id: str = "", chunk_count: int = 128, max_shards: int = 0) -> None:
    current_run_id = run_id or datetime.now(timezone.utc).strftime(
        "soc101_modal_%Y%m%d_%H%M%S"
    )
    manifest_result = build_manifest.remote(
        current_run_id, max_shards=max_shards or None
    )
    worker_results = list(
        run_chunk.starmap(
            [
                (current_run_id, chunk_index, chunk_count)
                for chunk_index in range(chunk_count)
            ]
        )
    )
    merge_result = merge_run.remote(current_run_id)
    print(
        json.dumps(
            {
                "run_id": current_run_id,
                "manifest": manifest_result,
                "workers_completed": len(worker_results),
                "merge": merge_result,
            },
            indent=2,
        )
    )