352 lines
17 KiB
Python
352 lines
17 KiB
Python
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())
|