from __future__ import annotations import json import os 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 unittest.mock import Mock from airfrans_frontier.training.config import load_training_config from airfrans_frontier.training.loop import train, 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, steps: int = 1000, log_interval: int = 250, checkpoint_interval_seconds: int = 1800, hidden_width: int = 128, ) -> 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 = {hidden_width} depth = 4 activation = "gelu" [optim] lr = 0.01 weight_decay = 0.0 steps = {steps} log_interval = {log_interval} [device] type = "{device_type}" allow_cpu_fallback = {str(allow_cpu_fallback).lower()} benchmark_kernels = true [loss] type = "normalized_mse" [checkpoint] interval_seconds = {checkpoint_interval_seconds} """.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, ) live_dir = tmp_path / "live" env = os.environ | { "AIRFRANS_OBSERVABILITY_DIR": str(live_dir), "AIRFRANS_REMOTE_RUN_ID": "test-run", } result = subprocess.run( [sys.executable, "-m", "airfrans_frontier.cli", "train", str(config_path)], text=True, capture_output=True, check=False, env=env, ) 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.assertFalse((run_dir / "checkpoint.pt").exists()) self.assertTrue((run_dir / "checkpoint_latest.pt").exists()) self.assertTrue((run_dir / "checkpoint_best.pt").exists()) self.assertTrue((run_dir / "checkpoint_final.pt").exists()) checkpoint = torch.load(run_dir / "checkpoint_latest.pt", map_location="cpu", weights_only=False) for key in ( "schema_version", "step", "model_state_dict", "optimizer_state_dict", "config_hash", "normalization", "feature_names", "target_names", "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") self.assertEqual(heartbeat["phase"], "completed") 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"]) def test_resume_uses_existing_run_dir_and_appends_metrics(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) write_training_config( config_path, data_root=data_root, artifact_dir=artifact_dir, device_type="cpu", steps=4, log_interval=1, checkpoint_interval_seconds=0, hidden_width=32, ) first = subprocess.run( [sys.executable, "-m", "airfrans_frontier.cli", "train", str(config_path)], text=True, capture_output=True, check=False, ) self.assertEqual(first.returncode, 0, msg=f"stdout={first.stdout}\nstderr={first.stderr}") run_dir_line = next(line for line in first.stdout.splitlines() if line.startswith("run_dir: ")) run_dir = Path(run_dir_line.removeprefix("run_dir: ")) second = subprocess.run( [ sys.executable, "-m", "airfrans_frontier.cli", "train", str(config_path), "--resume", str(run_dir / "checkpoint_latest.pt"), ], text=True, capture_output=True, check=False, ) self.assertEqual(second.returncode, 0, msg=f"stdout={second.stdout}\nstderr={second.stderr}") resumed_dir_line = next(line for line in second.stdout.splitlines() if line.startswith("run_dir: ")) self.assertEqual(Path(resumed_dir_line.removeprefix("run_dir: ")), run_dir) metrics = [json.loads(line) for line in (run_dir / "metrics.jsonl").read_text().splitlines()] self.assertTrue(any(metric["event"] == "resume" and metric["step"] == 4 for metric in metrics)) final_metrics = json.loads((run_dir / "final_metrics.json").read_text()) self.assertEqual(final_metrics["resumed_from"], str(run_dir / "checkpoint_latest.pt")) def test_nonfinite_loss_writes_failure_report_without_final_checkpoint(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) write_training_config( config_path, data_root=data_root, artifact_dir=artifact_dir, device_type="cpu", steps=4, log_interval=1, checkpoint_interval_seconds=0, hidden_width=32, ) config = load_training_config(config_path) bad_loss = torch.tensor(float("nan"), requires_grad=True) with patch("airfrans_frontier.training.loop.F.mse_loss", Mock(return_value=bad_loss)): with self.assertRaisesRegex(RuntimeError, "nonfinite loss"): train(config) run_dirs = sorted(path for path in artifact_dir.iterdir() if path.is_dir()) self.assertEqual(len(run_dirs), 1) run_dir = run_dirs[0] report = json.loads((run_dir / "failure_report.json").read_text()) self.assertEqual(report["error_type"], "NonFiniteLoss") self.assertEqual(report["step"], 1) checkpoint = torch.load(run_dir / "checkpoint_latest.pt", map_location="cpu", weights_only=False) self.assertEqual(checkpoint["step"], 0) self.assertFalse((run_dir / "checkpoint_final.pt").exists()) self.assertTrue((run_dir / "artifact_manifest.json").is_file()) self.assertTrue((run_dir / "checksums.txt").is_file()) def test_nonfinite_gradient_writes_failure_report(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) write_training_config( config_path, data_root=data_root, artifact_dir=artifact_dir, device_type="cpu", steps=4, log_interval=1, checkpoint_interval_seconds=0, hidden_width=32, ) config = load_training_config(config_path) with patch("airfrans_frontier.training.loop.torch.nn.utils.clip_grad_norm_", Mock(side_effect=RuntimeError("nonfinite"))): with self.assertRaisesRegex(RuntimeError, "nonfinite gradients"): train(config) run_dir = next(path for path in artifact_dir.iterdir() if path.is_dir()) report = json.loads((run_dir / "failure_report.json").read_text()) self.assertEqual(report["error_type"], "NonFiniteGradient") self.assertEqual(report["step"], 1) self.assertFalse((run_dir / "checkpoint_final.pt").exists()) self.assertTrue((run_dir / "artifact_manifest.json").is_file()) self.assertTrue((run_dir / "checksums.txt").is_file()) def test_film_fourier_training_smoke_learns_conditioned_point_field(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" data_root.mkdir() feature_names = np.array( [ "x", "y", "sdf", "u_inf", "log_re", "aoa_deg", "aoa_sin", "aoa_cos", "naca_param_0", "naca_param_1", "naca_param_2", "naca_param_3", "naca_param_0_mask", "naca_param_1_mask", "naca_param_2_mask", "naca_param_3_mask", ] ) rng = np.random.default_rng(42) for case_index in range(4): x = rng.uniform(-1.0, 1.0, size=96).astype(np.float32) y = rng.uniform(-1.0, 1.0, size=96).astype(np.float32) sdf = np.sqrt(x * x + y * y).astype(np.float32) aoa = np.float32(-5.0 + 5.0 * case_index) u_inf = np.float32(30.0 + case_index) condition = np.tile( np.array( [ u_inf, np.log(u_inf / np.float32(1.5e-5)), aoa, np.sin(np.deg2rad(aoa)), np.cos(np.deg2rad(aoa)), 2.0 + case_index, 4.0, 12.0, 0.0, 1.0, 1.0, 1.0, 0.0, ], dtype=np.float32, ), (x.shape[0], 1), ) features = np.concatenate((np.stack((x, y, sdf), axis=1), condition), axis=1).astype(np.float32) targets = np.stack( ( 0.1 * x + 0.01 * aoa, -0.2 * y + 0.001 * u_inf, x * y, 0.05 * sdf + 0.001 * case_index, ), axis=1, ).astype(np.float32) np.savez( data_root / f"case_{case_index:02d}.npz", features=features, targets=targets, feature_names=feature_names, target_names=TARGET_NAMES, ) config_path.write_text( f""" [run] name = "film_contract" seed = 0 artifact_dir = "{artifact_dir}" [data] root = "{data_root}" train_cases = 2 val_cases = 1 test_cases = 1 points_per_case = 96 batch_size = 64 [model] type = "film_fourier_mlp" hidden_width = 64 depth = 2 activation = "gelu" coordinate_features = ["x", "y", "sdf"] fourier_scales = [1.0, 2.0] condition_width = 32 condition_depth = 2 condition_dim = 32 [optim] lr = 0.003 weight_decay = 0.0 steps = 60 log_interval = 20 [device] type = "cpu" allow_cpu_fallback = false benchmark_kernels = false [loss] type = "normalized_mse" """.strip() + "\n" ) result = train(load_training_config(config_path)) self.assertTrue(np.isfinite(result.final_metrics["train_loss"])) self.assertEqual(result.final_metrics["model_type"], "film_fourier_mlp") self.assertEqual(result.final_metrics["precision"], "float32") self.assertTrue((result.run_dir / "checkpoint_final.pt").is_file()) if __name__ == "__main__": unittest.main()