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())