From 382baefca38cef08f018dcd7a8a53732e25dd15e Mon Sep 17 00:00:00 2001 From: Zachery Aaron Shores-Chmielewski Date: Sat, 25 Jul 2026 20:12:49 +0400 Subject: [PATCH] stash: failed model zoo and 70GB run --- .../runs/model/checkpoint_latest.pt | 1 + .skyignore | 7 +- configs/aggressive_smoke.toml | 19 +- configs/full_airfrans_incumbent_70gb.toml | 65 +++ .../model_zoo_7gb/deeponet_branch_trunk.toml | 61 +++ configs/model_zoo_7gb/film_fourier_inr.toml | 61 +++ ...shgraphnet_or_point_transformer_local.toml | 61 +++ configs/model_zoo_7gb/nerf_cfd_multires.toml | 62 +++ .../point_context_perceiver.toml | 63 +++ configs/model_zoo_7gb/raster_fno_unet.toml | 61 +++ .../model_zoo_7gb/siren_conditioned_inr.toml | 61 +++ configs/remote_full_70gb.toml | 86 ++++ .../deeponet_branch_trunk.toml | 89 ++++ .../film_fourier_inr.toml | 89 ++++ ...shgraphnet_or_point_transformer_local.toml | 89 ++++ .../nerf_cfd_multires.toml | 89 ++++ .../point_context_perceiver.toml | 89 ++++ .../remote_model_zoo_7gb/raster_fno_unet.toml | 89 ++++ .../siren_conditioned_inr.toml | 89 ++++ configs/remote_smoke.toml | 12 +- pyproject.toml | 4 + src/airfrans_frontier/cli.py | 96 ++++ src/airfrans_frontier/models/__init__.py | 19 +- src/airfrans_frontier/models/frontier.py | 404 +++++++++++++++++ src/airfrans_frontier/raw/process.py | 5 +- src/airfrans_frontier/raw/public.py | 253 +++++++++++ src/airfrans_frontier/remote/artifacts.py | 112 ++++- src/airfrans_frontier/remote/cli.py | 10 + src/airfrans_frontier/remote/smoke.py | 128 ++++-- src/airfrans_frontier/training/artifacts.py | 2 +- src/airfrans_frontier/training/calibration.py | 138 ++++++ src/airfrans_frontier/training/config.py | 79 +++- .../training/data_sources.py | 150 +++++++ src/airfrans_frontier/training/environment.py | 44 ++ src/airfrans_frontier/training/hf_upload.py | 215 +++++++++ src/airfrans_frontier/training/loop.py | 419 +++++++++++++++++- .../training/observability.py | 7 + src/airfrans_frontier/training/sanity.py | 167 +++++++ tests/test_data_sources.py | 98 ++++ tests/test_hf_upload.py | 77 ++++ tests/test_mlp.py | 83 +++- tests/test_public_data.py | 67 +++ tests/test_remote_run.py | 7 + tests/test_training_config.py | 11 + tests/test_training_loop.py | 11 + 45 files changed, 3793 insertions(+), 56 deletions(-) create mode 100644 .airfrans_hf_resume/2bc55a1e75e5f767/runs/model/checkpoint_latest.pt create mode 100644 configs/full_airfrans_incumbent_70gb.toml create mode 100644 configs/model_zoo_7gb/deeponet_branch_trunk.toml create mode 100644 configs/model_zoo_7gb/film_fourier_inr.toml create mode 100644 configs/model_zoo_7gb/meshgraphnet_or_point_transformer_local.toml create mode 100644 configs/model_zoo_7gb/nerf_cfd_multires.toml create mode 100644 configs/model_zoo_7gb/point_context_perceiver.toml create mode 100644 configs/model_zoo_7gb/raster_fno_unet.toml create mode 100644 configs/model_zoo_7gb/siren_conditioned_inr.toml create mode 100644 configs/remote_full_70gb.toml create mode 100644 configs/remote_model_zoo_7gb/deeponet_branch_trunk.toml create mode 100644 configs/remote_model_zoo_7gb/film_fourier_inr.toml create mode 100644 configs/remote_model_zoo_7gb/meshgraphnet_or_point_transformer_local.toml create mode 100644 configs/remote_model_zoo_7gb/nerf_cfd_multires.toml create mode 100644 configs/remote_model_zoo_7gb/point_context_perceiver.toml create mode 100644 configs/remote_model_zoo_7gb/raster_fno_unet.toml create mode 100644 configs/remote_model_zoo_7gb/siren_conditioned_inr.toml create mode 100644 src/airfrans_frontier/models/frontier.py create mode 100644 src/airfrans_frontier/raw/public.py create mode 100644 src/airfrans_frontier/training/calibration.py create mode 100644 src/airfrans_frontier/training/data_sources.py create mode 100644 src/airfrans_frontier/training/environment.py create mode 100644 src/airfrans_frontier/training/hf_upload.py create mode 100644 src/airfrans_frontier/training/sanity.py create mode 100644 tests/test_data_sources.py create mode 100644 tests/test_hf_upload.py create mode 100644 tests/test_public_data.py diff --git a/.airfrans_hf_resume/2bc55a1e75e5f767/runs/model/checkpoint_latest.pt b/.airfrans_hf_resume/2bc55a1e75e5f767/runs/model/checkpoint_latest.pt new file mode 100644 index 0000000..61f4290 --- /dev/null +++ b/.airfrans_hf_resume/2bc55a1e75e5f767/runs/model/checkpoint_latest.pt @@ -0,0 +1 @@ +checkpoint \ No newline at end of file diff --git a/.skyignore b/.skyignore index 3e29bf6..39de8f0 100644 --- a/.skyignore +++ b/.skyignore @@ -1,9 +1,10 @@ /artifacts /data/raw +/data/processed /.venv /notebooks -.env -HF_TOKEN -WANDB_API_KEY __pycache__ *.pyc +HF_TOKEN +WANDB_API_KEY +.env diff --git a/configs/aggressive_smoke.toml b/configs/aggressive_smoke.toml index 63aac64..c20cd31 100644 --- a/configs/aggressive_smoke.toml +++ b/configs/aggressive_smoke.toml @@ -1,5 +1,5 @@ [run] -name = "aggressive_smoke" +name = "model_class_frontier_7gb_01_film_fourier_inr" seed = 20260723 artifact_dir = "artifacts/current_run/training_runs" @@ -10,9 +10,14 @@ val_cases = 3 test_cases = 2 points_per_case = 999999999 batch_size = 4096 +source = "huggingface" +hf_repo_id = "zacheryasc/airfrans-processed" +hf_repo_type = "dataset" +hf_path_prefix = "processed/full" +cache_dir = "artifacts/data_cache/airfrans_processed" [model] -type = "film_fourier_mlp" +type = "film_fourier_inr" hidden_width = 4096 depth = 12 activation = "gelu" @@ -49,4 +54,12 @@ max_grad_norm = 1.0 backend = "wandb" entity = "zacheryasc-personal" project = "airfRANS-model-sweep" -tags = ["airfrans", "remote", "aggressive-smoke"] +group = "model_class_frontier_7gb_01" +tags = ["airfrans", "model-class-frontier", "7gb-subset", "flop-par", "hf-checkpoints"] + +[huggingface] +enabled = true +repo_id = "zacheryasc/airfrans-frontier-checkpoints" +repo_type = "model" +path_prefix = "model_class_frontier_7gb_01" +private = false diff --git a/configs/full_airfrans_incumbent_70gb.toml b/configs/full_airfrans_incumbent_70gb.toml new file mode 100644 index 0000000..05e3a1e --- /dev/null +++ b/configs/full_airfrans_incumbent_70gb.toml @@ -0,0 +1,65 @@ +[run] +name = "full_airfrans_incumbent_70gb_01" +seed = 20260723 +artifact_dir = "artifacts/current_run/training_runs" + +[data] +root = "data/processed/full" +train_cases = 900 +val_cases = 50 +test_cases = 50 +points_per_case = 999999999 +batch_size = 4096 +source = "huggingface" +hf_repo_id = "zacheryasc/airfrans-processed" +hf_repo_type = "dataset" +hf_path_prefix = "processed/full" +cache_dir = "artifacts/data_cache/airfrans_processed" + +[model] +type = "film_fourier_inr" +hidden_width = 4096 +depth = 12 +activation = "gelu" +coordinate_features = ["x", "y", "sdf"] +fourier_scales = [1.0, 2.0, 4.0, 8.0, 16.0, 32.0] +condition_width = 1024 +condition_depth = 3 +condition_dim = 512 + +[optim] +lr = 0.0001 +weight_decay = 0.0001 +steps = 5000 +log_interval = 500 + +[device] +type = "cuda" +allow_cpu_fallback = false +benchmark_kernels = true + +[loss] +type = "normalized_mse" + +[precision] +dtype = "bf16" + +[checkpoint] +interval_seconds = 900 + +[stability] +max_grad_norm = 1.0 + +[observability] +backend = "wandb" +entity = "zacheryasc-personal" +project = "airfRANS-model-sweep" +group = "full_airfrans_incumbent_70gb_01" +tags = ["airfrans", "full-data-frontier", "70gb", "incumbent", "hf-checkpoints"] + +[huggingface] +enabled = true +repo_id = "zacheryasc/airfrans-frontier-checkpoints" +repo_type = "model" +path_prefix = "full_airfrans_incumbent_70gb_01" +private = false diff --git a/configs/model_zoo_7gb/deeponet_branch_trunk.toml b/configs/model_zoo_7gb/deeponet_branch_trunk.toml new file mode 100644 index 0000000..0199bfb --- /dev/null +++ b/configs/model_zoo_7gb/deeponet_branch_trunk.toml @@ -0,0 +1,61 @@ +[run] +name = "model_class_frontier_7gb_01_deeponet_branch_trunk" +seed = 20260725 +artifact_dir = "artifacts/current_run/training_runs/deeponet_branch_trunk" + +[data] +root = "data/processed/full" +train_cases = 45 +val_cases = 3 +test_cases = 2 +points_per_case = 999999999 +batch_size = 4096 +source = "local" + +[model] +type = "deeponet_branch_trunk" +hidden_width = 1024 +depth = 8 +activation = "gelu" +coordinate_features = ["x", "y", "sdf"] +fourier_scales = [1.0, 2.0, 4.0, 8.0, 16.0, 32.0] +condition_width = 1024 +condition_depth = 3 +condition_dim = 512 + +[optim] +lr = 0.0001 +weight_decay = 0.0001 +steps = 5000 +log_interval = 500 + +[device] +type = "cuda" +allow_cpu_fallback = false +benchmark_kernels = true + +[loss] +type = "normalized_mse" + +[precision] +dtype = "bf16" + +[checkpoint] +interval_seconds = 1800 + +[stability] +max_grad_norm = 1.0 + +[observability] +backend = "wandb" +entity = "zacheryasc-personal" +project = "airfRANS-model-sweep" +group = "model_class_frontier_7gb_01" +tags = ["airfrans", "model-class-frontier", "7gb-subset", "tarball-data", "deeponet_branch_trunk"] + +[huggingface] +enabled = true +repo_id = "zacheryasc/airfrans-frontier-checkpoints" +repo_type = "model" +path_prefix = "model_class_frontier_7gb_01/deeponet_branch_trunk" +private = false diff --git a/configs/model_zoo_7gb/film_fourier_inr.toml b/configs/model_zoo_7gb/film_fourier_inr.toml new file mode 100644 index 0000000..a992043 --- /dev/null +++ b/configs/model_zoo_7gb/film_fourier_inr.toml @@ -0,0 +1,61 @@ +[run] +name = "model_class_frontier_7gb_01_film_fourier_inr" +seed = 20260725 +artifact_dir = "artifacts/current_run/training_runs/film_fourier_inr" + +[data] +root = "data/processed/full" +train_cases = 45 +val_cases = 3 +test_cases = 2 +points_per_case = 999999999 +batch_size = 4096 +source = "local" + +[model] +type = "film_fourier_inr" +hidden_width = 4096 +depth = 12 +activation = "gelu" +coordinate_features = ["x", "y", "sdf"] +fourier_scales = [1.0, 2.0, 4.0, 8.0, 16.0, 32.0] +condition_width = 1024 +condition_depth = 3 +condition_dim = 512 + +[optim] +lr = 0.0001 +weight_decay = 0.0001 +steps = 5000 +log_interval = 500 + +[device] +type = "cuda" +allow_cpu_fallback = false +benchmark_kernels = true + +[loss] +type = "normalized_mse" + +[precision] +dtype = "bf16" + +[checkpoint] +interval_seconds = 1800 + +[stability] +max_grad_norm = 1.0 + +[observability] +backend = "wandb" +entity = "zacheryasc-personal" +project = "airfRANS-model-sweep" +group = "model_class_frontier_7gb_01" +tags = ["airfrans", "model-class-frontier", "7gb-subset", "tarball-data", "film_fourier_inr"] + +[huggingface] +enabled = true +repo_id = "zacheryasc/airfrans-frontier-checkpoints" +repo_type = "model" +path_prefix = "model_class_frontier_7gb_01/film_fourier_inr" +private = false diff --git a/configs/model_zoo_7gb/meshgraphnet_or_point_transformer_local.toml b/configs/model_zoo_7gb/meshgraphnet_or_point_transformer_local.toml new file mode 100644 index 0000000..b80cd12 --- /dev/null +++ b/configs/model_zoo_7gb/meshgraphnet_or_point_transformer_local.toml @@ -0,0 +1,61 @@ +[run] +name = "model_class_frontier_7gb_01_meshgraphnet_or_point_transformer_local" +seed = 20260725 +artifact_dir = "artifacts/current_run/training_runs/meshgraphnet_or_point_transformer_local" + +[data] +root = "data/processed/full" +train_cases = 45 +val_cases = 3 +test_cases = 2 +points_per_case = 999999999 +batch_size = 4096 +source = "local" + +[model] +type = "meshgraphnet_or_point_transformer_local" +hidden_width = 512 +depth = 4 +activation = "gelu" +coordinate_features = ["x", "y", "sdf"] +condition_width = 512 +condition_depth = 3 +condition_dim = 512 +neighbors = 8 + +[optim] +lr = 0.0001 +weight_decay = 0.0001 +steps = 5000 +log_interval = 500 + +[device] +type = "cuda" +allow_cpu_fallback = false +benchmark_kernels = true + +[loss] +type = "normalized_mse" + +[precision] +dtype = "bf16" + +[checkpoint] +interval_seconds = 1800 + +[stability] +max_grad_norm = 1.0 + +[observability] +backend = "wandb" +entity = "zacheryasc-personal" +project = "airfRANS-model-sweep" +group = "model_class_frontier_7gb_01" +tags = ["airfrans", "model-class-frontier", "7gb-subset", "tarball-data", "meshgraphnet_or_point_transformer_local"] + +[huggingface] +enabled = true +repo_id = "zacheryasc/airfrans-frontier-checkpoints" +repo_type = "model" +path_prefix = "model_class_frontier_7gb_01/meshgraphnet_or_point_transformer_local" +private = false diff --git a/configs/model_zoo_7gb/nerf_cfd_multires.toml b/configs/model_zoo_7gb/nerf_cfd_multires.toml new file mode 100644 index 0000000..d2f668d --- /dev/null +++ b/configs/model_zoo_7gb/nerf_cfd_multires.toml @@ -0,0 +1,62 @@ +[run] +name = "model_class_frontier_7gb_01_nerf_cfd_multires" +seed = 20260725 +artifact_dir = "artifacts/current_run/training_runs/nerf_cfd_multires" + +[data] +root = "data/processed/full" +train_cases = 45 +val_cases = 3 +test_cases = 2 +points_per_case = 999999999 +batch_size = 4096 +source = "local" + +[model] +type = "nerf_cfd_multires" +hidden_width = 1024 +depth = 8 +activation = "gelu" +coordinate_features = ["x", "y", "sdf"] +condition_width = 512 +condition_depth = 3 +condition_dim = 512 +encoding_levels = 16 +features_per_level = 2 + +[optim] +lr = 0.0001 +weight_decay = 0.0001 +steps = 5000 +log_interval = 500 + +[device] +type = "cuda" +allow_cpu_fallback = false +benchmark_kernels = true + +[loss] +type = "normalized_mse" + +[precision] +dtype = "bf16" + +[checkpoint] +interval_seconds = 1800 + +[stability] +max_grad_norm = 1.0 + +[observability] +backend = "wandb" +entity = "zacheryasc-personal" +project = "airfRANS-model-sweep" +group = "model_class_frontier_7gb_01" +tags = ["airfrans", "model-class-frontier", "7gb-subset", "tarball-data", "nerf_cfd_multires"] + +[huggingface] +enabled = true +repo_id = "zacheryasc/airfrans-frontier-checkpoints" +repo_type = "model" +path_prefix = "model_class_frontier_7gb_01/nerf_cfd_multires" +private = false diff --git a/configs/model_zoo_7gb/point_context_perceiver.toml b/configs/model_zoo_7gb/point_context_perceiver.toml new file mode 100644 index 0000000..757bf43 --- /dev/null +++ b/configs/model_zoo_7gb/point_context_perceiver.toml @@ -0,0 +1,63 @@ +[run] +name = "model_class_frontier_7gb_01_point_context_perceiver" +seed = 20260725 +artifact_dir = "artifacts/current_run/training_runs/point_context_perceiver" + +[data] +root = "data/processed/full" +train_cases = 45 +val_cases = 3 +test_cases = 2 +points_per_case = 999999999 +batch_size = 4096 +source = "local" + +[model] +type = "point_context_perceiver" +hidden_width = 512 +depth = 4 +activation = "gelu" +coordinate_features = ["x", "y", "sdf"] +condition_width = 512 +condition_depth = 3 +condition_dim = 512 +context_points = 512 +latent_width = 512 +attention_depth = 4 + +[optim] +lr = 0.0001 +weight_decay = 0.0001 +steps = 5000 +log_interval = 500 + +[device] +type = "cuda" +allow_cpu_fallback = false +benchmark_kernels = true + +[loss] +type = "normalized_mse" + +[precision] +dtype = "bf16" + +[checkpoint] +interval_seconds = 1800 + +[stability] +max_grad_norm = 1.0 + +[observability] +backend = "wandb" +entity = "zacheryasc-personal" +project = "airfRANS-model-sweep" +group = "model_class_frontier_7gb_01" +tags = ["airfrans", "model-class-frontier", "7gb-subset", "tarball-data", "point_context_perceiver"] + +[huggingface] +enabled = true +repo_id = "zacheryasc/airfrans-frontier-checkpoints" +repo_type = "model" +path_prefix = "model_class_frontier_7gb_01/point_context_perceiver" +private = false diff --git a/configs/model_zoo_7gb/raster_fno_unet.toml b/configs/model_zoo_7gb/raster_fno_unet.toml new file mode 100644 index 0000000..8761256 --- /dev/null +++ b/configs/model_zoo_7gb/raster_fno_unet.toml @@ -0,0 +1,61 @@ +[run] +name = "model_class_frontier_7gb_01_raster_fno_unet" +seed = 20260725 +artifact_dir = "artifacts/current_run/training_runs/raster_fno_unet" + +[data] +root = "data/processed/full" +train_cases = 45 +val_cases = 3 +test_cases = 2 +points_per_case = 999999999 +batch_size = 4096 +source = "local" + +[model] +type = "raster_fno_unet" +hidden_width = 512 +depth = 6 +activation = "gelu" +coordinate_features = ["x", "y", "sdf"] +condition_width = 512 +condition_depth = 3 +condition_dim = 512 +grid_resolution = 128 + +[optim] +lr = 0.0001 +weight_decay = 0.0001 +steps = 5000 +log_interval = 500 + +[device] +type = "cuda" +allow_cpu_fallback = false +benchmark_kernels = true + +[loss] +type = "normalized_mse" + +[precision] +dtype = "bf16" + +[checkpoint] +interval_seconds = 1800 + +[stability] +max_grad_norm = 1.0 + +[observability] +backend = "wandb" +entity = "zacheryasc-personal" +project = "airfRANS-model-sweep" +group = "model_class_frontier_7gb_01" +tags = ["airfrans", "model-class-frontier", "7gb-subset", "tarball-data", "raster_fno_unet"] + +[huggingface] +enabled = true +repo_id = "zacheryasc/airfrans-frontier-checkpoints" +repo_type = "model" +path_prefix = "model_class_frontier_7gb_01/raster_fno_unet" +private = false diff --git a/configs/model_zoo_7gb/siren_conditioned_inr.toml b/configs/model_zoo_7gb/siren_conditioned_inr.toml new file mode 100644 index 0000000..ee97887 --- /dev/null +++ b/configs/model_zoo_7gb/siren_conditioned_inr.toml @@ -0,0 +1,61 @@ +[run] +name = "model_class_frontier_7gb_01_siren_conditioned_inr" +seed = 20260725 +artifact_dir = "artifacts/current_run/training_runs/siren_conditioned_inr" + +[data] +root = "data/processed/full" +train_cases = 45 +val_cases = 3 +test_cases = 2 +points_per_case = 999999999 +batch_size = 4096 +source = "local" + +[model] +type = "siren_conditioned_inr" +hidden_width = 1024 +depth = 6 +activation = "gelu" +coordinate_features = ["x", "y", "sdf"] +condition_width = 512 +condition_depth = 3 +condition_dim = 512 +siren_omega0 = 30.0 + +[optim] +lr = 0.0001 +weight_decay = 0.0001 +steps = 5000 +log_interval = 500 + +[device] +type = "cuda" +allow_cpu_fallback = false +benchmark_kernels = true + +[loss] +type = "normalized_mse" + +[precision] +dtype = "bf16" + +[checkpoint] +interval_seconds = 1800 + +[stability] +max_grad_norm = 1.0 + +[observability] +backend = "wandb" +entity = "zacheryasc-personal" +project = "airfRANS-model-sweep" +group = "model_class_frontier_7gb_01" +tags = ["airfrans", "model-class-frontier", "7gb-subset", "tarball-data", "siren_conditioned_inr"] + +[huggingface] +enabled = true +repo_id = "zacheryasc/airfrans-frontier-checkpoints" +repo_type = "model" +path_prefix = "model_class_frontier_7gb_01/siren_conditioned_inr" +private = false diff --git a/configs/remote_full_70gb.toml b/configs/remote_full_70gb.toml new file mode 100644 index 0000000..f027ab6 --- /dev/null +++ b/configs/remote_full_70gb.toml @@ -0,0 +1,86 @@ +[run] +name = "full_airfrans_incumbent_70gb_01" +timeout_minutes = 1440 +local_artifact_dir = "artifacts/remote_runs" +max_attempts = 5 + +[provider] +kind = "vastai" +disk_gb = 512 +max_price_per_hour = 0.80 +image = "vastai/base:0.0.2" + +[provider.gpu] +name = "RTX 4090" +count = 1 +min_vram_gb = 20 + +[selection] +min_reliability = 0.95 +min_down_mbps = 100 +min_up_mbps = 25 +require_verified = true +blocked_geos = ["CN"] +blacklist_hosts = [59017] +drop_cheap_frac = 0.30 +image_size_gb = 5.0 +base_url = "https://cloud.vast.ai" + +[workspace] +workdir = "." +exclude = [ + "/artifacts", + "/data/raw", + "/.venv", + "/notebooks", + "__pycache__", + "*.pyc", +] + +[bootstrap] +command = """ +uv sync --no-dev +uv run --no-dev python -c "import torch; print('torch_cuda_available=' + str(torch.cuda.is_available())); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(torch.cuda.device_count()))" +""" + +[data] +validation_command = """ +uv run --no-dev airfrans-frontier prepare-public-hf --repo-id zacheryasc/airfrans-processed --path-in-repo processed/full --work-dir artifacts/public_airfrans --output-dir artifacts/data_cache/airfrans_processed/processed/full --min-cases 1000 +uv run --no-dev python -c "from airfrans_frontier.training.config import load_training_config; c=load_training_config('configs/full_airfrans_incumbent_70gb.toml'); assert c.data.source == 'huggingface'; assert c.data.train_cases == 900; assert c.data.val_cases == 50; assert c.data.test_cases == 50; print('data_source=' + c.data.source + ' repo=' + str(c.data.hf_repo_id) + ' split=' + str((c.data.train_cases, c.data.val_cases, c.data.test_cases)))" +""" + +[job] +command = """ +uv run --no-dev remote-run smoke-train configs/full_airfrans_incumbent_70gb.toml --artifact-dir artifacts/current_run --run-id "$AIRFRANS_REMOTE_RUN_ID" +""" +artifact_dir = "artifacts/current_run" +heartbeat_file = "artifacts/current_run/heartbeat.json" +metrics_file = "artifacts/current_run/metrics.jsonl" + +[artifacts] +mode = "object_store_upload" +required = [ + "config.toml", + "metrics.jsonl", + "latest_metrics.json", + "heartbeat.json", + "checkpoint_latest.pt", + "checkpoint_best.pt", + "checkpoint_final.pt", + "final_metrics.json", + "split_manifest.json", + "data_manifest.json", + "normalization.json", + "environment_manifest.json", + "calibration_manifest.json", + "evaluation_protocol.json", + "hf_upload_manifest.json", + "run_manifest.json", + "artifact_manifest.json", + "checksums.txt", + "verification_report.json", +] + +[cleanup] +on_success = "sky_down" +on_failure = "collect_then_keep" diff --git a/configs/remote_model_zoo_7gb/deeponet_branch_trunk.toml b/configs/remote_model_zoo_7gb/deeponet_branch_trunk.toml new file mode 100644 index 0000000..40abc67 --- /dev/null +++ b/configs/remote_model_zoo_7gb/deeponet_branch_trunk.toml @@ -0,0 +1,89 @@ +[run] +name = "model_zoo_7gb_01_deeponet_branch_trunk" +timeout_minutes = 720 +local_artifact_dir = "artifacts/remote_runs" +max_attempts = 2 + +[provider] +kind = "vastai" +disk_gb = 128 +max_price_per_hour = 0.80 +image = "vastai/base:0.0.2" + +[provider.gpu] +name = "RTX 4090" +count = 1 +min_vram_gb = 20 + +[selection] +min_reliability = 0.95 +min_down_mbps = 100 +min_up_mbps = 25 +require_verified = true +blocked_geos = ["CN"] +blacklist_hosts = [59017] +drop_cheap_frac = 0.30 +image_size_gb = 5.0 +base_url = "https://cloud.vast.ai" + +[workspace] +workdir = "." +exclude = [ + "/artifacts", + "/data/raw", + "/data/processed", + "/.venv", + "/notebooks", + "__pycache__", + "*.pyc", +] + +[bootstrap] +command = """ +uv sync --no-dev +uv run --no-dev python -c "import torch; print('torch_cuda_available=' + str(torch.cuda.is_available())); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(torch.cuda.device_count()))" +""" + +[data] +validation_command = """ +mkdir -p data/processed +tar -xzf data/airfrans_processed_full_50cases.tar.gz -C data/processed +uv run --no-dev python -c "from pathlib import Path; files=list(Path('data/processed/full').glob('*.npz')); assert len(files) >= 50; print('tarball_processed_full_cases=' + str(len(files)))" +uv run --no-dev python -c "from airfrans_frontier.training.config import load_training_config; c=load_training_config('configs/model_zoo_7gb/deeponet_branch_trunk.toml'); assert c.data.source == 'local'; assert c.data.train_cases == 45; assert c.data.val_cases == 3; assert c.data.test_cases == 2; print('model_family=' + c.model.type + ' data_source=' + c.data.source)" +""" + +[job] +command = """ +uv run --no-dev remote-run smoke-train configs/model_zoo_7gb/deeponet_branch_trunk.toml --artifact-dir artifacts/current_run --run-id "$AIRFRANS_REMOTE_RUN_ID" +""" +artifact_dir = "artifacts/current_run" +heartbeat_file = "artifacts/current_run/heartbeat.json" +metrics_file = "artifacts/current_run/metrics.jsonl" + +[artifacts] +mode = "object_store_upload" +required = [ + "config.toml", + "metrics.jsonl", + "latest_metrics.json", + "heartbeat.json", + "checkpoint_latest.pt", + "checkpoint_best.pt", + "checkpoint_final.pt", + "final_metrics.json", + "split_manifest.json", + "data_manifest.json", + "normalization.json", + "environment_manifest.json", + "calibration_manifest.json", + "evaluation_protocol.json", + "hf_upload_manifest.json", + "run_manifest.json", + "artifact_manifest.json", + "checksums.txt", + "verification_report.json", +] + +[cleanup] +on_success = "sky_down" +on_failure = "collect_then_keep" diff --git a/configs/remote_model_zoo_7gb/film_fourier_inr.toml b/configs/remote_model_zoo_7gb/film_fourier_inr.toml new file mode 100644 index 0000000..3e76811 --- /dev/null +++ b/configs/remote_model_zoo_7gb/film_fourier_inr.toml @@ -0,0 +1,89 @@ +[run] +name = "model_zoo_7gb_01_film_fourier_inr" +timeout_minutes = 720 +local_artifact_dir = "artifacts/remote_runs" +max_attempts = 2 + +[provider] +kind = "vastai" +disk_gb = 128 +max_price_per_hour = 0.80 +image = "vastai/base:0.0.2" + +[provider.gpu] +name = "RTX 4090" +count = 1 +min_vram_gb = 20 + +[selection] +min_reliability = 0.95 +min_down_mbps = 100 +min_up_mbps = 25 +require_verified = true +blocked_geos = ["CN"] +blacklist_hosts = [59017] +drop_cheap_frac = 0.30 +image_size_gb = 5.0 +base_url = "https://cloud.vast.ai" + +[workspace] +workdir = "." +exclude = [ + "/artifacts", + "/data/raw", + "/data/processed", + "/.venv", + "/notebooks", + "__pycache__", + "*.pyc", +] + +[bootstrap] +command = """ +uv sync --no-dev +uv run --no-dev python -c "import torch; print('torch_cuda_available=' + str(torch.cuda.is_available())); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(torch.cuda.device_count()))" +""" + +[data] +validation_command = """ +mkdir -p data/processed +tar -xzf data/airfrans_processed_full_50cases.tar.gz -C data/processed +uv run --no-dev python -c "from pathlib import Path; files=list(Path('data/processed/full').glob('*.npz')); assert len(files) >= 50; print('tarball_processed_full_cases=' + str(len(files)))" +uv run --no-dev python -c "from airfrans_frontier.training.config import load_training_config; c=load_training_config('configs/model_zoo_7gb/film_fourier_inr.toml'); assert c.data.source == 'local'; assert c.data.train_cases == 45; assert c.data.val_cases == 3; assert c.data.test_cases == 2; print('model_family=' + c.model.type + ' data_source=' + c.data.source)" +""" + +[job] +command = """ +uv run --no-dev remote-run smoke-train configs/model_zoo_7gb/film_fourier_inr.toml --artifact-dir artifacts/current_run --run-id "$AIRFRANS_REMOTE_RUN_ID" +""" +artifact_dir = "artifacts/current_run" +heartbeat_file = "artifacts/current_run/heartbeat.json" +metrics_file = "artifacts/current_run/metrics.jsonl" + +[artifacts] +mode = "object_store_upload" +required = [ + "config.toml", + "metrics.jsonl", + "latest_metrics.json", + "heartbeat.json", + "checkpoint_latest.pt", + "checkpoint_best.pt", + "checkpoint_final.pt", + "final_metrics.json", + "split_manifest.json", + "data_manifest.json", + "normalization.json", + "environment_manifest.json", + "calibration_manifest.json", + "evaluation_protocol.json", + "hf_upload_manifest.json", + "run_manifest.json", + "artifact_manifest.json", + "checksums.txt", + "verification_report.json", +] + +[cleanup] +on_success = "sky_down" +on_failure = "collect_then_keep" diff --git a/configs/remote_model_zoo_7gb/meshgraphnet_or_point_transformer_local.toml b/configs/remote_model_zoo_7gb/meshgraphnet_or_point_transformer_local.toml new file mode 100644 index 0000000..f54e7b1 --- /dev/null +++ b/configs/remote_model_zoo_7gb/meshgraphnet_or_point_transformer_local.toml @@ -0,0 +1,89 @@ +[run] +name = "model_zoo_7gb_01_meshgraphnet_or_point_transformer_local" +timeout_minutes = 720 +local_artifact_dir = "artifacts/remote_runs" +max_attempts = 2 + +[provider] +kind = "vastai" +disk_gb = 128 +max_price_per_hour = 0.80 +image = "vastai/base:0.0.2" + +[provider.gpu] +name = "RTX 4090" +count = 1 +min_vram_gb = 20 + +[selection] +min_reliability = 0.95 +min_down_mbps = 100 +min_up_mbps = 25 +require_verified = true +blocked_geos = ["CN"] +blacklist_hosts = [59017] +drop_cheap_frac = 0.30 +image_size_gb = 5.0 +base_url = "https://cloud.vast.ai" + +[workspace] +workdir = "." +exclude = [ + "/artifacts", + "/data/raw", + "/data/processed", + "/.venv", + "/notebooks", + "__pycache__", + "*.pyc", +] + +[bootstrap] +command = """ +uv sync --no-dev +uv run --no-dev python -c "import torch; print('torch_cuda_available=' + str(torch.cuda.is_available())); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(torch.cuda.device_count()))" +""" + +[data] +validation_command = """ +mkdir -p data/processed +tar -xzf data/airfrans_processed_full_50cases.tar.gz -C data/processed +uv run --no-dev python -c "from pathlib import Path; files=list(Path('data/processed/full').glob('*.npz')); assert len(files) >= 50; print('tarball_processed_full_cases=' + str(len(files)))" +uv run --no-dev python -c "from airfrans_frontier.training.config import load_training_config; c=load_training_config('configs/model_zoo_7gb/meshgraphnet_or_point_transformer_local.toml'); assert c.data.source == 'local'; assert c.data.train_cases == 45; assert c.data.val_cases == 3; assert c.data.test_cases == 2; print('model_family=' + c.model.type + ' data_source=' + c.data.source)" +""" + +[job] +command = """ +uv run --no-dev remote-run smoke-train configs/model_zoo_7gb/meshgraphnet_or_point_transformer_local.toml --artifact-dir artifacts/current_run --run-id "$AIRFRANS_REMOTE_RUN_ID" +""" +artifact_dir = "artifacts/current_run" +heartbeat_file = "artifacts/current_run/heartbeat.json" +metrics_file = "artifacts/current_run/metrics.jsonl" + +[artifacts] +mode = "object_store_upload" +required = [ + "config.toml", + "metrics.jsonl", + "latest_metrics.json", + "heartbeat.json", + "checkpoint_latest.pt", + "checkpoint_best.pt", + "checkpoint_final.pt", + "final_metrics.json", + "split_manifest.json", + "data_manifest.json", + "normalization.json", + "environment_manifest.json", + "calibration_manifest.json", + "evaluation_protocol.json", + "hf_upload_manifest.json", + "run_manifest.json", + "artifact_manifest.json", + "checksums.txt", + "verification_report.json", +] + +[cleanup] +on_success = "sky_down" +on_failure = "collect_then_keep" diff --git a/configs/remote_model_zoo_7gb/nerf_cfd_multires.toml b/configs/remote_model_zoo_7gb/nerf_cfd_multires.toml new file mode 100644 index 0000000..8118085 --- /dev/null +++ b/configs/remote_model_zoo_7gb/nerf_cfd_multires.toml @@ -0,0 +1,89 @@ +[run] +name = "model_zoo_7gb_01_nerf_cfd_multires" +timeout_minutes = 720 +local_artifact_dir = "artifacts/remote_runs" +max_attempts = 2 + +[provider] +kind = "vastai" +disk_gb = 128 +max_price_per_hour = 0.80 +image = "vastai/base:0.0.2" + +[provider.gpu] +name = "RTX 4090" +count = 1 +min_vram_gb = 20 + +[selection] +min_reliability = 0.95 +min_down_mbps = 100 +min_up_mbps = 25 +require_verified = true +blocked_geos = ["CN"] +blacklist_hosts = [59017] +drop_cheap_frac = 0.30 +image_size_gb = 5.0 +base_url = "https://cloud.vast.ai" + +[workspace] +workdir = "." +exclude = [ + "/artifacts", + "/data/raw", + "/data/processed", + "/.venv", + "/notebooks", + "__pycache__", + "*.pyc", +] + +[bootstrap] +command = """ +uv sync --no-dev +uv run --no-dev python -c "import torch; print('torch_cuda_available=' + str(torch.cuda.is_available())); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(torch.cuda.device_count()))" +""" + +[data] +validation_command = """ +mkdir -p data/processed +tar -xzf data/airfrans_processed_full_50cases.tar.gz -C data/processed +uv run --no-dev python -c "from pathlib import Path; files=list(Path('data/processed/full').glob('*.npz')); assert len(files) >= 50; print('tarball_processed_full_cases=' + str(len(files)))" +uv run --no-dev python -c "from airfrans_frontier.training.config import load_training_config; c=load_training_config('configs/model_zoo_7gb/nerf_cfd_multires.toml'); assert c.data.source == 'local'; assert c.data.train_cases == 45; assert c.data.val_cases == 3; assert c.data.test_cases == 2; print('model_family=' + c.model.type + ' data_source=' + c.data.source)" +""" + +[job] +command = """ +uv run --no-dev remote-run smoke-train configs/model_zoo_7gb/nerf_cfd_multires.toml --artifact-dir artifacts/current_run --run-id "$AIRFRANS_REMOTE_RUN_ID" +""" +artifact_dir = "artifacts/current_run" +heartbeat_file = "artifacts/current_run/heartbeat.json" +metrics_file = "artifacts/current_run/metrics.jsonl" + +[artifacts] +mode = "object_store_upload" +required = [ + "config.toml", + "metrics.jsonl", + "latest_metrics.json", + "heartbeat.json", + "checkpoint_latest.pt", + "checkpoint_best.pt", + "checkpoint_final.pt", + "final_metrics.json", + "split_manifest.json", + "data_manifest.json", + "normalization.json", + "environment_manifest.json", + "calibration_manifest.json", + "evaluation_protocol.json", + "hf_upload_manifest.json", + "run_manifest.json", + "artifact_manifest.json", + "checksums.txt", + "verification_report.json", +] + +[cleanup] +on_success = "sky_down" +on_failure = "collect_then_keep" diff --git a/configs/remote_model_zoo_7gb/point_context_perceiver.toml b/configs/remote_model_zoo_7gb/point_context_perceiver.toml new file mode 100644 index 0000000..26e3055 --- /dev/null +++ b/configs/remote_model_zoo_7gb/point_context_perceiver.toml @@ -0,0 +1,89 @@ +[run] +name = "model_zoo_7gb_01_point_context_perceiver" +timeout_minutes = 720 +local_artifact_dir = "artifacts/remote_runs" +max_attempts = 2 + +[provider] +kind = "vastai" +disk_gb = 128 +max_price_per_hour = 0.80 +image = "vastai/base:0.0.2" + +[provider.gpu] +name = "RTX 4090" +count = 1 +min_vram_gb = 20 + +[selection] +min_reliability = 0.95 +min_down_mbps = 100 +min_up_mbps = 25 +require_verified = true +blocked_geos = ["CN"] +blacklist_hosts = [59017] +drop_cheap_frac = 0.30 +image_size_gb = 5.0 +base_url = "https://cloud.vast.ai" + +[workspace] +workdir = "." +exclude = [ + "/artifacts", + "/data/raw", + "/data/processed", + "/.venv", + "/notebooks", + "__pycache__", + "*.pyc", +] + +[bootstrap] +command = """ +uv sync --no-dev +uv run --no-dev python -c "import torch; print('torch_cuda_available=' + str(torch.cuda.is_available())); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(torch.cuda.device_count()))" +""" + +[data] +validation_command = """ +mkdir -p data/processed +tar -xzf data/airfrans_processed_full_50cases.tar.gz -C data/processed +uv run --no-dev python -c "from pathlib import Path; files=list(Path('data/processed/full').glob('*.npz')); assert len(files) >= 50; print('tarball_processed_full_cases=' + str(len(files)))" +uv run --no-dev python -c "from airfrans_frontier.training.config import load_training_config; c=load_training_config('configs/model_zoo_7gb/point_context_perceiver.toml'); assert c.data.source == 'local'; assert c.data.train_cases == 45; assert c.data.val_cases == 3; assert c.data.test_cases == 2; print('model_family=' + c.model.type + ' data_source=' + c.data.source)" +""" + +[job] +command = """ +uv run --no-dev remote-run smoke-train configs/model_zoo_7gb/point_context_perceiver.toml --artifact-dir artifacts/current_run --run-id "$AIRFRANS_REMOTE_RUN_ID" +""" +artifact_dir = "artifacts/current_run" +heartbeat_file = "artifacts/current_run/heartbeat.json" +metrics_file = "artifacts/current_run/metrics.jsonl" + +[artifacts] +mode = "object_store_upload" +required = [ + "config.toml", + "metrics.jsonl", + "latest_metrics.json", + "heartbeat.json", + "checkpoint_latest.pt", + "checkpoint_best.pt", + "checkpoint_final.pt", + "final_metrics.json", + "split_manifest.json", + "data_manifest.json", + "normalization.json", + "environment_manifest.json", + "calibration_manifest.json", + "evaluation_protocol.json", + "hf_upload_manifest.json", + "run_manifest.json", + "artifact_manifest.json", + "checksums.txt", + "verification_report.json", +] + +[cleanup] +on_success = "sky_down" +on_failure = "collect_then_keep" diff --git a/configs/remote_model_zoo_7gb/raster_fno_unet.toml b/configs/remote_model_zoo_7gb/raster_fno_unet.toml new file mode 100644 index 0000000..af814ef --- /dev/null +++ b/configs/remote_model_zoo_7gb/raster_fno_unet.toml @@ -0,0 +1,89 @@ +[run] +name = "model_zoo_7gb_01_raster_fno_unet" +timeout_minutes = 720 +local_artifact_dir = "artifacts/remote_runs" +max_attempts = 2 + +[provider] +kind = "vastai" +disk_gb = 128 +max_price_per_hour = 0.80 +image = "vastai/base:0.0.2" + +[provider.gpu] +name = "RTX 4090" +count = 1 +min_vram_gb = 20 + +[selection] +min_reliability = 0.95 +min_down_mbps = 100 +min_up_mbps = 25 +require_verified = true +blocked_geos = ["CN"] +blacklist_hosts = [59017] +drop_cheap_frac = 0.30 +image_size_gb = 5.0 +base_url = "https://cloud.vast.ai" + +[workspace] +workdir = "." +exclude = [ + "/artifacts", + "/data/raw", + "/data/processed", + "/.venv", + "/notebooks", + "__pycache__", + "*.pyc", +] + +[bootstrap] +command = """ +uv sync --no-dev +uv run --no-dev python -c "import torch; print('torch_cuda_available=' + str(torch.cuda.is_available())); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(torch.cuda.device_count()))" +""" + +[data] +validation_command = """ +mkdir -p data/processed +tar -xzf data/airfrans_processed_full_50cases.tar.gz -C data/processed +uv run --no-dev python -c "from pathlib import Path; files=list(Path('data/processed/full').glob('*.npz')); assert len(files) >= 50; print('tarball_processed_full_cases=' + str(len(files)))" +uv run --no-dev python -c "from airfrans_frontier.training.config import load_training_config; c=load_training_config('configs/model_zoo_7gb/raster_fno_unet.toml'); assert c.data.source == 'local'; assert c.data.train_cases == 45; assert c.data.val_cases == 3; assert c.data.test_cases == 2; print('model_family=' + c.model.type + ' data_source=' + c.data.source)" +""" + +[job] +command = """ +uv run --no-dev remote-run smoke-train configs/model_zoo_7gb/raster_fno_unet.toml --artifact-dir artifacts/current_run --run-id "$AIRFRANS_REMOTE_RUN_ID" +""" +artifact_dir = "artifacts/current_run" +heartbeat_file = "artifacts/current_run/heartbeat.json" +metrics_file = "artifacts/current_run/metrics.jsonl" + +[artifacts] +mode = "object_store_upload" +required = [ + "config.toml", + "metrics.jsonl", + "latest_metrics.json", + "heartbeat.json", + "checkpoint_latest.pt", + "checkpoint_best.pt", + "checkpoint_final.pt", + "final_metrics.json", + "split_manifest.json", + "data_manifest.json", + "normalization.json", + "environment_manifest.json", + "calibration_manifest.json", + "evaluation_protocol.json", + "hf_upload_manifest.json", + "run_manifest.json", + "artifact_manifest.json", + "checksums.txt", + "verification_report.json", +] + +[cleanup] +on_success = "sky_down" +on_failure = "collect_then_keep" diff --git a/configs/remote_model_zoo_7gb/siren_conditioned_inr.toml b/configs/remote_model_zoo_7gb/siren_conditioned_inr.toml new file mode 100644 index 0000000..3d6a474 --- /dev/null +++ b/configs/remote_model_zoo_7gb/siren_conditioned_inr.toml @@ -0,0 +1,89 @@ +[run] +name = "model_zoo_7gb_01_siren_conditioned_inr" +timeout_minutes = 720 +local_artifact_dir = "artifacts/remote_runs" +max_attempts = 2 + +[provider] +kind = "vastai" +disk_gb = 128 +max_price_per_hour = 0.80 +image = "vastai/base:0.0.2" + +[provider.gpu] +name = "RTX 4090" +count = 1 +min_vram_gb = 20 + +[selection] +min_reliability = 0.95 +min_down_mbps = 100 +min_up_mbps = 25 +require_verified = true +blocked_geos = ["CN"] +blacklist_hosts = [59017] +drop_cheap_frac = 0.30 +image_size_gb = 5.0 +base_url = "https://cloud.vast.ai" + +[workspace] +workdir = "." +exclude = [ + "/artifacts", + "/data/raw", + "/data/processed", + "/.venv", + "/notebooks", + "__pycache__", + "*.pyc", +] + +[bootstrap] +command = """ +uv sync --no-dev +uv run --no-dev python -c "import torch; print('torch_cuda_available=' + str(torch.cuda.is_available())); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(torch.cuda.device_count()))" +""" + +[data] +validation_command = """ +mkdir -p data/processed +tar -xzf data/airfrans_processed_full_50cases.tar.gz -C data/processed +uv run --no-dev python -c "from pathlib import Path; files=list(Path('data/processed/full').glob('*.npz')); assert len(files) >= 50; print('tarball_processed_full_cases=' + str(len(files)))" +uv run --no-dev python -c "from airfrans_frontier.training.config import load_training_config; c=load_training_config('configs/model_zoo_7gb/siren_conditioned_inr.toml'); assert c.data.source == 'local'; assert c.data.train_cases == 45; assert c.data.val_cases == 3; assert c.data.test_cases == 2; print('model_family=' + c.model.type + ' data_source=' + c.data.source)" +""" + +[job] +command = """ +uv run --no-dev remote-run smoke-train configs/model_zoo_7gb/siren_conditioned_inr.toml --artifact-dir artifacts/current_run --run-id "$AIRFRANS_REMOTE_RUN_ID" +""" +artifact_dir = "artifacts/current_run" +heartbeat_file = "artifacts/current_run/heartbeat.json" +metrics_file = "artifacts/current_run/metrics.jsonl" + +[artifacts] +mode = "object_store_upload" +required = [ + "config.toml", + "metrics.jsonl", + "latest_metrics.json", + "heartbeat.json", + "checkpoint_latest.pt", + "checkpoint_best.pt", + "checkpoint_final.pt", + "final_metrics.json", + "split_manifest.json", + "data_manifest.json", + "normalization.json", + "environment_manifest.json", + "calibration_manifest.json", + "evaluation_protocol.json", + "hf_upload_manifest.json", + "run_manifest.json", + "artifact_manifest.json", + "checksums.txt", + "verification_report.json", +] + +[cleanup] +on_success = "sky_down" +on_failure = "collect_then_keep" diff --git a/configs/remote_smoke.toml b/configs/remote_smoke.toml index 6d825de..40c9373 100644 --- a/configs/remote_smoke.toml +++ b/configs/remote_smoke.toml @@ -1,5 +1,5 @@ [run] -name = "airfrans-aggressive-smoke" +name = "model_class_frontier_7gb_01_film_fourier_inr" timeout_minutes = 360 local_artifact_dir = "artifacts/remote_runs" max_attempts = 2 @@ -45,7 +45,7 @@ uv run --no-dev python -c "import torch; assert torch.cuda.is_available(); print [data] validation_command = """ -uv run --no-dev python -c "from pathlib import Path; files=sorted(Path('data/processed/full').glob('*.npz')); assert len(files) >= 50; print(f'processed_full_cases={len(files)}')" +uv run --no-dev python -c "from airfrans_frontier.training.config import load_training_config; c=load_training_config('configs/aggressive_smoke.toml'); assert c.data.source == 'huggingface'; print('data_source=' + c.data.source + ' repo=' + str(c.data.hf_repo_id))" """ [job] @@ -57,7 +57,7 @@ heartbeat_file = "artifacts/current_run/heartbeat.json" metrics_file = "artifacts/current_run/metrics.jsonl" [artifacts] -mode = "rsync" +mode = "object_store_upload" required = [ "config.toml", "metrics.jsonl", @@ -70,10 +70,14 @@ required = [ "split_manifest.json", "data_manifest.json", "normalization.json", - "run_manifest.json", "environment_manifest.json", + "calibration_manifest.json", + "evaluation_protocol.json", + "hf_upload_manifest.json", + "run_manifest.json", "artifact_manifest.json", "checksums.txt", + "verification_report.json", ] [cleanup] diff --git a/pyproject.toml b/pyproject.toml index 638418b..f40b930 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,3 +31,7 @@ dev = [ "skypilot[vast]>=0.12.3.post1", "pytest>=9.1.1", ] + + +[tool.pytest.ini_options] +testpaths = ["tests"] \ No newline at end of file diff --git a/src/airfrans_frontier/cli.py b/src/airfrans_frontier/cli.py index 35a6a8b..84c9bb4 100644 --- a/src/airfrans_frontier/cli.py +++ b/src/airfrans_frontier/cli.py @@ -1,6 +1,7 @@ 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 @@ -24,11 +25,40 @@ def build_parser() -> argparse.ArgumentParser: 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 @@ -74,6 +104,48 @@ def main(argv: list[str] | None = None) -> int: 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 @@ -90,6 +162,30 @@ def main(argv: list[str] | None = None) -> int: 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 diff --git a/src/airfrans_frontier/models/__init__.py b/src/airfrans_frontier/models/__init__.py index 9da57ab..1e99d11 100644 --- a/src/airfrans_frontier/models/__init__.py +++ b/src/airfrans_frontier/models/__init__.py @@ -1,6 +1,23 @@ """Baseline model definitions.""" from airfrans_frontier.models.film import FourierFiLMMLP +from airfrans_frontier.models.frontier import ( + DeepONetBranchTrunk, + LocalPointTransformer, + NeRFCFDMultiRes, + PointContextPerceiver, + RasterFNOUNet, + SirenConditionedINR, +) from airfrans_frontier.models.mlp import PointwiseMLP -__all__ = ["FourierFiLMMLP", "PointwiseMLP"] +__all__ = [ + "DeepONetBranchTrunk", + "FourierFiLMMLP", + "LocalPointTransformer", + "NeRFCFDMultiRes", + "PointContextPerceiver", + "PointwiseMLP", + "RasterFNOUNet", + "SirenConditionedINR", +] diff --git a/src/airfrans_frontier/models/frontier.py b/src/airfrans_frontier/models/frontier.py new file mode 100644 index 0000000..2b3f5c5 --- /dev/null +++ b/src/airfrans_frontier/models/frontier.py @@ -0,0 +1,404 @@ +from __future__ import annotations + +import math +from collections.abc import Sequence + +import torch +from torch import nn +from torch.nn import functional as F + + +class NeRFCFDMultiRes(nn.Module): + def __init__( + self, + *, + feature_names: Sequence[str], + output_dim: int, + coordinate_features: Sequence[str], + encoding_levels: int, + hidden_width: int, + depth: int, + condition_width: int, + condition_depth: int, + activation: str, + ) -> None: + super().__init__() + coordinate_indices = _indices(feature_names, coordinate_features) + condition_indices = _complement_indices(len(feature_names), coordinate_indices) + self.register_buffer("coordinate_indices", torch.tensor(coordinate_indices, dtype=torch.long), persistent=False) + self.register_buffer("condition_indices", torch.tensor(condition_indices, dtype=torch.long), persistent=False) + self.encoding_levels = int(encoding_levels) + encoded_dim = len(coordinate_indices) * (1 + 2 * self.encoding_levels) + condition_input_dim = len(condition_indices) if condition_indices else 1 + self.condition_encoder = _mlp( + input_dim=condition_input_dim, + hidden_width=condition_width, + output_dim=condition_width, + depth=condition_depth, + activation=activation, + ) + self.decoder = _mlp( + input_dim=encoded_dim + condition_width, + hidden_width=hidden_width, + output_dim=output_dim, + depth=depth, + activation=activation, + activate_output=False, + ) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + coordinates = features.index_select(dim=1, index=self.coordinate_indices) + condition = _gather_or_zeros(features, self.condition_indices) + encoded = _multires_encode(coordinates, self.encoding_levels) + condition_embedding = self.condition_encoder(condition) + return self.decoder(torch.cat((encoded, condition_embedding), dim=1)) + + +class DeepONetBranchTrunk(nn.Module): + def __init__( + self, + *, + feature_names: Sequence[str], + output_dim: int, + coordinate_features: Sequence[str], + fourier_scales: Sequence[float], + hidden_width: int, + depth: int, + condition_width: int, + condition_depth: int, + activation: str, + ) -> None: + super().__init__() + coordinate_indices = _indices(feature_names, coordinate_features) + condition_indices = _complement_indices(len(feature_names), coordinate_indices) + self.register_buffer("coordinate_indices", torch.tensor(coordinate_indices, dtype=torch.long), persistent=False) + self.register_buffer("condition_indices", torch.tensor(condition_indices, dtype=torch.long), persistent=False) + self.register_buffer("fourier_scales", torch.tensor(tuple(float(scale) for scale in fourier_scales), dtype=torch.float32), persistent=False) + trunk_input_dim = len(coordinate_indices) * (1 + 2 * len(fourier_scales)) + condition_input_dim = len(condition_indices) if condition_indices else 1 + self.branch = _mlp( + input_dim=condition_input_dim, + hidden_width=condition_width, + output_dim=hidden_width, + depth=condition_depth, + activation=activation, + ) + self.trunk = _mlp( + input_dim=trunk_input_dim, + hidden_width=hidden_width, + output_dim=hidden_width, + depth=depth, + activation=activation, + ) + self.head = nn.Linear(hidden_width, output_dim) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + coordinates = features.index_select(dim=1, index=self.coordinate_indices) + condition = _gather_or_zeros(features, self.condition_indices) + trunk = self.trunk(_fourier_features(coordinates, self.fourier_scales)) + branch = self.branch(condition) + return self.head(trunk * branch) + + +class PointContextPerceiver(nn.Module): + def __init__( + self, + *, + input_dim: int, + output_dim: int, + hidden_width: int, + latent_width: int, + context_points: int, + attention_depth: int, + activation: str, + ) -> None: + super().__init__() + self.context_tokens = nn.Parameter(torch.empty(context_points, latent_width)) + nn.init.normal_(self.context_tokens, std=latent_width ** -0.5) + self.input_projection = nn.Linear(input_dim, latent_width) + self.blocks = nn.ModuleList( + [_PerceiverPointBlock(latent_width=latent_width, hidden_width=hidden_width, activation=activation) for _ in range(attention_depth)] + ) + self.head = _mlp( + input_dim=latent_width, + hidden_width=hidden_width, + output_dim=output_dim, + depth=max(1, attention_depth), + activation=activation, + activate_output=False, + ) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + hidden = self.input_projection(features) + tokens = self.context_tokens.to(dtype=hidden.dtype, device=hidden.device) + for block in self.blocks: + hidden = block(hidden, tokens) + return self.head(hidden) + + +class LocalPointTransformer(nn.Module): + def __init__( + self, + *, + feature_names: Sequence[str], + output_dim: int, + coordinate_features: Sequence[str], + hidden_width: int, + depth: int, + neighbors: int, + activation: str, + ) -> None: + super().__init__() + coordinate_indices = _indices(feature_names, coordinate_features) + self.register_buffer("coordinate_indices", torch.tensor(coordinate_indices, dtype=torch.long), persistent=False) + self.neighbors = int(neighbors) + self.input_projection = nn.Linear(len(feature_names), hidden_width) + self.blocks = nn.ModuleList([_LocalPointBlock(hidden_width=hidden_width, activation=activation) for _ in range(depth)]) + self.head = nn.Linear(hidden_width, output_dim) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + coordinates = features.index_select(dim=1, index=self.coordinate_indices) + hidden = self.input_projection(features) + for block in self.blocks: + hidden = block(hidden, coordinates, self.neighbors) + return self.head(hidden) + + +class RasterFNOUNet(nn.Module): + def __init__( + self, + *, + feature_names: Sequence[str], + output_dim: int, + coordinate_features: Sequence[str], + grid_resolution: int, + hidden_width: int, + depth: int, + condition_width: int, + condition_depth: int, + activation: str, + ) -> None: + super().__init__() + coordinate_indices = _indices(feature_names, coordinate_features[:2]) + condition_indices = _complement_indices(len(feature_names), coordinate_indices) + self.register_buffer("coordinate_indices", torch.tensor(coordinate_indices, dtype=torch.long), persistent=False) + self.register_buffer("condition_indices", torch.tensor(condition_indices, dtype=torch.long), persistent=False) + self.grid_resolution = int(grid_resolution) + self.grid = nn.Parameter(torch.empty(self.grid_resolution, self.grid_resolution, hidden_width)) + nn.init.normal_(self.grid, std=hidden_width ** -0.5) + condition_input_dim = len(condition_indices) if condition_indices else 1 + self.condition_encoder = _mlp( + input_dim=condition_input_dim, + hidden_width=condition_width, + output_dim=condition_width, + depth=condition_depth, + activation=activation, + ) + self.decoder = _mlp( + input_dim=hidden_width + condition_width, + hidden_width=hidden_width, + output_dim=output_dim, + depth=depth, + activation=activation, + activate_output=False, + ) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + xy = features.index_select(dim=1, index=self.coordinate_indices) + sampled = _sample_grid(self.grid, xy) + condition = self.condition_encoder(_gather_or_zeros(features, self.condition_indices)) + return self.decoder(torch.cat((sampled, condition), dim=1)) + + +class SirenConditionedINR(nn.Module): + def __init__( + self, + *, + feature_names: Sequence[str], + output_dim: int, + coordinate_features: Sequence[str], + hidden_width: int, + depth: int, + condition_width: int, + condition_depth: int, + omega0: float, + activation: str, + ) -> None: + super().__init__() + coordinate_indices = _indices(feature_names, coordinate_features) + condition_indices = _complement_indices(len(feature_names), coordinate_indices) + self.register_buffer("coordinate_indices", torch.tensor(coordinate_indices, dtype=torch.long), persistent=False) + self.register_buffer("condition_indices", torch.tensor(condition_indices, dtype=torch.long), persistent=False) + condition_input_dim = len(condition_indices) if condition_indices else 1 + self.condition_encoder = _mlp( + input_dim=condition_input_dim, + hidden_width=condition_width, + output_dim=condition_width, + depth=condition_depth, + activation=activation, + ) + layers: list[nn.Module] = [] + input_dim = len(coordinate_indices) + condition_width + for layer_index in range(depth): + layers.append(_SineLayer(input_dim if layer_index == 0 else hidden_width, hidden_width, omega0=omega0, first=layer_index == 0)) + self.net = nn.Sequential(*layers) + self.head = nn.Linear(hidden_width, output_dim) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + coordinates = features.index_select(dim=1, index=self.coordinate_indices) + condition = self.condition_encoder(_gather_or_zeros(features, self.condition_indices)) + hidden = self.net(torch.cat((coordinates, condition), dim=1)) + return self.head(hidden) + + +class _PerceiverPointBlock(nn.Module): + def __init__(self, *, latent_width: int, hidden_width: int, activation: str) -> None: + super().__init__() + self.norm = nn.LayerNorm(latent_width) + self.ffn = _mlp( + input_dim=latent_width, + hidden_width=hidden_width, + output_dim=latent_width, + depth=2, + activation=activation, + activate_output=False, + ) + + def forward(self, hidden: torch.Tensor, tokens: torch.Tensor) -> torch.Tensor: + scale = hidden.shape[1] ** -0.5 + attention = torch.softmax(hidden @ tokens.T * scale, dim=1) + context = attention @ tokens + return hidden + self.ffn(self.norm(hidden + context)) + + +class _LocalPointBlock(nn.Module): + def __init__(self, *, hidden_width: int, activation: str) -> None: + super().__init__() + self.norm = nn.LayerNorm(hidden_width) + self.update = _mlp( + input_dim=hidden_width * 2, + hidden_width=hidden_width, + output_dim=hidden_width, + depth=2, + activation=activation, + activate_output=False, + ) + + def forward(self, hidden: torch.Tensor, coordinates: torch.Tensor, neighbors: int) -> torch.Tensor: + if hidden.shape[0] <= 1 or neighbors <= 0: + neighborhood = hidden + else: + k = min(neighbors + 1, hidden.shape[0]) + distances = torch.cdist(coordinates.float(), coordinates.float()) + indices = distances.topk(k=k, largest=False).indices[:, 1:] if k > 1 else distances.topk(k=k, largest=False).indices + gathered = hidden.index_select(dim=0, index=indices.reshape(-1)).reshape(hidden.shape[0], -1, hidden.shape[1]) + neighborhood = gathered.mean(dim=1) + return hidden + self.update(torch.cat((self.norm(hidden), neighborhood), dim=1)) + + +class _SineLayer(nn.Module): + def __init__(self, input_dim: int, output_dim: int, *, omega0: float, first: bool) -> None: + super().__init__() + self.linear = nn.Linear(input_dim, output_dim) + self.omega0 = float(omega0) + with torch.no_grad(): + bound = 1.0 / input_dim if first else math.sqrt(6.0 / input_dim) / self.omega0 + self.linear.weight.uniform_(-bound, bound) + + def forward(self, values: torch.Tensor) -> torch.Tensor: + return torch.sin(self.omega0 * self.linear(values)) + + +def _indices(feature_names: Sequence[str], selected_names: Sequence[str]) -> tuple[int, ...]: + indices: list[int] = [] + for name in selected_names: + try: + indices.append(tuple(feature_names).index(name)) + except ValueError as exc: + raise ValueError(f"Coordinate feature {name!r} is not present in dataset features") from exc + if not indices: + raise ValueError("At least one coordinate feature is required") + return tuple(indices) + + +def _complement_indices(size: int, excluded: Sequence[int]) -> tuple[int, ...]: + excluded_set = set(excluded) + return tuple(index for index in range(size) if index not in excluded_set) + + +def _gather_or_zeros(features: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: + if indices.numel() == 0: + return features.new_zeros((features.shape[0], 1)) + return features.index_select(dim=1, index=indices) + + +def _multires_encode(coordinates: torch.Tensor, levels: int) -> torch.Tensor: + if levels <= 0: + return coordinates + pieces = [coordinates] + for level in range(levels): + scale = float(2**level) * math.pi + pieces.append(torch.sin(coordinates * scale)) + pieces.append(torch.cos(coordinates * scale)) + return torch.cat(pieces, dim=1) + + +def _fourier_features(coordinates: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + if scales.numel() == 0: + return coordinates + phases = coordinates.unsqueeze(-1) * scales.to(device=coordinates.device, dtype=coordinates.dtype) * math.pi + return torch.cat((coordinates, torch.sin(phases).flatten(1), torch.cos(phases).flatten(1)), dim=1) + + +def _sample_grid(grid: torch.Tensor, xy: torch.Tensor) -> torch.Tensor: + resolution = grid.shape[0] + if xy.shape[1] < 2: + raise ValueError("Raster model requires at least x and y coordinate features") + scaled = ((xy[:, :2].clamp(-1.0, 1.0) + 1.0) * 0.5) * float(resolution - 1) + x = scaled[:, 0] + y = scaled[:, 1] + x0 = torch.floor(x).long().clamp(0, resolution - 1) + y0 = torch.floor(y).long().clamp(0, resolution - 1) + x1 = (x0 + 1).clamp(0, resolution - 1) + y1 = (y0 + 1).clamp(0, resolution - 1) + wx = (x - x0.to(x.dtype)).unsqueeze(1) + wy = (y - y0.to(y.dtype)).unsqueeze(1) + g00 = grid[y0, x0] + g10 = grid[y0, x1] + g01 = grid[y1, x0] + g11 = grid[y1, x1] + return (1 - wx) * (1 - wy) * g00 + wx * (1 - wy) * g10 + (1 - wx) * wy * g01 + wx * wy * g11 + + +def _mlp( + *, + input_dim: int, + hidden_width: int, + output_dim: int, + depth: int, + activation: str, + activate_output: bool = True, +) -> nn.Sequential: + layers: list[nn.Module] = [] + current_dim = input_dim + for _ in range(max(depth - 1, 0)): + layers.append(nn.Linear(current_dim, hidden_width)) + layers.append(_activation(activation)) + current_dim = hidden_width + layers.append(nn.Linear(current_dim, output_dim)) + if activate_output: + layers.append(_activation(activation)) + return nn.Sequential(*layers) + + +def _activation(name: str) -> nn.Module: + normalized = name.lower() + if normalized == "gelu": + return nn.GELU() + if normalized == "relu": + return nn.ReLU() + if normalized == "silu": + return nn.SiLU() + if normalized == "tanh": + return nn.Tanh() + raise ValueError(f"Unsupported activation: {name}") diff --git a/src/airfrans_frontier/raw/process.py b/src/airfrans_frontier/raw/process.py index ebc9269..bcce47a 100644 --- a/src/airfrans_frontier/raw/process.py +++ b/src/airfrans_frontier/raw/process.py @@ -74,6 +74,7 @@ def process_raw_dataset( *, limit: int | None = None, force: bool = False, + progress_every: int | None = None, ) -> ProcessingResult: raw_root = Path(raw_dir).expanduser() if not raw_root.is_dir(): @@ -90,7 +91,7 @@ def process_raw_dataset( records: list[dict[str, object]] = [] started = time.perf_counter() total_points = 0 - for case_dir in case_dirs: + for index, case_dir in enumerate(case_dirs, start=1): target_path = out_root / f"{case_dir.name}.npz" if target_path.exists() and not force: with np.load(target_path, allow_pickle=False) as npz: @@ -110,6 +111,8 @@ def process_raw_dataset( points = int(features.shape[0]) total_points += points records.append({"case_id": case_dir.name, "path": str(target_path), "points": points, "metadata": _metadata_json(metadata)}) + if progress_every is not None and progress_every > 0 and (index % progress_every == 0 or index == len(case_dirs)): + print(f"processed_airfrans_cases={index}/{len(case_dirs)} total_points={total_points}", flush=True) manifest = { "raw_dir": str(raw_root), diff --git a/src/airfrans_frontier/raw/public.py b/src/airfrans_frontier/raw/public.py new file mode 100644 index 0000000..c54fb67 --- /dev/null +++ b/src/airfrans_frontier/raw/public.py @@ -0,0 +1,253 @@ +from __future__ import annotations + +import json +import os +import shutil +import time +import urllib.error +import urllib.request +import zipfile +from pathlib import Path +from typing import Any + +from airfrans_frontier.training.data_sources import publish_processed_dataset + +PUBLIC_OF_DATASET_URL = "https://data.isir.upmc.fr/extrality/NeurIPS_2022/OF_dataset.zip" +DEFAULT_PUBLIC_WORK_DIR = Path("artifacts/public_airfrans") +DEFAULT_PUBLIC_OUTPUT_DIR = Path("artifacts/data_cache/airfrans_processed/processed/full") + +def ensure_public_airfrans_processed_hf( + *, + repo_id: str, + path_in_repo: str = "processed/full", + work_dir: str | Path = DEFAULT_PUBLIC_WORK_DIR, + output_dir: str | Path = DEFAULT_PUBLIC_OUTPUT_DIR, + source_url: str = PUBLIC_OF_DATASET_URL, + min_cases: int = 1000, + private: bool = False, + force: bool = False, +) -> dict[str, Any]: + if min_cases <= 0: + raise ValueError("min_cases must be positive") + prefix = path_in_repo.strip("/") + started = time.time() + existing = _hf_dataset_status(repo_id=repo_id, path_in_repo=prefix) + if not force and existing["npz_file_count"] >= min_cases and existing["has_manifest"]: + return { + "ok": True, + "phase": "already_published", + "repo_id": repo_id, + "repo_type": "dataset", + "repo_url": f"https://huggingface.co/datasets/{repo_id}", + "path_in_repo": prefix, + "min_cases": min_cases, + "elapsed_seconds": time.time() - started, + **existing, + } + + work_root = Path(work_dir).expanduser() + output_root = Path(output_dir).expanduser() + work_root.mkdir(parents=True, exist_ok=True) + output_root.mkdir(parents=True, exist_ok=True) + + archive_path = work_root / "OF_dataset.zip" + download = download_file(source_url, archive_path) + extract_root = work_root / "raw" + raw_root = extract_of_dataset(archive_path, extract_root, min_cases=min_cases) + from airfrans_frontier.raw.process import process_raw_dataset + + print(f"process_airfrans_raw raw_root={raw_root} output_dir={output_root}", flush=True) + processed = process_raw_dataset(raw_root, output_root, force=force, progress_every=25) + if processed.case_count < min_cases: + raise RuntimeError(f"Processed only {processed.case_count} cases from public AirfRANS archive; expected at least {min_cases}") + print(f"publish_airfrans_processed_hf repo={repo_id} path_in_repo={prefix}", flush=True) + publish = publish_processed_dataset( + data_root=output_root, + repo_id=repo_id, + path_in_repo=prefix, + private=private, + manifest_out=output_root / "hf_dataset_manifest.json", + ) + final = _hf_dataset_status(repo_id=repo_id, path_in_repo=prefix) + if final["npz_file_count"] < min_cases: + raise RuntimeError(f"Published dataset has {final['npz_file_count']} .npz files under {prefix}; expected at least {min_cases}") + if not final["has_manifest"]: + raise RuntimeError(f"Published dataset is missing hf_dataset_manifest.json under {prefix}") + return { + "ok": True, + "phase": "published", + "repo_id": repo_id, + "repo_type": "dataset", + "repo_url": f"https://huggingface.co/datasets/{repo_id}", + "path_in_repo": prefix, + "source_url": source_url, + "archive_path": str(archive_path), + "archive_bytes": archive_path.stat().st_size, + "raw_root": str(raw_root), + "output_dir": str(output_root), + "processed_case_count": processed.case_count, + "processed_total_points": processed.total_points, + "processed_manifest_path": str(processed.manifest_path), + "download": download, + "publish": publish, + "elapsed_seconds": time.time() - started, + **final, + } + + +def download_file(url: str, destination: str | Path, *, chunk_size: int = 16 * 1024 * 1024) -> dict[str, Any]: + path = Path(destination).expanduser() + path.parent.mkdir(parents=True, exist_ok=True) + expected_size = _remote_content_length(url) + existing_size = path.stat().st_size if path.exists() else 0 + if expected_size is not None and existing_size == expected_size: + return {"url": url, "path": str(path), "bytes": existing_size, "resumed": False, "skipped": True} + + headers: dict[str, str] = {} + mode = "wb" + resumed = False + if expected_size is not None and 0 < existing_size < expected_size: + headers["Range"] = f"bytes={existing_size}-" + mode = "ab" + resumed = True + + print( + f"download_airfrans_zip url={url} path={path} existing_bytes={existing_size} expected_bytes={expected_size}", + flush=True, + ) + request = urllib.request.Request(url, headers=headers) + try: + response = urllib.request.urlopen(request, timeout=60) + except urllib.error.HTTPError as exc: + if exc.code == 416 and expected_size is not None and existing_size >= expected_size: + return {"url": url, "path": str(path), "bytes": existing_size, "resumed": False, "skipped": True} + raise + with response: + if resumed and getattr(response, "status", None) != 206: + mode = "wb" + resumed = False + existing_size = 0 + written = existing_size + next_report = ((written // 1_000_000_000) + 1) * 1_000_000_000 + with path.open(mode) as handle: + while True: + chunk = response.read(chunk_size) + if not chunk: + break + handle.write(chunk) + written += len(chunk) + if written >= next_report: + print(f"downloaded_airfrans_zip_bytes={written}", flush=True) + next_report += 1_000_000_000 + final_size = path.stat().st_size + if expected_size is not None and final_size != expected_size: + raise RuntimeError(f"Downloaded {final_size} bytes from {url}, expected {expected_size}") + return {"url": url, "path": str(path), "bytes": final_size, "resumed": resumed, "skipped": False} + + +def extract_of_dataset(archive_path: str | Path, extract_root: str | Path, *, min_cases: int = 1000) -> Path: + archive = Path(archive_path).expanduser() + root = Path(extract_root).expanduser() + root.mkdir(parents=True, exist_ok=True) + existing = _find_of_dataset_root(root) + if existing is not None and _case_count(existing) >= min_cases: + return existing + print(f"extract_airfrans_zip archive={archive} root={root}", flush=True) + with zipfile.ZipFile(archive) as zf: + members = zf.infolist() + for index, member in enumerate(members, start=1): + _safe_extract_member(zf, member, root) + if index % 1000 == 0 or index == len(members): + print(f"extracted_airfrans_members={index}/{len(members)}", flush=True) + found = _find_of_dataset_root(root) + if found is None: + raise RuntimeError(f"OF_dataset directory not found after extracting {archive}") + case_count = _case_count(found) + if case_count < min_cases: + raise RuntimeError(f"Extracted AirfRANS OF_dataset has {case_count} cases; expected at least {min_cases}") + return found + + +def _hf_dataset_status(*, repo_id: str, path_in_repo: str) -> dict[str, Any]: + try: + from huggingface_hub import HfApi + except ModuleNotFoundError as exc: + raise RuntimeError("huggingface_hub is required for AirfRANS public data preparation") from exc + token = _optional_secret("HF_TOKEN") + api = HfApi(token=token) + try: + files = api.list_repo_files(repo_id=repo_id, repo_type="dataset") + except Exception: + files = [] + prefix = path_in_repo.strip("/") + base = f"{prefix}/" if prefix else "" + npz_count = sum(1 for item in files if item.startswith(base) and item.endswith(".npz")) + has_manifest = any(item == f"{base}hf_dataset_manifest.json" for item in files) + return { + "file_count": len(files), + "npz_file_count": npz_count, + "has_manifest": has_manifest, + } + + +def _remote_content_length(url: str) -> int | None: + request = urllib.request.Request(url, method="HEAD") + try: + with urllib.request.urlopen(request, timeout=60) as response: + raw = response.headers.get("Content-Length") + except Exception: + return None + if raw is None: + return None + try: + return int(raw) + except ValueError: + return None + + +def _safe_extract_member(zf: zipfile.ZipFile, member: zipfile.ZipInfo, root: Path) -> None: + target = root / member.filename + resolved_root = root.resolve() + resolved_target = target.resolve() + if resolved_root != resolved_target and resolved_root not in resolved_target.parents: + raise RuntimeError(f"Unsafe path in AirfRANS archive: {member.filename}") + if member.is_dir(): + target.mkdir(parents=True, exist_ok=True) + return + target.parent.mkdir(parents=True, exist_ok=True) + with zf.open(member) as source, target.open("wb") as destination: + shutil.copyfileobj(source, destination, length=16 * 1024 * 1024) + + +def _find_of_dataset_root(root: Path) -> Path | None: + direct = root / "OF_dataset" + if direct.is_dir(): + return direct + for candidate in root.glob("*/OF_dataset"): + if candidate.is_dir(): + return candidate + if _case_count(root) > 0: + return root + return None + + +def _case_count(root: Path) -> int: + return sum(1 for path in root.iterdir() if path.is_dir() and path.name.startswith("airFoil2D_")) if root.is_dir() else 0 + + +def _optional_secret(name: str) -> str | None: + value = os.environ.get(name) + if value: + return value + for path in (Path(".env") / name, Path(".env") / f"{name}.txt"): + if path.is_file(): + text = path.read_text().strip() + if text: + return text + return None + + +def write_json_report(path: str | Path, payload: dict[str, Any]) -> None: + report_path = Path(path).expanduser() + report_path.parent.mkdir(parents=True, exist_ok=True) + report_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") diff --git a/src/airfrans_frontier/remote/artifacts.py b/src/airfrans_frontier/remote/artifacts.py index 61a70cd..96dbc8a 100644 --- a/src/airfrans_frontier/remote/artifacts.py +++ b/src/airfrans_frontier/remote/artifacts.py @@ -5,7 +5,6 @@ import json from pathlib import Path from typing import Any, Iterable -import torch BASE_REQUIRED = ( "config.toml", @@ -25,6 +24,7 @@ def verify_artifacts( required: Iterable[str] = DEFAULT_REQUIRED, *, require_terminal: bool = True, + verify_hf_remote: bool = False, ) -> dict[str, Any]: root = Path(artifact_dir) if not root.exists(): @@ -50,6 +50,10 @@ def verify_artifacts( if missing_failure: raise ValueError(f"Failed artifact directory missing files: {', '.join(missing_failure)}") + checks: dict[str, Any] = { + "required_files": {name: True for name in required_names}, + "terminal_artifact": "final_metrics.json" if has_final else "failure_report.json" if has_failure else None, + } for json_name in ( "latest_metrics.json", "heartbeat.json", @@ -59,20 +63,33 @@ def verify_artifacts( "split_manifest.json", "data_manifest.json", "normalization.json", + "environment_manifest.json", + "evaluation_protocol.json", + "artifact_manifest.json", + "hf_upload_manifest.json", + "verification_report.json", ): path = root / json_name if path.is_file(): _validate_json(path) + checks[f"json:{json_name}"] = True _validate_jsonl(root / "metrics.jsonl") + checks["jsonl:metrics.jsonl"] = True for checkpoint_name in ("checkpoint_latest.pt", "checkpoint_best.pt", "checkpoint_final.pt"): path = root / checkpoint_name if path.is_file(): _validate_checkpoint_metadata(path) + checks[f"checkpoint:{checkpoint_name}"] = True + + if (root / "hf_upload_manifest.json").is_file(): + checks["hf_upload_manifest.json"] = _validate_hf_upload_manifest(root / "hf_upload_manifest.json") + if verify_hf_remote: + checks["hf_remote_paths"] = _verify_hf_remote_paths(root / "hf_upload_manifest.json") files = sorted( path for path in root.rglob("*") - if not path.is_symlink() and path.is_file() and path.name not in {"artifact_manifest.json", "checksums.txt"} + if not path.is_symlink() and path.is_file() and path.name not in {"artifact_manifest.json", "checksums.txt", "verification_report.json"} ) manifest = { "artifact_dir": str(root), @@ -90,6 +107,17 @@ def verify_artifacts( (root / "checksums.txt").write_text( "".join(f"{item['sha256']} {item['path']}\n" for item in manifest["files"]) ) + checks["artifact_manifest.json"] = _validate_artifact_manifest(root / "artifact_manifest.json", root) + checks["checksums.txt"] = _validate_checksums(root / "checksums.txt", root) + report = { + "ok": True, + "artifact_dir": str(root), + "required": list(required_names), + "checked_at": __import__("time").time(), + "checks": checks, + "manifest_file_count": manifest["file_count"], + } + (root / "verification_report.json").write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") return manifest @@ -121,6 +149,7 @@ def _validate_jsonl(path: Path) -> None: def _validate_checkpoint_metadata(path: Path) -> None: + import torch try: checkpoint = torch.load(path, map_location="cpu", weights_only=False) except Exception as exc: @@ -131,3 +160,82 @@ def _validate_checkpoint_metadata(path: Path) -> None: missing = [name for name in required if name not in checkpoint] if missing: raise ValueError(f"Checkpoint artifact {path} missing keys: {', '.join(missing)}") + + +def _validate_checksums(path: Path, root: Path) -> dict[str, int]: + checked = 0 + for line_number, raw_line in enumerate(path.read_text().splitlines(), start=1): + if not raw_line.strip(): + continue + try: + expected, relative = raw_line.split(" ", 1) + except ValueError as exc: + raise ValueError(f"Invalid checksum line {path}:{line_number}") from exc + target = root / relative + if not target.is_file(): + raise ValueError(f"Checksum references missing artifact: {relative}") + actual = sha256_file(target) + if actual != expected: + raise ValueError(f"Checksum mismatch for artifact: {relative}") + checked += 1 + return {"checked": checked} + + +def _validate_artifact_manifest(path: Path, root: Path) -> dict[str, int]: + data = json.loads(path.read_text()) + if not isinstance(data, dict): + raise ValueError(f"Artifact manifest is not a mapping: {path}") + files = data.get("files") + if not isinstance(files, list): + raise ValueError(f"Artifact manifest missing files list: {path}") + checked = 0 + for item in files: + if not isinstance(item, dict): + raise ValueError(f"Artifact manifest file entry is not a mapping: {path}") + relative = item.get("path") + expected = item.get("sha256") + if not isinstance(relative, str) or not isinstance(expected, str): + raise ValueError(f"Artifact manifest file entry missing path or sha256: {path}") + target = root / relative + if not target.is_file(): + raise ValueError(f"Artifact manifest references missing artifact: {relative}") + if sha256_file(target) != expected: + raise ValueError(f"Artifact manifest checksum mismatch: {relative}") + checked += 1 + return {"checked": checked} + + +def _validate_hf_upload_manifest(path: Path) -> dict[str, Any]: + data = json.loads(path.read_text()) + if not isinstance(data, dict): + raise ValueError(f"HF upload manifest is not a mapping: {path}") + enabled = data.get("enabled") + if enabled is False: + return {"enabled": False, "uploaded_paths": 0} + uploaded_paths = data.get("uploaded_paths") + if uploaded_paths is None: + uploaded_paths = [] + if not isinstance(uploaded_paths, list) or not all(isinstance(item, str) for item in uploaded_paths): + raise ValueError(f"HF upload manifest uploaded_paths must be a list of strings: {path}") + return {"enabled": bool(enabled), "uploaded_paths": len(uploaded_paths)} + + +def _verify_hf_remote_paths(path: Path) -> dict[str, Any]: + data = json.loads(path.read_text()) + if not data.get("enabled"): + return {"enabled": False} + repo_id = data.get("repo_id") + repo_type = data.get("repo_type", "model") + uploaded_paths = data.get("uploaded_paths", []) + if not isinstance(repo_id, str) or not isinstance(uploaded_paths, list): + raise ValueError(f"HF upload manifest cannot be remote-verified: {path}") + try: + from huggingface_hub import HfApi + except ModuleNotFoundError as exc: + raise RuntimeError("huggingface_hub is required for remote HF artifact verification") from exc + api = HfApi() + remote_files = set(api.list_repo_files(repo_id=repo_id, repo_type=repo_type)) + missing = [item for item in uploaded_paths if item not in remote_files] + if missing: + raise ValueError(f"HF repo is missing uploaded artifact paths: {', '.join(missing)}") + return {"enabled": True, "checked": len(uploaded_paths)} diff --git a/src/airfrans_frontier/remote/cli.py b/src/airfrans_frontier/remote/cli.py index 04c5913..d3cf728 100644 --- a/src/airfrans_frontier/remote/cli.py +++ b/src/airfrans_frontier/remote/cli.py @@ -79,6 +79,10 @@ def main(argv: list[str] | None = None) -> int: parser = build_parser() args = parser.parse_args(argv) try: + if args.command in {"doctor", "run"}: + from airfrans_frontier.runtime import remove_pythonpath_entries + + remove_pythonpath_entries() if args.command == "doctor": return _doctor(apply=args.apply_skypilot_patch) if args.command == "select": @@ -96,6 +100,9 @@ def main(argv: list[str] | None = None) -> int: print(text) return 0 if args.command == "verify-artifacts": + from airfrans_frontier.runtime import remove_pythonpath_entries + + remove_pythonpath_entries() manifest = verify_artifacts(args.artifact_dir) print(json.dumps({"status": "ok", "file_count": manifest["file_count"]}, sort_keys=True)) return 0 @@ -422,8 +429,10 @@ _TERMINAL_ARTIFACT_NAMES = ( "artifact_manifest.json", "checksums.txt", "config.toml", + "calibration_manifest.json", "data_manifest.json", "environment_manifest.json", + "evaluation_protocol.json", "failure_report.json", "final_metrics.json", "heartbeat.json", @@ -434,6 +443,7 @@ _TERMINAL_ARTIFACT_NAMES = ( "run_manifest.json", "split_manifest.json", "wandb_smoke_manifest.json", + "verification_report.json", ) diff --git a/src/airfrans_frontier/remote/smoke.py b/src/airfrans_frontier/remote/smoke.py index 919b427..1b9f380 100644 --- a/src/airfrans_frontier/remote/smoke.py +++ b/src/airfrans_frontier/remote/smoke.py @@ -52,8 +52,25 @@ def run_smoke_training( previous_run_id = os.environ.get("AIRFRANS_REMOTE_RUN_ID") os.environ["AIRFRANS_OBSERVABILITY_DIR"] = str(output_dir) os.environ["AIRFRANS_REMOTE_RUN_ID"] = run_id + training_dir: Path | None = None + error: Exception | None = None try: result = train_from_config_path(config_path, resume_path=resume_path or os.environ.get("AIRFRANS_RESUME_CHECKPOINT")) + training_dir = result.run_dir + except Exception as exc: + error = exc + training_dir = _latest_training_run_dir(config_path) + if training_dir is None: + _write_json( + output_dir / "failure_report.json", + { + "run_id": run_id, + "phase": "training", + "error_type": type(exc).__name__, + "error_message": str(exc), + "timestamp": time.time(), + }, + ) finally: if previous_observability_dir is None: os.environ.pop("AIRFRANS_OBSERVABILITY_DIR", None) @@ -65,47 +82,33 @@ def run_smoke_training( os.environ["AIRFRANS_REMOTE_RUN_ID"] = previous_run_id finished = time.time() - training_dir = result.run_dir - required_from_training = [ - "final_metrics.json", - "metrics.jsonl", - "latest_metrics.json", - "heartbeat.json", - "checkpoint_latest.pt", - "checkpoint_best.pt", - "checkpoint_final.pt", - "config.toml", - "normalization.json", - "split_manifest.json", - ] - for name in required_from_training: - source = training_dir / name - if source.is_file(): - shutil.copy2(source, output_dir / name) + if training_dir is not None: + _copy_training_artifacts(training_dir, output_dir) + latest_metrics = _read_json(output_dir / "latest_metrics.json") run_manifest: dict[str, Any] = { "run_id": run_id, "command": f"remote-run smoke-train {config_path}", "started_at": started, "finished_at": finished, "elapsed_seconds": finished - started, - "exit_code": 0, - "training_run_dir": str(training_dir), + "exit_code": 0 if error is None else 1, + "training_run_dir": str(training_dir) if training_dir is not None else None, "artifact_dir": str(output_dir), - "final_metrics_path": str(output_dir / "final_metrics.json"), - "checkpoint_path": str(output_dir / "checkpoint_latest.pt"), + "final_metrics_path": str(output_dir / "final_metrics.json") if (output_dir / "final_metrics.json").is_file() else None, + "failure_report_path": str(output_dir / "failure_report.json") if (output_dir / "failure_report.json").is_file() else None, + "checkpoint_path": str(output_dir / "checkpoint_latest.pt") if (output_dir / "checkpoint_latest.pt").is_file() else None, "resume_path": str(resume_path) if resume_path is not None else None, } _write_json(output_dir / "run_manifest.json", run_manifest) - latest_metrics = _read_json(output_dir / "latest_metrics.json") _write_json( heartbeat_path, { "run_id": run_id, - "phase": "completed", + "phase": "completed" if error is None else "failed", "epoch": latest_metrics.get("epoch"), "step": latest_metrics.get("step"), - "latest_checkpoint": "checkpoint_final.pt", + "latest_checkpoint": "checkpoint_final.pt" if error is None else "checkpoint_latest.pt", "latest_metrics": latest_metrics, "started_at": started, "finished_at": finished, @@ -113,8 +116,81 @@ def run_smoke_training( "timestamp": time.time(), }, ) - verify_artifacts(output_dir) - return output_dir + if error is None: + verify_artifacts(output_dir, required=_smoke_required(success=True)) + return output_dir + verify_artifacts(output_dir, required=_smoke_required(success=False)) + raise error + + +def _copy_training_artifacts(training_dir: Path, output_dir: Path) -> None: + names = ( + "config.toml", + "metrics.jsonl", + "latest_metrics.json", + "heartbeat.json", + "checkpoint_latest.pt", + "checkpoint_best.pt", + "checkpoint_final.pt", + "final_metrics.json", + "failure_report.json", + "normalization.json", + "split_manifest.json", + "data_manifest.json", + "environment_manifest.json", + "calibration_manifest.json", + "evaluation_protocol.json", + "hf_upload_manifest.json", + "artifact_manifest.json", + "checksums.txt", + "verification_report.json", + ) + for name in names: + source = training_dir / name + if source.is_file(): + shutil.copy2(source, output_dir / name) + + +def _latest_training_run_dir(config_path: str | Path) -> Path | None: + try: + from airfrans_frontier.training.config import load_training_config + + config = load_training_config(config_path) + except Exception: + return None + root = config.run.artifact_dir + if not root.is_dir(): + return None + candidates = [path for path in root.iterdir() if path.is_dir()] + if not candidates: + return None + return max(candidates, key=lambda path: path.stat().st_mtime) + + +def _smoke_required(*, success: bool) -> tuple[str, ...]: + required = [ + "config.toml", + "metrics.jsonl", + "latest_metrics.json", + "heartbeat.json", + "checkpoint_latest.pt", + "checkpoint_best.pt", + "normalization.json", + "split_manifest.json", + "data_manifest.json", + "environment_manifest.json", + "calibration_manifest.json", + "evaluation_protocol.json", + "hf_upload_manifest.json", + "run_manifest.json", + "artifact_manifest.json", + "checksums.txt", + ] + if success: + required.extend(("checkpoint_final.pt", "final_metrics.json")) + else: + required.append("failure_report.json") + return tuple(required) def run_hf_upload_smoke( *, diff --git a/src/airfrans_frontier/training/artifacts.py b/src/airfrans_frontier/training/artifacts.py index bafbc02..cb5c20e 100644 --- a/src/airfrans_frontier/training/artifacts.py +++ b/src/airfrans_frontier/training/artifacts.py @@ -144,7 +144,7 @@ def _artifact_manifest_text(root: Path) -> tuple[str, str]: files = sorted( path for path in root.rglob("*") - if path.is_file() and path.name not in {"artifact_manifest.json", "checksums.txt"} + if path.is_file() and path.name not in {"artifact_manifest.json", "checksums.txt", "verification_report.json"} ) manifest = { "artifact_dir": str(root), diff --git a/src/airfrans_frontier/training/calibration.py b/src/airfrans_frontier/training/calibration.py new file mode 100644 index 0000000..ad49d32 --- /dev/null +++ b/src/airfrans_frontier/training/calibration.py @@ -0,0 +1,138 @@ +from __future__ import annotations + +import json +import tempfile +import time +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch import nn +from torch.nn import functional as F + +from airfrans_frontier.training.config import TrainingConfig + + +def estimate_forward_flops_per_item(model: nn.Module) -> int: + total = 0 + for module in model.modules(): + if isinstance(module, nn.Linear): + total += 2 * module.in_features * module.out_features + if module.bias is not None: + total += module.out_features + return int(total) + + +def estimate_training_compute(*, steps: int, batch_size: int, forward_flops_per_item: int) -> int: + return int(steps * batch_size * forward_flops_per_item * 3) + + +def checkpoint_size_bytes(run_dir: Path, name: str) -> int | None: + path = run_dir / name + if not path.is_file(): + return None + return int(path.stat().st_size) + + +def gpu_memory_metrics(device: torch.device) -> dict[str, int | None]: + if device.type != "cuda": + return { + "gpu_memory_allocated_mb": None, + "gpu_memory_reserved_mb": None, + "gpu_memory_peak_allocated_mb": None, + } + index = device.index if device.index is not None else torch.cuda.current_device() + return { + "gpu_memory_allocated_mb": int(torch.cuda.memory_allocated(index) // (1024 * 1024)), + "gpu_memory_reserved_mb": int(torch.cuda.memory_reserved(index) // (1024 * 1024)), + "gpu_memory_peak_allocated_mb": int(torch.cuda.max_memory_allocated(index) // (1024 * 1024)), + } + + +def measure_training_step( + model: nn.Module, + optimizer: torch.optim.Optimizer, + features: np.ndarray, + targets: np.ndarray, + *, + batch_size: int, + steps: int, + device: torch.device, +) -> dict[str, Any]: + if steps <= 0: + raise ValueError("Calibration steps must be positive") + rng = np.random.default_rng(1729) + model.train() + if device.type == "cuda": + torch.cuda.reset_peak_memory_stats(device) + torch.cuda.synchronize(device) + started = time.perf_counter() + last_loss = 0.0 + for _ in range(steps): + indices = rng.integers(0, features.shape[0], size=batch_size) + batch_features = torch.from_numpy(np.ascontiguousarray(features[indices], dtype=np.float32)).to(device) + batch_targets = torch.from_numpy(np.ascontiguousarray(targets[indices], dtype=np.float32)).to(device) + optimizer.zero_grad(set_to_none=True) + predictions = model(batch_features) + loss = F.mse_loss(predictions, batch_targets) + loss.backward() + optimizer.step() + last_loss = float(loss.detach().cpu().item()) + if device.type == "cuda": + torch.cuda.synchronize(device) + elapsed = time.perf_counter() - started + return { + "calibration_steps": steps, + "step_time_seconds": elapsed / steps, + "points_per_sec": steps * batch_size / max(elapsed, 1e-12), + "last_calibration_loss": last_loss, + **gpu_memory_metrics(device), + } + + +def measure_validation_runtime( + model: nn.Module, + features: np.ndarray, + *, + batch_size: int, + device: torch.device, +) -> float: + model.eval() + if device.type == "cuda": + torch.cuda.synchronize(device) + started = time.perf_counter() + with torch.no_grad(): + for start in range(0, features.shape[0], batch_size): + stop = min(start + batch_size, features.shape[0]) + batch_features = torch.from_numpy(np.ascontiguousarray(features[start:stop], dtype=np.float32)).to(device) + model(batch_features) + if device.type == "cuda": + torch.cuda.synchronize(device) + return time.perf_counter() - started + + +def measure_checkpoint_size(payload: dict[str, Any]) -> int: + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "checkpoint.pt" + torch.save(payload, path) + return int(path.stat().st_size) + + +def write_calibration_report(path: str | Path, data: dict[str, Any]) -> Path: + report_path = Path(path) + report_path.parent.mkdir(parents=True, exist_ok=True) + report_path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n") + return report_path + + +def static_calibration_fields(config: TrainingConfig, model: nn.Module) -> dict[str, Any]: + forward_flops = estimate_forward_flops_per_item(model) + return { + "estimated_forward_flops_per_item": forward_flops, + "estimated_train_flops": estimate_training_compute( + steps=config.optim.steps, + batch_size=config.data.batch_size, + forward_flops_per_item=forward_flops, + ), + } diff --git a/src/airfrans_frontier/training/config.py b/src/airfrans_frontier/training/config.py index bb50747..1a29b9f 100644 --- a/src/airfrans_frontier/training/config.py +++ b/src/airfrans_frontier/training/config.py @@ -6,6 +6,18 @@ from pathlib import Path from typing import Any +_MODEL_TYPES = { + "mlp", + "film_fourier_mlp", + "film_fourier_inr", + "nerf_cfd_multires", + "deeponet_branch_trunk", + "point_context_perceiver", + "meshgraphnet_or_point_transformer_local", + "raster_fno_unet", + "siren_conditioned_inr", +} + @dataclass(frozen=True) class RunConfig: name: str @@ -21,6 +33,11 @@ class DataConfig: test_cases: int points_per_case: int batch_size: int + source: str + hf_repo_id: str | None + hf_repo_type: str + hf_path_prefix: str + cache_dir: Path | None @dataclass(frozen=True) @@ -34,6 +51,14 @@ class ModelConfig: condition_width: int condition_depth: int condition_dim: int + encoding_levels: int + features_per_level: int + context_points: int + latent_width: int + attention_depth: int + neighbors: int + grid_resolution: int + siren_omega0: float @dataclass(frozen=True) @@ -76,10 +101,20 @@ class ObservabilityConfig: backend: str project: str entity: str | None + group: str | None mode: str tags: tuple[str, ...] +@dataclass(frozen=True) +class HuggingFaceConfig: + enabled: bool + repo_id: str + repo_type: str + path_prefix: str + private: bool + + @dataclass(frozen=True) class TrainingConfig: path: Path @@ -94,6 +129,7 @@ class TrainingConfig: stability: StabilityConfig precision: PrecisionConfig observability: ObservabilityConfig + huggingface: HuggingFaceConfig _REQUIRED_SECTIONS = ("run", "data", "model", "optim", "device", "loss") @@ -141,6 +177,12 @@ def load_training_config(path: str | Path) -> TrainingConfig: observability_raw = {} if not isinstance(observability_raw, dict): raise ValueError("Training config [observability] section must be a table") + huggingface_raw = raw.get("huggingface", {}) + if huggingface_raw is None: + huggingface_raw = {} + if not isinstance(huggingface_raw, dict): + raise ValueError("Training config [huggingface] section must be a table") + run = RunConfig( @@ -155,9 +197,14 @@ def load_training_config(path: str | Path) -> TrainingConfig: test_cases=_integer(data_raw, "test_cases", minimum=0), points_per_case=_integer(data_raw, "points_per_case", minimum=1), batch_size=_integer(data_raw, "batch_size", minimum=1), + source=_choice(_string(data_raw, "source", default="local").lower(), {"local", "huggingface"}, "data.source"), + hf_repo_id=_optional_string(data_raw, "hf_repo_id"), + hf_repo_type=_choice(_string(data_raw, "hf_repo_type", default="dataset"), {"dataset"}, "data.hf_repo_type"), + hf_path_prefix=_string(data_raw, "hf_path_prefix", default=""), + cache_dir=_path(data_raw, "cache_dir") if "cache_dir" in data_raw else None, ) model = ModelConfig( - type=_choice(_string(model_raw, "type"), {"mlp", "film_fourier_mlp"}, "model.type"), + type=_choice(_string(model_raw, "type"), _MODEL_TYPES, "model.type"), hidden_width=_integer(model_raw, "hidden_width", minimum=1), depth=_integer(model_raw, "depth", minimum=1), activation=_choice(_string(model_raw, "activation").lower(), {"gelu", "relu", "silu", "tanh"}, "model.activation"), @@ -166,6 +213,14 @@ def load_training_config(path: str | Path) -> TrainingConfig: condition_width=_integer(model_raw, "condition_width", minimum=1, default=_integer(model_raw, "hidden_width", minimum=1)), condition_depth=_integer(model_raw, "condition_depth", minimum=1, default=2), condition_dim=_integer(model_raw, "condition_dim", minimum=1, default=_integer(model_raw, "hidden_width", minimum=1)), + encoding_levels=_integer(model_raw, "encoding_levels", minimum=0, default=8), + features_per_level=_integer(model_raw, "features_per_level", minimum=1, default=2), + context_points=_integer(model_raw, "context_points", minimum=1, default=512), + latent_width=_integer(model_raw, "latent_width", minimum=1, default=_integer(model_raw, "hidden_width", minimum=1)), + attention_depth=_integer(model_raw, "attention_depth", minimum=1, default=2), + neighbors=_integer(model_raw, "neighbors", minimum=0, default=8), + grid_resolution=_integer(model_raw, "grid_resolution", minimum=2, default=32), + siren_omega0=_number(model_raw, "siren_omega0", minimum=0.0, exclusive_minimum=True, default=30.0), ) optim = OptimConfig( lr=_number(optim_raw, "lr", minimum=0.0, exclusive_minimum=True), @@ -190,10 +245,18 @@ def load_training_config(path: str | Path) -> TrainingConfig: ) observability = ObservabilityConfig( backend=_choice(_string(observability_raw, "backend", default="none").lower(), {"none", "wandb"}, "observability.backend"), - project=_string(observability_raw, "project", default="airfrans"), + project=_string(observability_raw, "project", default="airfRANS-model-sweep"), entity=_optional_string(observability_raw, "entity"), + group=_optional_string(observability_raw, "group"), mode=_choice(_string(observability_raw, "mode", default="online").lower(), {"online", "offline", "disabled"}, "observability.mode"), - tags=_string_tuple(observability_raw, "tags", default=()), + tags=_string_tuple(observability_raw, "tags", default=("airfrans",)), + ) + huggingface = HuggingFaceConfig( + enabled=_boolean(huggingface_raw, "enabled") if "enabled" in huggingface_raw else False, + repo_id=_string(huggingface_raw, "repo_id", default="zacheryasc/airfrans-frontier-checkpoints"), + repo_type=_choice(_string(huggingface_raw, "repo_type", default="model"), {"model", "dataset", "space"}, "huggingface.repo_type"), + path_prefix=_string(huggingface_raw, "path_prefix", default="training_runs"), + private=_boolean(huggingface_raw, "private") if "private" in huggingface_raw else False, ) @@ -214,6 +277,7 @@ def load_training_config(path: str | Path) -> TrainingConfig: stability=stability, precision=precision, observability=observability, + huggingface=huggingface, ) @@ -267,8 +331,15 @@ def _number( *, minimum: float | None = None, exclusive_minimum: bool = False, + default: float | None = None, ) -> float: - value = _required(section, key) + if key not in section: + if default is None: + value = _required(section, key) + else: + value = default + else: + value = section[key] if isinstance(value, bool) or not isinstance(value, (int, float)): raise ValueError(f"Expected number for {key}") result = float(value) diff --git a/src/airfrans_frontier/training/data_sources.py b/src/airfrans_frontier/training/data_sources.py new file mode 100644 index 0000000..155267c --- /dev/null +++ b/src/airfrans_frontier/training/data_sources.py @@ -0,0 +1,150 @@ +from __future__ import annotations + +import hashlib +import json +import os +import time +from pathlib import Path +from typing import Any + +from airfrans_frontier.training.config import DataConfig + + +def resolve_training_data_root(data: DataConfig) -> Path: + if data.source == "local": + return data.root + if data.source != "huggingface": + raise ValueError(f"Unsupported data source: {data.source}") + if not data.hf_repo_id: + raise ValueError("data.hf_repo_id is required when data.source = 'huggingface'") + try: + from huggingface_hub import snapshot_download + except ModuleNotFoundError as exc: + raise RuntimeError("huggingface_hub is required when data.source = 'huggingface'") from exc + token = _resolve_optional_token("HF_TOKEN") + prefix = data.hf_path_prefix.strip("/") + allow_patterns = [f"{prefix}/**"] if prefix else ["*.npz", "*.json", "*.jsonl", "*.txt"] + local_dir = data.cache_dir or data.root + local_dir.mkdir(parents=True, exist_ok=True) + downloaded = Path( + snapshot_download( + repo_id=data.hf_repo_id, + repo_type=data.hf_repo_type, + allow_patterns=allow_patterns, + local_dir=str(local_dir), + token=token, + ) + ) + resolved = downloaded / prefix if prefix else downloaded + if not resolved.is_dir(): + raise FileNotFoundError(f"Downloaded Hugging Face data path is missing: {resolved}") + return resolved + + +def publish_processed_dataset( + *, + data_root: str | Path, + repo_id: str, + path_in_repo: str, + private: bool = False, + manifest_out: str | Path | None = None, +) -> dict[str, Any]: + root = Path(data_root).expanduser() + if not root.is_dir(): + raise NotADirectoryError(f"Processed data path is not a directory: {root}") + files = sorted(path for path in root.rglob("*") if path.is_file() and not path.is_symlink()) + npz_files = [path for path in files if path.suffix == ".npz"] + if not npz_files: + raise ValueError(f"No .npz simulation files found under: {root}") + path_prefix = path_in_repo.strip("/") + manifest = { + "repo_id": repo_id, + "repo_type": "dataset", + "repo_url": f"https://huggingface.co/datasets/{repo_id}", + "path_in_repo": path_prefix, + "created_at": time.time(), + "source_root": str(root), + "file_count": len(files), + "npz_file_count": len(npz_files), + "total_bytes": sum(path.stat().st_size for path in files), + "files": [ + { + "path": str(path.relative_to(root)), + "bytes": path.stat().st_size, + "sha256": _sha256_file(path), + } + for path in files + ], + } + manifest_path = Path(manifest_out).expanduser() if manifest_out is not None else root / "hf_dataset_manifest.json" + manifest_path.parent.mkdir(parents=True, exist_ok=True) + manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n") + + try: + from huggingface_hub import HfApi + except ModuleNotFoundError as exc: + raise RuntimeError("huggingface_hub is required to publish processed data") from exc + token = _resolve_required_token("HF_TOKEN", "Hugging Face dataset publishing") + api = HfApi(token=token) + api.create_repo(repo_id=repo_id, repo_type="dataset", private=private, exist_ok=True) + commit = api.upload_folder( + repo_id=repo_id, + repo_type="dataset", + folder_path=str(root), + path_in_repo=path_prefix, + commit_message=f"Publish AirfRANS processed dataset {path_prefix or 'root'}", + ) + manifest_repo_path = f"{path_prefix}/hf_dataset_manifest.json" if path_prefix else "hf_dataset_manifest.json" + api.upload_file( + repo_id=repo_id, + repo_type="dataset", + path_or_fileobj=str(manifest_path), + path_in_repo=manifest_repo_path, + commit_message=f"Add AirfRANS dataset manifest {path_prefix or 'root'}", + ) + manifest["uploaded_manifest_path"] = manifest_repo_path + manifest["commit"] = { + "commit_url": getattr(commit, "commit_url", None), + "commit_hash": getattr(commit, "oid", None) or getattr(commit, "commit_hash", None), + "pr_url": getattr(commit, "pr_url", None), + } + manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n") + return manifest + + +def _resolve_optional_token(name: str) -> str | None: + try: + return _resolve_required_token(name, "optional Hugging Face access") + except RuntimeError: + return None + + +def _resolve_required_token(name: str, purpose: str) -> str: + value = os.environ.get(name) + if value and value.strip(): + return value.strip() + for path in (Path(name), Path(".env") / name): + if path.is_file(): + value = path.read_text().strip() + if value: + return value + env_file = Path(".env") + if env_file.is_file(): + for raw_line in env_file.read_text().splitlines(): + line = raw_line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, value = line.split("=", 1) + if key.strip() == name: + value = value.strip().strip("\"'") + if value: + return value + raise RuntimeError(f"{name} env var or local secret file is required for {purpose}") + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() diff --git a/src/airfrans_frontier/training/environment.py b/src/airfrans_frontier/training/environment.py new file mode 100644 index 0000000..ec6dfc7 --- /dev/null +++ b/src/airfrans_frontier/training/environment.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import platform +import subprocess +import sys +from typing import Any + + +def environment_manifest() -> dict[str, Any]: + manifest: dict[str, Any] = { + "python": sys.version, + "platform": platform.platform(), + "executable": sys.executable, + } + try: + import torch + + manifest.update( + { + "torch_version": torch.__version__, + "cuda_available": torch.cuda.is_available(), + "cuda_version": torch.version.cuda, + "gpu_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else None, + } + ) + except Exception as exc: + manifest["torch_error"] = repr(exc) + try: + result = subprocess.run( + ["nvidia-smi", "--query-gpu=name,memory.total,driver_version", "--format=csv,noheader"], + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + timeout=15, + check=False, + ) + manifest["nvidia_smi"] = { + "returncode": result.returncode, + "stdout": result.stdout.strip(), + "stderr": result.stderr.strip(), + } + except OSError as exc: + manifest["nvidia_smi"] = {"error": str(exc)} + return manifest diff --git a/src/airfrans_frontier/training/hf_upload.py b/src/airfrans_frontier/training/hf_upload.py new file mode 100644 index 0000000..77422a2 --- /dev/null +++ b/src/airfrans_frontier/training/hf_upload.py @@ -0,0 +1,215 @@ +from __future__ import annotations + +import hashlib +import json +import os +import time +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Iterable + +from airfrans_frontier.training.config import TrainingConfig + + +@dataclass +class UploadRecord: + local_path: str + repo_path: str + bytes: int + sha256: str + uploaded_at: float + + +@dataclass +class UploadManifest: + enabled: bool + repo_id: str | None + repo_type: str | None + repo_url: str | None + path_in_repo: str | None + uploaded_paths: list[str] = field(default_factory=list) + uploaded_files: list[dict[str, Any]] = field(default_factory=list) + commits: list[dict[str, Any]] = field(default_factory=list) + last_error: str | None = None + + +class HfArtifactUploader: + def __init__( + self, + *, + enabled: bool, + run_dir: Path, + repo_id: str | None = None, + repo_type: str | None = None, + path_in_repo: str | None = None, + private: bool = False, + ) -> None: + self.enabled = enabled + self.run_dir = run_dir + self.repo_id = repo_id + self.repo_type = repo_type + self.path_in_repo = path_in_repo.strip("/") if path_in_repo else None + self.private = private + self._api: Any | None = None + self._manifest = UploadManifest( + enabled=enabled, + repo_id=repo_id, + repo_type=repo_type, + repo_url=f"https://huggingface.co/{repo_id}" if repo_id else None, + path_in_repo=self.path_in_repo, + ) + self.write_manifest() + + @classmethod + def from_config(cls, config: TrainingConfig, *, run_dir: Path, run_id: str) -> HfArtifactUploader: + path_prefix = config.huggingface.path_prefix.strip("/") + path_parts = [part for part in (path_prefix, config.model.type, run_id) if part] + return cls( + enabled=config.huggingface.enabled, + run_dir=run_dir, + repo_id=config.huggingface.repo_id, + repo_type=config.huggingface.repo_type, + path_in_repo="/".join(path_parts), + private=config.huggingface.private, + ) + + @property + def repo_url(self) -> str | None: + return self._manifest.repo_url + + def upload_files(self, names: Iterable[str], *, commit_message: str) -> dict[str, Any]: + names = tuple(dict.fromkeys(names)) + if not self.enabled: + return {"enabled": False, "uploaded": [], "missing": []} + missing = [name for name in names if not (self.run_dir / name).is_file()] + if missing: + raise FileNotFoundError(f"Cannot upload missing Hugging Face artifacts: {', '.join(missing)}") + api = self._ensure_api() + uploaded: list[str] = [] + try: + for name in names: + local_path = self.run_dir / name + repo_path = f"{self.path_in_repo}/{name}" if self.path_in_repo else name + commit = api.upload_file( + repo_id=self.repo_id, + repo_type=self.repo_type, + path_or_fileobj=str(local_path), + path_in_repo=repo_path, + commit_message=commit_message, + ) + uploaded.append(repo_path) + record = UploadRecord( + local_path=str(local_path), + repo_path=repo_path, + bytes=local_path.stat().st_size, + sha256=_sha256_file(local_path), + uploaded_at=time.time(), + ) + self._manifest.uploaded_paths.append(repo_path) + self._manifest.uploaded_files.append(record.__dict__) + self._manifest.commits.append(_commit_payload(commit)) + self._manifest.uploaded_paths = sorted(set(self._manifest.uploaded_paths)) + self._manifest.last_error = None + self.write_manifest() + return {"enabled": True, "uploaded": uploaded, "missing": []} + except Exception as exc: + self._manifest.last_error = str(exc) + self.write_manifest() + raise + + def write_manifest(self) -> Path: + path = self.run_dir / "hf_upload_manifest.json" + path.write_text(json.dumps(self._manifest.__dict__, indent=2, sort_keys=True) + "\n") + return path + + def _ensure_api(self) -> Any: + if self._api is not None: + return self._api + try: + from huggingface_hub import HfApi + except ModuleNotFoundError as exc: + raise RuntimeError("huggingface_hub is required when [huggingface].enabled = true") from exc + token = _resolve_secret("HF_TOKEN", "Hugging Face checkpoint uploads") + api = HfApi(token=token) + assert self.repo_id is not None + assert self.repo_type is not None + api.create_repo(repo_id=self.repo_id, repo_type=self.repo_type, private=self.private, exist_ok=True) + self._api = api + return api + + +def resolve_resume_checkpoint(resume_path: str | Path | None) -> tuple[Path | None, dict[str, Any]]: + if resume_path is None: + return None, {"resume_source": None, "resume_downloaded": False, "resume_downloaded_path": None} + raw = str(resume_path) + if not raw.startswith("hf://"): + return Path(resume_path).expanduser(), {"resume_source": raw, "resume_downloaded": False, "resume_downloaded_path": None} + repo_id, filename = _parse_hf_checkpoint_uri(raw) + try: + from huggingface_hub import hf_hub_download + except ModuleNotFoundError as exc: + raise RuntimeError("huggingface_hub is required to resume from hf:// checkpoints") from exc + token = _resolve_secret("HF_TOKEN", "Hugging Face checkpoint download") + cache_dir = Path(".airfrans_hf_resume") / hashlib.sha256(raw.encode()).hexdigest()[:16] + cache_dir.mkdir(parents=True, exist_ok=True) + downloaded = Path( + hf_hub_download( + repo_id=repo_id, + repo_type="model", + filename=filename, + token=token, + local_dir=str(cache_dir), + ) + ) + return downloaded, {"resume_source": raw, "resume_downloaded": True, "resume_downloaded_path": str(downloaded)} + + +def _parse_hf_checkpoint_uri(uri: str) -> tuple[str, str]: + rest = uri.removeprefix("hf://") + parts = rest.split("/") + if len(parts) < 3: + raise ValueError("HF checkpoint URI must be hf://namespace/repo/path/to/checkpoint.pt") + repo_id = "/".join(parts[:2]) + filename = "/".join(parts[2:]) + if not filename: + raise ValueError("HF checkpoint URI is missing checkpoint path") + return repo_id, filename + + +def _resolve_secret(name: str, purpose: str) -> str: + value = os.environ.get(name) + if value and value.strip(): + return value.strip() + for path in (Path(name), Path(".env") / name): + if path.is_file(): + value = path.read_text().strip() + if value: + return value + env_file = Path(".env") + if env_file.is_file(): + for raw_line in env_file.read_text().splitlines(): + line = raw_line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, value = line.split("=", 1) + if key.strip() == name: + value = value.strip().strip("\"'") + if value: + return value + raise RuntimeError(f"{name} env var or local secret file is required for {purpose}") + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _commit_payload(commit: Any) -> dict[str, Any]: + return { + "commit_url": getattr(commit, "commit_url", None), + "commit_hash": getattr(commit, "oid", None) or getattr(commit, "commit_hash", None), + "pr_url": getattr(commit, "pr_url", None), + } diff --git a/src/airfrans_frontier/training/loop.py b/src/airfrans_frontier/training/loop.py index 30b5b79..47c517a 100644 --- a/src/airfrans_frontier/training/loop.py +++ b/src/airfrans_frontier/training/loop.py @@ -12,9 +12,23 @@ import numpy as np import torch from torch.nn import functional as F -from airfrans_frontier.models import FourierFiLMMLP, PointwiseMLP +from airfrans_frontier.models import ( + DeepONetBranchTrunk, + FourierFiLMMLP, + LocalPointTransformer, + NeRFCFDMultiRes, + PointContextPerceiver, + PointwiseMLP, + RasterFNOUNet, + SirenConditionedINR, +) +from airfrans_frontier.remote.artifacts import verify_artifacts from airfrans_frontier.training.artifacts import ArtifactWriter +from airfrans_frontier.training.calibration import checkpoint_size_bytes, static_calibration_fields from airfrans_frontier.training.config import TrainingConfig, load_training_config +from airfrans_frontier.training.data_sources import resolve_training_data_root +from airfrans_frontier.training.environment import environment_manifest +from airfrans_frontier.training.hf_upload import HfArtifactUploader, resolve_resume_checkpoint from airfrans_frontier.training.observability import start_observer from airfrans_frontier.training.data import DatasetBundle, build_dataset_bundle, load_processed_dataset from airfrans_frontier.training.metrics import count_parameters, device_metrics, overall_mse, per_channel_mse @@ -45,17 +59,57 @@ def train_from_config_path(path: str | Path, resume_path: str | Path | None = No def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> TrainingResult: _seed_all(config.run.seed) device = select_device(config) - resume = Path(resume_path).expanduser() if resume_path else None + resume, resume_info = resolve_resume_checkpoint(resume_path) writer = ArtifactWriter.resume_or_create(config.run.artifact_dir, config.run.name, resume) + run_id = os.environ.get("AIRFRANS_REMOTE_RUN_ID", config.run.name) writer.write_config(config.config_text) + writer.write_json("environment_manifest.json", environment_manifest()) observer = start_observer(config, run_dir=writer.run_dir) + uploader = HfArtifactUploader.from_config(config, run_dir=writer.run_dir, run_id=run_id) + started_wall = time.time() + run_manifest = _run_manifest( + config=config, + run_id=run_id, + run_dir=writer.run_dir, + phase="starting", + started_at=started_wall, + resume_info=resume_info, + wandb_run_url=observer.url, + hf_repo_url=uploader.repo_url, + hf_path_in_repo=uploader.path_in_repo, + ) + writer.write_json("run_manifest.json", run_manifest) + observer.update_config( + { + "run_id": run_id, + "model_family": config.model.type, + "hf_repo_url": uploader.repo_url, + "hf_path_in_repo": uploader.path_in_repo, + } + ) def record_metrics(metrics: dict[str, Any]) -> None: writer.append_metrics(metrics) observer.log(metrics) + def publish_artifacts(names: tuple[str, ...], *, event: str, step: int) -> None: + if not config.huggingface.enabled: + return + upload_result = uploader.upload_files(names, commit_message=f"{run_id}: {event} step {step}") + observer.log( + { + "event": "artifact_upload", + "phase": "artifacts", + "step": step, + "artifact_event": event, + "hf_uploaded_count": len(upload_result["uploaded"]), + "hf_repo_url": uploader.repo_url, + "hf_path_in_repo": uploader.path_in_repo, + } + ) - samples = load_processed_dataset(config.data.root) + data_root = resolve_training_data_root(config.data) + samples = load_processed_dataset(data_root) bundle = build_dataset_bundle( samples, train_cases=config.data.train_cases, @@ -74,7 +128,13 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T writer.write_json( "data_manifest.json", { - "root": str(config.data.root), + "root": str(data_root), + "configured_root": str(config.data.root), + "source": config.data.source, + "hf_repo_id": config.data.hf_repo_id, + "hf_repo_type": config.data.hf_repo_type, + "hf_path_prefix": config.data.hf_path_prefix, + "cache_dir": str(config.data.cache_dir) if config.data.cache_dir is not None else None, "case_count": len(samples), "total_points": sum(sample.num_points for sample in samples), "feature_names": list(bundle.feature_names), @@ -105,6 +165,38 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T weight_decay=config.optim.weight_decay, ) + calibration_fields = static_calibration_fields(config, model) + protocol_fields = _evaluation_protocol(config.model.type) + writer.write_json("calibration_manifest.json", calibration_fields) + writer.write_json("evaluation_protocol.json", protocol_fields) + observer.update_config({**calibration_fields, **protocol_fields}) + run_manifest.update( + { + "phase": "initialized", + "parameter_count": count_parameters(model), + **calibration_fields, + **protocol_fields, + } + ) + writer.write_json("run_manifest.json", run_manifest) + writer.write_artifact_manifest() + publish_artifacts( + ( + "config.toml", + "environment_manifest.json", + "split_manifest.json", + "data_manifest.json", + "normalization.json", + "calibration_manifest.json", + "evaluation_protocol.json", + "run_manifest.json", + "artifact_manifest.json", + "checksums.txt", + ), + event="initialized", + step=0, + ) + rng = np.random.default_rng(config.run.seed + 404) started = time.perf_counter() start_step = 0 @@ -211,12 +303,27 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T best_val_loss=best_val_loss, initial_train_loss=initial_train_loss, ) + writer.write_artifact_manifest() + publish_artifacts( + ( + "metrics.jsonl", + "latest_metrics.json", + "heartbeat.json", + LATEST_CHECKPOINT, + BEST_CHECKPOINT, + "artifact_manifest.json", + "checksums.txt", + ), + event="initial_checkpoint", + step=start_step, + ) log_interval = config.optim.log_interval or max(1, config.optim.steps // 10) last_checkpoint_at = time.monotonic() last_log_at = time.perf_counter() last_log_step = start_step last_grad_norm: float | None = None + last_points_per_sec: float | None = None model.train() try: for step in range(start_step + 1, config.optim.steps + 1): @@ -288,6 +395,12 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T initial_train_loss=initial_train_loss, ) last_checkpoint_at = now + writer.write_artifact_manifest() + publish_artifacts( + (LATEST_CHECKPOINT, "metrics.jsonl", "latest_metrics.json", "heartbeat.json", "artifact_manifest.json", "checksums.txt"), + event="latest_checkpoint", + step=step, + ) if step % log_interval == 0 or step == config.optim.steps: train_eval = evaluate_arrays( @@ -326,9 +439,16 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T best_val_loss=best_val_loss, initial_train_loss=initial_train_loss, ) + writer.write_artifact_manifest() + publish_artifacts( + (BEST_CHECKPOINT, "metrics.jsonl", "latest_metrics.json", "heartbeat.json", "artifact_manifest.json", "checksums.txt"), + event="best_checkpoint", + step=step, + ) elapsed = time.perf_counter() - started interval_elapsed = max(time.perf_counter() - last_log_at, 1e-9) points_per_sec = (step - last_log_step) * config.data.batch_size / interval_elapsed + last_points_per_sec = points_per_sec record_metrics( _log_metrics( event="train_eval", @@ -346,7 +466,48 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T last_log_at = time.perf_counter() last_log_step = step model.train() - except Exception: + except Exception as exc: + failure_step = int(locals().get("step", start_step)) + if not (writer.run_dir / "failure_report.json").is_file(): + _write_failure( + writer, + phase="training", + step=failure_step, + error_type=type(exc).__name__, + error_message=str(exc), + latest_grad_norm=last_grad_norm, + latest_checkpoint=LATEST_CHECKPOINT, + ) + run_manifest.update({"phase": "failed", "finished_at": time.time(), "exit_code": 1}) + writer.write_json("run_manifest.json", run_manifest) + writer.write_artifact_manifest() + try: + verify_artifacts(writer.run_dir, required=_verification_required(success=False)) + except Exception as verification_exc: + writer.write_json( + "verification_report.json", + { + "ok": False, + "error_type": type(verification_exc).__name__, + "error_message": str(verification_exc), + "checked_at": time.time(), + }, + ) + publish_artifacts( + ( + "failure_report.json", + "metrics.jsonl", + "latest_metrics.json", + "heartbeat.json", + "run_manifest.json", + "artifact_manifest.json", + "checksums.txt", + "verification_report.json", + ), + event="failure", + step=failure_step, + ) + observer.finish(exit_code=1) raise final_train = evaluate_arrays( @@ -357,6 +518,7 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T device=device, target_names=bundle.target_names, ) + validation_started = time.perf_counter() final_val = ( evaluate_arrays( model, @@ -369,6 +531,7 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T if val_features is not None and val_targets is not None else None ) + validation_runtime_seconds = time.perf_counter() - validation_started if final_val is not None else None final_test = ( evaluate_arrays( model, @@ -394,6 +557,7 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T "best_val_loss": best_val_loss, "parameter_count": count_parameters(model), "model_type": config.model.type, + "model_family": config.model.type, "precision": config.precision.dtype, "train_cases": len(bundle.split.train_ids), "val_cases": len(bundle.split.val_ids), @@ -401,13 +565,25 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T "points_per_case": config.data.points_per_case, "steps": config.optim.steps, "elapsed_seconds": elapsed, + "points_per_sec": last_points_per_sec, + "step_time_seconds": (config.data.batch_size / last_points_per_sec) if last_points_per_sec else None, + "validation_runtime_seconds": validation_runtime_seconds, "checkpoint_interval_seconds": config.checkpoint.interval_seconds, + "data_source": config.data.source, + "data_hf_repo_id": config.data.hf_repo_id, + "data_hf_path_prefix": config.data.hf_path_prefix, + "data_cache_dir": str(config.data.cache_dir) if config.data.cache_dir is not None else None, "resumed_from": str(resume) if resume is not None else None, + **resume_info, + **calibration_fields, + **protocol_fields, **device_metrics(device), } if observer.url is not None: final_metrics["wandb_run_url"] = observer.url - observer.update_summary(final_metrics) + if uploader.repo_url is not None: + final_metrics["hf_repo_url"] = uploader.repo_url + final_metrics["hf_path_in_repo"] = uploader.path_in_repo writer.write_final_metrics(final_metrics) _save_training_checkpoint( writer, @@ -423,6 +599,15 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T initial_train_loss=initial_train_loss, final_metrics=final_metrics, ) + final_metrics.update( + { + "checkpoint_latest_bytes": checkpoint_size_bytes(writer.run_dir, LATEST_CHECKPOINT), + "checkpoint_best_bytes": checkpoint_size_bytes(writer.run_dir, BEST_CHECKPOINT), + "checkpoint_final_bytes": checkpoint_size_bytes(writer.run_dir, FINAL_CHECKPOINT), + } + ) + writer.write_final_metrics(final_metrics) + observer.update_summary(final_metrics) record_metrics( _log_metrics( event="completed", @@ -438,7 +623,33 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T latest_checkpoint=FINAL_CHECKPOINT, ) ) + run_manifest.update( + { + "phase": "completed", + "finished_at": time.time(), + "exit_code": 0, + "final_metrics_path": str(writer.run_dir / "final_metrics.json"), + "checkpoint_latest_path": str(writer.run_dir / LATEST_CHECKPOINT), + "checkpoint_best_path": str(writer.run_dir / BEST_CHECKPOINT), + "checkpoint_final_path": str(writer.run_dir / FINAL_CHECKPOINT), + "wandb_run_url": observer.url, + "hf_repo_url": uploader.repo_url, + "hf_path_in_repo": uploader.path_in_repo, + **calibration_fields, + **protocol_fields, + } + ) + writer.write_json("run_manifest.json", run_manifest) writer.write_artifact_manifest() + verify_artifacts(writer.run_dir, required=_verification_required(success=True)) + publish_artifacts(_final_upload_names(), event="completed", step=config.optim.steps) + writer.write_artifact_manifest() + verify_artifacts(writer.run_dir, required=_verification_required(success=True)) + publish_artifacts( + ("hf_upload_manifest.json", "artifact_manifest.json", "checksums.txt", "verification_report.json"), + event="verification", + step=config.optim.steps, + ) observer.finish(exit_code=0) return TrainingResult(run_dir=writer.run_dir, final_metrics=final_metrics) @@ -450,15 +661,17 @@ def _autocast_context(config: TrainingConfig, device: torch.device): def _build_model(config: TrainingConfig, bundle: DatasetBundle, *, output_dim: int) -> torch.nn.Module: - if config.model.type == "mlp": + input_dim = bundle.train.features.shape[1] + model_type = config.model.type + if model_type == "mlp": return PointwiseMLP( - input_dim=bundle.train.features.shape[1], + input_dim=input_dim, output_dim=output_dim, hidden_width=config.model.hidden_width, depth=config.model.depth, activation=config.model.activation, ) - if config.model.type == "film_fourier_mlp": + if model_type in {"film_fourier_mlp", "film_fourier_inr"}: return FourierFiLMMLP( feature_names=bundle.feature_names, output_dim=output_dim, @@ -471,7 +684,75 @@ def _build_model(config: TrainingConfig, bundle: DatasetBundle, *, output_dim: i condition_dim=config.model.condition_dim, activation=config.model.activation, ) - raise ValueError(f"Unsupported model type: {config.model.type}") + if model_type == "nerf_cfd_multires": + return NeRFCFDMultiRes( + feature_names=bundle.feature_names, + output_dim=output_dim, + coordinate_features=config.model.coordinate_features, + encoding_levels=config.model.encoding_levels, + hidden_width=config.model.hidden_width, + depth=config.model.depth, + condition_width=config.model.condition_width, + condition_depth=config.model.condition_depth, + activation=config.model.activation, + ) + if model_type == "deeponet_branch_trunk": + return DeepONetBranchTrunk( + feature_names=bundle.feature_names, + output_dim=output_dim, + coordinate_features=config.model.coordinate_features, + fourier_scales=config.model.fourier_scales, + hidden_width=config.model.hidden_width, + depth=config.model.depth, + condition_width=config.model.condition_width, + condition_depth=config.model.condition_depth, + activation=config.model.activation, + ) + if model_type == "point_context_perceiver": + return PointContextPerceiver( + input_dim=input_dim, + output_dim=output_dim, + hidden_width=config.model.hidden_width, + latent_width=config.model.latent_width, + context_points=config.model.context_points, + attention_depth=config.model.attention_depth, + activation=config.model.activation, + ) + if model_type == "meshgraphnet_or_point_transformer_local": + return LocalPointTransformer( + feature_names=bundle.feature_names, + output_dim=output_dim, + coordinate_features=config.model.coordinate_features, + hidden_width=config.model.hidden_width, + depth=config.model.depth, + neighbors=config.model.neighbors, + activation=config.model.activation, + ) + if model_type == "raster_fno_unet": + return RasterFNOUNet( + feature_names=bundle.feature_names, + output_dim=output_dim, + coordinate_features=config.model.coordinate_features, + grid_resolution=config.model.grid_resolution, + hidden_width=config.model.hidden_width, + depth=config.model.depth, + condition_width=config.model.condition_width, + condition_depth=config.model.condition_depth, + activation=config.model.activation, + ) + if model_type == "siren_conditioned_inr": + return SirenConditionedINR( + feature_names=bundle.feature_names, + output_dim=output_dim, + coordinate_features=config.model.coordinate_features, + hidden_width=config.model.hidden_width, + depth=config.model.depth, + condition_width=config.model.condition_width, + condition_depth=config.model.condition_depth, + omega0=config.model.siren_omega0, + activation=config.model.activation, + ) + raise ValueError(f"Unsupported model type: {model_type}") def select_device(config: TrainingConfig) -> torch.device: @@ -578,6 +859,123 @@ def _log_metrics( } +def _run_manifest( + *, + config: TrainingConfig, + run_id: str, + run_dir: Path, + phase: str, + started_at: float, + resume_info: dict[str, Any], + wandb_run_url: str | None, + hf_repo_url: str | None, + hf_path_in_repo: str | None, +) -> dict[str, Any]: + return { + "run_id": run_id, + "run_name": config.run.name, + "model_family": config.model.type, + "phase": phase, + "started_at": started_at, + "artifact_dir": str(run_dir), + "config_path": str(config.path), + "config_hash": _config_hash(config), + "data_source": config.data.source, + "data_hf_repo_id": config.data.hf_repo_id, + "data_hf_path_prefix": config.data.hf_path_prefix, + "data_cache_dir": str(config.data.cache_dir) if config.data.cache_dir is not None else None, + "wandb_project": config.observability.project, + "wandb_group": config.observability.group, + "wandb_tags": list(config.observability.tags), + "wandb_run_url": wandb_run_url, + "hf_repo_url": hf_repo_url, + "hf_path_in_repo": hf_path_in_repo, + "checkpoint_latest": LATEST_CHECKPOINT, + "checkpoint_best": BEST_CHECKPOINT, + "checkpoint_final": FINAL_CHECKPOINT, + **resume_info, + } + + +def _evaluation_protocol(model_type: str) -> dict[str, Any]: + protocol = { + "context_target_values_allowed": False, + "target_context_policy": "no_target_context", + "shared_metric_space": "validation_test_points", + } + if model_type == "point_context_perceiver": + protocol.update( + { + "context_protocol": "geometry_and_condition_only", + "target_leakage_policy": "target values are not used as context inputs", + } + ) + if model_type == "meshgraphnet_or_point_transformer_local": + protocol.update( + { + "locality_protocol": "point_knn_fallback", + "mesh_adjacency_required": False, + } + ) + if model_type == "raster_fno_unet": + protocol.update( + { + "raster_protocol": "interpolate_grid_predictions_to_points", + "raster_metric_space": "sampled_validation_test_points", + } + ) + return protocol + + +def _verification_required(*, success: bool) -> tuple[str, ...]: + required = [ + "config.toml", + "metrics.jsonl", + "latest_metrics.json", + "heartbeat.json", + LATEST_CHECKPOINT, + BEST_CHECKPOINT, + "split_manifest.json", + "data_manifest.json", + "normalization.json", + "environment_manifest.json", + "calibration_manifest.json", + "evaluation_protocol.json", + "run_manifest.json", + "hf_upload_manifest.json", + "artifact_manifest.json", + "checksums.txt", + ] + if success: + required.extend(("final_metrics.json", FINAL_CHECKPOINT)) + else: + required.append("failure_report.json") + return tuple(required) + + +def _final_upload_names() -> tuple[str, ...]: + return ( + "config.toml", + "metrics.jsonl", + "latest_metrics.json", + "heartbeat.json", + LATEST_CHECKPOINT, + BEST_CHECKPOINT, + FINAL_CHECKPOINT, + "final_metrics.json", + "split_manifest.json", + "data_manifest.json", + "normalization.json", + "environment_manifest.json", + "calibration_manifest.json", + "evaluation_protocol.json", + "run_manifest.json", + "artifact_manifest.json", + "checksums.txt", + "verification_report.json", + ) + + def _memory_metrics(device: torch.device) -> dict[str, int | None]: if device.type != "cuda": return { @@ -655,6 +1053,7 @@ def _checkpoint_payload( "model_config": asdict(config.model), "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), + "scheduler_state_dict": None, "config": config.config_text, "config_hash": _config_hash(config), "normalization": stats.to_dict(), diff --git a/src/airfrans_frontier/training/observability.py b/src/airfrans_frontier/training/observability.py index 4cf24c4..7d62494 100644 --- a/src/airfrans_frontier/training/observability.py +++ b/src/airfrans_frontier/training/observability.py @@ -18,6 +18,9 @@ class TrainingObserver: def update_summary(self, metrics: Mapping[str, Any]) -> None: return None + def update_config(self, values: Mapping[str, Any]) -> None: + return None + def finish(self, *, exit_code: int = 0) -> None: return None @@ -47,6 +50,9 @@ class WandbObserver(TrainingObserver): for key, value in _json_safe(dict(metrics)).items(): self._run.summary[key] = value + def update_config(self, values: Mapping[str, Any]) -> None: + self._run.config.update(_json_safe(dict(values)), allow_val_change=True) + def finish(self, *, exit_code: int = 0) -> None: self._wandb.finish(exit_code=exit_code) @@ -66,6 +72,7 @@ def start_observer(config: TrainingConfig, *, run_dir: Path) -> TrainingObserver run = wandb.init( entity=observability.entity, project=observability.project, + group=observability.group, name=config.run.name, tags=list(observability.tags), mode=observability.mode, diff --git a/src/airfrans_frontier/training/sanity.py b/src/airfrans_frontier/training/sanity.py new file mode 100644 index 0000000..ab0156b --- /dev/null +++ b/src/airfrans_frontier/training/sanity.py @@ -0,0 +1,167 @@ +from __future__ import annotations + +import json +import tempfile +from pathlib import Path +from typing import Any, Iterable + +import numpy as np +import torch + +from airfrans_frontier.training.config import load_training_config +from airfrans_frontier.training.loop import train + +MODEL_FAMILIES = ( + "film_fourier_inr", + "nerf_cfd_multires", + "deeponet_branch_trunk", + "point_context_perceiver", + "meshgraphnet_or_point_transformer_local", + "raster_fno_unet", + "siren_conditioned_inr", +) + +FEATURE_NAMES = np.array(["x", "y", "sdf", "u_inf", "log_re", "aoa_deg", "aoa_sin", "aoa_cos"], dtype="U16") +TARGET_NAMES = np.array(["velocity_x", "velocity_y", "pressure", "turbulent_viscosity"], dtype="U32") + + +def run_model_sanity( + *, + artifact_dir: str | Path, + device_type: str = "auto", + families: Iterable[str] = MODEL_FAMILIES, + steps: int = 80, +) -> dict[str, Any]: + output_dir = Path(artifact_dir) + output_dir.mkdir(parents=True, exist_ok=True) + selected = tuple(families) + results: dict[str, Any] = { + "device_requested": device_type, + "cuda_available": torch.cuda.is_available(), + "families": {}, + } + with tempfile.TemporaryDirectory() as tmp: + tmp_path = Path(tmp) + data_root = tmp_path / "toy_data" + _write_toy_dataset(data_root) + for family in selected: + config_path = tmp_path / f"{family}.toml" + family_artifacts = output_dir / family + config_path.write_text(_config_text(family, data_root=data_root, artifact_dir=family_artifacts, device_type=device_type, steps=steps)) + result = train(load_training_config(config_path)) + final_metrics = result.final_metrics + initial = float(final_metrics["initial_train_loss"]) + final = float(final_metrics["train_loss"]) + decreased = final < initial + results["families"][family] = { + "run_dir": str(result.run_dir), + "initial_train_loss": initial, + "final_train_loss": final, + "loss_decreased": decreased, + "device": final_metrics.get("device"), + "parameter_count": final_metrics.get("parameter_count"), + "points_per_sec": final_metrics.get("points_per_sec"), + "step_time_seconds": final_metrics.get("step_time_seconds"), + "validation_runtime_seconds": final_metrics.get("validation_runtime_seconds"), + "checkpoint_latest_bytes": final_metrics.get("checkpoint_latest_bytes"), + "checkpoint_best_bytes": final_metrics.get("checkpoint_best_bytes"), + "checkpoint_final_bytes": final_metrics.get("checkpoint_final_bytes"), + "gpu_memory_peak_allocated_mb": final_metrics.get("gpu_memory_peak_allocated_mb"), + "estimated_forward_flops_per_item": final_metrics.get("estimated_forward_flops_per_item"), + "estimated_train_flops": final_metrics.get("estimated_train_flops"), + "target_context_policy": final_metrics.get("target_context_policy"), + "locality_protocol": final_metrics.get("locality_protocol"), + "raster_protocol": final_metrics.get("raster_protocol"), + } + if not decreased: + raise RuntimeError(f"Toy sanity loss did not decrease for {family}: initial={initial}, final={final}") + results["ok"] = True + report_path = output_dir / "model_sanity_results.json" + report_path.write_text(json.dumps(results, indent=2, sort_keys=True) + "\n") + return results + + +def _write_toy_dataset(root: Path, *, cases: int = 4, points: int = 64) -> None: + root.mkdir(parents=True, exist_ok=True) + rng = np.random.default_rng(8675309) + for case_index in range(cases): + x = rng.uniform(-1.0, 1.0, size=points).astype(np.float32) + y = rng.uniform(-1.0, 1.0, size=points).astype(np.float32) + sdf = (np.sqrt(x * x + y * y) - 0.35).astype(np.float32) + aoa = np.float32(-6.0 + 4.0 * case_index) + u_inf = np.float32(20.0 + 2.0 * case_index) + log_re = np.log(u_inf / np.float32(1.5e-5)).astype(np.float32) + condition = np.tile( + np.array([u_inf / 40.0, log_re / 20.0, aoa / 10.0, np.sin(np.deg2rad(aoa)), np.cos(np.deg2rad(aoa))], dtype=np.float32), + (points, 1), + ) + features = np.concatenate((np.stack((x, y, sdf), axis=1), condition), axis=1).astype(np.float32) + targets = np.stack( + ( + 0.35 * x + 0.10 * y + 0.04 * aoa, + -0.25 * y + 0.02 * u_inf / 40.0, + x * y + 0.05 * sdf, + sdf * sdf + 0.03 * np.sin(np.deg2rad(aoa)) + 0.01 * x, + ), + axis=1, + ).astype(np.float32) + np.savez(root / f"case_{case_index:02d}.npz", features=features, targets=targets, feature_names=FEATURE_NAMES, target_names=TARGET_NAMES) + + +def _config_text(family: str, *, data_root: Path, artifact_dir: Path, device_type: str, steps: int) -> str: + return f""" +[run] +name = "sanity_{family}" +seed = 7 +artifact_dir = "{artifact_dir}" + +[data] +root = "{data_root}" +train_cases = 2 +val_cases = 1 +test_cases = 1 +points_per_case = 64 +batch_size = 32 + +[model] +type = "{family}" +hidden_width = 24 +depth = 2 +activation = "gelu" +coordinate_features = ["x", "y", "sdf"] +fourier_scales = [1.0, 2.0] +condition_width = 24 +condition_depth = 2 +condition_dim = 24 +encoding_levels = 3 +features_per_level = 2 +context_points = 16 +latent_width = 24 +attention_depth = 2 +neighbors = 4 +grid_resolution = 8 +siren_omega0 = 10.0 + +[optim] +lr = 0.01 +weight_decay = 0.0 +steps = {steps} +log_interval = {max(1, steps // 4)} + +[device] +type = "{device_type}" +allow_cpu_fallback = true +benchmark_kernels = false + +[loss] +type = "normalized_mse" + +[checkpoint] +interval_seconds = 0 + +[observability] +backend = "none" + +[huggingface] +enabled = false +""".strip() + "\n" diff --git a/tests/test_data_sources.py b/tests/test_data_sources.py new file mode 100644 index 0000000..07875d9 --- /dev/null +++ b/tests/test_data_sources.py @@ -0,0 +1,98 @@ +from __future__ import annotations + +import json +import os +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +from airfrans_frontier.runtime import remove_pythonpath_entries +from airfrans_frontier.training.config import DataConfig +from airfrans_frontier.training.data_sources import publish_processed_dataset, resolve_training_data_root + +remove_pythonpath_entries() +import numpy as np + + +class DataSourceTests(unittest.TestCase): + def test_huggingface_source_downloads_prefix_to_cache(self) -> None: + calls: list[dict[str, object]] = [] + + def fake_snapshot_download(**kwargs): + calls.append(kwargs) + local_dir = Path(str(kwargs["local_dir"])) + target = local_dir / "processed" / "full" + target.mkdir(parents=True) + return str(local_dir) + + fake_module = types.SimpleNamespace(snapshot_download=fake_snapshot_download) + with tempfile.TemporaryDirectory() as tmp, patch.dict(sys.modules, {"huggingface_hub": fake_module}), patch.dict(os.environ, {}, clear=False): + tmp_path = Path(tmp) + config = DataConfig( + root=tmp_path / "configured-root", + train_cases=1, + val_cases=0, + test_cases=0, + points_per_case=1, + batch_size=1, + source="huggingface", + hf_repo_id="owner/airfrans-processed", + hf_repo_type="dataset", + hf_path_prefix="processed/full", + cache_dir=tmp_path / "cache", + ) + + resolved = resolve_training_data_root(config) + + self.assertEqual(resolved, tmp_path / "cache" / "processed" / "full") + self.assertEqual(calls[0]["repo_id"], "owner/airfrans-processed") + self.assertEqual(calls[0]["repo_type"], "dataset") + self.assertEqual(calls[0]["allow_patterns"], ["processed/full/**"]) + + def test_publish_processed_dataset_uploads_folder_and_manifest(self) -> None: + created: list[tuple[str, str, bool]] = [] + uploaded_folders: list[tuple[str, str]] = [] + uploaded_files: list[str] = [] + + class FakeApi: + def __init__(self, token: str) -> None: + self.token = token + + def create_repo(self, *, repo_id: str, repo_type: str, private: bool, exist_ok: bool) -> None: + created.append((repo_id, repo_type, private)) + + def upload_folder(self, *, repo_id: str, repo_type: str, folder_path: str, path_in_repo: str, commit_message: str): + uploaded_folders.append((folder_path, path_in_repo)) + return types.SimpleNamespace(commit_url="https://huggingface.co/datasets/owner/repo/commit/abc", oid="abc") + + def upload_file(self, *, repo_id: str, repo_type: str, path_or_fileobj: str, path_in_repo: str, commit_message: str): + uploaded_files.append(path_in_repo) + return types.SimpleNamespace(commit_url="https://huggingface.co/datasets/owner/repo/commit/def", oid="def") + + fake_module = types.SimpleNamespace(HfApi=FakeApi) + with tempfile.TemporaryDirectory() as tmp, patch.dict(sys.modules, {"huggingface_hub": fake_module}), patch.dict(os.environ, {"HF_TOKEN": "token"}): + root = Path(tmp) / "processed" + root.mkdir() + np.savez(root / "case_00.npz", features=np.zeros((2, 2), dtype=np.float32), targets=np.zeros((2, 1), dtype=np.float32)) + manifest_path = Path(tmp) / "manifest.json" + + manifest = publish_processed_dataset( + data_root=root, + repo_id="owner/repo", + path_in_repo="processed/full", + manifest_out=manifest_path, + ) + + self.assertEqual(created, [("owner/repo", "dataset", False)]) + self.assertEqual(uploaded_folders, [(str(root), "processed/full")]) + self.assertEqual(uploaded_files, ["processed/full/hf_dataset_manifest.json"]) + self.assertEqual(manifest["npz_file_count"], 1) + self.assertTrue(manifest_path.is_file()) + self.assertEqual(json.loads(manifest_path.read_text())["uploaded_manifest_path"], "processed/full/hf_dataset_manifest.json") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_hf_upload.py b/tests/test_hf_upload.py new file mode 100644 index 0000000..f3a902f --- /dev/null +++ b/tests/test_hf_upload.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import json +import os +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +from airfrans_frontier.training.hf_upload import HfArtifactUploader, resolve_resume_checkpoint + + +class HuggingFaceUploadTests(unittest.TestCase): + def test_uploader_creates_repo_uploads_file_and_writes_manifest(self) -> None: + created: list[tuple[str, str, bool]] = [] + uploaded: list[tuple[str, str, str]] = [] + + class FakeApi: + def __init__(self, token: str) -> None: + self.token = token + + def create_repo(self, *, repo_id: str, repo_type: str, private: bool, exist_ok: bool) -> None: + created.append((repo_id, repo_type, private)) + + def upload_file(self, *, repo_id: str, repo_type: str, path_or_fileobj: str, path_in_repo: str, commit_message: str): + uploaded.append((repo_id, repo_type, path_in_repo)) + return types.SimpleNamespace(commit_url="https://huggingface.co/repo/commit/abc", oid="abc") + + fake_module = types.SimpleNamespace(HfApi=FakeApi) + with tempfile.TemporaryDirectory() as tmp, patch.dict(sys.modules, {"huggingface_hub": fake_module}), patch.dict(os.environ, {"HF_TOKEN": "token"}): + root = Path(tmp) + (root / "checkpoint_latest.pt").write_bytes(b"checkpoint") + uploader = HfArtifactUploader( + enabled=True, + run_dir=root, + repo_id="owner/repo", + repo_type="model", + path_in_repo="runs/model/run-1", + private=False, + ) + + result = uploader.upload_files(("checkpoint_latest.pt",), commit_message="upload checkpoint") + + self.assertEqual(result["uploaded"], ["runs/model/run-1/checkpoint_latest.pt"]) + self.assertEqual(created, [("owner/repo", "model", False)]) + self.assertEqual(uploaded, [("owner/repo", "model", "runs/model/run-1/checkpoint_latest.pt")]) + manifest = json.loads((root / "hf_upload_manifest.json").read_text()) + self.assertTrue(manifest["enabled"]) + self.assertEqual(manifest["repo_id"], "owner/repo") + self.assertIn("runs/model/run-1/checkpoint_latest.pt", manifest["uploaded_paths"]) + + def test_resolve_resume_checkpoint_downloads_hf_uri(self) -> None: + calls: list[tuple[str, str]] = [] + + def fake_download(*, repo_id: str, repo_type: str, filename: str, token: str, local_dir: str) -> str: + calls.append((repo_id, filename)) + path = Path(local_dir) / filename + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"checkpoint") + return str(path) + + fake_module = types.SimpleNamespace(hf_hub_download=fake_download) + with patch.dict(sys.modules, {"huggingface_hub": fake_module}), patch.dict(os.environ, {"HF_TOKEN": "token"}): + path, info = resolve_resume_checkpoint("hf://owner/repo/runs/model/checkpoint_latest.pt") + + self.assertIsNotNone(path) + assert path is not None + self.assertTrue(path.is_file()) + self.assertEqual(calls, [("owner/repo", "runs/model/checkpoint_latest.pt")]) + self.assertTrue(info["resume_downloaded"]) + self.assertEqual(info["resume_source"], "hf://owner/repo/runs/model/checkpoint_latest.pt") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_mlp.py b/tests/test_mlp.py index 2ed553e..a837475 100644 --- a/tests/test_mlp.py +++ b/tests/test_mlp.py @@ -8,7 +8,15 @@ remove_pythonpath_entries() import torch -from airfrans_frontier.models import PointwiseMLP +from airfrans_frontier.models import ( + DeepONetBranchTrunk, + LocalPointTransformer, + NeRFCFDMultiRes, + PointContextPerceiver, + PointwiseMLP, + RasterFNOUNet, + SirenConditionedINR, +) class PointwiseMLPTests(unittest.TestCase): @@ -20,6 +28,79 @@ class PointwiseMLPTests(unittest.TestCase): self.assertEqual(tuple(output.shape), (7, 4)) + def test_frontier_models_return_batch_by_target_dim(self) -> None: + feature_names = ("x", "y", "sdf", "u_inf", "log_re", "aoa_deg") + batch = torch.randn(8, len(feature_names)) + models = [ + NeRFCFDMultiRes( + feature_names=feature_names, + output_dim=4, + coordinate_features=("x", "y", "sdf"), + encoding_levels=2, + hidden_width=16, + depth=2, + condition_width=12, + condition_depth=2, + activation="gelu", + ), + DeepONetBranchTrunk( + feature_names=feature_names, + output_dim=4, + coordinate_features=("x", "y", "sdf"), + fourier_scales=(1.0, 2.0), + hidden_width=16, + depth=2, + condition_width=12, + condition_depth=2, + activation="gelu", + ), + PointContextPerceiver( + input_dim=len(feature_names), + output_dim=4, + hidden_width=16, + latent_width=12, + context_points=6, + attention_depth=2, + activation="gelu", + ), + LocalPointTransformer( + feature_names=feature_names, + output_dim=4, + coordinate_features=("x", "y", "sdf"), + hidden_width=16, + depth=2, + neighbors=3, + activation="gelu", + ), + RasterFNOUNet( + feature_names=feature_names, + output_dim=4, + coordinate_features=("x", "y", "sdf"), + grid_resolution=4, + hidden_width=16, + depth=2, + condition_width=12, + condition_depth=2, + activation="gelu", + ), + SirenConditionedINR( + feature_names=feature_names, + output_dim=4, + coordinate_features=("x", "y", "sdf"), + hidden_width=16, + depth=2, + condition_width=12, + condition_depth=2, + omega0=10.0, + activation="gelu", + ), + ] + + for model in models: + with self.subTest(model=type(model).__name__): + output = model(batch) + self.assertEqual(tuple(output.shape), (8, 4)) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_public_data.py b/tests/test_public_data.py new file mode 100644 index 0000000..c6c5250 --- /dev/null +++ b/tests/test_public_data.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +import sys +import tempfile +import types +import unittest +import zipfile +from pathlib import Path +from unittest.mock import patch + +from airfrans_frontier.raw.public import ensure_public_airfrans_processed_hf, extract_of_dataset + + +class PublicAirfransDataTests(unittest.TestCase): + def test_prepare_public_hf_skips_when_dataset_already_published(self) -> None: + class FakeApi: + def __init__(self, token=None): + self.token = token + + def list_repo_files(self, *, repo_id: str, repo_type: str): + assert repo_id == "owner/airfrans-processed" + assert repo_type == "dataset" + return [ + "processed/full/case_000.npz", + "processed/full/case_001.npz", + "processed/full/hf_dataset_manifest.json", + ] + + fake_module = types.SimpleNamespace(HfApi=FakeApi) + with tempfile.TemporaryDirectory() as tmp, patch.dict(sys.modules, {"huggingface_hub": fake_module}), patch.dict( + "os.environ", {"HF_TOKEN": "token"} + ): + report = ensure_public_airfrans_processed_hf( + repo_id="owner/airfrans-processed", + path_in_repo="processed/full", + work_dir=Path(tmp) / "work", + output_dir=Path(tmp) / "out", + min_cases=2, + ) + + self.assertTrue(report["ok"]) + self.assertEqual(report["phase"], "already_published") + self.assertEqual(report["npz_file_count"], 2) + self.assertTrue(report["has_manifest"]) + + def test_extract_of_dataset_finds_public_archive_root(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + tmp_path = Path(tmp) + archive = tmp_path / "OF_dataset.zip" + with zipfile.ZipFile(archive, "w") as zf: + zf.writestr("OF_dataset/airFoil2D_SST_demo/system/controlDict", "ok") + root = extract_of_dataset(archive, tmp_path / "raw", min_cases=1) + + self.assertEqual(root.name, "OF_dataset") + + def test_extract_of_dataset_rejects_zip_slip_paths(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + tmp_path = Path(tmp) + archive = tmp_path / "bad.zip" + with zipfile.ZipFile(archive, "w") as zf: + zf.writestr("../escape.txt", "bad") + with self.assertRaisesRegex(RuntimeError, "Unsafe path"): + extract_of_dataset(archive, tmp_path / "raw", min_cases=1) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_remote_run.py b/tests/test_remote_run.py index 5e621c3..b68b0e4 100644 --- a/tests/test_remote_run.py +++ b/tests/test_remote_run.py @@ -6,6 +6,10 @@ import shutil import unittest from pathlib import Path +from airfrans_frontier.runtime import remove_pythonpath_entries + +remove_pythonpath_entries() + import torch from airfrans_frontier.remote.artifacts import verify_artifacts @@ -101,6 +105,9 @@ class ArtifactVerificationTests(unittest.TestCase): self.assertGreaterEqual(manifest["file_count"], 8) self.assertTrue((root / "artifact_manifest.json").is_file()) self.assertTrue((root / "checksums.txt").is_file()) + self.assertTrue((root / "verification_report.json").is_file()) + report = json.loads((root / "verification_report.json").read_text()) + self.assertTrue(report["ok"]) def test_verify_artifacts_accepts_failure_report_terminal_state(self) -> None: with tempfile.TemporaryDirectory() as tmp: diff --git a/tests/test_training_config.py b/tests/test_training_config.py index 7e7fd0c..7b8ecc2 100644 --- a/tests/test_training_config.py +++ b/tests/test_training_config.py @@ -16,6 +16,17 @@ class TrainingConfigTests(unittest.TestCase): self.assertEqual(config.loss.type, "normalized_mse") self.assertEqual(config.device.type, "cuda") self.assertTrue(config.data.root.is_absolute()) + self.assertEqual(config.data.source, "local") + self.assertIsNone(config.data.hf_repo_id) + self.assertIsNone(config.data.cache_dir) + + def test_config_loader_accepts_huggingface_data_source(self) -> None: + config = load_training_config("configs/aggressive_smoke.toml") + + self.assertEqual(config.data.source, "huggingface") + self.assertEqual(config.data.hf_repo_id, "zacheryasc/airfrans-processed") + self.assertEqual(config.data.hf_path_prefix, "processed/full") + self.assertTrue(config.data.cache_dir is not None) def test_config_loader_rejects_missing_section(self) -> None: with tempfile.TemporaryDirectory() as tmp: diff --git a/tests/test_training_loop.py b/tests/test_training_loop.py index 0203be5..8bd8051 100644 --- a/tests/test_training_loop.py +++ b/tests/test_training_loop.py @@ -176,11 +176,18 @@ class TrainingLoopTests(unittest.TestCase): "rng_state", "torch_rng_state", "batch_rng_state", + "scheduler_state_dict", ): self.assertIn(key, checkpoint) self.assertEqual(list(run_dir.glob("*.tmp")), []) self.assertTrue((run_dir / "artifact_manifest.json").is_file()) self.assertTrue((run_dir / "checksums.txt").is_file()) + self.assertTrue((run_dir / "verification_report.json").is_file()) + self.assertTrue((run_dir / "hf_upload_manifest.json").is_file()) + self.assertTrue((run_dir / "calibration_manifest.json").is_file()) + self.assertTrue((run_dir / "environment_manifest.json").is_file()) + self.assertTrue((run_dir / "evaluation_protocol.json").is_file()) + self.assertTrue((run_dir / "run_manifest.json").is_file()) self.assertTrue((run_dir / "metrics.jsonl").exists()) heartbeat = json.loads((live_dir / "heartbeat.json").read_text()) self.assertEqual(heartbeat["run_id"], "test-run") @@ -188,6 +195,10 @@ class TrainingLoopTests(unittest.TestCase): self.assertTrue(np.isfinite(heartbeat["latest_metrics"]["train_loss"])) self.assertTrue((live_dir / "latest_metrics.json").is_file()) self.assertEqual(final_metrics["device"].startswith("cuda"), torch.cuda.is_available()) + self.assertIn("estimated_forward_flops_per_item", final_metrics) + self.assertIn("estimated_train_flops", final_metrics) + self.assertIn("checkpoint_final_bytes", final_metrics) + self.assertFalse(final_metrics["context_target_values_allowed"]) if torch.cuda.is_available(): self.assertIn("T550", final_metrics["gpu_name"])