airfRANS-model-exploration/src/airfrans_frontier/cli.py

from __future__ import annotations

import argparse
import json
import sys

from airfrans_frontier.paths import DEFAULT_RAW_DATA_DIR, DEFAULT_RAW_MANIFEST_PATH, resolve_path
from airfrans_frontier.raw.inspect import format_raw_inspection, inspect_raw_subset


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="airfrans-frontier")
    subparsers = parser.add_subparsers(required=True)

    inspect_raw = subparsers.add_parser("inspect-raw", help="inspect the local raw AirfRANS subset")
    inspect_raw.add_argument("--data-dir", default=str(DEFAULT_RAW_DATA_DIR))
    inspect_raw.add_argument("--manifest", default=str(DEFAULT_RAW_MANIFEST_PATH))
    inspect_raw.add_argument("--sample-limit", type=int, default=5)
    inspect_raw.set_defaults(command="inspect-raw")

    process_raw = subparsers.add_parser("process-raw", help="convert raw OpenFOAM cases into training tensors")
    process_raw.add_argument("--raw-dir", default=str(DEFAULT_RAW_DATA_DIR))
    process_raw.add_argument("--output-dir", default="data/processed/full")
    process_raw.add_argument("--limit", type=int)
    process_raw.add_argument("--force", action="store_true")
    process_raw.set_defaults(command="process-raw")

    publish_processed = subparsers.add_parser("publish-processed-hf", help="publish processed .npz data to a Hugging Face dataset repo")
    publish_processed.add_argument("--data-root", required=True)
    publish_processed.add_argument("--repo-id", required=True)
    publish_processed.add_argument("--path-in-repo", default="processed/full")
    publish_processed.add_argument("--private", action="store_true")
    publish_processed.add_argument("--manifest-out")
    publish_processed.set_defaults(command="publish-processed-hf")

    prepare_public = subparsers.add_parser(
        "prepare-public-hf",
        help="download public AirfRANS OF_dataset.zip, process it, and publish processed .npz files to HF",
    )
    prepare_public.add_argument("--repo-id", required=True)
    prepare_public.add_argument("--path-in-repo", default="processed/full")
    prepare_public.add_argument("--work-dir", default="artifacts/public_airfrans")
    prepare_public.add_argument("--output-dir", default="artifacts/data_cache/airfrans_processed/processed/full")
    prepare_public.add_argument("--source-url", default="https://data.isir.upmc.fr/extrality/NeurIPS_2022/OF_dataset.zip")
    prepare_public.add_argument("--min-cases", type=int, default=1000)
    prepare_public.add_argument("--private", action="store_true")
    prepare_public.add_argument("--force", action="store_true")
    prepare_public.set_defaults(command="prepare-public-hf")

    train = subparsers.add_parser("train", help="train a configured baseline model")
    train.add_argument("config", help="path to a training config TOML file")
    train.add_argument("--resume", help="path to checkpoint_latest.pt to resume from")
    train.set_defaults(command="train")

    sanity = subparsers.add_parser("model-sanity", help="run toy loss-decrease checks for frontier model families")
    sanity.add_argument("--artifact-dir", default="artifacts/model_sanity")
    sanity.add_argument("--device", choices=("auto", "cuda", "cpu"), default="auto")
    sanity.add_argument("--steps", type=int, default=80)
    sanity.add_argument("--families", nargs="*", help="model families to check; defaults to every frontier family")
    sanity.set_defaults(command="model-sanity")

    return parser


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

    if args.command == "inspect-raw":
        if args.sample_limit < 0:
            print("error: --sample-limit must be non-negative", file=sys.stderr)
            return 1

        data_dir = resolve_path(args.data_dir)
        manifest_path = resolve_path(args.manifest)
        try:
            report = inspect_raw_subset(data_dir, manifest_path)
        except (FileNotFoundError, NotADirectoryError, ValueError) as exc:
            print(f"error: {exc}", file=sys.stderr)
            return 1

        print(format_raw_inspection(report, sample_limit=args.sample_limit))
        return 0 if report.matches_manifest else 1

    if args.command == "process-raw":
        if args.limit is not None and args.limit <= 0:
            print("error: --limit must be positive", file=sys.stderr)
            return 1
        from airfrans_frontier.raw.process import process_raw_dataset

        try:
            result = process_raw_dataset(
                resolve_path(args.raw_dir),
                resolve_path(args.output_dir),
                limit=args.limit,
                force=args.force,
            )
        except (FileNotFoundError, NotADirectoryError, ValueError) as exc:
            print(f"error: {exc}", file=sys.stderr)
            return 1
        print(f"output_dir: {result.output_dir}")
        print(f"case_count: {result.case_count}")
        print(f"total_points: {result.total_points}")
        print(f"manifest: {result.manifest_path}")
        return 0

    if args.command == "publish-processed-hf":
        from airfrans_frontier.training.data_sources import publish_processed_dataset

        try:
            manifest = publish_processed_dataset(
                data_root=resolve_path(args.data_root),
                repo_id=args.repo_id,
                path_in_repo=args.path_in_repo,
                private=args.private,
                manifest_out=resolve_path(args.manifest_out) if args.manifest_out else None,
            )
        except (FileNotFoundError, NotADirectoryError, ValueError, RuntimeError) as exc:
            print(f"error: {exc}", file=sys.stderr)
            return 1
        print(f"repo_url: {manifest['repo_url']}")
        print(f"path_in_repo: {manifest['path_in_repo']}")
        print(f"npz_files: {manifest['npz_file_count']}")
        return 0

    if args.command == "prepare-public-hf":
        if args.min_cases <= 0:
            print("error: --min-cases must be positive", file=sys.stderr)
            return 1
        from airfrans_frontier.raw.public import ensure_public_airfrans_processed_hf

        try:
            report = ensure_public_airfrans_processed_hf(
                repo_id=args.repo_id,
                path_in_repo=args.path_in_repo,
                work_dir=resolve_path(args.work_dir),
                output_dir=resolve_path(args.output_dir),
                source_url=args.source_url,
                min_cases=args.min_cases,
                private=args.private,
                force=args.force,
            )
        except (FileNotFoundError, NotADirectoryError, ValueError, RuntimeError) as exc:
            print(f"error: {exc}", file=sys.stderr)
            return 1
        print(json.dumps(report, indent=2, sort_keys=True))
        return 0

    if args.command == "train":
        from airfrans_frontier.runtime import remove_pythonpath_entries

        remove_pythonpath_entries()
        from airfrans_frontier.training.loop import train_from_config_path

        try:
            result = train_from_config_path(resolve_path(args.config), resume_path=resolve_path(args.resume) if args.resume else None)
        except (FileNotFoundError, NotADirectoryError, ValueError, RuntimeError) as exc:
            print(f"error: {exc}", file=sys.stderr)
            return 1

        print(f"run_dir: {result.run_dir}")
        print(f"final_metrics: {result.run_dir / 'final_metrics.json'}")
        return 0

    if args.command == "model-sanity":
        if args.steps <= 0:
            print("error: --steps must be positive", file=sys.stderr)
            return 1
        from airfrans_frontier.runtime import remove_pythonpath_entries

        remove_pythonpath_entries()
        from airfrans_frontier.training.sanity import MODEL_FAMILIES, run_model_sanity

        families = tuple(args.families) if args.families else MODEL_FAMILIES
        try:
            result = run_model_sanity(
                artifact_dir=resolve_path(args.artifact_dir),
                device_type=args.device,
                families=families,
                steps=args.steps,
            )
        except (FileNotFoundError, NotADirectoryError, ValueError, RuntimeError) as exc:
            print(f"error: {exc}", file=sys.stderr)
            return 1
        print(f"report: {resolve_path(args.artifact_dir) / 'model_sanity_results.json'}")
        print(f"families: {len(result['families'])}")
        return 0

    parser.error(f"unknown command: {args.command}")
    return 2


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