155 lines
4.9 KiB
Python
155 lines
4.9 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from airfrans_frontier.runtime import remove_pythonpath_entries
|
|
|
|
remove_pythonpath_entries()
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from airfrans_frontier.training.config import load_training_config
|
|
from airfrans_frontier.training.loop import select_device
|
|
|
|
|
|
FEATURE_NAMES = np.array(["re_norm", "aoa_norm", "x", "y", "sdf"])
|
|
TARGET_NAMES = np.array(["velocity_x", "velocity_y", "pressure", "turbulent_viscosity"])
|
|
|
|
|
|
def write_toy_simulator_dataset(root: Path, *, cases: int = 4, points: int = 64) -> None:
|
|
root.mkdir(parents=True)
|
|
rng = np.random.default_rng(123)
|
|
for case_index in range(cases):
|
|
re_norm = np.full(points, -0.5 + 0.25 * case_index, dtype=np.float32)
|
|
aoa_norm = np.full(points, -0.2 + 0.15 * case_index, dtype=np.float32)
|
|
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.5).astype(np.float32)
|
|
features = np.stack([re_norm, aoa_norm, x, y, sdf], axis=1).astype(np.float32)
|
|
targets = np.stack(
|
|
[
|
|
0.5 * x + 0.2 * y + 0.1 * aoa_norm,
|
|
-0.3 * x + 0.7 * sdf,
|
|
x * y + 0.05 * re_norm,
|
|
sdf**2 + 0.1 * y,
|
|
],
|
|
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 write_training_config(
|
|
path: Path,
|
|
*,
|
|
data_root: Path,
|
|
artifact_dir: Path,
|
|
device_type: str,
|
|
allow_cpu_fallback: bool = False,
|
|
) -> None:
|
|
path.write_text(
|
|
f"""
|
|
[run]
|
|
name = "test_mlp"
|
|
seed = 0
|
|
artifact_dir = "{artifact_dir}"
|
|
|
|
[data]
|
|
root = "{data_root}"
|
|
train_cases = 2
|
|
val_cases = 1
|
|
test_cases = 1
|
|
points_per_case = 64
|
|
batch_size = 64
|
|
|
|
[model]
|
|
type = "mlp"
|
|
hidden_width = 128
|
|
depth = 4
|
|
activation = "gelu"
|
|
|
|
[optim]
|
|
lr = 0.01
|
|
weight_decay = 0.0
|
|
steps = 1000
|
|
log_interval = 250
|
|
|
|
[device]
|
|
type = "{device_type}"
|
|
allow_cpu_fallback = {str(allow_cpu_fallback).lower()}
|
|
benchmark_kernels = true
|
|
|
|
[loss]
|
|
type = "normalized_mse"
|
|
""".strip()
|
|
+ "\n"
|
|
)
|
|
|
|
|
|
class TrainingLoopTests(unittest.TestCase):
|
|
def test_cuda_config_fails_clearly_when_cuda_unavailable(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
tmp_path = Path(tmp)
|
|
config_path = tmp_path / "config.toml"
|
|
write_training_config(
|
|
config_path,
|
|
data_root=tmp_path / "data",
|
|
artifact_dir=tmp_path / "artifacts",
|
|
device_type="cuda",
|
|
)
|
|
config = load_training_config(config_path)
|
|
|
|
with patch("torch.cuda.is_available", return_value=False):
|
|
with self.assertRaisesRegex(RuntimeError, "CUDA requested"):
|
|
select_device(config)
|
|
|
|
def test_mlp_training_smoke_learns_tiny_deterministic_simulator(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
tmp_path = Path(tmp)
|
|
data_root = tmp_path / "data"
|
|
artifact_dir = tmp_path / "artifacts"
|
|
config_path = tmp_path / "config.toml"
|
|
write_toy_simulator_dataset(data_root, cases=4, points=64)
|
|
device_type = "cuda" if torch.cuda.is_available() else "cpu"
|
|
write_training_config(
|
|
config_path,
|
|
data_root=data_root,
|
|
artifact_dir=artifact_dir,
|
|
device_type=device_type,
|
|
)
|
|
|
|
result = subprocess.run(
|
|
[sys.executable, "-m", "airfrans_frontier.cli", "train", str(config_path)],
|
|
text=True,
|
|
capture_output=True,
|
|
check=False,
|
|
)
|
|
self.assertEqual(result.returncode, 0, msg=f"stdout={result.stdout}\nstderr={result.stderr}")
|
|
run_dir_line = next(line for line in result.stdout.splitlines() if line.startswith("run_dir: "))
|
|
run_dir = Path(run_dir_line.removeprefix("run_dir: "))
|
|
final_metrics = json.loads((run_dir / "final_metrics.json").read_text())
|
|
|
|
self.assertTrue(np.isfinite(final_metrics["train_loss"]))
|
|
self.assertLess(final_metrics["train_loss"], final_metrics["initial_train_loss"] * 0.1)
|
|
self.assertLess(final_metrics["train_loss"], 1e-2)
|
|
self.assertTrue((run_dir / "checkpoint.pt").exists())
|
|
self.assertTrue((run_dir / "metrics.jsonl").exists())
|
|
self.assertEqual(final_metrics["device"].startswith("cuda"), torch.cuda.is_available())
|
|
if torch.cuda.is_available():
|
|
self.assertIn("T550", final_metrics["gpu_name"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|