stash: changes up to the first deploy test of the agressive run

This commit is contained in:
Zachery Aaron Shores-Chmielewski 2026-07-23 12:36:59 +04:00
parent 05476f1724
commit bdbbc2458e
20 changed files with 2110 additions and 143 deletions

1
.gitignore vendored
View file

@ -2,6 +2,7 @@
data/
outputs/
artifacts/
.airfrans_resume/
# Python
.venv/

View 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

View 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

View file

@ -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",

View file

@ -27,4 +27,5 @@ dev = [
"nbclient>=0.10.2",
"nbformat>=5.10.4",
"skypilot[vast]>=0.12.3.post1",
"pytest>=9.1.1",
]

View file

@ -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

View file

@ -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"]

View 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}")

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

View file

@ -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)}")

View file

@ -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)
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)
except Exception as exc:
state("FAILED", selected_offer_id=selection.selected_offer_id, error=str(exc))
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)
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)
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)
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)

View file

@ -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,

View file

@ -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),

View file

@ -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")

View file

@ -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:

View file

@ -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):

View file

@ -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,19 +160,60 @@ 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):
try:
for step in range(start_step + 1, config.optim.steps + 1):
batch_features, batch_targets = _sample_batch(
train_features,
train_targets,
@ -121,11 +224,65 @@ def train(config: TrainingConfig) -> TrainingResult:
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,
@ -147,15 +304,44 @@ def train(config: TrainingConfig) -> TrainingResult:
if val_features is not None and val_targets is not None
else None
)
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=time.perf_counter() - started,
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()

View file

@ -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,

View file

@ -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
View file

@ -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"