diff --git a/.gitignore b/.gitignore index 2dd5719..4734818 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ data/ outputs/ artifacts/ +.airfrans_resume/ # Python .venv/ diff --git a/configs/aggressive_smoke.toml b/configs/aggressive_smoke.toml new file mode 100644 index 0000000..f7fd9d2 --- /dev/null +++ b/configs/aggressive_smoke.toml @@ -0,0 +1,46 @@ +[run] +name = "aggressive_smoke" +seed = 20260723 +artifact_dir = "artifacts/current_run/training_runs" + +[data] +root = "data/processed/full" +train_cases = 45 +val_cases = 3 +test_cases = 2 +points_per_case = 999999999 +batch_size = 4096 + +[model] +type = "film_fourier_mlp" +hidden_width = 4096 +depth = 12 +activation = "gelu" +coordinate_features = ["x", "y", "sdf"] +fourier_scales = [1.0, 2.0, 4.0, 8.0, 16.0, 32.0] +condition_width = 1024 +condition_depth = 3 +condition_dim = 512 + +[optim] +lr = 0.0001 +weight_decay = 0.0001 +steps = 5000 +log_interval = 500 + +[device] +type = "cuda" +allow_cpu_fallback = false +benchmark_kernels = true + +[loss] +type = "normalized_mse" + +[precision] +dtype = "bf16" + +[checkpoint] +interval_seconds = 1800 + +[stability] +max_grad_norm = 1.0 diff --git a/configs/aggressive_smoke_local.toml b/configs/aggressive_smoke_local.toml new file mode 100644 index 0000000..4649b46 --- /dev/null +++ b/configs/aggressive_smoke_local.toml @@ -0,0 +1,46 @@ +[run] +name = "aggressive_smoke_local" +seed = 20260723 +artifact_dir = "artifacts/runs" + +[data] +root = "artifacts/processed_preflight_one" +train_cases = 1 +val_cases = 0 +test_cases = 0 +points_per_case = 8192 +batch_size = 512 + +[model] +type = "film_fourier_mlp" +hidden_width = 256 +depth = 4 +activation = "gelu" +coordinate_features = ["x", "y", "sdf"] +fourier_scales = [1.0, 2.0, 4.0, 8.0] +condition_width = 128 +condition_depth = 2 +condition_dim = 128 + +[optim] +lr = 0.0003 +weight_decay = 0.0001 +steps = 20 +log_interval = 10 + +[device] +type = "cuda" +allow_cpu_fallback = false +benchmark_kernels = true + +[loss] +type = "normalized_mse" + +[precision] +dtype = "bf16" + +[checkpoint] +interval_seconds = 0 + +[stability] +max_grad_norm = 1.0 diff --git a/configs/remote_smoke.toml b/configs/remote_smoke.toml index c05edf4..6d825de 100644 --- a/configs/remote_smoke.toml +++ b/configs/remote_smoke.toml @@ -1,19 +1,19 @@ [run] -name = "airfrans-smoke" -timeout_minutes = 45 +name = "airfrans-aggressive-smoke" +timeout_minutes = 360 local_artifact_dir = "artifacts/remote_runs" max_attempts = 2 [provider] kind = "vastai" -disk_gb = 64 -max_price_per_hour = 0.60 +disk_gb = 128 +max_price_per_hour = 0.80 image = "vastai/base:0.0.2" [provider.gpu] name = "RTX 4090" count = 1 -min_vram_gb = 16 +min_vram_gb = 20 [selection] min_reliability = 0.95 @@ -45,12 +45,12 @@ uv run --no-dev python -c "import torch; assert torch.cuda.is_available(); print [data] validation_command = """ -uv run --no-dev python -c "from pathlib import Path; files=sorted(Path('data/processed/minimal').glob('*.npz')); assert len(files) >= 6; print(f'processed_minimal_cases={len(files)}')" +uv run --no-dev python -c "from pathlib import Path; files=sorted(Path('data/processed/full').glob('*.npz')); assert len(files) >= 50; print(f'processed_full_cases={len(files)}')" """ [job] command = """ -uv run --no-dev remote-run smoke-train configs/remote_mlp_tiny.toml --artifact-dir artifacts/current_run --run-id "$AIRFRANS_REMOTE_RUN_ID" +uv run --no-dev remote-run smoke-train configs/aggressive_smoke.toml --artifact-dir artifacts/current_run --run-id "$AIRFRANS_REMOTE_RUN_ID" """ artifact_dir = "artifacts/current_run" heartbeat_file = "artifacts/current_run/heartbeat.json" @@ -59,9 +59,17 @@ metrics_file = "artifacts/current_run/metrics.jsonl" [artifacts] mode = "rsync" required = [ - "final_metrics.json", + "config.toml", "metrics.jsonl", - "checkpoint.pt", + "latest_metrics.json", + "heartbeat.json", + "checkpoint_latest.pt", + "checkpoint_best.pt", + "checkpoint_final.pt", + "final_metrics.json", + "split_manifest.json", + "data_manifest.json", + "normalization.json", "run_manifest.json", "environment_manifest.json", "artifact_manifest.json", diff --git a/pyproject.toml b/pyproject.toml index fd545a1..36abd9f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,4 +27,5 @@ dev = [ "nbclient>=0.10.2", "nbformat>=5.10.4", "skypilot[vast]>=0.12.3.post1", + "pytest>=9.1.1", ] diff --git a/src/airfrans_frontier/cli.py b/src/airfrans_frontier/cli.py index e3db888..35a6a8b 100644 --- a/src/airfrans_frontier/cli.py +++ b/src/airfrans_frontier/cli.py @@ -17,8 +17,16 @@ def build_parser() -> argparse.ArgumentParser: inspect_raw.add_argument("--sample-limit", type=int, default=5) inspect_raw.set_defaults(command="inspect-raw") + process_raw = subparsers.add_parser("process-raw", help="convert raw OpenFOAM cases into training tensors") + process_raw.add_argument("--raw-dir", default=str(DEFAULT_RAW_DATA_DIR)) + process_raw.add_argument("--output-dir", default="data/processed/full") + process_raw.add_argument("--limit", type=int) + process_raw.add_argument("--force", action="store_true") + process_raw.set_defaults(command="process-raw") + train = subparsers.add_parser("train", help="train a configured baseline model") train.add_argument("config", help="path to a training config TOML file") + train.add_argument("--resume", help="path to checkpoint_latest.pt to resume from") train.set_defaults(command="train") return parser @@ -44,6 +52,28 @@ def main(argv: list[str] | None = None) -> int: print(format_raw_inspection(report, sample_limit=args.sample_limit)) return 0 if report.matches_manifest else 1 + if args.command == "process-raw": + if args.limit is not None and args.limit <= 0: + print("error: --limit must be positive", file=sys.stderr) + return 1 + from airfrans_frontier.raw.process import process_raw_dataset + + try: + result = process_raw_dataset( + resolve_path(args.raw_dir), + resolve_path(args.output_dir), + limit=args.limit, + force=args.force, + ) + except (FileNotFoundError, NotADirectoryError, ValueError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + print(f"output_dir: {result.output_dir}") + print(f"case_count: {result.case_count}") + print(f"total_points: {result.total_points}") + print(f"manifest: {result.manifest_path}") + return 0 + if args.command == "train": from airfrans_frontier.runtime import remove_pythonpath_entries @@ -51,7 +81,7 @@ def main(argv: list[str] | None = None) -> int: from airfrans_frontier.training.loop import train_from_config_path try: - result = train_from_config_path(resolve_path(args.config)) + result = train_from_config_path(resolve_path(args.config), resume_path=resolve_path(args.resume) if args.resume else None) except (FileNotFoundError, NotADirectoryError, ValueError, RuntimeError) as exc: print(f"error: {exc}", file=sys.stderr) return 1 diff --git a/src/airfrans_frontier/models/__init__.py b/src/airfrans_frontier/models/__init__.py index 7a24280..9da57ab 100644 --- a/src/airfrans_frontier/models/__init__.py +++ b/src/airfrans_frontier/models/__init__.py @@ -1,5 +1,6 @@ """Baseline model definitions.""" +from airfrans_frontier.models.film import FourierFiLMMLP from airfrans_frontier.models.mlp import PointwiseMLP -__all__ = ["PointwiseMLP"] +__all__ = ["FourierFiLMMLP", "PointwiseMLP"] diff --git a/src/airfrans_frontier/models/film.py b/src/airfrans_frontier/models/film.py new file mode 100644 index 0000000..5c071ff --- /dev/null +++ b/src/airfrans_frontier/models/film.py @@ -0,0 +1,144 @@ +from __future__ import annotations + +import math +from collections.abc import Sequence + +import torch +from torch import nn + + +class FourierFiLMMLP(nn.Module): + """Coordinate trunk modulated by per-simulation condition features. + + `feature_names` identifies coordinate columns. Every non-coordinate column is + treated as a condition feature; repeated per-point condition rows are + collapsed inside `forward` so FiLM parameters are generated once per unique + simulation condition in a batch. + """ + + def __init__( + self, + *, + feature_names: Sequence[str], + output_dim: int, + coordinate_features: Sequence[str] = ("x", "y", "sdf"), + fourier_scales: Sequence[float] = (1.0, 2.0, 4.0, 8.0, 16.0), + trunk_width: int = 1024, + trunk_depth: int = 8, + condition_width: int = 512, + condition_depth: int = 3, + condition_dim: int = 512, + activation: str = "gelu", + ) -> None: + super().__init__() + if output_dim <= 0: + raise ValueError("output_dim must be positive") + if trunk_width <= 0: + raise ValueError("trunk_width must be positive") + if trunk_depth <= 0: + raise ValueError("trunk_depth must be positive") + if condition_width <= 0: + raise ValueError("condition_width must be positive") + if condition_depth <= 0: + raise ValueError("condition_depth must be positive") + if condition_dim <= 0: + raise ValueError("condition_dim must be positive") + + names = tuple(feature_names) + if len(set(names)) != len(names): + raise ValueError("feature_names must be unique") + coordinate_names = tuple(coordinate_features) + missing = [name for name in coordinate_names if name not in names] + if missing: + raise ValueError(f"Missing coordinate features for FourierFiLMMLP: {missing}") + condition_names = tuple(name for name in names if name not in set(coordinate_names)) + if not condition_names: + raise ValueError("FourierFiLMMLP requires at least one condition feature") + + self.feature_names = names + self.coordinate_names = coordinate_names + self.condition_names = condition_names + self.register_buffer("coordinate_indices", torch.tensor([names.index(name) for name in coordinate_names], dtype=torch.long), persistent=False) + self.register_buffer("condition_indices", torch.tensor([names.index(name) for name in condition_names], dtype=torch.long), persistent=False) + self.register_buffer("fourier_scales", torch.tensor(tuple(float(scale) for scale in fourier_scales), dtype=torch.float32), persistent=False) + + coordinate_dim = len(coordinate_names) + fourier_dim = coordinate_dim * (1 + 2 * len(tuple(fourier_scales))) + self.input = nn.Linear(fourier_dim, trunk_width) + self.condition_encoder = _mlp( + input_dim=len(condition_names), + hidden_width=condition_width, + output_dim=condition_dim, + depth=condition_depth, + activation=activation, + ) + self.blocks = nn.ModuleList( + _FiLMResidualBlock(width=trunk_width, condition_dim=condition_dim, activation=activation) + for _ in range(trunk_depth) + ) + self.output_norm = nn.LayerNorm(trunk_width) + self.output = nn.Linear(trunk_width, output_dim) + self.activation = _activation(activation) + + def forward(self, features: torch.Tensor) -> torch.Tensor: + coordinates = features.index_select(dim=1, index=self.coordinate_indices) + conditions = features.index_select(dim=1, index=self.condition_indices) + coordinate_embedding = _fourier_features(coordinates, self.fourier_scales) + hidden = self.input(coordinate_embedding) + + unique_conditions, inverse = torch.unique(conditions, dim=0, return_inverse=True) + condition_embedding = self.condition_encoder(unique_conditions) + for block in self.blocks: + hidden = block(hidden, condition_embedding, inverse) + hidden = self.output_norm(hidden) + hidden = self.activation(hidden) + return self.output(hidden) + + +class _FiLMResidualBlock(nn.Module): + def __init__(self, *, width: int, condition_dim: int, activation: str) -> None: + super().__init__() + self.norm = nn.LayerNorm(width) + self.linear = nn.Linear(width, width) + self.film = nn.Linear(condition_dim, 2 * width) + self.activation = _activation(activation) + + def forward(self, hidden: torch.Tensor, condition_embedding: torch.Tensor, inverse: torch.Tensor) -> torch.Tensor: + gamma_beta = self.film(condition_embedding).index_select(dim=0, index=inverse) + gamma, beta = gamma_beta.chunk(2, dim=1) + update = self.linear(self.activation(self.norm(hidden))) + update = update * (1.0 + gamma) + beta + return hidden + update + + +def _fourier_features(coordinates: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + if scales.numel() == 0: + return coordinates + phases = coordinates.unsqueeze(-1) * scales.to(device=coordinates.device, dtype=coordinates.dtype) * math.pi + encoded = torch.cat((coordinates, torch.sin(phases).flatten(1), torch.cos(phases).flatten(1)), dim=1) + return encoded + + +def _mlp(*, input_dim: int, hidden_width: int, output_dim: int, depth: int, activation: str) -> nn.Sequential: + layers: list[nn.Module] = [] + current_dim = input_dim + for _ in range(depth - 1): + layers.append(nn.Linear(current_dim, hidden_width)) + layers.append(_activation(activation)) + current_dim = hidden_width + layers.append(nn.Linear(current_dim, output_dim)) + layers.append(_activation(activation)) + return nn.Sequential(*layers) + + +def _activation(name: str) -> nn.Module: + normalized = name.lower() + if normalized == "gelu": + return nn.GELU() + if normalized == "relu": + return nn.ReLU() + if normalized == "silu": + return nn.SiLU() + if normalized == "tanh": + return nn.Tanh() + raise ValueError(f"Unsupported activation: {name}") diff --git a/src/airfrans_frontier/raw/process.py b/src/airfrans_frontier/raw/process.py new file mode 100644 index 0000000..ebc9269 --- /dev/null +++ b/src/airfrans_frontier/raw/process.py @@ -0,0 +1,414 @@ +from __future__ import annotations + +import gzip +import json +import math +import re +import time +from collections.abc import Iterator +from dataclasses import dataclass, asdict +from pathlib import Path +from typing import TextIO + +import numpy as np +from numpy.typing import NDArray + +from airfrans_frontier.paths import DEFAULT_RAW_DATA_DIR + +FloatArray = NDArray[np.float32] +INT_RE = re.compile(r"-?\d+") +FLOAT_RE = re.compile(r"[-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?") +SIM_RE = re.compile( + r"^airFoil2D_(?P[^_]+)_" + r"(?P-?\d+(?:\.\d+)?)_" + r"(?P-?\d+(?:\.\d+)?)_" + r"(?P.+)$" +) +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", + ], + dtype="U32", +) +TARGET_NAMES = np.array(["velocity_x", "velocity_y", "pressure", "turbulent_viscosity"], dtype="U32") + + +@dataclass(frozen=True) +class ProcessingResult: + output_dir: Path + case_count: int + total_points: int + manifest_path: Path + + +@dataclass(frozen=True) +class CaseMetadata: + case_id: str + turbulence: str + u_inf: float + alpha_deg: float + naca_params: tuple[float, ...] + nu: float + reynolds: float + timestep: str + + +def process_raw_dataset( + raw_dir: str | Path = DEFAULT_RAW_DATA_DIR, + output_dir: str | Path = "data/processed/full", + *, + limit: int | None = None, + force: bool = False, +) -> ProcessingResult: + raw_root = Path(raw_dir).expanduser() + if not raw_root.is_dir(): + raise NotADirectoryError(f"Raw AirfRANS directory not found: {raw_root}") + out_root = Path(output_dir).expanduser() + out_root.mkdir(parents=True, exist_ok=True) + + case_dirs = sorted(path for path in raw_root.iterdir() if path.is_dir() and path.name.startswith("airFoil2D_")) + if limit is not None: + case_dirs = case_dirs[:limit] + if not case_dirs: + raise ValueError(f"No raw AirfRANS case directories found under: {raw_root}") + + records: list[dict[str, object]] = [] + started = time.perf_counter() + total_points = 0 + for case_dir in case_dirs: + target_path = out_root / f"{case_dir.name}.npz" + if target_path.exists() and not force: + with np.load(target_path, allow_pickle=False) as npz: + points = int(npz["features"].shape[0]) + records.append({"case_id": case_dir.name, "path": str(target_path), "points": points, "skipped_existing": True}) + total_points += points + continue + metadata, features, targets = process_raw_case(case_dir) + _atomic_save_npz( + target_path, + features=features, + targets=targets, + feature_names=FEATURE_NAMES, + target_names=TARGET_NAMES, + metadata=json.dumps(_metadata_json(metadata), sort_keys=True), + ) + points = int(features.shape[0]) + total_points += points + records.append({"case_id": case_dir.name, "path": str(target_path), "points": points, "metadata": _metadata_json(metadata)}) + + manifest = { + "raw_dir": str(raw_root), + "output_dir": str(out_root), + "case_count": len(records), + "total_points": total_points, + "feature_names": FEATURE_NAMES.tolist(), + "target_names": TARGET_NAMES.tolist(), + "elapsed_seconds": time.perf_counter() - started, + "cases": records, + } + manifest_path = out_root / "manifest.json" + manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n") + return ProcessingResult(output_dir=out_root, case_count=len(records), total_points=total_points, manifest_path=manifest_path) + + +def process_raw_case(case_dir: str | Path) -> tuple[CaseMetadata, FloatArray, FloatArray]: + root = Path(case_dir).expanduser() + if not root.is_dir(): + raise NotADirectoryError(f"Raw case directory not found: {root}") + timestep = _latest_timestep(root) + metadata = _case_metadata(root, timestep=timestep) + + points = _parse_vector_list(root / "constant" / "polyMesh" / "points.gz", columns=3).astype(np.float64, copy=False) + owner = _parse_label_list(root / "constant" / "polyMesh" / "owner.gz") + neighbour = _parse_label_list(root / "constant" / "polyMesh" / "neighbour.gz") + U = _parse_vector_list(root / timestep / "U.gz", columns=3) + p = _parse_scalar_list(root / timestep / "p.gz") + nut = _parse_scalar_list(_field_path(root / timestep, "turbulenceProperties:nut.gz", "nut.gz")) + n_cells = int(U.shape[0]) + if p.shape[0] != n_cells or nut.shape[0] != n_cells: + raise ValueError(f"Field row count mismatch in {root}") + + centers, airfoil_centers = _cell_centers_and_airfoil(points, owner, neighbour, root / "constant" / "polyMesh") + if centers.shape[0] != n_cells: + raise ValueError(f"Mesh cell count mismatch in {root}: centers={centers.shape[0]} field={n_cells}") + sdf = _distance_to_airfoil(centers, airfoil_centers) + condition = _condition_features(metadata, n_cells) + features = np.concatenate((centers.astype(np.float32), sdf[:, None], condition), axis=1) + targets = np.stack((U[:, 0], U[:, 1], p, nut), axis=1).astype(np.float32, copy=False) + _validate_finite(root, features, targets) + return metadata, features.astype(np.float32, copy=False), targets + + +def _metadata_json(metadata: CaseMetadata) -> dict[str, object]: + payload = asdict(metadata) + payload["naca_params"] = list(metadata.naca_params) + return payload + + +def _latest_timestep(root: Path) -> str: + candidates: list[tuple[float, str]] = [] + for path in root.iterdir(): + if not path.is_dir(): + continue + try: + value = float(path.name) + except ValueError: + continue + if (path / "U.gz").is_file() and (path / "p.gz").is_file(): + candidates.append((value, path.name)) + if not candidates: + raise ValueError(f"No numeric timestep with U.gz/p.gz found in {root}") + return max(candidates)[1] + + +def _case_metadata(root: Path, *, timestep: str) -> CaseMetadata: + match = SIM_RE.match(root.name) + if match is None: + raise ValueError(f"Unexpected AirfRANS case name: {root.name}") + params = tuple(float(value) for value in match.group("params").split("_")) + u_inf = float(match.group("u_inf")) + alpha = float(match.group("alpha")) + nu = _transport_nu(root / "constant" / "transportProperties") + return CaseMetadata( + case_id=root.name, + turbulence=match.group("turbulence"), + u_inf=u_inf, + alpha_deg=alpha, + naca_params=params, + nu=nu, + reynolds=u_inf / nu, + timestep=timestep, + ) + + +def _transport_nu(path: Path) -> float: + text = _read_text(path, max_bytes=20_000) + match = re.search(r"^\s*nu\s+([^;]+);", text, flags=re.MULTILINE) + if match is None: + raise ValueError(f"Could not parse nu from {path}") + values = FLOAT_RE.findall(match.group(1)) + if not values: + raise ValueError(f"Could not parse numeric nu from {path}") + return float(values[-1]) + + +def _condition_features(metadata: CaseMetadata, n_cells: int) -> FloatArray: + params = np.zeros(4, dtype=np.float32) + mask = np.zeros(4, dtype=np.float32) + for index, value in enumerate(metadata.naca_params[:4]): + params[index] = float(value) + mask[index] = 1.0 + alpha_rad = math.radians(metadata.alpha_deg) + row = np.array( + [ + metadata.u_inf, + math.log(metadata.reynolds), + metadata.alpha_deg, + math.sin(alpha_rad), + math.cos(alpha_rad), + *params.tolist(), + *mask.tolist(), + ], + dtype=np.float32, + ) + return np.repeat(row[None, :], n_cells, axis=0) + + +def _cell_centers_and_airfoil( + points: NDArray[np.float64], + owner: NDArray[np.int64], + neighbour: NDArray[np.int64], + poly_mesh: Path, +) -> tuple[FloatArray, FloatArray]: + n_cells = int(max(int(owner.max(initial=0)), int(neighbour.max(initial=0))) + 1) + sums = np.zeros((n_cells, 2), dtype=np.float64) + counts = np.zeros(n_cells, dtype=np.int32) + boundary = _parse_boundary(poly_mesh / "boundary") + if "aerofoil" not in boundary: + raise ValueError(f"Missing aerofoil boundary patch in {poly_mesh / 'boundary'}") + aerofoil_start = int(boundary["aerofoil"]["startFace"]) + aerofoil_stop = aerofoil_start + int(boundary["aerofoil"]["nFaces"]) + airfoil_faces: list[NDArray[np.float64]] = [] + + for face_index, vertices in enumerate(_iter_faces(poly_mesh / "faces.gz")): + face_center = points[vertices, :2].mean(axis=0) + owner_cell = int(owner[face_index]) + sums[owner_cell] += face_center + counts[owner_cell] += 1 + if face_index < len(neighbour): + neighbour_cell = int(neighbour[face_index]) + sums[neighbour_cell] += face_center + counts[neighbour_cell] += 1 + if aerofoil_start <= face_index < aerofoil_stop: + airfoil_faces.append(face_center) + if np.any(counts == 0): + raise ValueError(f"Mesh contains {int(np.sum(counts == 0))} cells without faces in {poly_mesh}") + if not airfoil_faces: + raise ValueError(f"No aerofoil face centers parsed from {poly_mesh}") + centers = (sums / counts[:, None]).astype(np.float32) + return centers, np.asarray(airfoil_faces, dtype=np.float32) + + +def _distance_to_airfoil(centers: FloatArray, airfoil_centers: FloatArray, *, chunk_size: int = 8192) -> FloatArray: + result = np.empty(centers.shape[0], dtype=np.float32) + airfoil = airfoil_centers.astype(np.float32, copy=False) + for start in range(0, centers.shape[0], chunk_size): + stop = min(start + chunk_size, centers.shape[0]) + diff = centers[start:stop, None, :] - airfoil[None, :, :] + dist2 = np.sum(diff * diff, axis=2) + result[start:stop] = np.sqrt(np.min(dist2, axis=1)).astype(np.float32) + return result + + +def _parse_boundary(path: Path) -> dict[str, dict[str, int | str]]: + text = _read_text(path, max_bytes=500_000) + patches: dict[str, dict[str, int | str]] = {} + for name, body in re.findall(r"\n\s*([A-Za-z][A-Za-z0-9_]*)\s*\n\s*\{(.*?)\n\s*\}", text, flags=re.DOTALL): + n_faces = _assignment(body, "nFaces") + start_face = _assignment(body, "startFace") + patch_type = _assignment(body, "type") or "" + if n_faces is not None and start_face is not None: + patches[name] = {"type": patch_type, "nFaces": int(n_faces), "startFace": int(start_face)} + return patches + + +def _assignment(text: str, key: str) -> str | None: + match = re.search(rf"^\s*{re.escape(key)}\s+([^;]+);", text, flags=re.MULTILINE) + return match.group(1).strip() if match else None + + +def _parse_vector_list(path: Path, *, columns: int) -> NDArray[np.float32]: + rows: list[list[float]] = [] + expected = None + in_values = False + with _open_text(path) as stream: + for line in stream: + stripped = line.strip() + if not in_values: + if expected is None and stripped.isdigit(): + expected = int(stripped) + continue + if expected is not None and stripped == "(": + in_values = True + continue + continue + if stripped == ")": + break + values = [float(item) for item in FLOAT_RE.findall(stripped)] + if len(values) >= columns: + rows.append(values[:columns]) + array = np.asarray(rows, dtype=np.float32) + if expected is not None and array.shape[0] != expected: + raise ValueError(f"{path}: parsed {array.shape[0]} rows, expected {expected}") + return array + + +def _parse_scalar_list(path: Path) -> FloatArray: + return _parse_vector_list(path, columns=1).reshape(-1) + + +def _parse_label_list(path: Path) -> NDArray[np.int64]: + values: list[int] = [] + expected = None + in_values = False + with _open_text(path) as stream: + for line in stream: + stripped = line.strip() + if not in_values: + if expected is None and stripped.isdigit(): + expected = int(stripped) + continue + if expected is not None and stripped == "(": + in_values = True + continue + continue + if stripped == ")": + break + if stripped: + values.append(int(stripped)) + array = np.asarray(values, dtype=np.int64) + if expected is not None and array.shape[0] != expected: + raise ValueError(f"{path}: parsed {array.shape[0]} labels, expected {expected}") + return array + + +def _iter_faces(path: Path) -> Iterator[NDArray[np.int64]]: + expected = None + in_values = False + count = 0 + with _open_text(path) as stream: + for line in stream: + stripped = line.strip() + if not in_values: + if expected is None and stripped.isdigit(): + expected = int(stripped) + continue + if expected is not None and stripped == "(": + in_values = True + continue + continue + if stripped == ")": + break + values = [int(item) for item in INT_RE.findall(stripped)] + if values: + count += 1 + yield np.asarray(values[1:], dtype=np.int64) + if expected is not None and count != expected: + raise ValueError(f"{path}: parsed {count} faces, expected {expected}") + + +def _field_path(root: Path, preferred: str, fallback: str) -> Path: + preferred_path = root / preferred + if preferred_path.is_file(): + return preferred_path + fallback_path = root / fallback + if fallback_path.is_file(): + return fallback_path + raise FileNotFoundError(f"Missing field {preferred} or {fallback} under {root}") + + +def _open_text(path: Path) -> TextIO: + if path.suffix == ".gz": + return gzip.open(path, "rt", errors="replace") + return path.open("rt", errors="replace") + + +def _read_text(path: Path, *, max_bytes: int) -> str: + if path.suffix == ".gz": + with gzip.open(path, "rb") as stream: + data = stream.read(max_bytes) + else: + data = path.read_bytes()[:max_bytes] + return data.decode("utf-8", errors="replace") + + +def _validate_finite(root: Path, features: FloatArray, targets: FloatArray) -> None: + if features.shape[1] != len(FEATURE_NAMES): + raise ValueError(f"Feature width mismatch in {root}: {features.shape[1]} != {len(FEATURE_NAMES)}") + if targets.shape[1] != len(TARGET_NAMES): + raise ValueError(f"Target width mismatch in {root}: {targets.shape[1]} != {len(TARGET_NAMES)}") + if not np.isfinite(features).all(): + raise ValueError(f"Non-finite features in {root}") + if not np.isfinite(targets).all(): + raise ValueError(f"Non-finite targets in {root}") + + +def _atomic_save_npz(path: Path, **arrays: object) -> None: + tmp = path.with_name(path.name + ".tmp.npz") + np.savez_compressed(tmp, **arrays) + tmp.replace(path) diff --git a/src/airfrans_frontier/remote/artifacts.py b/src/airfrans_frontier/remote/artifacts.py index b4a1108..0bc4ba3 100644 --- a/src/airfrans_frontier/remote/artifacts.py +++ b/src/airfrans_frontier/remote/artifacts.py @@ -5,26 +5,71 @@ import json from pathlib import Path from typing import Any, Iterable +import torch -DEFAULT_REQUIRED = ("final_metrics.json", "metrics.jsonl", "checkpoint.pt", "run_manifest.json") +BASE_REQUIRED = ( + "config.toml", + "metrics.jsonl", + "latest_metrics.json", + "heartbeat.json", + "checkpoint_latest.pt", + "checkpoint_best.pt", +) +SUCCESS_REQUIRED = ("final_metrics.json", "checkpoint_final.pt") +FAILURE_REQUIRED = ("failure_report.json",) +DEFAULT_REQUIRED = BASE_REQUIRED -def verify_artifacts(artifact_dir: str | Path, required: Iterable[str] = DEFAULT_REQUIRED) -> dict[str, Any]: +def verify_artifacts( + artifact_dir: str | Path, + required: Iterable[str] = DEFAULT_REQUIRED, + *, + require_terminal: bool = True, +) -> dict[str, Any]: root = Path(artifact_dir) if not root.exists(): raise FileNotFoundError(f"Artifact directory not found: {root}") if not root.is_dir(): raise ValueError(f"Artifact path is not a directory: {root}") - missing = [name for name in required if not (root / name).is_file()] + required_names = tuple(required) + missing = [name for name in required_names if not (root / name).is_file()] if missing: raise ValueError(f"Artifact directory missing required files: {', '.join(missing)}") - _validate_json(root / "final_metrics.json") - _validate_json(root / "run_manifest.json") - _validate_jsonl(root / "metrics.jsonl") + has_final = (root / "final_metrics.json").is_file() + has_failure = (root / "failure_report.json").is_file() + if require_terminal and not has_final and not has_failure: + raise ValueError("Artifact directory has no terminal artifact: final_metrics.json or failure_report.json") + if has_final: + missing_success = [name for name in SUCCESS_REQUIRED if not (root / name).is_file()] + if missing_success: + raise ValueError(f"Successful artifact directory missing files: {', '.join(missing_success)}") + if has_failure: + missing_failure = [name for name in FAILURE_REQUIRED if not (root / name).is_file()] + if missing_failure: + raise ValueError(f"Failed artifact directory missing files: {', '.join(missing_failure)}") - files = sorted(path for path in root.rglob("*") if path.is_file()) + for json_name in ( + "latest_metrics.json", + "heartbeat.json", + "final_metrics.json", + "failure_report.json", + "run_manifest.json", + "split_manifest.json", + "data_manifest.json", + "normalization.json", + ): + path = root / json_name + if path.is_file(): + _validate_json(path) + _validate_jsonl(root / "metrics.jsonl") + for checkpoint_name in ("checkpoint_latest.pt", "checkpoint_best.pt", "checkpoint_final.pt"): + path = root / checkpoint_name + if path.is_file(): + _validate_checkpoint_metadata(path) + + files = sorted(path for path in root.rglob("*") if path.is_file() and path.name not in {"artifact_manifest.json", "checksums.txt"}) manifest = { "artifact_dir": str(root), "file_count": len(files), @@ -69,3 +114,16 @@ def _validate_jsonl(path: Path) -> None: json.loads(stripped) except json.JSONDecodeError as exc: raise ValueError(f"Invalid JSONL artifact {path}:{line_number}: {exc}") from exc + + +def _validate_checkpoint_metadata(path: Path) -> None: + try: + checkpoint = torch.load(path, map_location="cpu", weights_only=False) + except Exception as exc: + raise ValueError(f"Invalid checkpoint artifact {path}: {exc}") from exc + if not isinstance(checkpoint, dict): + raise ValueError(f"Checkpoint artifact is not a mapping: {path}") + required = ("schema_version", "step", "model_state_dict", "optimizer_state_dict", "normalization") + missing = [name for name in required if name not in checkpoint] + if missing: + raise ValueError(f"Checkpoint artifact {path} missing keys: {', '.join(missing)}") diff --git a/src/airfrans_frontier/remote/cli.py b/src/airfrans_frontier/remote/cli.py index b1d7188..a4ee77f 100644 --- a/src/airfrans_frontier/remote/cli.py +++ b/src/airfrans_frontier/remote/cli.py @@ -46,6 +46,7 @@ def build_parser() -> argparse.ArgumentParser: smoke.add_argument("training_config") smoke.add_argument("--artifact-dir", required=True) smoke.add_argument("--run-id", required=True) + smoke.add_argument("--resume", help="path to checkpoint_latest.pt to resume from") smoke.set_defaults(command="smoke-train") run = subparsers.add_parser("run", help="select a Vast offer and execute a SkyPilot run") @@ -82,7 +83,12 @@ def main(argv: list[str] | None = None) -> int: print(json.dumps({"status": "ok", "file_count": manifest["file_count"]}, sort_keys=True)) return 0 if args.command == "smoke-train": - output = run_smoke_training(args.training_config, artifact_dir=args.artifact_dir, run_id=args.run_id) + output = run_smoke_training( + args.training_config, + artifact_dir=args.artifact_dir, + run_id=args.run_id, + resume_path=args.resume or os.environ.get("AIRFRANS_RESUME_CHECKPOINT"), + ) print(f"artifact_dir: {output}") return 0 if args.command == "run": @@ -147,15 +153,18 @@ def _run(config_path: str | Path, *, dry_run: bool, skip_down: bool) -> int: selection = select_offer(config) selection_path = local_run_dir / "selection_manifest.json" selection_path.write_text(json.dumps(selection.to_manifest(), indent=2, sort_keys=True) + "\n") - - state("RENDERING_SKYPILOT_CONFIG", selected_offer_id=selection.selected_offer_id) - sky_yaml = render_skypilot_yaml(config, selection, run_id=run_id) - sky_yaml_path = local_run_dir / "sky.yaml" - sky_yaml_path.write_text(sky_yaml) (local_run_dir / "config.toml").write_text(config.raw_text) write_skyignore(config) if dry_run: + sky_yaml_path = _write_attempt_sky_yaml( + config=config, + selection=selection, + run_id=run_id, + local_run_dir=local_run_dir, + resume_checkpoint=None, + attempt=1, + ) state("DRY_RUN", selected_offer_id=selection.selected_offer_id, sky_yaml=str(sky_yaml_path)) print(f"run_id: {run_id}") print(f"selection: {selection_path}") @@ -163,28 +172,122 @@ def _run(config_path: str | Path, *, dry_run: bool, skip_down: bool) -> int: return 0 env = _subprocess_env() - try: - state("PROVISIONING", selected_offer_id=selection.selected_offer_id) - _run_checked(["sky", "launch", "-c", run_id, str(sky_yaml_path), "-y"], env=env, timeout=config.run.timeout_minutes * 60) - - state("COLLECTING", selected_offer_id=selection.selected_offer_id) - _collect_with_rsync(cluster=run_id, remote_dir=config.job.artifact_dir, local_dir=local_run_dir, env=env) - - state("VERIFYING_ARTIFACTS", selected_offer_id=selection.selected_offer_id) - verify_artifacts(local_run_dir, required=config.artifacts.required) - except Exception as exc: - state("FAILED", selected_offer_id=selection.selected_offer_id, error=str(exc)) + last_error: str | None = None + for attempt in range(1, config.run.max_attempts + 1): + resume_checkpoint = _stage_resume_checkpoint(local_run_dir, run_id) + sky_yaml_path = _write_attempt_sky_yaml( + config=config, + selection=selection, + run_id=run_id, + local_run_dir=local_run_dir, + resume_checkpoint=resume_checkpoint, + attempt=attempt, + ) + state( + "PROVISIONING", + selected_offer_id=selection.selected_offer_id, + attempt=attempt, + resume_checkpoint=str(resume_checkpoint) if resume_checkpoint is not None else None, + ) + return_code = _run_sky_with_periodic_collection( + cluster=run_id, + sky_yaml_path=sky_yaml_path, + config=config, + local_run_dir=local_run_dir, + env=env, + ) + state("COLLECTING", selected_offer_id=selection.selected_offer_id, attempt=attempt, return_code=return_code) + _collect_best_effort(cluster=run_id, remote_dir=config.job.artifact_dir, local_dir=local_run_dir, env=env) + status = _classify_artifacts(local_run_dir) + if status == "success": + state("VERIFYING_ARTIFACTS", selected_offer_id=selection.selected_offer_id, attempt=attempt) + verify_artifacts(local_run_dir, required=config.artifacts.required) + if config.cleanup.on_success == "sky_down" and not skip_down: + state("CLEANING_UP", selected_offer_id=selection.selected_offer_id, attempt=attempt) + _run_checked(["sky", "down", run_id, "-y"], env=env, timeout=300) + state("SUCCEEDED", selected_offer_id=selection.selected_offer_id, attempt=attempt) + print(f"run_id: {run_id}") + print(f"artifacts: {local_run_dir}") + return 0 + if status == "failure_report": + last_error = "training wrote failure_report.json" + state("FAILED_TRAINING", selected_offer_id=selection.selected_offer_id, attempt=attempt, error=last_error) + if config.cleanup.on_failure == "sky_down" and not skip_down: + _run_best_effort(["sky", "down", run_id, "-y"], env=env) + raise RuntimeError(last_error) + last_error = f"SkyPilot job ended without terminal artifacts, return_code={return_code}" + state("RETRYING", selected_offer_id=selection.selected_offer_id, attempt=attempt, error=last_error) if config.cleanup.on_failure == "sky_down" and not skip_down: _run_best_effort(["sky", "down", run_id, "-y"], env=env) - raise - else: - if config.cleanup.on_success == "sky_down" and not skip_down: - state("CLEANING_UP", selected_offer_id=selection.selected_offer_id) - _run_checked(["sky", "down", run_id, "-y"], env=env, timeout=300) - state("SUCCEEDED", selected_offer_id=selection.selected_offer_id) - print(f"run_id: {run_id}") - print(f"artifacts: {local_run_dir}") - return 0 + + state("FAILED", selected_offer_id=selection.selected_offer_id, error=last_error or "max attempts exhausted") + raise RuntimeError(last_error or "max attempts exhausted") + + +def _write_attempt_sky_yaml( + *, + config: RemoteRunConfig, + selection: SelectionResult, + run_id: str, + local_run_dir: Path, + resume_checkpoint: Path | None, + attempt: int, +) -> Path: + sky_yaml = render_skypilot_yaml( + config, + selection, + run_id=run_id, + resume_checkpoint=resume_checkpoint, + ) + sky_yaml_path = local_run_dir / f"sky_attempt_{attempt}.yaml" + sky_yaml_path.write_text(sky_yaml) + if attempt == 1: + (local_run_dir / "sky.yaml").write_text(sky_yaml) + return sky_yaml_path + + +def _stage_resume_checkpoint(local_run_dir: Path, run_id: str) -> Path | None: + latest = local_run_dir / "checkpoint_latest.pt" + if not latest.is_file(): + return None + resume_dir = Path(".airfrans_resume") / run_id + resume_dir.mkdir(parents=True, exist_ok=True) + destination = resume_dir / "checkpoint_latest.pt" + shutil.copy2(latest, destination) + return destination + + +def _run_sky_with_periodic_collection( + *, + cluster: str, + sky_yaml_path: Path, + config: RemoteRunConfig, + local_run_dir: Path, + env: dict[str, str], +) -> int: + process = subprocess.Popen(["sky", "launch", "-c", cluster, str(sky_yaml_path), "-y"], env=env) + deadline = time.monotonic() + config.run.timeout_minutes * 60 + next_collect = time.monotonic() + config.run.artifact_sync_interval_seconds + while True: + return_code = process.poll() + if return_code is not None: + return int(return_code) + now = time.monotonic() + if now >= deadline: + process.terminate() + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + process.kill() + return int(process.returncode or 124) + if config.run.artifact_sync_interval_seconds == 0 or now >= next_collect: + _collect_best_effort(cluster=cluster, remote_dir=config.job.artifact_dir, local_dir=local_run_dir, env=env) + if _classify_artifacts(local_run_dir) in {"success", "failure_report"}: + return_code = process.poll() + if return_code is not None: + return int(return_code) + next_collect = now + max(1, config.run.artifact_sync_interval_seconds) + time.sleep(min(5.0, max(0.1, next_collect - now))) def _collect_with_rsync(*, cluster: str, remote_dir: Path, local_dir: Path, env: dict[str, str]) -> None: @@ -193,6 +296,23 @@ def _collect_with_rsync(*, cluster: str, remote_dir: Path, local_dir: Path, env: _run_checked(["rsync", "-Pavz", source, f"{local_dir}/"], env=env, timeout=600) +def _collect_best_effort(*, cluster: str, remote_dir: Path, local_dir: Path, env: dict[str, str]) -> None: + try: + _collect_with_rsync(cluster=cluster, remote_dir=remote_dir, local_dir=local_dir, env=env) + except Exception: + pass + + +def _classify_artifacts(local_run_dir: Path) -> str: + if (local_run_dir / "final_metrics.json").is_file() and (local_run_dir / "checkpoint_final.pt").is_file(): + return "success" + if (local_run_dir / "failure_report.json").is_file(): + return "failure_report" + if (local_run_dir / "checkpoint_latest.pt").is_file(): + return "restartable" + return "incomplete" + + def _run_checked(argv: list[str], *, env: dict[str, str], timeout: int) -> None: subprocess.run(argv, check=True, env=env, timeout=timeout) diff --git a/src/airfrans_frontier/remote/config.py b/src/airfrans_frontier/remote/config.py index 69cc305..dfcfd63 100644 --- a/src/airfrans_frontier/remote/config.py +++ b/src/airfrans_frontier/remote/config.py @@ -12,6 +12,7 @@ class RunConfig: timeout_minutes: int local_artifact_dir: Path max_attempts: int + artifact_sync_interval_seconds: int @dataclass(frozen=True) @@ -190,6 +191,7 @@ def load_remote_run_config(path: str | Path) -> RemoteRunConfig: timeout_minutes=_integer(run_raw, "timeout_minutes", minimum=1, default=60), local_artifact_dir=_path(run_raw, "local_artifact_dir", default="artifacts/remote_runs"), max_attempts=_integer(run_raw, "max_attempts", minimum=1, default=2), + artifact_sync_interval_seconds=_integer(run_raw, "artifact_sync_interval_seconds", minimum=0, default=1800), ), provider=provider, selection=selection, diff --git a/src/airfrans_frontier/remote/skypilot.py b/src/airfrans_frontier/remote/skypilot.py index 4e9c9b6..2d74f15 100644 --- a/src/airfrans_frontier/remote/skypilot.py +++ b/src/airfrans_frontier/remote/skypilot.py @@ -6,7 +6,13 @@ from airfrans_frontier.remote.config import RemoteRunConfig from airfrans_frontier.remote.vast import SelectionResult -def render_skypilot_yaml(config: RemoteRunConfig, selection: SelectionResult, *, run_id: str) -> str: +def render_skypilot_yaml( + config: RemoteRunConfig, + selection: SelectionResult, + *, + run_id: str, + resume_checkpoint: str | Path | None = None, +) -> str: setup = _compose_setup(config) run = _compose_run(config, run_id=run_id) accelerator = _accelerator(config) @@ -25,13 +31,17 @@ def render_skypilot_yaml(config: RemoteRunConfig, selection: SelectionResult, *, if not image.startswith("docker:"): image = f"docker:{image}" lines.append(f" image_id: {image}") + env_lines = [ + "envs:", + f" AIRFRANS_REMOTE_RUN_ID: {run_id}", + ] + if resume_checkpoint is not None: + env_lines.append(f" AIRFRANS_RESUME_CHECKPOINT: {_yaml_scalar(str(resume_checkpoint))}") lines.extend( [ "", f"workdir: {_yaml_scalar(str(config.workspace.workdir))}", - "", - "envs:", - f" AIRFRANS_REMOTE_RUN_ID: {run_id}", + *env_lines, "", "setup: |", *_indent_block(setup), diff --git a/src/airfrans_frontier/remote/smoke.py b/src/airfrans_frontier/remote/smoke.py index c9b5370..bc62b43 100644 --- a/src/airfrans_frontier/remote/smoke.py +++ b/src/airfrans_frontier/remote/smoke.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import os import platform import shutil import subprocess @@ -13,7 +14,13 @@ from airfrans_frontier.remote.artifacts import verify_artifacts from airfrans_frontier.runtime import remove_pythonpath_entries -def run_smoke_training(config_path: str | Path, *, artifact_dir: str | Path, run_id: str) -> Path: +def run_smoke_training( + config_path: str | Path, + *, + artifact_dir: str | Path, + run_id: str, + resume_path: str | Path | None = None, +) -> Path: remove_pythonpath_entries() from airfrans_frontier.training.loop import train_from_config_path @@ -40,13 +47,32 @@ def run_smoke_training(config_path: str | Path, *, artifact_dir: str | Path, run "timestamp": time.time(), }, ) - result = train_from_config_path(config_path) + previous_observability_dir = os.environ.get("AIRFRANS_OBSERVABILITY_DIR") + previous_run_id = os.environ.get("AIRFRANS_REMOTE_RUN_ID") + os.environ["AIRFRANS_OBSERVABILITY_DIR"] = str(output_dir) + os.environ["AIRFRANS_REMOTE_RUN_ID"] = run_id + try: + result = train_from_config_path(config_path, resume_path=resume_path or os.environ.get("AIRFRANS_RESUME_CHECKPOINT")) + finally: + if previous_observability_dir is None: + os.environ.pop("AIRFRANS_OBSERVABILITY_DIR", None) + else: + os.environ["AIRFRANS_OBSERVABILITY_DIR"] = previous_observability_dir + if previous_run_id is None: + os.environ.pop("AIRFRANS_REMOTE_RUN_ID", None) + else: + os.environ["AIRFRANS_REMOTE_RUN_ID"] = previous_run_id finished = time.time() training_dir = result.run_dir required_from_training = [ "final_metrics.json", "metrics.jsonl", + "latest_metrics.json", + "heartbeat.json", + "checkpoint_latest.pt", + "checkpoint_best.pt", + "checkpoint_final.pt", "checkpoint.pt", "config.toml", "normalization.json", @@ -67,16 +93,23 @@ def run_smoke_training(config_path: str | Path, *, artifact_dir: str | Path, run "training_run_dir": str(training_dir), "artifact_dir": str(output_dir), "final_metrics_path": str(output_dir / "final_metrics.json"), - "checkpoint_path": str(output_dir / "checkpoint.pt"), + "checkpoint_path": str(output_dir / "checkpoint_latest.pt"), + "resume_path": str(resume_path) if resume_path is not None else None, } _write_json(output_dir / "run_manifest.json", run_manifest) + latest_metrics = _read_json(output_dir / "latest_metrics.json") _write_json( heartbeat_path, { "run_id": run_id, "phase": "completed", + "epoch": latest_metrics.get("epoch"), + "step": latest_metrics.get("step"), + "latest_checkpoint": "checkpoint_final.pt", + "latest_metrics": latest_metrics, "started_at": started, "finished_at": finished, + "updated_at": time.time(), "timestamp": time.time(), }, ) @@ -123,5 +156,12 @@ def environment_manifest() -> dict[str, Any]: return manifest +def _read_json(path: Path) -> dict[str, Any]: + if not path.is_file(): + return {} + data = json.loads(path.read_text()) + return data if isinstance(data, dict) else {} + + def _write_json(path: Path, data: Any) -> None: path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n") diff --git a/src/airfrans_frontier/training/artifacts.py b/src/airfrans_frontier/training/artifacts.py index 481df23..bafbc02 100644 --- a/src/airfrans_frontier/training/artifacts.py +++ b/src/airfrans_frontier/training/artifacts.py @@ -1,6 +1,9 @@ from __future__ import annotations +import hashlib import json +import os +import time from datetime import datetime, timezone from pathlib import Path from typing import Any, Mapping @@ -9,8 +12,9 @@ import torch class ArtifactWriter: - def __init__(self, run_dir: Path) -> None: + def __init__(self, run_dir: Path, *, mirror_dir: Path | None = None) -> None: self.run_dir = run_dir + self.mirror_dir = mirror_dir @classmethod def create(cls, base_dir: str | Path, run_name: str) -> ArtifactWriter: @@ -24,10 +28,20 @@ class ArtifactWriter: candidate = base_path / f"{timestamp}_{safe_name}_{suffix}" suffix += 1 candidate.mkdir(parents=True) - return cls(candidate) + mirror_dir = _observability_dir() + return cls(candidate, mirror_dir=mirror_dir) + + @classmethod + def resume_or_create(cls, base_dir: str | Path, run_name: str, resume_path: str | Path | None) -> ArtifactWriter: + mirror_dir = _observability_dir() + if resume_path is not None: + candidate = Path(resume_path).expanduser().parent + if (candidate / "metrics.jsonl").is_file(): + return cls(candidate, mirror_dir=mirror_dir) + return cls.create(base_dir, run_name) def write_config(self, config_text: str) -> None: - (self.run_dir / "config.toml").write_text(config_text) + self._write_text_file("config.toml", config_text) def write_split_manifest(self, split_manifest: Mapping[str, Any]) -> None: self.write_json("split_manifest.json", dict(split_manifest)) @@ -36,17 +50,124 @@ class ArtifactWriter: self.write_json("normalization.json", dict(normalization)) def append_metrics(self, metrics: Mapping[str, Any]) -> None: - with (self.run_dir / "metrics.jsonl").open("a") as file: - file.write(json.dumps(dict(metrics), sort_keys=True) + "\n") + payload = dict(metrics) + line = json.dumps(payload, sort_keys=True) + "\n" + self._append_text("metrics.jsonl", line) + self._write_json_file("latest_metrics.json", payload) + self._write_json_file("heartbeat.json", _heartbeat_payload(payload, self.run_dir)) def write_final_metrics(self, metrics: Mapping[str, Any]) -> None: self.write_json("final_metrics.json", dict(metrics)) def write_json(self, name: str, data: Mapping[str, Any]) -> None: - (self.run_dir / name).write_text(json.dumps(dict(data), indent=2, sort_keys=True) + "\n") + self._write_json_file(name, dict(data)) - def save_checkpoint(self, payload: Mapping[str, Any]) -> None: - torch.save(_to_cpu(payload), self.run_dir / "checkpoint.pt") + def write_failure_report(self, report: Mapping[str, Any]) -> None: + self.write_json("failure_report.json", dict(report)) + + def write_artifact_manifest(self) -> None: + manifest_text, checksums_text = _artifact_manifest_text(self.run_dir) + self._write_text_file("artifact_manifest.json", manifest_text) + self._write_text_file("checksums.txt", checksums_text) + + def save_checkpoint(self, payload: Mapping[str, Any], *, name: str) -> None: + cpu_payload = _to_cpu(payload) + self._save_checkpoint_file(name, cpu_payload) + + def _save_checkpoint_file(self, name: str, payload: Mapping[str, Any]) -> None: + _atomic_torch_save(payload, self.run_dir / name) + if self.mirror_dir is not None: + self.mirror_dir.mkdir(parents=True, exist_ok=True) + _atomic_torch_save(payload, self.mirror_dir / name) + + def _append_text(self, name: str, text: str) -> None: + _append_text(self.run_dir / name, text) + if self.mirror_dir is not None: + self.mirror_dir.mkdir(parents=True, exist_ok=True) + _append_text(self.mirror_dir / name, text) + + def _write_json_file(self, name: str, data: Mapping[str, Any]) -> None: + text = json.dumps(dict(data), indent=2, sort_keys=True) + "\n" + _write_text(self.run_dir / name, text) + if self.mirror_dir is not None: + self.mirror_dir.mkdir(parents=True, exist_ok=True) + _write_text(self.mirror_dir / name, text) + + def _write_text_file(self, name: str, text: str) -> None: + _write_text(self.run_dir / name, text) + if self.mirror_dir is not None: + self.mirror_dir.mkdir(parents=True, exist_ok=True) + _write_text(self.mirror_dir / name, text) + + +def _heartbeat_payload(metrics: Mapping[str, Any], run_dir: Path) -> dict[str, Any]: + return { + "run_id": os.environ.get("AIRFRANS_REMOTE_RUN_ID"), + "phase": metrics.get("phase", "training"), + "epoch": metrics.get("epoch"), + "step": metrics.get("step"), + "updated_at": time.time(), + "run_dir": str(run_dir), + "latest_checkpoint": metrics.get("latest_checkpoint"), + "latest_metrics": dict(metrics), + } + + +def _observability_dir() -> Path | None: + raw = os.environ.get("AIRFRANS_OBSERVABILITY_DIR") + if not raw: + return None + return Path(raw).expanduser() + + +def _append_text(path: Path, text: str) -> None: + with path.open("a") as file: + file.write(text) + file.flush() + os.fsync(file.fileno()) + + +def _write_text(path: Path, text: str) -> None: + with path.open("w") as file: + file.write(text) + file.flush() + os.fsync(file.fileno()) + + +def _atomic_torch_save(payload: Mapping[str, Any], path: Path) -> None: + tmp_path = path.with_name(f"{path.name}.tmp") + torch.save(payload, tmp_path) + tmp_path.replace(path) + + +def _artifact_manifest_text(root: Path) -> tuple[str, str]: + files = sorted( + path + for path in root.rglob("*") + if path.is_file() and path.name not in {"artifact_manifest.json", "checksums.txt"} + ) + manifest = { + "artifact_dir": str(root), + "file_count": len(files), + "files": [ + { + "path": str(path.relative_to(root)), + "bytes": path.stat().st_size, + "sha256": _sha256_file(path), + } + for path in files + ], + } + checksums = "".join(f"{item['sha256']} {item['path']}\n" for item in manifest["files"]) + return json.dumps(manifest, indent=2, sort_keys=True) + "\n", checksums + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() def _safe_name(value: str) -> str: diff --git a/src/airfrans_frontier/training/config.py b/src/airfrans_frontier/training/config.py index 62cbbac..8702ede 100644 --- a/src/airfrans_frontier/training/config.py +++ b/src/airfrans_frontier/training/config.py @@ -29,6 +29,11 @@ class ModelConfig: hidden_width: int depth: int activation: str + coordinate_features: tuple[str, ...] + fourier_scales: tuple[float, ...] + condition_width: int + condition_depth: int + condition_dim: int @dataclass(frozen=True) @@ -39,6 +44,21 @@ class OptimConfig: log_interval: int | None = None +@dataclass(frozen=True) +class CheckpointConfig: + interval_seconds: int + + +@dataclass(frozen=True) +class StabilityConfig: + max_grad_norm: float | None + + +@dataclass(frozen=True) +class PrecisionConfig: + dtype: str + + @dataclass(frozen=True) class DeviceConfig: type: str @@ -61,6 +81,9 @@ class TrainingConfig: optim: OptimConfig device: DeviceConfig loss: LossConfig + checkpoint: CheckpointConfig + stability: StabilityConfig + precision: PrecisionConfig _REQUIRED_SECTIONS = ("run", "data", "model", "optim", "device", "loss") @@ -89,6 +112,21 @@ def load_training_config(path: str | Path) -> TrainingConfig: optim_raw = raw["optim"] device_raw = raw["device"] loss_raw = raw["loss"] + checkpoint_raw = raw.get("checkpoint", {}) + if checkpoint_raw is None: + checkpoint_raw = {} + if not isinstance(checkpoint_raw, dict): + raise ValueError("Training config [checkpoint] section must be a table") + stability_raw = raw.get("stability", {}) + if stability_raw is None: + stability_raw = {} + if not isinstance(stability_raw, dict): + raise ValueError("Training config [stability] section must be a table") + precision_raw = raw.get("precision", {}) + if precision_raw is None: + precision_raw = {} + if not isinstance(precision_raw, dict): + raise ValueError("Training config [precision] section must be a table") run = RunConfig( name=_string(run_raw, "name"), @@ -104,10 +142,15 @@ def load_training_config(path: str | Path) -> TrainingConfig: batch_size=_integer(data_raw, "batch_size", minimum=1), ) model = ModelConfig( - type=_choice(_string(model_raw, "type"), {"mlp"}, "model.type"), + type=_choice(_string(model_raw, "type"), {"mlp", "film_fourier_mlp"}, "model.type"), hidden_width=_integer(model_raw, "hidden_width", minimum=1), depth=_integer(model_raw, "depth", minimum=1), activation=_choice(_string(model_raw, "activation").lower(), {"gelu", "relu", "silu", "tanh"}, "model.activation"), + coordinate_features=_string_tuple(model_raw, "coordinate_features", default=("x", "y", "sdf")), + fourier_scales=_number_tuple(model_raw, "fourier_scales", default=(1.0, 2.0, 4.0, 8.0, 16.0)), + condition_width=_integer(model_raw, "condition_width", minimum=1, default=_integer(model_raw, "hidden_width", minimum=1)), + condition_depth=_integer(model_raw, "condition_depth", minimum=1, default=2), + condition_dim=_integer(model_raw, "condition_dim", minimum=1, default=_integer(model_raw, "hidden_width", minimum=1)), ) optim = OptimConfig( lr=_number(optim_raw, "lr", minimum=0.0, exclusive_minimum=True), @@ -121,6 +164,15 @@ def load_training_config(path: str | Path) -> TrainingConfig: benchmark_kernels=_boolean(device_raw, "benchmark_kernels"), ) loss = LossConfig(type=_choice(_string(loss_raw, "type"), {"normalized_mse"}, "loss.type")) + checkpoint = CheckpointConfig( + interval_seconds=_integer(checkpoint_raw, "interval_seconds", minimum=0, default=1800), + ) + stability = StabilityConfig( + max_grad_norm=_optional_number(stability_raw, "max_grad_norm", minimum=0.0, exclusive_minimum=True), + ) + precision = PrecisionConfig( + dtype=_choice(_string(precision_raw, "dtype", default="float32").lower(), {"float32", "bf16"}, "precision.dtype"), + ) requested_cases = data.train_cases + data.val_cases + data.test_cases if requested_cases <= 0: @@ -135,6 +187,9 @@ def load_training_config(path: str | Path) -> TrainingConfig: optim=optim, device=device, loss=loss, + checkpoint=checkpoint, + stability=stability, + precision=precision, ) @@ -146,15 +201,23 @@ def _path(section: dict[str, Any], key: str) -> Path: return Path.cwd() / path -def _string(section: dict[str, Any], key: str) -> str: - value = _required(section, key) +def _string(section: dict[str, Any], key: str, *, default: str | None = None) -> str: + if key not in section: + if default is not None: + return default + raise ValueError(f"Training config missing key: {key}") + value = section[key] if not isinstance(value, str) or not value: raise ValueError(f"Expected non-empty string for {key}") return value -def _integer(section: dict[str, Any], key: str, *, minimum: int | None = None) -> int: - value = _required(section, key) +def _integer(section: dict[str, Any], key: str, *, minimum: int | None = None, default: int | None = None) -> int: + if key not in section: + if default is not None: + return default + raise ValueError(f"Training config missing key: {key}") + value = section[key] if isinstance(value, bool) or not isinstance(value, int): raise ValueError(f"Expected integer for {key}") if minimum is not None and value < minimum: @@ -187,6 +250,48 @@ def _number( return result +def _optional_number( + section: dict[str, Any], + key: str, + *, + minimum: float | None = None, + exclusive_minimum: bool = False, +) -> float | None: + if key not in section: + return None + return _number(section, key, minimum=minimum, exclusive_minimum=exclusive_minimum) + + +def _number_tuple(section: dict[str, Any], key: str, *, default: tuple[float, ...]) -> tuple[float, ...]: + if key not in section: + return default + value = section[key] + if not isinstance(value, list) or not value: + raise ValueError(f"Expected non-empty number array for {key}") + result: list[float] = [] + for item in value: + if isinstance(item, bool) or not isinstance(item, (int, float)): + raise ValueError(f"Expected number array for {key}") + result.append(float(item)) + return tuple(result) + + +def _string_tuple(section: dict[str, Any], key: str, *, default: tuple[str, ...]) -> tuple[str, ...]: + if key not in section: + return default + value = section[key] + if not isinstance(value, list) or not value: + raise ValueError(f"Expected non-empty string array for {key}") + result: list[str] = [] + for item in value: + if not isinstance(item, str) or not item: + raise ValueError(f"Expected non-empty string array for {key}") + result.append(item) + if len(set(result)) != len(result): + raise ValueError(f"Expected unique strings for {key}") + return tuple(result) + + def _boolean(section: dict[str, Any], key: str) -> bool: value = _required(section, key) if not isinstance(value, bool): diff --git a/src/airfrans_frontier/training/loop.py b/src/airfrans_frontier/training/loop.py index f47e0ec..31a4fb9 100644 --- a/src/airfrans_frontier/training/loop.py +++ b/src/airfrans_frontier/training/loop.py @@ -1,5 +1,7 @@ from __future__ import annotations +import hashlib +import os import random import time from dataclasses import asdict, dataclass @@ -10,7 +12,7 @@ import numpy as np import torch from torch.nn import functional as F -from airfrans_frontier.models import PointwiseMLP +from airfrans_frontier.models import FourierFiLMMLP, PointwiseMLP from airfrans_frontier.training.artifacts import ArtifactWriter from airfrans_frontier.training.config import TrainingConfig, load_training_config from airfrans_frontier.training.data import DatasetBundle, build_dataset_bundle, load_processed_dataset @@ -22,6 +24,12 @@ from airfrans_frontier.training.normalize import ( normalize_targets, ) +CHECKPOINT_SCHEMA_VERSION = 1 +LATEST_CHECKPOINT = "checkpoint_latest.pt" +BEST_CHECKPOINT = "checkpoint_best.pt" +FINAL_CHECKPOINT = "checkpoint_final.pt" +LEGACY_CHECKPOINT = "checkpoint.pt" + @dataclass(frozen=True) class TrainingResult: @@ -29,14 +37,17 @@ class TrainingResult: final_metrics: dict[str, Any] -def train_from_config_path(path: str | Path) -> TrainingResult: +def train_from_config_path(path: str | Path, resume_path: str | Path | None = None) -> TrainingResult: config = load_training_config(path) - return train(config) + return train(config, resume_path=resume_path or os.environ.get("AIRFRANS_RESUME_CHECKPOINT")) -def train(config: TrainingConfig) -> TrainingResult: +def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> TrainingResult: _seed_all(config.run.seed) device = select_device(config) + resume = Path(resume_path).expanduser() if resume_path else None + writer = ArtifactWriter.resume_or_create(config.run.artifact_dir, config.run.name, resume) + writer.write_config(config.config_text) samples = load_processed_dataset(config.data.root) bundle = build_dataset_bundle( @@ -53,6 +64,27 @@ def train(config: TrainingConfig) -> TrainingResult: feature_names=bundle.feature_names, target_names=bundle.target_names, ) + writer.write_split_manifest(bundle.split.to_dict()) + writer.write_json( + "data_manifest.json", + { + "root": str(config.data.root), + "case_count": len(samples), + "total_points": sum(sample.num_points for sample in samples), + "feature_names": list(bundle.feature_names), + "target_names": list(bundle.target_names), + "cases": [ + { + "case_id": sample.case_id, + "points": sample.num_points, + "path": str(sample.source_path), + } + for sample in samples + ], + }, + ) + writer.write_normalization(stats.to_dict()) + train_features = normalize_features(bundle.train.features, stats) train_targets = normalize_targets(bundle.train.targets, stats) val_features = normalize_features(bundle.val.features, stats) if bundle.val is not None else None @@ -60,24 +92,54 @@ def train(config: TrainingConfig) -> TrainingResult: test_features = normalize_features(bundle.test.features, stats) if bundle.test is not None else None test_targets = normalize_targets(bundle.test.targets, stats) if bundle.test is not None else None - model = PointwiseMLP( - input_dim=train_features.shape[1], - output_dim=train_targets.shape[1], - hidden_width=config.model.hidden_width, - depth=config.model.depth, - activation=config.model.activation, - ).to(device) + model = _build_model(config, bundle, output_dim=train_targets.shape[1]).to(device) optimizer = torch.optim.AdamW( model.parameters(), lr=config.optim.lr, weight_decay=config.optim.weight_decay, ) - writer = ArtifactWriter.create(config.run.artifact_dir, config.run.name) - writer.write_config(config.config_text) - writer.write_split_manifest(bundle.split.to_dict()) - writer.write_normalization(stats.to_dict()) + rng = np.random.default_rng(config.run.seed + 404) started = time.perf_counter() + start_step = 0 + best_val_loss: float | None = None + initial_train_loss: float | None = None + + if resume is not None: + try: + checkpoint = _load_checkpoint(resume, device) + _validate_resume_checkpoint(checkpoint, config, bundle, stats) + model.load_state_dict(checkpoint["model_state_dict"]) + optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) + start_step = int(checkpoint["step"]) + best_val_loss = _optional_float(checkpoint.get("best_val_loss")) + initial_train_loss = _optional_float(checkpoint.get("initial_train_loss")) + _restore_rng_state(checkpoint, rng) + except Exception as exc: + _write_failure( + writer, + phase="resume", + step=0, + error_type=type(exc).__name__, + error_message=str(exc), + latest_checkpoint=str(resume), + ) + raise + writer.append_metrics( + _log_metrics( + event="resume", + step=start_step, + train_loss=None, + val_loss=best_val_loss, + elapsed_seconds=0.0, + lr=_learning_rate(optimizer), + grad_norm=None, + points_per_sec=None, + device=device, + latest_checkpoint=LATEST_CHECKPOINT, + ) + ) + initial_train = evaluate_arrays( model, train_features, @@ -98,64 +160,188 @@ def train(config: TrainingConfig) -> TrainingResult: if val_features is not None and val_targets is not None else None ) + if initial_train_loss is None: + initial_train_loss = initial_train["loss"] + if best_val_loss is None and initial_val is not None: + best_val_loss = initial_val["loss"] + writer.append_metrics( _log_metrics( - step=0, + event="initial_eval" if start_step == 0 else "resume_eval", + step=start_step, train_loss=initial_train["loss"], val_loss=initial_val["loss"] if initial_val is not None else None, elapsed_seconds=0.0, + lr=_learning_rate(optimizer), + grad_norm=None, + points_per_sec=None, + device=device, + latest_checkpoint=LATEST_CHECKPOINT, ) ) + _save_training_checkpoint( + writer, + LATEST_CHECKPOINT, + config=config, + bundle=bundle, + stats=stats, + model=model, + optimizer=optimizer, + rng=rng, + step=start_step, + best_val_loss=best_val_loss, + initial_train_loss=initial_train_loss, + ) + _save_training_checkpoint( + writer, + BEST_CHECKPOINT, + config=config, + bundle=bundle, + stats=stats, + model=model, + optimizer=optimizer, + rng=rng, + step=start_step, + best_val_loss=best_val_loss, + initial_train_loss=initial_train_loss, + ) log_interval = config.optim.log_interval or max(1, config.optim.steps // 10) - rng = np.random.default_rng(config.run.seed + 404) + last_checkpoint_at = time.monotonic() + last_log_at = time.perf_counter() + last_log_step = start_step + last_grad_norm: float | None = None model.train() - for step in range(1, config.optim.steps + 1): - batch_features, batch_targets = _sample_batch( - train_features, - train_targets, - batch_size=config.data.batch_size, - rng=rng, - ) - features_tensor = _to_device(batch_features, device) - targets_tensor = _to_device(batch_targets, device) - - optimizer.zero_grad(set_to_none=True) - predictions = model(features_tensor) - loss = F.mse_loss(predictions, targets_tensor) - loss.backward() - optimizer.step() - - if step % log_interval == 0 or step == config.optim.steps: - train_eval = evaluate_arrays( - model, + try: + for step in range(start_step + 1, config.optim.steps + 1): + batch_features, batch_targets = _sample_batch( train_features, train_targets, batch_size=config.data.batch_size, - device=device, - target_names=bundle.target_names, + rng=rng, ) - val_eval = ( - evaluate_arrays( + features_tensor = _to_device(batch_features, device) + targets_tensor = _to_device(batch_targets, device) + + optimizer.zero_grad(set_to_none=True) + with _autocast_context(config, device): + predictions = model(features_tensor) + loss = F.mse_loss(predictions, targets_tensor) + if not torch.isfinite(loss): + _write_failure( + writer, + phase="training", + step=step, + error_type="NonFiniteLoss", + error_message="loss is NaN or Inf", + latest_loss=float(loss.detach().cpu().item()), + latest_grad_norm=last_grad_norm, + latest_checkpoint=LATEST_CHECKPOINT, + ) + raise RuntimeError("nonfinite loss") + loss.backward() + try: + grad_norm_tensor = torch.nn.utils.clip_grad_norm_( + model.parameters(), + config.stability.max_grad_norm if config.stability.max_grad_norm is not None else float("inf"), + error_if_nonfinite=True, + ) + except RuntimeError as exc: + _write_failure( + writer, + phase="training", + step=step, + error_type="NonFiniteGradient", + error_message=str(exc), + latest_loss=float(loss.detach().cpu().item()), + latest_grad_norm=last_grad_norm, + latest_checkpoint=LATEST_CHECKPOINT, + ) + raise RuntimeError("nonfinite gradients") from exc + last_grad_norm = float(grad_norm_tensor.detach().cpu().item()) + optimizer.step() + + now = time.monotonic() + should_checkpoint = ( + config.checkpoint.interval_seconds == 0 + or now - last_checkpoint_at >= config.checkpoint.interval_seconds + or step == config.optim.steps + ) + if should_checkpoint: + _save_training_checkpoint( + writer, + LATEST_CHECKPOINT, + config=config, + bundle=bundle, + stats=stats, + model=model, + optimizer=optimizer, + rng=rng, + step=step, + best_val_loss=best_val_loss, + initial_train_loss=initial_train_loss, + ) + last_checkpoint_at = now + + if step % log_interval == 0 or step == config.optim.steps: + train_eval = evaluate_arrays( model, - val_features, - val_targets, + train_features, + train_targets, batch_size=config.data.batch_size, device=device, target_names=bundle.target_names, ) - if val_features is not None and val_targets is not None - else None - ) - writer.append_metrics( - _log_metrics( - step=step, - train_loss=train_eval["loss"], - val_loss=val_eval["loss"] if val_eval is not None else None, - elapsed_seconds=time.perf_counter() - started, + val_eval = ( + evaluate_arrays( + model, + val_features, + val_targets, + batch_size=config.data.batch_size, + device=device, + target_names=bundle.target_names, + ) + if val_features is not None and val_targets is not None + else None ) - ) - model.train() + current_metric = val_eval["loss"] if val_eval is not None else train_eval["loss"] + if best_val_loss is None or current_metric < best_val_loss: + best_val_loss = current_metric + _save_training_checkpoint( + writer, + BEST_CHECKPOINT, + config=config, + bundle=bundle, + stats=stats, + model=model, + optimizer=optimizer, + rng=rng, + step=step, + best_val_loss=best_val_loss, + initial_train_loss=initial_train_loss, + ) + elapsed = time.perf_counter() - started + interval_elapsed = max(time.perf_counter() - last_log_at, 1e-9) + points_per_sec = (step - last_log_step) * config.data.batch_size / interval_elapsed + writer.append_metrics( + _log_metrics( + event="train_eval", + step=step, + train_loss=train_eval["loss"], + val_loss=val_eval["loss"] if val_eval is not None else None, + elapsed_seconds=elapsed, + lr=_learning_rate(optimizer), + grad_norm=last_grad_norm, + points_per_sec=points_per_sec, + device=device, + latest_checkpoint=LATEST_CHECKPOINT if should_checkpoint else None, + ) + ) + last_log_at = time.perf_counter() + last_log_step = step + model.train() + except Exception: + raise final_train = evaluate_arrays( model, @@ -192,40 +378,105 @@ def train(config: TrainingConfig) -> TrainingResult: elapsed = time.perf_counter() - started final_metrics: dict[str, Any] = { - "initial_train_loss": initial_train["loss"], + "initial_train_loss": initial_train_loss, "train_loss": final_train["loss"], "train_mse_per_channel": final_train["per_channel_mse"], "val_loss": final_val["loss"] if final_val is not None else None, "val_mse_per_channel": final_val["per_channel_mse"] if final_val is not None else None, "test_loss": final_test["loss"] if final_test is not None else None, "test_mse_per_channel": final_test["per_channel_mse"] if final_test is not None else None, + "best_val_loss": best_val_loss, "parameter_count": count_parameters(model), + "model_type": config.model.type, + "precision": config.precision.dtype, "train_cases": len(bundle.split.train_ids), "val_cases": len(bundle.split.val_ids), "test_cases": len(bundle.split.test_ids), "points_per_case": config.data.points_per_case, "steps": config.optim.steps, "elapsed_seconds": elapsed, + "checkpoint_interval_seconds": config.checkpoint.interval_seconds, + "resumed_from": str(resume) if resume is not None else None, **device_metrics(device), } writer.write_final_metrics(final_metrics) - writer.save_checkpoint( - { - "step": config.optim.steps, - "model_type": config.model.type, - "input_dim": train_features.shape[1], - "output_dim": train_targets.shape[1], - "model_config": asdict(config.model), - "model_state_dict": model.state_dict(), - "optimizer_state_dict": optimizer.state_dict(), - "normalization": stats.to_dict(), - "target_names": bundle.target_names, - "feature_names": bundle.feature_names, - "final_metrics": final_metrics, - } + _save_training_checkpoint( + writer, + FINAL_CHECKPOINT, + config=config, + bundle=bundle, + stats=stats, + model=model, + optimizer=optimizer, + rng=rng, + step=config.optim.steps, + best_val_loss=best_val_loss, + initial_train_loss=initial_train_loss, + final_metrics=final_metrics, ) + _save_training_checkpoint( + writer, + LEGACY_CHECKPOINT, + config=config, + bundle=bundle, + stats=stats, + model=model, + optimizer=optimizer, + rng=rng, + step=config.optim.steps, + best_val_loss=best_val_loss, + initial_train_loss=initial_train_loss, + final_metrics=final_metrics, + ) + writer.append_metrics( + _log_metrics( + event="completed", + phase="completed", + step=config.optim.steps, + train_loss=final_train["loss"], + val_loss=final_val["loss"] if final_val is not None else None, + elapsed_seconds=elapsed, + lr=_learning_rate(optimizer), + grad_norm=last_grad_norm, + points_per_sec=None, + device=device, + latest_checkpoint=FINAL_CHECKPOINT, + ) + ) + writer.write_artifact_manifest() return TrainingResult(run_dir=writer.run_dir, final_metrics=final_metrics) +def _autocast_context(config: TrainingConfig, device: torch.device): + if config.precision.dtype == "float32" or device.type != "cuda": + return torch.autocast(device_type=device.type, enabled=False) + return torch.autocast(device_type=device.type, dtype=torch.bfloat16) + + + +def _build_model(config: TrainingConfig, bundle: DatasetBundle, *, output_dim: int) -> torch.nn.Module: + if config.model.type == "mlp": + return PointwiseMLP( + input_dim=bundle.train.features.shape[1], + output_dim=output_dim, + hidden_width=config.model.hidden_width, + depth=config.model.depth, + activation=config.model.activation, + ) + if config.model.type == "film_fourier_mlp": + return FourierFiLMMLP( + feature_names=bundle.feature_names, + output_dim=output_dim, + coordinate_features=config.model.coordinate_features, + fourier_scales=config.model.fourier_scales, + trunk_width=config.model.hidden_width, + trunk_depth=config.model.depth, + condition_width=config.model.condition_width, + condition_depth=config.model.condition_depth, + condition_dim=config.model.condition_dim, + activation=config.model.activation, + ) + raise ValueError(f"Unsupported model type: {config.model.type}") + def select_device(config: TrainingConfig) -> torch.device: requested = config.device.type @@ -303,14 +554,210 @@ def _seed_all(seed: int) -> None: def _log_metrics( *, + event: str, step: int, - train_loss: float, + train_loss: float | None, val_loss: float | None, elapsed_seconds: float, + lr: float, + grad_norm: float | None, + points_per_sec: float | None, + device: torch.device, + latest_checkpoint: str | None, + phase: str = "training", ) -> dict[str, Any]: return { + "event": event, + "phase": phase, + "epoch": 0, "step": step, "train_loss": train_loss, "val_loss": val_loss, + "lr": lr, + "grad_norm": grad_norm, "elapsed_seconds": elapsed_seconds, + "points_per_sec": points_per_sec, + "latest_checkpoint": latest_checkpoint, + **_memory_metrics(device), } + + +def _memory_metrics(device: torch.device) -> dict[str, int | None]: + if device.type != "cuda": + return { + "gpu_memory_allocated_mb": None, + "gpu_memory_reserved_mb": None, + "gpu_memory_peak_allocated_mb": None, + } + index = device.index if device.index is not None else torch.cuda.current_device() + return { + "gpu_memory_allocated_mb": int(torch.cuda.memory_allocated(index) // (1024 * 1024)), + "gpu_memory_reserved_mb": int(torch.cuda.memory_reserved(index) // (1024 * 1024)), + "gpu_memory_peak_allocated_mb": int(torch.cuda.max_memory_allocated(index) // (1024 * 1024)), + } + + +def _learning_rate(optimizer: torch.optim.Optimizer) -> float: + return float(optimizer.param_groups[0]["lr"]) + + +def _save_training_checkpoint( + writer: ArtifactWriter, + name: str, + *, + config: TrainingConfig, + bundle: DatasetBundle, + stats: NormalizationStats, + model: torch.nn.Module, + optimizer: torch.optim.Optimizer, + rng: np.random.Generator, + step: int, + best_val_loss: float | None, + initial_train_loss: float | None, + final_metrics: dict[str, Any] | None = None, +) -> None: + writer.save_checkpoint( + _checkpoint_payload( + config=config, + bundle=bundle, + stats=stats, + model=model, + optimizer=optimizer, + rng=rng, + step=step, + best_val_loss=best_val_loss, + initial_train_loss=initial_train_loss, + final_metrics=final_metrics, + ), + name=name, + ) + + +def _checkpoint_payload( + *, + config: TrainingConfig, + bundle: DatasetBundle, + stats: NormalizationStats, + model: torch.nn.Module, + optimizer: torch.optim.Optimizer, + rng: np.random.Generator, + step: int, + best_val_loss: float | None, + initial_train_loss: float | None, + final_metrics: dict[str, Any] | None, +) -> dict[str, Any]: + return { + "schema_version": CHECKPOINT_SCHEMA_VERSION, + "run_id": os.environ.get("AIRFRANS_REMOTE_RUN_ID", config.run.name), + "epoch": 0, + "step": step, + "best_val_loss": best_val_loss, + "initial_train_loss": initial_train_loss, + "model_type": config.model.type, + "input_dim": bundle.train.features.shape[1], + "output_dim": bundle.train.targets.shape[1], + "model_config": asdict(config.model), + "model_state_dict": model.state_dict(), + "optimizer_state_dict": optimizer.state_dict(), + "config": config.config_text, + "config_hash": _config_hash(config), + "normalization": stats.to_dict(), + "target_names": bundle.target_names, + "feature_names": bundle.feature_names, + "rng_state": random.getstate(), + "numpy_rng_state": np.random.get_state(), + "torch_rng_state": torch.get_rng_state(), + "cuda_rng_state": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None, + "batch_rng_state": rng.bit_generator.state, + "final_metrics": final_metrics, + } + + +def _load_checkpoint(path: Path, device: torch.device) -> dict[str, Any]: + if not path.is_file(): + raise FileNotFoundError(f"Resume checkpoint not found: {path}") + payload = torch.load(path, map_location=device, weights_only=False) + if not isinstance(payload, dict): + raise ValueError(f"Resume checkpoint is not a mapping: {path}") + return payload + + +def _validate_resume_checkpoint( + checkpoint: dict[str, Any], + config: TrainingConfig, + bundle: DatasetBundle, + stats: NormalizationStats, +) -> None: + if checkpoint.get("schema_version") != CHECKPOINT_SCHEMA_VERSION: + raise ValueError("unsupported checkpoint schema_version") + if checkpoint.get("model_type") != config.model.type: + raise ValueError("checkpoint model type does not match config") + if checkpoint.get("model_config") != asdict(config.model): + raise ValueError("checkpoint model config does not match config") + if checkpoint.get("config_hash") != _config_hash(config): + raise ValueError("checkpoint config hash does not match config") + if tuple(checkpoint.get("feature_names", ())) != bundle.feature_names: + raise ValueError("checkpoint feature names do not match dataset") + if tuple(checkpoint.get("target_names", ())) != bundle.target_names: + raise ValueError("checkpoint target names do not match dataset") + if "normalization" not in checkpoint: + raise ValueError("checkpoint missing normalization") + if "optimizer_state_dict" not in checkpoint: + raise ValueError("checkpoint missing optimizer state") + if checkpoint.get("normalization") != stats.to_dict(): + raise ValueError("checkpoint normalization does not match dataset") + + +def _restore_rng_state(checkpoint: dict[str, Any], rng: np.random.Generator) -> None: + if "rng_state" in checkpoint: + random.setstate(checkpoint["rng_state"]) + if "numpy_rng_state" in checkpoint: + np.random.set_state(checkpoint["numpy_rng_state"]) + if "torch_rng_state" in checkpoint: + torch_state = checkpoint["torch_rng_state"] + if isinstance(torch_state, torch.Tensor): + torch_state = torch_state.cpu() + torch.set_rng_state(torch_state) + cuda_rng_state = checkpoint.get("cuda_rng_state") + if cuda_rng_state is not None and torch.cuda.is_available(): + torch.cuda.set_rng_state_all([state.cpu() if isinstance(state, torch.Tensor) else state for state in cuda_rng_state]) + if "batch_rng_state" in checkpoint: + rng.bit_generator.state = checkpoint["batch_rng_state"] + + +def _config_hash(config: TrainingConfig) -> str: + return hashlib.sha256(config.config_text.encode()).hexdigest() + + +def _optional_float(value: Any) -> float | None: + if value is None: + return None + return float(value) + + +def _write_failure( + writer: ArtifactWriter, + *, + phase: str, + step: int, + error_type: str, + error_message: str, + latest_loss: float | None = None, + latest_grad_norm: float | None = None, + latest_checkpoint: str | None = None, +) -> None: + writer.write_failure_report( + { + "run_id": os.environ.get("AIRFRANS_REMOTE_RUN_ID"), + "phase": phase, + "epoch": 0, + "step": step, + "error_type": error_type, + "error_message": error_message, + "latest_loss": latest_loss, + "latest_grad_norm": latest_grad_norm, + "latest_checkpoint": latest_checkpoint, + "timestamp": time.time(), + } + ) + writer.write_artifact_manifest() diff --git a/tests/test_remote_run.py b/tests/test_remote_run.py index 53fdd8e..be316e5 100644 --- a/tests/test_remote_run.py +++ b/tests/test_remote_run.py @@ -2,10 +2,14 @@ from __future__ import annotations import json import tempfile +import shutil import unittest from pathlib import Path +import torch + from airfrans_frontier.remote.artifacts import verify_artifacts +from airfrans_frontier.remote.cli import _classify_artifacts, _stage_resume_checkpoint from airfrans_frontier.remote.config import load_remote_run_config from airfrans_frontier.remote.skypilot import render_skypilot_yaml from airfrans_frontier.remote.vast import VastOffer, choose_offer @@ -18,7 +22,7 @@ class RemoteRunConfigTests(unittest.TestCase): self.assertEqual(config.provider.kind, "vastai") self.assertEqual(config.provider.gpu.name, "RTX 4090") self.assertEqual(config.job.artifact_dir.as_posix(), "artifacts/current_run") - self.assertIn("checkpoint.pt", config.artifacts.required) + self.assertIn("checkpoint_latest.pt", config.artifacts.required) class VastSelectionTests(unittest.TestCase): @@ -50,23 +54,78 @@ class VastSelectionTests(unittest.TestCase): self.assertIn("selected_offer_id: 123", yaml) self.assertNotIn("sky launch", yaml) self.assertIn("remote-run smoke-train", yaml) + self.assertIn("configs/aggressive_smoke.toml", yaml) + + def test_rendered_yaml_can_pass_resume_checkpoint(self) -> None: + config = load_remote_run_config("configs/remote_smoke.toml") + result = choose_offer([offer(123, price=0.30, host=22), offer(124, price=0.40, host=23)], config, query={}) + + yaml = render_skypilot_yaml( + config, + result, + run_id="airfrans-test", + resume_checkpoint=".airfrans_resume/airfrans-test/checkpoint_latest.pt", + ) + + self.assertIn("AIRFRANS_RESUME_CHECKPOINT: .airfrans_resume/airfrans-test/checkpoint_latest.pt", yaml) class ArtifactVerificationTests(unittest.TestCase): def test_verify_artifacts_requires_contract_files_and_writes_manifest(self) -> None: with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) - (root / "final_metrics.json").write_text(json.dumps({"loss": 1.0}) + "\n") - (root / "metrics.jsonl").write_text(json.dumps({"step": 0}) + "\n") - (root / "checkpoint.pt").write_bytes(b"weights") - (root / "run_manifest.json").write_text(json.dumps({"exit_code": 0}) + "\n") + _write_contract_artifacts(root, success=True) manifest = verify_artifacts(root) - self.assertEqual(manifest["file_count"], 4) + self.assertGreaterEqual(manifest["file_count"], 8) self.assertTrue((root / "artifact_manifest.json").is_file()) self.assertTrue((root / "checksums.txt").is_file()) + def test_verify_artifacts_accepts_failure_report_terminal_state(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + _write_contract_artifacts(root, success=False) + (root / "failure_report.json").write_text(json.dumps({"error_type": "NonFiniteLoss"}) + "\n") + + manifest = verify_artifacts(root) + + self.assertGreaterEqual(manifest["file_count"], 7) + + def test_classifies_and_stages_restart_checkpoint(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + self.assertEqual(_classify_artifacts(root), "incomplete") + _write_contract_artifacts(root, success=False) + self.assertEqual(_classify_artifacts(root), "restartable") + + staged = _stage_resume_checkpoint(root, "test-run") + + self.assertIsNotNone(staged) + assert staged is not None + self.assertTrue(staged.is_file()) + self.assertEqual(staged.as_posix(), ".airfrans_resume/test-run/checkpoint_latest.pt") + shutil.rmtree(".airfrans_resume") + + +def _write_contract_artifacts(root: Path, *, success: bool) -> None: + (root / "config.toml").write_text("[run]\nname = 'test'\n") + (root / "metrics.jsonl").write_text(json.dumps({"step": 0}) + "\n") + (root / "latest_metrics.json").write_text(json.dumps({"step": 0}) + "\n") + (root / "heartbeat.json").write_text(json.dumps({"phase": "training"}) + "\n") + checkpoint = { + "schema_version": 1, + "step": 0, + "model_state_dict": {}, + "optimizer_state_dict": {}, + "normalization": {}, + } + torch.save(checkpoint, root / "checkpoint_latest.pt") + torch.save(checkpoint, root / "checkpoint_best.pt") + if success: + torch.save(checkpoint, root / "checkpoint_final.pt") + (root / "final_metrics.json").write_text(json.dumps({"loss": 1.0}) + "\n") + def offer( offer_id: int, diff --git a/tests/test_training_loop.py b/tests/test_training_loop.py index c9c58a7..f641644 100644 --- a/tests/test_training_loop.py +++ b/tests/test_training_loop.py @@ -1,6 +1,7 @@ from __future__ import annotations import json +import os import subprocess import sys import tempfile @@ -14,9 +15,10 @@ 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 select_device +from airfrans_frontier.training.loop import train, select_device FEATURE_NAMES = np.array(["re_norm", "aoa_norm", "x", "y", "sdf"]) @@ -58,6 +60,10 @@ def write_training_config( 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""" @@ -76,15 +82,15 @@ batch_size = 64 [model] type = "mlp" -hidden_width = 128 +hidden_width = {hidden_width} depth = 4 activation = "gelu" [optim] lr = 0.01 weight_decay = 0.0 -steps = 1000 -log_interval = 250 +steps = {steps} +log_interval = {log_interval} [device] type = "{device_type}" @@ -93,6 +99,9 @@ benchmark_kernels = true [loss] type = "normalized_mse" + +[checkpoint] +interval_seconds = {checkpoint_interval_seconds} """.strip() + "\n" ) @@ -130,11 +139,17 @@ class TrainingLoopTests(unittest.TestCase): 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: ")) @@ -145,11 +160,274 @@ class TrainingLoopTests(unittest.TestCase): 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 / "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", + ): + 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 / "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()) 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() diff --git a/uv.lock b/uv.lock index 994bb1d..31a9174 100644 --- a/uv.lock +++ b/uv.lock @@ -189,6 +189,7 @@ dev = [ { name = "nbformat" }, { name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "numpy", version = "2.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, + { name = "pytest" }, { name = "skypilot", extra = ["vast"] }, ] @@ -205,6 +206,7 @@ dev = [ { name = "nbclient", specifier = ">=0.10.2" }, { name = "nbformat", specifier = ">=5.10.4" }, { name = "numpy", specifier = ">=2.4.0" }, + { name = "pytest", specifier = ">=9.1.1" }, { name = "skypilot", extras = ["vast"], specifier = ">=0.12.3.post1" }, ] @@ -1396,6 +1398,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ea/02/aafbf0c3e1468c7c0f607065363b49c381de7e4bb43ae6674684a3fafe92/ijson-3.5.1-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:4b75b6bf4b0dbb0df24947db6722cd5723ce8d6e6b13fddbfc98db312ba82237", size = 54922, upload-time = "2026-07-06T17:37:41.879Z" }, ] +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + [[package]] name = "invoke" version = "3.0.3" @@ -2752,6 +2763,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ec/73/6fd0bb9ce84138c3857f12e9de63bc901852975a092d545f18087a204aa2/platformdirs-4.10.1-py3-none-any.whl", hash = "sha256:0e4eff26be2d75293977f7cddc153fd9b8eaa7fb0c7b64ffe4076cb443117443", size = 22906, upload-time = "2026-07-18T03:53:42.576Z" }, ] +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + [[package]] name = "prettytable" version = "3.18.0" @@ -3219,6 +3239,22 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/10/bd/c038d7cc38edc1aa5bf91ab8068b63d4308c66c4c8bb3cbba7dfbc049f9c/pyparsing-3.3.2-py3-none-any.whl", hash = "sha256:850ba148bd908d7e2411587e247a1e4f0327839c40e2e5e6d05a007ecc69911d", size = 122781, upload-time = "2026-01-21T03:57:55.912Z" }, ] +[[package]] +name = "pytest" +version = "9.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369, upload-time = "2026-06-19T10:58:32.857Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, +] + [[package]] name = "python-barcode" version = "0.16.1"