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="range-stream public AirfRANS, process bounded chunks, and publish verified .npz files to HF",
    )
    prepare_public.add_argument("--repo-id", default="zacheryasc/airfrans-processed")
    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("--chunk-max-bytes", type=int, default=10 * 1024**3)
    prepare_public.add_argument("--state-path")
    prepare_public.add_argument("--train-cases", type=int, default=900)
    prepare_public.add_argument("--val-cases", type=int, default=50)
    prepare_public.add_argument("--test-cases", type=int, default=50)
    prepare_public.add_argument("--split-seed", type=int, default=20260726)
    prepare_public.add_argument("--limit-cases", type=int)
    prepare_public.add_argument("--verify-download-limit-bytes", type=int, default=64 * 1024 * 1024)
    prepare_public.add_argument("--workers", type=int, default=4)
    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")

    sweep_generate = subparsers.add_parser("sweep-generate", help="generate aggressive OOM sweep pool.toml, jobs.jsonl, and configs")
    sweep_generate.add_argument("--output-dir", default="artifacts/aggressive_oom_sweep")
    sweep_generate.add_argument("--data-root", default="artifacts/data_cache/airfrans_processed/processed/full")
    sweep_generate.add_argument("--artifact-dir", default="artifacts/runs")
    sweep_generate.add_argument("--hf-repo-id", default="zacheryasc/airfrans-frontier-checkpoints")
    sweep_generate.add_argument("--group", default="aggressive_oom_sweep_01")
    sweep_generate.add_argument("--bands", nargs="*", default=["100m", "700m"])
    sweep_generate.add_argument("--families", nargs="*")
    sweep_generate.add_argument("--encodings", nargs="*")
    sweep_generate.set_defaults(command="sweep-generate")

    sweep_collect = subparsers.add_parser("sweep-collect", help="reconstruct sweep job status from artifact files")
    sweep_collect.add_argument("--jobs", default="artifacts/aggressive_oom_sweep/jobs.jsonl")
    sweep_collect.set_defaults(command="sweep-collect")

    sweep_gate = subparsers.add_parser("sweep-gate", help="evaluate GPU pool expansion gate")
    sweep_gate.add_argument("--pool", default="artifacts/aggressive_oom_sweep/pool.toml")
    sweep_gate.add_argument("--utilization", required=True)
    sweep_gate.add_argument("--backlog", type=int, required=True)
    sweep_gate.add_argument("--recent-failures", type=int, default=0)
    sweep_gate.add_argument("--spent-usd", type=float, default=0.0)
    sweep_gate.add_argument("--reserved-usd", type=float, default=0.0)
    sweep_gate.add_argument("--next-pool-cost-usd", type=float, required=True)
    sweep_gate.set_defaults(command="sweep-gate")

    sweep_sample_util = subparsers.add_parser("sweep-sample-util", help="append one nvidia-smi utilization sample as JSONL")
    sweep_sample_util.add_argument("--output", default="artifacts/aggressive_oom_sweep/utilization.jsonl")
    sweep_sample_util.set_defaults(command="sweep-sample-util")

    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    if args.command in {"process-raw", "prepare-public-hf", "train", "model-sanity"}:
        from airfrans_frontier.runtime import remove_pythonpath_entries

        remove_pythonpath_entries()

    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
        if args.chunk_max_bytes <= 0:
            print("error: --chunk-max-bytes must be positive", file=sys.stderr)
            return 1
        if args.workers <= 0:
            print("error: --workers must be positive", file=sys.stderr)
            return 1
        if args.limit_cases is not None and args.limit_cases <= 0:
            print("error: --limit-cases must be positive", file=sys.stderr)
            return 1
        if args.train_cases + args.val_cases + args.test_cases <= 0:
            print("error: at least one split case is required", file=sys.stderr)
            return 1
        from airfrans_frontier.raw.bounded_public import prepare_public_airfrans_processed_hf_bounded

        try:
            report = prepare_public_airfrans_processed_hf_bounded(
                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,
                chunk_max_bytes=args.chunk_max_bytes,
                state_path=resolve_path(args.state_path) if args.state_path else None,
                train_cases=args.train_cases,
                val_cases=args.val_cases,
                test_cases=args.test_cases,
                split_seed=args.split_seed,
                private=args.private,
                force=args.force,
                limit_cases=args.limit_cases,
                verify_download_limit_bytes=args.verify_download_limit_bytes,
                workers=args.workers,
            )
        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.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.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

    if args.command == "sweep-generate":
        from airfrans_frontier.sweep import DEFAULT_ENCODINGS, DEFAULT_FAMILIES, generate_jobs

        try:
            jobs = generate_jobs(
                output_dir=resolve_path(args.output_dir),
                data_root=args.data_root,
                artifact_dir=args.artifact_dir,
                hf_repo_id=args.hf_repo_id,
                group=args.group,
                bands=tuple(args.bands),
                families=tuple(args.families) if args.families else DEFAULT_FAMILIES,
                encodings=tuple(args.encodings) if args.encodings else DEFAULT_ENCODINGS,
            )
        except (FileNotFoundError, NotADirectoryError, ValueError, RuntimeError) as exc:
            print(f"error: {exc}", file=sys.stderr)
            return 1
        print(f"output_dir: {resolve_path(args.output_dir)}")
        print(f"jobs: {len(jobs)}")
        print(f"jobs_jsonl: {resolve_path(args.output_dir) / 'jobs.jsonl'}")
        return 0

    if args.command == "sweep-collect":
        from airfrans_frontier.sweep import collect_job_status

        try:
            status = collect_job_status(jobs_path=resolve_path(args.jobs))
        except (FileNotFoundError, NotADirectoryError, ValueError, RuntimeError) as exc:
            print(f"error: {exc}", file=sys.stderr)
            return 1
        print(json.dumps(status, indent=2, sort_keys=True))
        return 0

    if args.command == "sweep-gate":
        from airfrans_frontier.sweep import BudgetLedger, can_expand_pool, load_pool_config, utilization_summary

        if args.backlog < 0:
            print("error: --backlog must be non-negative", file=sys.stderr)
            return 1
        if args.recent_failures < 0:
            print("error: --recent-failures must be non-negative", file=sys.stderr)
            return 1
        try:
            pool = load_pool_config(resolve_path(args.pool))
            utilization = utilization_summary(resolve_path(args.utilization))
            ledger = BudgetLedger(budget_usd=pool.budget_usd, spent_usd=args.spent_usd, reserved_usd=args.reserved_usd)
            allowed, reasons = can_expand_pool(
                pool=pool,
                utilization=utilization,
                backlog=args.backlog,
                recent_failures=args.recent_failures,
                ledger=ledger,
                next_pool_cost_usd=args.next_pool_cost_usd,
            )
        except (FileNotFoundError, NotADirectoryError, ValueError, RuntimeError) as exc:
            print(f"error: {exc}", file=sys.stderr)
            return 1
        print(json.dumps({"allowed": allowed, "reasons": reasons, "utilization": utilization}, indent=2, sort_keys=True))
        return 0 if allowed else 1

    if args.command == "sweep-sample-util":
        from airfrans_frontier.sweep import append_utilization_sample

        sample = append_utilization_sample(resolve_path(args.output))
        print(json.dumps(sample, indent=2, sort_keys=True))
        return 0

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


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