100 lines
3 KiB
Python
100 lines
3 KiB
Python
from __future__ import annotations
|
|
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from airfrans_frontier.training.config import load_training_config
|
|
|
|
|
|
class TrainingConfigTests(unittest.TestCase):
|
|
def test_config_loader_accepts_mlp_tiny(self) -> None:
|
|
config = load_training_config("configs/mlp_tiny.toml")
|
|
|
|
self.assertEqual(config.run.name, "mlp_tiny")
|
|
self.assertEqual(config.model.type, "mlp")
|
|
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_parses_model_metadata_and_defaults_old_configs(self) -> None:
|
|
config = load_training_config("configs/mlp_tiny.toml")
|
|
self.assertEqual(config.model_metadata.reported_family, "mlp")
|
|
self.assertFalse(config.model_metadata.is_proxy)
|
|
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
config_path = Path(tmp) / "proxy.toml"
|
|
config_path.write_text(
|
|
f"""
|
|
[run]
|
|
name = "proxy"
|
|
seed = 0
|
|
artifact_dir = "{Path(tmp) / "runs"}"
|
|
|
|
[data]
|
|
root = "{Path(tmp) / "data"}"
|
|
train_cases = 1
|
|
val_cases = 0
|
|
test_cases = 0
|
|
points_per_case = 1
|
|
batch_size = 1
|
|
|
|
[model]
|
|
type = "raster_fno_unet"
|
|
hidden_width = 8
|
|
depth = 1
|
|
activation = "gelu"
|
|
|
|
[model_metadata]
|
|
requested_family = "raster_fno_unet"
|
|
implementation_family = "raster_fno_unet"
|
|
reported_family = "raster_fno_unet_proxy"
|
|
is_proxy = true
|
|
proxy_for = "fno_or_raster_field_model"
|
|
proxy_notes = "proxy"
|
|
coordinate_encoding_compatibility = "raw_only"
|
|
|
|
[optim]
|
|
lr = 0.001
|
|
weight_decay = 0.0
|
|
steps = 1
|
|
|
|
[device]
|
|
type = "cpu"
|
|
allow_cpu_fallback = false
|
|
benchmark_kernels = false
|
|
|
|
[loss]
|
|
type = "normalized_mse"
|
|
""".strip()
|
|
+ "\n"
|
|
)
|
|
|
|
parsed = load_training_config(config_path)
|
|
|
|
self.assertEqual(parsed.model_metadata.reported_family, "raster_fno_unet_proxy")
|
|
self.assertTrue(parsed.model_metadata.is_proxy)
|
|
self.assertEqual(parsed.model_metadata.coordinate_encoding_compatibility, "raw_only")
|
|
|
|
def test_config_loader_rejects_missing_section(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
config_path = Path(tmp) / "bad.toml"
|
|
config_path.write_text("[run]\nname = 'bad'\n")
|
|
|
|
with self.assertRaisesRegex(ValueError, r"missing \[data\] section"):
|
|
load_training_config(config_path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|