social-data-attribution-w2/src/dolma/quality/cli.py

"""CLI for Modal CPU quality sidecars."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

from data_attribution.cli.config import configure_logging
from dolma.quality.artifact import write_json_artifact
from dolma.quality.benchmark import select_recommended_shape
from dolma.quality.fasttext import DEFAULT_QUALITY_MODEL_REPO
from dolma.quality.r2 import (
    R2Config,
    create_r2_client,
    group_source_keys,
    pending_source_keys,
)


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Dolma quality sidecar CLI")
    subparsers = parser.add_subparsers(dest="command", required=True)

    manifest = subparsers.add_parser("manifest")
    benchmark = subparsers.add_parser("benchmark")
    launch = subparsers.add_parser("launch")

    for subparser in (manifest, benchmark, launch):
        subparser.add_argument("--output-prefix", required=True)
        subparser.add_argument(
            "--input-prefixes",
            nargs="+",
            default=["soc127/phase1_pool_shared", "soc127/phase2_nonpool_final"],
        )
        subparser.add_argument("--group-size", type=int, default=20)
        subparser.add_argument(
            "--verbose", action=argparse.BooleanOptionalAction, default=False
        )

    manifest.add_argument("--manifest-output", type=Path, default=None)
    manifest.add_argument("--summary-output", type=Path, default=None)

    benchmark.add_argument("--pilot-groups", type=int, default=5)
    benchmark.add_argument("--cpu-shapes", nargs="+", type=int, default=[4, 8, 16])
    benchmark.add_argument("--batch-size", type=int, default=2048)
    benchmark.add_argument("--model-repo", default=DEFAULT_QUALITY_MODEL_REPO)
    benchmark.add_argument("--summary-output", type=Path, required=True)

    launch.add_argument("--cpu", type=int, default=None)
    launch.add_argument("--wave-size", type=int, default=1000)
    launch.add_argument("--max-groups", type=int, default=None)
    launch.add_argument("--batch-size", type=int, default=2048)
    launch.add_argument("--model-repo", default=DEFAULT_QUALITY_MODEL_REPO)
    launch.add_argument("--benchmark-output", type=Path, default=None)
    launch.add_argument("--manifest-summary-output", type=Path, required=True)
    return parser


def _manifest_payload(
    grouped: list[list[str]], pending: list[str], done_count: int, config: R2Config
) -> dict[str, object]:
    return {
        "bucket": config.bucket,
        "input_prefixes": list(config.input_prefixes),
        "output_prefix": config.output_prefix,
        "source_key_count": len(pending) + done_count,
        "pending_source_key_count": len(pending),
        "completed_source_key_count": done_count,
        "group_count": len(grouped),
    }


def _config_from_args(args: argparse.Namespace) -> R2Config:
    return R2Config.from_env(
        output_prefix=args.output_prefix,
        input_prefixes=args.input_prefixes,
    )


def _grouped_sources(
    args: argparse.Namespace,
) -> tuple[R2Config, list[str], list[list[str]], int]:
    config = _config_from_args(args)
    client = create_r2_client(config)
    pending, done_count = pending_source_keys(client, config)
    grouped = group_source_keys(pending, args.group_size)
    return config, pending, grouped, done_count


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    configure_logging(args.verbose)

    config, pending, grouped, done_count = _grouped_sources(args)
    manifest_payload = _manifest_payload(grouped, pending, done_count, config)

    if args.command == "manifest":
        if args.manifest_output is not None:
            args.manifest_output.parent.mkdir(parents=True, exist_ok=True)
            args.manifest_output.write_text(
                "\n".join(" ".join(group) for group in grouped)
                + ("\n" if grouped else ""),
                encoding="utf-8",
            )
        if args.summary_output is not None:
            write_json_artifact(args.summary_output, manifest_payload)
        return 0

    job = {
        "r2": config.__dict__,
        "model_repo": args.model_repo,
        "batch_size": args.batch_size,
    }

    if args.command == "benchmark":
        from dolma.quality.modal import benchmark_modal_groups

        benchmark_summary = benchmark_modal_groups(
            grouped,
            job,
            cpu_shapes=args.cpu_shapes,
            pilot_groups=args.pilot_groups,
            total_pending_groups=len(grouped),
        )
        benchmark_summary["recommended"] = select_recommended_shape(
            benchmark_summary["cpu_summaries"]
        )
        write_json_artifact(args.summary_output, benchmark_summary)
        return 0

    if args.command == "launch":
        from dolma.quality.modal import benchmark_modal_groups, submit_modal_groups

        launch_groups = (
            grouped[: args.max_groups] if args.max_groups is not None else grouped
        )
        if args.cpu is None:
            if args.benchmark_output is None:
                raise ValueError("--benchmark-output is required when --cpu is omitted")
            benchmark_summary = benchmark_modal_groups(
                launch_groups,
                job,
                cpu_shapes=[4, 8, 16],
                pilot_groups=min(5, len(launch_groups)),
                total_pending_groups=len(launch_groups),
            )
            recommended = select_recommended_shape(benchmark_summary["cpu_summaries"])
            benchmark_summary["recommended"] = recommended
            write_json_artifact(args.benchmark_output, benchmark_summary)
            cpu = int(recommended["cpu"])
        else:
            cpu = args.cpu
        submission = submit_modal_groups(
            launch_groups,
            job,
            cpu=cpu,
            wave_size=args.wave_size,
        )
        write_json_artifact(
            args.manifest_summary_output,
            {**manifest_payload, **submission},
        )
        return 0

    raise ValueError(f"Unknown command: {args.command}")


if __name__ == "__main__":
    sys.exit(main())