stash: failed model zoo and 70GB run
This commit is contained in:
parent
1da467d1d6
commit
382baefca3
45 changed files with 3793 additions and 56 deletions
|
|
@ -0,0 +1 @@
|
|||
checkpoint
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
65
configs/full_airfrans_incumbent_70gb.toml
Normal file
65
configs/full_airfrans_incumbent_70gb.toml
Normal file
|
|
@ -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
|
||||
61
configs/model_zoo_7gb/deeponet_branch_trunk.toml
Normal file
61
configs/model_zoo_7gb/deeponet_branch_trunk.toml
Normal file
|
|
@ -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
|
||||
61
configs/model_zoo_7gb/film_fourier_inr.toml
Normal file
61
configs/model_zoo_7gb/film_fourier_inr.toml
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
62
configs/model_zoo_7gb/nerf_cfd_multires.toml
Normal file
62
configs/model_zoo_7gb/nerf_cfd_multires.toml
Normal file
|
|
@ -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
|
||||
63
configs/model_zoo_7gb/point_context_perceiver.toml
Normal file
63
configs/model_zoo_7gb/point_context_perceiver.toml
Normal file
|
|
@ -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
|
||||
61
configs/model_zoo_7gb/raster_fno_unet.toml
Normal file
61
configs/model_zoo_7gb/raster_fno_unet.toml
Normal file
|
|
@ -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
|
||||
61
configs/model_zoo_7gb/siren_conditioned_inr.toml
Normal file
61
configs/model_zoo_7gb/siren_conditioned_inr.toml
Normal file
|
|
@ -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
|
||||
86
configs/remote_full_70gb.toml
Normal file
86
configs/remote_full_70gb.toml
Normal file
|
|
@ -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"
|
||||
89
configs/remote_model_zoo_7gb/deeponet_branch_trunk.toml
Normal file
89
configs/remote_model_zoo_7gb/deeponet_branch_trunk.toml
Normal file
|
|
@ -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"
|
||||
89
configs/remote_model_zoo_7gb/film_fourier_inr.toml
Normal file
89
configs/remote_model_zoo_7gb/film_fourier_inr.toml
Normal file
|
|
@ -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"
|
||||
|
|
@ -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"
|
||||
89
configs/remote_model_zoo_7gb/nerf_cfd_multires.toml
Normal file
89
configs/remote_model_zoo_7gb/nerf_cfd_multires.toml
Normal file
|
|
@ -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"
|
||||
89
configs/remote_model_zoo_7gb/point_context_perceiver.toml
Normal file
89
configs/remote_model_zoo_7gb/point_context_perceiver.toml
Normal file
|
|
@ -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"
|
||||
89
configs/remote_model_zoo_7gb/raster_fno_unet.toml
Normal file
89
configs/remote_model_zoo_7gb/raster_fno_unet.toml
Normal file
|
|
@ -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"
|
||||
89
configs/remote_model_zoo_7gb/siren_conditioned_inr.toml
Normal file
89
configs/remote_model_zoo_7gb/siren_conditioned_inr.toml
Normal file
|
|
@ -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"
|
||||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -31,3 +31,7 @@ dev = [
|
|||
"skypilot[vast]>=0.12.3.post1",
|
||||
"pytest>=9.1.1",
|
||||
]
|
||||
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
testpaths = ["tests"]
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
404
src/airfrans_frontier/models/frontier.py
Normal file
404
src/airfrans_frontier/models/frontier.py
Normal file
|
|
@ -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}")
|
||||
|
|
@ -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),
|
||||
|
|
|
|||
253
src/airfrans_frontier/raw/public.py
Normal file
253
src/airfrans_frontier/raw/public.py
Normal file
|
|
@ -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")
|
||||
|
|
@ -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)}
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
*,
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
138
src/airfrans_frontier/training/calibration.py
Normal file
138
src/airfrans_frontier/training/calibration.py
Normal file
|
|
@ -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,
|
||||
),
|
||||
}
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
150
src/airfrans_frontier/training/data_sources.py
Normal file
150
src/airfrans_frontier/training/data_sources.py
Normal file
|
|
@ -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()
|
||||
44
src/airfrans_frontier/training/environment.py
Normal file
44
src/airfrans_frontier/training/environment.py
Normal file
|
|
@ -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
|
||||
215
src/airfrans_frontier/training/hf_upload.py
Normal file
215
src/airfrans_frontier/training/hf_upload.py
Normal file
|
|
@ -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),
|
||||
}
|
||||
|
|
@ -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(),
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
167
src/airfrans_frontier/training/sanity.py
Normal file
167
src/airfrans_frontier/training/sanity.py
Normal file
|
|
@ -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"
|
||||
98
tests/test_data_sources.py
Normal file
98
tests/test_data_sources.py
Normal file
|
|
@ -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()
|
||||
77
tests/test_hf_upload.py
Normal file
77
tests/test_hf_upload.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
67
tests/test_public_data.py
Normal file
67
tests/test_public_data.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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"])
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue