airfRANS-model-exploration/tests/test_training_loop.py

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()