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

"""CLI for quality sidecar validation on actual data."""

from __future__ import annotations

import argparse
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, list_keys
from dolma.quality.validation.analysis import analyze_validation_sources
from dolma.quality.validation.manifest import (
    build_manifest_summary,
    flatten_groups,
    group_source_keys,
    read_group_manifest,
    select_validation_source_keys,
    write_group_manifest,
    write_manifest_summary,
)
from dolma.quality.validation.outputs import write_quality_validation_outputs
from dolma.quality.validation.run import run_local_smoke


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="command", required=True)
    manifest = subparsers.add_parser("manifest")
    smoke = subparsers.add_parser("smoke")
    pilot = subparsers.add_parser("pilot")
    analyze = subparsers.add_parser("analyze")
    for subparser in (manifest, smoke, pilot, analyze):
        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(
            "--verbose", action=argparse.BooleanOptionalAction, default=False
        )
    manifest.add_argument("--manifest-output", type=Path, required=True)
    manifest.add_argument("--summary-output", type=Path, required=True)
    manifest.add_argument("--total-shards", type=int, default=500)
    manifest.add_argument("--group-size", type=int, default=20)
    smoke.add_argument("--manifest-path", type=Path, required=True)
    smoke.add_argument("--summary-output", type=Path, required=True)
    smoke.add_argument("--smoke-shards", type=int, default=1)
    smoke.add_argument("--batch-size", type=int, default=2048)
    smoke.add_argument("--model-repo", default=DEFAULT_QUALITY_MODEL_REPO)
    pilot.add_argument("--manifest-path", type=Path, required=True)
    pilot.add_argument("--benchmark-output", type=Path, required=True)
    pilot.add_argument("--submission-output", type=Path, required=True)
    pilot.add_argument("--cpu", type=int, default=None)
    pilot.add_argument("--batch-size", type=int, default=2048)
    pilot.add_argument("--pilot-groups", type=int, default=5)
    pilot.add_argument("--cpu-shapes", nargs="+", type=int, default=[4, 8, 16])
    pilot.add_argument("--wave-size", type=int, default=1000)
    pilot.add_argument("--model-repo", default=DEFAULT_QUALITY_MODEL_REPO)
    analyze.add_argument("--manifest-path", type=Path, required=True)
    analyze.add_argument("--analysis-output-dir", type=Path, required=True)
    analyze.add_argument("--soc91-prefix", default="soc91-labels")
    return parser


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


def _source_keys_from_r2(args: argparse.Namespace) -> tuple[object, object, list[str]]:
    config = _config_from_args(args)
    client = create_r2_client(config)
    source_keys: list[str] = []
    for prefix in config.input_prefixes:
        source_keys.extend(
            list_keys(client, bucket=config.bucket, prefix=prefix, suffix=".jsonl.zst")
        )
    return config, client, sorted(set(source_keys))


def _manifest_groups(path: Path) -> list[list[str]]:
    groups = read_group_manifest(path)
    if not groups:
        raise ValueError(f"Empty manifest: {path}")
    return groups


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    configure_logging(args.verbose)
    if args.command == "manifest":
        config, _, source_keys = _source_keys_from_r2(args)
        selected = select_validation_source_keys(
            source_keys,
            input_prefixes=config.input_prefixes,
            total_shards=args.total_shards,
        )
        groups = group_source_keys(selected, args.group_size)
        write_group_manifest(args.manifest_output, groups)
        write_manifest_summary(
            args.summary_output,
            {
                **build_manifest_summary(
                    selected,
                    input_prefixes=config.input_prefixes,
                    group_size=args.group_size,
                ),
                "bucket": config.bucket,
                "input_prefixes": list(config.input_prefixes),
                "output_prefix": config.output_prefix,
            },
        )
        return 0
    config = _config_from_args(args)
    client = create_r2_client(config)
    groups = _manifest_groups(args.manifest_path)
    source_keys = flatten_groups(groups)
    if args.command == "smoke":
        summaries = run_local_smoke(
            client,
            bucket=config.bucket,
            output_prefix=config.output_prefix,
            source_keys=source_keys[: args.smoke_shards],
            model_repo=args.model_repo,
            batch_size=args.batch_size,
        )
        write_json_artifact(
            args.summary_output,
            {
                "smoke_shards": len(summaries),
                "output_prefix": config.output_prefix,
                "sources": summaries,
            },
        )
        return 0
    if args.command == "pilot":
        from dolma.quality.modal import benchmark_modal_groups, submit_modal_groups

        job = {
            "r2": config.__dict__,
            "model_repo": args.model_repo,
            "batch_size": args.batch_size,
        }
        if args.cpu is None:
            benchmark = benchmark_modal_groups(
                groups,
                job,
                cpu_shapes=args.cpu_shapes,
                pilot_groups=min(args.pilot_groups, len(groups)),
                total_pending_groups=len(groups),
            )
            recommended = select_recommended_shape(benchmark["cpu_summaries"])
            benchmark["recommended"] = recommended
            cpu = int(recommended["cpu"])
        else:
            benchmark = None
            cpu = args.cpu
        if benchmark is not None:
            write_json_artifact(args.benchmark_output, benchmark)
        submission = submit_modal_groups(groups, job, cpu=cpu, wave_size=args.wave_size)
        write_json_artifact(
            args.submission_output, {**submission, "source_key_count": len(source_keys)}
        )
        return 0
    payload = analyze_validation_sources(
        client,
        bucket=config.bucket,
        source_keys=source_keys,
        output_prefix=config.output_prefix,
        soc91_prefix=args.soc91_prefix,
    )
    write_quality_validation_outputs(args.analysis_output_dir, payload)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())