649 lines
27 KiB
Python
649 lines
27 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
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.hf_upload import HfArtifactUploader
|
|
from airfrans_frontier.training.observability import WandbObserver
|
|
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,
|
|
checkpoint_include_optimizer_state: bool = True,
|
|
checkpoint_include_rng_state: bool = True,
|
|
dead_curve_patience_evals: int | None = None,
|
|
) -> 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}
|
|
include_optimizer_state = {str(checkpoint_include_optimizer_state).lower()}
|
|
include_rng_state = {str(checkpoint_include_rng_state).lower()}
|
|
{f'''
|
|
[stability]
|
|
dead_curve_patience_evals = {dead_curve_patience_evals}
|
|
dead_curve_min_relative_improvement = 0.01
|
|
dead_curve_warmup_steps = 0
|
|
''' if dead_curve_patience_evals is not None else ''}
|
|
""".strip()
|
|
+ "\n"
|
|
)
|
|
|
|
|
|
class HfArtifactUploaderTests(unittest.TestCase):
|
|
def test_upload_files_commits_batch_once_to_reduce_hub_rate_limit(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
run_dir = Path(tmp)
|
|
(run_dir / "metrics.jsonl").write_text("{}\n")
|
|
(run_dir / "latest_metrics.json").write_text("{}\n")
|
|
uploader = HfArtifactUploader(
|
|
enabled=True,
|
|
run_dir=run_dir,
|
|
repo_id="owner/repo",
|
|
repo_type="model",
|
|
path_in_repo="runs/model",
|
|
)
|
|
fake_api = Mock()
|
|
fake_api.create_commit.return_value = Mock(oid="abc123", commit_url="https://hf/commit/abc123", pr_url=None)
|
|
uploader._api = fake_api
|
|
|
|
result = uploader.upload_files(("metrics.jsonl", "latest_metrics.json"), commit_message="batch artifacts")
|
|
|
|
self.assertEqual(result["uploaded"], ["runs/model/metrics.jsonl", "runs/model/latest_metrics.json"])
|
|
fake_api.create_commit.assert_called_once()
|
|
call_kwargs = fake_api.create_commit.call_args.kwargs
|
|
self.assertEqual(call_kwargs["repo_id"], "owner/repo")
|
|
self.assertEqual(call_kwargs["repo_type"], "model")
|
|
self.assertEqual(call_kwargs["commit_message"], "batch artifacts")
|
|
self.assertEqual(
|
|
[operation.path_in_repo for operation in call_kwargs["operations"]],
|
|
["runs/model/metrics.jsonl", "runs/model/latest_metrics.json"],
|
|
)
|
|
manifest = json.loads((run_dir / "hf_upload_manifest.json").read_text())
|
|
self.assertEqual(len(manifest["commits"]), 1)
|
|
self.assertEqual(set(manifest["uploaded_paths"]), set(result["uploaded"]))
|
|
|
|
def test_upload_files_suppresses_large_artifacts_and_uploads_metadata(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
run_dir = Path(tmp)
|
|
(run_dir / "metrics.jsonl").write_text("{}\n")
|
|
(run_dir / "checkpoint_large.pt").write_text("checkpoint")
|
|
uploader = HfArtifactUploader(
|
|
enabled=True,
|
|
run_dir=run_dir,
|
|
repo_id="owner/repo",
|
|
repo_type="model",
|
|
path_in_repo="runs/model",
|
|
)
|
|
fake_api = Mock()
|
|
fake_api.create_commit.return_value = Mock(oid="abc123", commit_url="https://hf/commit/abc123", pr_url=None)
|
|
uploader._api = fake_api
|
|
|
|
with patch.dict(os.environ, {"AIRFRANS_HF_MAX_FILE_BYTES": "4"}):
|
|
result = uploader.upload_files(
|
|
("metrics.jsonl", "checkpoint_large.pt"),
|
|
commit_message="skip oversized checkpoints",
|
|
)
|
|
|
|
self.assertEqual(result["uploaded"], ["runs/model/metrics.jsonl"])
|
|
self.assertEqual(result["suppressed_large_files"], ["runs/model/checkpoint_large.pt"])
|
|
fake_api.create_commit.assert_called_once()
|
|
self.assertEqual(
|
|
[operation.path_in_repo for operation in fake_api.create_commit.call_args.kwargs["operations"]],
|
|
["runs/model/metrics.jsonl"],
|
|
)
|
|
manifest = json.loads((run_dir / "hf_upload_manifest.json").read_text())
|
|
self.assertEqual(manifest["suppressed_uploads"][-1]["event"], "file_too_large")
|
|
self.assertEqual(manifest["suppressed_uploads"][-1]["repo_path"], "runs/model/checkpoint_large.pt")
|
|
|
|
class WandbObserverTests(unittest.TestCase):
|
|
def test_logs_existing_run_artifacts_with_event_and_step_aliases(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
run_dir = Path(tmp)
|
|
(run_dir / "metrics.jsonl").write_text("{}\n")
|
|
(run_dir / "checkpoint_latest.pt").write_text("checkpoint")
|
|
fake_artifact = Mock()
|
|
fake_wandb = Mock()
|
|
fake_wandb.Artifact.return_value = fake_artifact
|
|
fake_run = Mock()
|
|
fake_run.name = "run with spaces"
|
|
observer = WandbObserver(fake_run, fake_wandb)
|
|
|
|
observer.log_artifact_files(
|
|
run_dir,
|
|
("metrics.jsonl", "checkpoint_latest.pt", "missing.json"),
|
|
event="latest_checkpoint",
|
|
step=5,
|
|
aliases=("latest_checkpoint", "step-5", "latest"),
|
|
)
|
|
|
|
fake_wandb.Artifact.assert_called_once()
|
|
artifact_kwargs = fake_wandb.Artifact.call_args.kwargs
|
|
self.assertEqual(artifact_kwargs["type"], "airfrans-run-artifacts")
|
|
self.assertEqual(artifact_kwargs["metadata"]["event"], "latest_checkpoint")
|
|
self.assertEqual(artifact_kwargs["metadata"]["step"], 5)
|
|
self.assertEqual(fake_artifact.add_file.call_count, 2)
|
|
self.assertEqual(
|
|
[call.kwargs["name"] for call in fake_artifact.add_file.call_args_list],
|
|
["metrics.jsonl", "checkpoint_latest.pt"],
|
|
)
|
|
fake_run.log_artifact.assert_called_once_with(
|
|
fake_artifact,
|
|
aliases=["latest_checkpoint", "step-5", "latest"],
|
|
)
|
|
|
|
def test_finish_timeout_does_not_block_completed_training(self) -> None:
|
|
fake_wandb = Mock()
|
|
|
|
def slow_finish(*, exit_code: int) -> None:
|
|
time.sleep(1.0)
|
|
|
|
fake_wandb.finish.side_effect = slow_finish
|
|
observer = WandbObserver(Mock(), fake_wandb)
|
|
|
|
started = time.perf_counter()
|
|
with patch.dict(os.environ, {"AIRFRANS_WANDB_FINISH_TIMEOUT_SECONDS": "0.01"}):
|
|
observer.finish(exit_code=0)
|
|
|
|
self.assertLess(time.perf_counter() - started, 0.5)
|
|
|
|
|
|
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",
|
|
"normalization_policy",
|
|
):
|
|
self.assertIn(key, checkpoint)
|
|
self.assertEqual(checkpoint["normalization_policy"]["raw_feature_names"], ["x", "y", "sdf"])
|
|
normalization_manifest = json.loads((run_dir / "normalization.json").read_text())
|
|
self.assertEqual(normalization_manifest["input_policy"]["raw_feature_names"], ["x", "y", "sdf"])
|
|
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())
|
|
utilization_lines = [json.loads(line) for line in (run_dir / "utilization.jsonl").read_text().splitlines()]
|
|
self.assertGreaterEqual(len(utilization_lines), 2)
|
|
self.assertIn("gpu_util_percent", utilization_lines[-1])
|
|
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"])
|
|
self.assertEqual(final_metrics["reported_family"], "mlp")
|
|
self.assertFalse(final_metrics["is_proxy"])
|
|
self.assertEqual(final_metrics["coordinate_encoding_type"], "fixed_fourier")
|
|
run_manifest = json.loads((run_dir / "run_manifest.json").read_text())
|
|
self.assertEqual(run_manifest["reported_family"], "mlp")
|
|
job_manifest = json.loads((run_dir / "job_manifest.json").read_text())
|
|
self.assertEqual(job_manifest["model_metadata"]["reported_family"], "mlp")
|
|
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["failure_category"], "nonfinite")
|
|
self.assertEqual(report["reported_family"], "mlp")
|
|
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_dead_curve_interruption_writes_plateau_failure(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=6,
|
|
log_interval=1,
|
|
checkpoint_interval_seconds=0,
|
|
hidden_width=32,
|
|
dead_curve_patience_evals=2,
|
|
)
|
|
config = load_training_config(config_path)
|
|
with patch("airfrans_frontier.training.loop.torch.optim.AdamW.step", lambda self: None):
|
|
with self.assertRaisesRegex(RuntimeError, "loss did not improve"):
|
|
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["failure_category"], "plateau_or_dead_recipe")
|
|
self.assertEqual(report["error_type"], "DeadCurveError")
|
|
self.assertIn("dead_curve", report)
|
|
|
|
def test_lightweight_checkpoint_omits_optimizer_payload(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=2,
|
|
log_interval=1,
|
|
checkpoint_interval_seconds=0,
|
|
hidden_width=32,
|
|
checkpoint_include_optimizer_state=False,
|
|
checkpoint_include_rng_state=False,
|
|
)
|
|
|
|
result = train(load_training_config(config_path))
|
|
|
|
checkpoint = torch.load(result.run_dir / "checkpoint_final.pt", map_location="cpu", weights_only=False)
|
|
self.assertIsNone(checkpoint["optimizer_state_dict"])
|
|
self.assertIsNone(checkpoint["rng_state"])
|
|
self.assertFalse(result.final_metrics["checkpoint_include_optimizer_state"])
|
|
|
|
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["failure_category"], "nonfinite")
|
|
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()
|