stash: changes up to the first deploy test of the agressive run
This commit is contained in:
parent
05476f1724
commit
bdbbc2458e
20 changed files with 2110 additions and 143 deletions
1
.gitignore
vendored
1
.gitignore
vendored
|
|
@ -2,6 +2,7 @@
|
|||
data/
|
||||
outputs/
|
||||
artifacts/
|
||||
.airfrans_resume/
|
||||
|
||||
# Python
|
||||
.venv/
|
||||
|
|
|
|||
46
configs/aggressive_smoke.toml
Normal file
46
configs/aggressive_smoke.toml
Normal file
|
|
@ -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
|
||||
46
configs/aggressive_smoke_local.toml
Normal file
46
configs/aggressive_smoke_local.toml
Normal file
|
|
@ -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
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -27,4 +27,5 @@ dev = [
|
|||
"nbclient>=0.10.2",
|
||||
"nbformat>=5.10.4",
|
||||
"skypilot[vast]>=0.12.3.post1",
|
||||
"pytest>=9.1.1",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
144
src/airfrans_frontier/models/film.py
Normal file
144
src/airfrans_frontier/models/film.py
Normal file
|
|
@ -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}")
|
||||
414
src/airfrans_frontier/raw/process.py
Normal file
414
src/airfrans_frontier/raw/process.py
Normal file
|
|
@ -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<turbulence>[^_]+)_"
|
||||
r"(?P<u_inf>-?\d+(?:\.\d+)?)_"
|
||||
r"(?P<alpha>-?\d+(?:\.\d+)?)_"
|
||||
r"(?P<params>.+)$"
|
||||
)
|
||||
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)
|
||||
|
|
@ -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)}")
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
36
uv.lock
36
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"
|
||||
|
|
|
|||
Loading…
Reference in a new issue