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_run_node = subparsers.add_parser("sweep-run-node", help="drain sweep jobs sequentially for one named node") sweep_run_node.add_argument("--jobs", default="artifacts/aggressive_oom_sweep/jobs.jsonl") sweep_run_node.add_argument("--node", required=True) sweep_run_node.add_argument("--max-jobs", type=int) sweep_run_node.add_argument("--stale-after-seconds", type=float, default=21600.0) sweep_run_node.add_argument("--max-attempts", type=int, default=2) sweep_run_node.set_defaults(command="sweep-run-node") 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-run-node": from airfrans_frontier.sweep import run_node if args.max_jobs is not None and args.max_jobs < 0: print("error: --max-jobs must be non-negative", file=sys.stderr) return 1 if args.stale_after_seconds < 0: print("error: --stale-after-seconds must be non-negative", file=sys.stderr) return 1 if args.max_attempts <= 0: print("error: --max-attempts must be positive", file=sys.stderr) return 1 try: summary = run_node( jobs_path=resolve_path(args.jobs), node_name=args.node, max_jobs=args.max_jobs, stale_after_seconds=args.stale_after_seconds, max_attempts=args.max_attempts, ) except (FileNotFoundError, NotADirectoryError, ValueError, RuntimeError) as exc: print(f"error: {exc}", file=sys.stderr) return 1 print(json.dumps(summary, indent=2, sort_keys=True)) 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())