zacheryasc/airfRANS-model-exploration
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: use pre-processed dataset

Browse source
This commit is contained in:
Zachery Aaron Shores-Chmielewski 2026-07-29 13:42:52 +04:00
parent 2071670a4d
commit d2b102cafb
18 changed files with 1571 additions and 114 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.skyignore
View file

@ -1,5 +1,6 @@
/artifacts/remote_runs
/artifacts/runs
/artifacts/runs_minimal_scaling_20260728
/artifacts/data_cache
/artifacts/model_sanity
/artifacts/current_run

119
configs/remote_minimal_scaling_node0.toml Normal file
View file

@ -0,0 +1,119 @@
[run]
name = "minimal_scaling_sweep_node0"
timeout_minutes = 1440
local_artifact_dir = "artifacts/remote_runs"
max_attempts = 1
artifact_sync_interval_seconds = 60
[provider]
kind = "vastai"
disk_gb = 192
max_price_per_hour = 0.80
image = "vastai/base:0.0.2"
[provider.gpu]
name = "RTX 4090"
count = 1
min_vram_gb = 20
[selection]
min_reliability = 0.95
min_down_mbps = 100
min_up_mbps = 25
require_verified = true
blocked_geos = ["CN"]
blacklist_hosts = [59017, 1647, 92578, 1276, 75481, 1256, 85323, 34031, 528237]
drop_cheap_frac = 0.30
image_size_gb = 5.0
base_url = "https://cloud.vast.ai"
[workspace]
workdir = "."
exclude = [
"/artifacts/remote_runs",
"/artifacts/runs",
"/artifacts/runs_minimal_scaling_20260728",
"/artifacts/data_cache",
"/artifacts/model_sanity",
"/artifacts/current_run",
"/artifacts/public_airfrans",
"/artifacts/notes",
"/artifacts/preflight_20260725",
"/artifacts/remote_state_20260725",
"/artifacts/sweep_20260725",
"/artifacts/model_sanity_decisions",
"/artifacts/model_sanity_verify",
"/artifacts/*.json",
"/artifacts/*.yaml",
"/artifacts/*.md",
"/data/raw",
"/data/processed",
"/.venv",
"/.git",
"/.airfrans_hf_resume",
"/sky_logs",
"/notebooks",
"__pycache__",
"*.pyc",
]
[bootstrap]
command = """
uv sync --no-dev
uv run --no-dev python -c "import torch; ok=torch.cuda.is_available(); print('torch_cuda_available=' + str(ok)); print(torch.cuda.get_device_name(0) if ok else 'no_cuda_device'); assert ok, 'torch CUDA unavailable'"
uv run --no-dev python -c "import huggingface_hub, wandb; print('hf_wandb_import_ok')"
"""
[data]
validation_command = """
uv run --no-dev python - <<'PY'
from airfrans_frontier.sweep import read_jobs
jobs = read_jobs('artifacts/minimal_scaling_sweep_20260728/jobs.jsonl')
pending = sum(1 for job in jobs if job['status'] == 'pending')
phases = [job.get('phase') for job in jobs]
transitions = []
for phase in phases:
if not transitions or transitions[-1] != phase:
transitions.append(phase)
assert len(jobs) == 23, len(jobs)
assert pending > 0
assert transitions == ['A_measurement_canary', 'B_family_viability', 'C_scaling_ladder'], transitions
print('minimal_scaling_pending_jobs=' + str(pending))
print('minimal_scaling_phase_order=' + ','.join(transitions))
PY
"""
[job]
command = """
uv run --no-dev python scripts/aggressive_oom_node_wrapper.py --jobs artifacts/minimal_scaling_sweep_20260728/jobs.jsonl --node node-0 --artifact-dir artifacts/current_run --utilization artifacts/minimal_scaling_sweep_20260728/utilization_node0.jsonl --sample-interval-seconds 30 --stale-after-seconds 21600 --max-attempts 2 --require-success
"""
artifact_dir = "artifacts/current_run"
heartbeat_file = "artifacts/current_run/heartbeat.json"
metrics_file = "artifacts/current_run/metrics.jsonl"
[artifacts]
mode = "object_store_upload"
required = [
"config.toml",
"job_manifest.json",
"metrics.jsonl",
"latest_metrics.json",
"heartbeat.json",
"checkpoint_latest.pt",
"checkpoint_best.pt",
"checkpoint_final.pt",
"final_metrics.json",
"run_manifest.json",
"environment_manifest.json",
"utilization.jsonl",
"node_summary.json",
"sweep_collect.json",
"sweep_jobs.jsonl",
"sweep_attempts.jsonl",
"sweep_utilization.jsonl",
]
[cleanup]
on_success = "sky_down"
on_failure = "collect_then_keep"

37
src/airfrans_frontier/cli.py
View file

@ -73,7 +73,9 @@ def build_parser() -> argparse.ArgumentParser:
sweep_generate.add_argument("--data-root", default="artifacts/data_cache/airfrans_processed/processed/full")
sweep_generate.add_argument("--artifact-dir", default="artifacts/runs")
sweep_generate.add_argument("--hf-repo-id", default="zacheryasc/airfrans-frontier-checkpoints")
sweep_generate.add_argument("--group", default="aggressive_oom_sweep_01")
sweep_generate.add_argument("--group", default="minimal_scaling_sweep_20260728")
sweep_generate.add_argument("--plan", choices=("minimal-scaling", "legacy-aggressive"), default="minimal-scaling")
sweep_generate.add_argument("--include-proxies", action="store_true")
sweep_generate.add_argument("--bands", nargs="*", default=["100m", "700m"])
sweep_generate.add_argument("--families", nargs="*")
sweep_generate.add_argument("--encodings", nargs="*")
@ -250,21 +252,30 @@ def main(argv: list[str] | None = None) -> int:
print(f"report: {resolve_path(args.artifact_dir) / 'model_sanity_results.json'}")
print(f"families: {len(result['families'])}")
return 0
if args.command == "sweep-generate":
from airfrans_frontier.sweep import DEFAULT_ENCODINGS, DEFAULT_FAMILIES, generate_jobs
from airfrans_frontier.sweep import DEFAULT_ENCODINGS, DEFAULT_FAMILIES, generate_jobs, generate_minimal_scaling_jobs
try:
jobs = generate_jobs(
output_dir=resolve_path(args.output_dir),
data_root=args.data_root,
artifact_dir=args.artifact_dir,
hf_repo_id=args.hf_repo_id,
group=args.group,
bands=tuple(args.bands),
families=tuple(args.families) if args.families else DEFAULT_FAMILIES,
encodings=tuple(args.encodings) if args.encodings else DEFAULT_ENCODINGS,
)
if args.plan == "minimal-scaling":
jobs = generate_minimal_scaling_jobs(
output_dir=resolve_path(args.output_dir),
data_root=args.data_root,
artifact_dir=args.artifact_dir,
hf_repo_id=args.hf_repo_id,
group=args.group,
include_proxies=args.include_proxies,
)
else:
jobs = generate_jobs(
output_dir=resolve_path(args.output_dir),
data_root=args.data_root,
artifact_dir=args.artifact_dir,
hf_repo_id=args.hf_repo_id,
group=args.group,
bands=tuple(args.bands),
families=tuple(args.families) if args.families else DEFAULT_FAMILIES,
encodings=tuple(args.encodings) if args.encodings else DEFAULT_ENCODINGS,
)
except (FileNotFoundError, NotADirectoryError, ValueError, RuntimeError) as exc:
print(f"error: {exc}", file=sys.stderr)
return 1

3
src/airfrans_frontier/models/__init__.py
View file

@ -10,12 +10,13 @@ from airfrans_frontier.models.frontier import (
RasterFNOUNet,
SirenConditionedINR,
)
from airfrans_frontier.models.mlp import PointwiseMLP
from airfrans_frontier.models.mlp import EncodedPointwiseMLP, PointwiseMLP
__all__ = [
"CoordinateEncoder",
"CoordinateEncodingSpec",
"DeepONetBranchTrunk",
"EncodedPointwiseMLP",
"FourierFiLMMLP",
"LocalPointTransformer",
"NeRFCFDMultiRes",

46
src/airfrans_frontier/models/mlp.py
View file

@ -3,6 +3,52 @@ from __future__ import annotations
import torch
from torch import nn
from airfrans_frontier.models.coordinate_encoding import CoordinateEncoder, CoordinateEncodingSpec
class EncodedPointwiseMLP(nn.Module):
def __init__(
self,
*,
feature_names: tuple[str, ...],
output_dim: int,
coordinate_features: tuple[str, ...],
coordinate_encoding_spec: CoordinateEncodingSpec,
hidden_width: int = 128,
depth: int = 4,
activation: str = "gelu",
) -> None:
super().__init__()
if output_dim <= 0:
raise ValueError("output_dim must be positive")
if hidden_width <= 0:
raise ValueError("hidden_width must be positive")
if depth <= 0:
raise ValueError("depth must be positive")
name_to_index = {name: index for index, name in enumerate(feature_names)}
missing = [name for name in coordinate_features if name not in name_to_index]
if missing:
raise ValueError(f"coordinate features are missing from input schema: {missing}")
coord_indices = [name_to_index[name] for name in coordinate_features]
condition_indices = [index for index, name in enumerate(feature_names) if name not in set(coordinate_features)]
self.register_buffer("coord_indices", torch.tensor(coord_indices, dtype=torch.long), persistent=False)
self.register_buffer("condition_indices", torch.tensor(condition_indices, dtype=torch.long), persistent=False)
self.encoder = CoordinateEncoder(coordinate_encoding_spec, coordinate_dim=len(coord_indices))
input_dim = self.encoder.output_dim + len(condition_indices)
layers: list[nn.Module] = [nn.Linear(input_dim, hidden_width), _activation(activation)]
for _ in range(depth - 1):
layers.extend((nn.Linear(hidden_width, hidden_width), _activation(activation)))
layers.append(nn.Linear(hidden_width, output_dim))
self.network = nn.Sequential(*layers)
def forward(self, features: torch.Tensor) -> torch.Tensor:
coordinates = features.index_select(dim=1, index=self.coord_indices)
encoded = self.encoder(coordinates)
if self.condition_indices.numel():
conditions = features.index_select(dim=1, index=self.condition_indices)
encoded = torch.cat((encoded, conditions), dim=-1)
return self.network(encoded)
class PointwiseMLP(nn.Module):
def __init__(

336
src/airfrans_frontier/sweep.py
View file

@ -30,28 +30,37 @@ SUCCESS_ARTIFACTS = REQUIRED_JOB_ARTIFACTS + ("final_metrics.json",)
FAILURE_ARTIFACTS = REQUIRED_JOB_ARTIFACTS + ("failure_report.json",)
OOM_BANDS: dict[str, dict[str, int]] = {
"10m": {"hidden_width": 1024, "depth": 8, "steps": 300, "batch_size": 4096},
"100m": {"hidden_width": 2048, "depth": 16, "steps": 1500, "batch_size": 4096},
"700m": {"hidden_width": 4096, "depth": 40, "steps": 500, "batch_size": 2048},
"10m": {"hidden_width": 1024, "depth": 8, "steps": 5000, "batch_size": 4096},
"100m": {"hidden_width": 2048, "depth": 16, "steps": 2500, "batch_size": 4096},
"700m": {"hidden_width": 4096, "depth": 40, "steps": 1000, "batch_size": 2048},
}
FILM_DEPTH_OVERRIDES = {"700m": 32}
DEFAULT_FAMILIES = (
"mlp",
MAIN_SCALING_FAMILIES = (
"mlp_encoded_baseline",
"nerf_cfd_multires",
"deeponet_branch_trunk",
"film_fourier_inr",
"siren_conditioned_inr",
)
PROXY_SCALING_FAMILIES = (
"raster_fno_unet",
"point_context_perceiver",
"meshgraphnet_or_point_transformer_local",
)
DEFAULT_ENCODINGS = (
DEFAULT_FAMILIES = MAIN_SCALING_FAMILIES
SUPPORTED_ENCODINGS = (
"raw",
"fixed_fourier",
"nerf_multires",
"random_fourier",
)
DEFAULT_ENCODINGS = (
"fixed_fourier",
"nerf_multires",
"random_fourier",
)
ENCODING_COMPATIBLE_FAMILIES = {
"mlp_encoded_baseline",
"nerf_cfd_multires",
"deeponet_branch_trunk",
"film_fourier_mlp",
@ -59,8 +68,8 @@ ENCODING_COMPATIBLE_FAMILIES = {
}
COORDINATE_ENCODING_EXCLUSION_REASONS = {
"mlp": "architecture_uses_full_raw_point_features",
"siren_conditioned_inr": "siren_uses_raw_coordinates_with_omega0",
"mlp": "legacy pointwise MLP uses full raw point features; use mlp_encoded_baseline for serious sweeps",
"siren_conditioned_inr": "siren_uses_sine_coordinate_basis_with_omega0",
"raster_fno_unet": "raster_proxy_uses_xy_sampling_grid",
"point_context_perceiver": "perceiver_proxy_uses_full_feature_projection",
"meshgraphnet_or_point_transformer_local": "local_point_proxy_uses_knn_coordinates",
@ -143,7 +152,7 @@ def write_default_pool(path: str | Path) -> Path:
return target
def coordinate_encoding_decision(family: str, encoding: str) -> dict[str, Any]:
if encoding not in DEFAULT_ENCODINGS:
if encoding not in SUPPORTED_ENCODINGS:
raise ValueError(f"unsupported coordinate encoding: {encoding}")
if family in ENCODING_COMPATIBLE_FAMILIES:
return {
@ -240,6 +249,235 @@ def generate_jobs(
return jobs
MINIMAL_SWEEP_RECIPES: dict[str, dict[str, Any]] = {
"mlp_encoded_baseline": {"encoding": "nerf_multires", "lr": 0.0001},
"nerf_cfd_multires": {"encoding": "nerf_multires", "lr": 0.0001},
"deeponet_branch_trunk": {"encoding": "nerf_multires", "lr": 0.0001},
"film_fourier_inr": {"encoding": "nerf_multires", "lr": 0.0001},
"siren_conditioned_inr": {"encoding": "raw", "lr": 0.00003},
}
def generate_minimal_scaling_jobs(
*,
output_dir: str | Path,
data_root: str = "artifacts/data_cache/airfrans_processed/processed/full",
artifact_dir: str = "artifacts/runs_minimal_scaling",
hf_repo_id: str = "zacheryasc/airfrans-frontier-checkpoints",
group: str = "minimal_scaling_sweep_20260728",
include_proxies: bool = False,
seed: int = 20260728,
) -> list[dict[str, Any]]:
root = Path(output_dir)
config_dir = root / "configs"
config_dir.mkdir(parents=True, exist_ok=True)
jobs: list[dict[str, Any]] = []
def add_job(
*,
job_id: str,
phase: str,
purpose: str,
family: str,
band: str,
encoding: str,
lr: float,
steps: int,
batch_size: int,
train_cases: int,
val_cases: int,
test_cases: int,
log_interval: int,
checkpoint_policy: str,
include_optimizer_state: bool,
include_rng_state: bool,
dead_curve_patience_evals: int | None,
dead_curve_warmup_steps: int,
expected_terminal_status: str = "succeeded",
run_seed: int | None = None,
) -> None:
decision = coordinate_encoding_decision(family, encoding)
metadata = model_metadata_for_family(family, coordinate_encoding_compatibility=str(decision["compatibility"]))
config_path = config_dir / f"{job_id}.toml"
config_path.write_text(
training_config_text(
run_name=job_id,
seed=seed + len(jobs) if run_seed is None else run_seed,
data_root=data_root,
artifact_dir=artifact_dir,
model_family=family,
band=band,
coordinate_encoding=encoding,
hf_repo_id=hf_repo_id,
group=group,
train_cases=train_cases,
val_cases=val_cases,
test_cases=test_cases,
streaming_normalization_cases=min(4, train_cases),
data_source="huggingface_streaming",
all_points_per_case=True,
points_per_case=None,
steps=steps,
batch_size=batch_size,
lr=lr,
log_interval=log_interval,
checkpoint_policy=checkpoint_policy,
checkpoint_include_optimizer_state=include_optimizer_state,
checkpoint_include_rng_state=include_rng_state,
dead_curve_patience_evals=dead_curve_patience_evals,
dead_curve_warmup_steps=dead_curve_warmup_steps,
)
)
jobs.append(
{
"job_id": job_id,
"status": "pending",
"phase": phase,
"purpose": purpose,
"band": band,
"model_family": family,
"requested_family": metadata["requested_family"],
"implementation_family": metadata["implementation_family"],
"reported_family": metadata["reported_family"],
"is_proxy": metadata["is_proxy"],
"proxy_for": metadata["proxy_for"],
"proxy_notes": metadata["proxy_notes"],
"coordinate_encoding": encoding,
"coordinate_encoding_decision": decision,
"coordinate_encoding_compatibility": decision["compatibility"],
"all_points_per_case": True,
"expected_terminal_status": expected_terminal_status,
"checkpoint_policy": checkpoint_policy,
"include_optimizer_state": include_optimizer_state,
"dead_curve_patience_evals": dead_curve_patience_evals,
"min_post_warmup_eval_points": 8,
"config_path": str(config_path),
"attempts": 0,
}
)
add_job(
job_id="A_data_all_points_smoke",
phase="A_measurement_canary",
purpose="all_points_data_path_smoke",
family="mlp_encoded_baseline",
band="10m",
encoding="fixed_fourier",
lr=0.0001,
steps=50,
batch_size=512,
train_cases=2,
val_cases=1,
test_cases=1,
log_interval=5,
checkpoint_policy="full",
include_optimizer_state=True,
include_rng_state=True,
dead_curve_patience_evals=None,
dead_curve_warmup_steps=0,
)
add_job(
job_id="A_checkpoint_700m_footprint_smoke",
phase="A_measurement_canary",
purpose="full_700m_checkpoint_smoke",
family="film_fourier_inr",
band="700m",
encoding="nerf_multires",
lr=0.0001,
steps=2,
batch_size=64,
train_cases=2,
val_cases=1,
test_cases=1,
log_interval=1,
checkpoint_policy="full",
include_optimizer_state=True,
include_rng_state=True,
dead_curve_patience_evals=None,
dead_curve_warmup_steps=0,
)
add_job(
job_id="A_plateau_interrupt_smoke",
phase="A_measurement_canary",
purpose="dead_curve_interrupt_smoke",
family="mlp_encoded_baseline",
band="10m",
encoding="fixed_fourier",
lr=1e-12,
steps=40,
batch_size=512,
train_cases=2,
val_cases=1,
test_cases=1,
log_interval=5,
checkpoint_policy="full",
include_optimizer_state=True,
include_rng_state=True,
dead_curve_patience_evals=3,
dead_curve_warmup_steps=5,
expected_terminal_status="failed",
)
families = tuple(MINIMAL_SWEEP_RECIPES)
if include_proxies:
families = families + PROXY_SCALING_FAMILIES
family_recipes = [
(family, MINIMAL_SWEEP_RECIPES.get(family, {"encoding": "fixed_fourier", "lr": 0.0001}))
for family in families
]
for family, recipe in family_recipes:
add_job(
job_id=f"B_10m_{family}_{recipe['encoding']}_viability",
phase="B_family_viability",
purpose="viability",
family=family,
band="10m",
encoding=str(recipe["encoding"]),
lr=float(recipe["lr"]),
steps=300,
batch_size=2048,
train_cases=900,
val_cases=50,
test_cases=50,
log_interval=25,
checkpoint_policy="full",
include_optimizer_state=True,
include_rng_state=True,
dead_curve_patience_evals=8,
dead_curve_warmup_steps=50,
)
for family_index, (family, recipe) in enumerate(family_recipes):
for band in ("10m", "100m", "700m"):
shape = OOM_BANDS[band]
add_job(
job_id=f"C_{band}_{family}_{recipe['encoding']}_scale",
phase="C_scaling_ladder",
purpose="scaling_curve",
family=family,
band=band,
encoding=str(recipe["encoding"]),
lr=float(recipe["lr"]),
steps=shape["steps"],
batch_size=shape["batch_size"],
train_cases=900,
val_cases=50,
test_cases=50,
log_interval=50 if band == "10m" else 25,
checkpoint_policy="full",
include_optimizer_state=True,
include_rng_state=True,
dead_curve_patience_evals=8,
dead_curve_warmup_steps=100 if band == "10m" else 50,
run_seed=seed + 1000 + family_index,
)
write_jobs(root / "jobs.jsonl", jobs)
write_default_pool(root / "pool.toml")
(root / "budget_ledger.json").write_text(json.dumps({"budget_usd": 25.0, "spent_usd": 0.0, "reserved_usd": 0.0}, indent=2) + "\n")
write_rescue_templates(root / RESCUE_TEMPLATES_JSON)
return jobs
def training_config_text(
*,
run_name: str,
@ -251,11 +489,63 @@ def training_config_text(
coordinate_encoding: str,
hf_repo_id: str,
group: str,
train_cases: int = 900,
val_cases: int = 50,
test_cases: int = 50,
streaming_normalization_cases: int | None = None,
data_source: str = "huggingface_streaming",
all_points_per_case: bool = False,
points_per_case: int | None = 8192,
steps: int | None = None,
batch_size: int | None = None,
lr: float = 0.0001,
log_interval: int = 25,
checkpoint_policy: str = "full",
checkpoint_include_optimizer_state: bool = True,
checkpoint_include_rng_state: bool = True,
dead_curve_patience_evals: int | None = None,
dead_curve_min_relative_improvement: float = 0.005,
dead_curve_warmup_steps: int = 0,
) -> str:
shape = dict(OOM_BANDS[band])
if model_family == "film_fourier_inr" and band in FILM_DEPTH_OVERRIDES:
shape["depth"] = FILM_DEPTH_OVERRIDES[band]
condition_width = max(512, shape["hidden_width"] // 2)
if steps is not None:
shape["steps"] = int(steps)
if batch_size is not None:
shape["batch_size"] = int(batch_size)
data_source_lines = {
"huggingface": (
"source = \"huggingface\"",
"hf_repo_id = \"zacheryasc/airfrans-processed\"",
"hf_repo_type = \"dataset\"",
"hf_path_prefix = \"processed/full\"",
"cache_dir = \"artifacts/data_cache/airfrans_processed\"",
),
"huggingface_streaming": (
"source = \"huggingface_streaming\"",
"hf_repo_id = \"zacheryasc/airfrans-processed\"",
"hf_repo_type = \"dataset\"",
"hf_path_prefix = \"processed/full\"",
"cache_dir = \"artifacts/data_cache/airfrans_processed\"",
"streaming_queue_max_cases = 8",
),
"public_zip_streaming": (
"source = \"public_zip_streaming\"",
"hf_repo_type = \"dataset\"",
"cache_dir = \"artifacts/data_cache/airfrans_public_streaming\"",
"streaming_scratch_dir = \"artifacts/data_cache/airfrans_public_streaming/_raw\"",
"streaming_cache_max_bytes = 34359738368",
"streaming_cache_high_water_bytes = 30064771072",
"streaming_cache_low_water_bytes = 21474836480",
"streaming_queue_max_cases = 8",
),
}
if data_source not in data_source_lines:
raise ValueError(f"unsupported data source for sweep config: {data_source}")
if all_points_per_case and points_per_case is not None:
raise ValueError("all_points_per_case cannot be combined with points_per_case")
condition_depth = 4 if band in {"100m", "700m"} else 2
condition_dim = min(shape["hidden_width"], 1024)
decision = coordinate_encoding_decision(model_family, coordinate_encoding)
@ -272,17 +562,13 @@ def training_config_text(
"",
"[data]",
f"root = {json.dumps(data_root)}",
"source = \"huggingface\"",
"hf_repo_id = \"zacheryasc/airfrans-processed\"",
"hf_repo_type = \"dataset\"",
"hf_path_prefix = \"processed/full\"",
"cache_dir = \"artifacts/data_cache/airfrans_processed\"",
"train_cases = 900",
"val_cases = 50",
"test_cases = 50",
"points_per_case = 8192",
*data_source_lines[data_source],
f"train_cases = {train_cases}",
f"val_cases = {val_cases}",
f"test_cases = {test_cases}",
"all_points_per_case = true" if all_points_per_case else f"points_per_case = {points_per_case}",
f"batch_size = {shape['batch_size']}",
"streaming_normalization_cases = 16" if band == "700m" else None,
f"streaming_normalization_cases = {streaming_normalization_cases}" if streaming_normalization_cases is not None else None,
"",
"[coordinate_encoding]",
_coordinate_encoding_toml(coordinate_encoding),
@ -313,10 +599,10 @@ def training_config_text(
f"coordinate_encoding_compatibility = {json.dumps(metadata['coordinate_encoding_compatibility'])}",
"",
"[optim]",
"lr = 0.0001",
f"lr = {lr}",
"weight_decay = 0.0001",
f"steps = {shape['steps']}",
"log_interval = 25",
f"log_interval = {log_interval}",
"",
"[device]",
"type = \"cuda\"",
@ -331,16 +617,22 @@ def training_config_text(
"",
"[checkpoint]",
"interval_seconds = 900",
f"policy = {json.dumps(checkpoint_policy)}",
f"include_optimizer_state = {str(checkpoint_include_optimizer_state).lower()}",
f"include_rng_state = {str(checkpoint_include_rng_state).lower()}",
"",
"[stability]",
"max_grad_norm = 1.0",
f"dead_curve_patience_evals = {dead_curve_patience_evals}" if dead_curve_patience_evals is not None else None,
f"dead_curve_min_relative_improvement = {dead_curve_min_relative_improvement}",
f"dead_curve_warmup_steps = {dead_curve_warmup_steps}",
"",
"[observability]",
"backend = \"wandb\"",
"entity = \"zacheryasc-personal\"",
"project = \"airfRANS-model-sweep\"",
f"group = {json.dumps(group)}",
f"tags = [\"airfrans\", \"aggressive-oom-sweep\", {json.dumps(band)}, {json.dumps(model_family)}, {json.dumps(coordinate_encoding)}]",
f"tags = [\"airfrans\", {json.dumps(group)}, {json.dumps(band)}, {json.dumps(model_family)}, {json.dumps(coordinate_encoding)}]",
"",
"[huggingface]",
"enabled = true",

2
src/airfrans_frontier/training/artifacts.py
View file

@ -147,7 +147,7 @@ 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", "verification_report.json"}
if path.is_file() and path.name not in {"artifact_manifest.json", "checksums.txt", "verification_report.json"} and not path.name.endswith(".tmp")
)
manifest = {
"artifact_dir": str(root),

26
src/airfrans_frontier/training/config.py
View file

@ -9,6 +9,7 @@ from airfrans_frontier.model_metadata import model_metadata_for_family
_MODEL_TYPES = {
"mlp",
"mlp_encoded_baseline",
"film_fourier_mlp",
"film_fourier_inr",
"nerf_cfd_multires",
@ -32,7 +33,8 @@ class DataConfig:
train_cases: int
val_cases: int
test_cases: int
points_per_case: int
points_per_case: int | None
all_points_per_case: bool
batch_size: int
source: str
hf_repo_id: str | None
@ -49,7 +51,6 @@ class DataConfig:
streaming_upload_batch_size: int = 8
streaming_normalization_cases: int | None = None
@dataclass(frozen=True)
class ModelConfig:
type: str
@ -104,12 +105,17 @@ class OptimConfig:
@dataclass(frozen=True)
class CheckpointConfig:
interval_seconds: int
include_optimizer_state: bool
include_rng_state: bool
policy: str
@dataclass(frozen=True)
class StabilityConfig:
max_grad_norm: float | None
dead_curve_patience_evals: int | None
dead_curve_min_relative_improvement: float
dead_curve_warmup_steps: int
@dataclass(frozen=True)
class PrecisionConfig:
@ -279,15 +285,19 @@ def load_training_config(path: str | Path) -> TrainingConfig:
raise ValueError("data.streaming_cache_high_water_bytes must be <= data.streaming_cache_max_bytes")
if streaming_low_water_bytes >= streaming_high_water_bytes:
raise ValueError("data.streaming_cache_low_water_bytes must be < data.streaming_cache_high_water_bytes")
all_points_per_case = _boolean(data_raw, "all_points_per_case") if "all_points_per_case" in data_raw else False
points_per_case = None if all_points_per_case else _integer(data_raw, "points_per_case", minimum=1)
data = DataConfig(
root=_path(data_raw, "root"),
train_cases=_integer(data_raw, "train_cases", minimum=1),
val_cases=_integer(data_raw, "val_cases", minimum=0),
test_cases=_integer(data_raw, "test_cases", minimum=0),
points_per_case=_integer(data_raw, "points_per_case", minimum=1),
points_per_case=points_per_case,
all_points_per_case=all_points_per_case,
batch_size=_integer(data_raw, "batch_size", minimum=1),
source=_choice(_string(data_raw, "source", default="local").lower(), {"local", "huggingface", "public_zip_streaming"}, "data.source"),
source=_choice(_string(data_raw, "source", default="local").lower(), {"local", "huggingface", "huggingface_streaming", "public_zip_streaming"}, "data.source"),
hf_repo_id=_optional_string(data_raw, "hf_repo_id"),
hf_repo_type=_choice(_string(data_raw, "hf_repo_type", default="dataset"), {"dataset"}, "data.hf_repo_type"),
hf_path_prefix=_string(data_raw, "hf_path_prefix", default=""),
@ -354,9 +364,15 @@ def load_training_config(path: str | Path) -> TrainingConfig:
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),
include_optimizer_state=_boolean(checkpoint_raw, "include_optimizer_state") if "include_optimizer_state" in checkpoint_raw else True,
include_rng_state=_boolean(checkpoint_raw, "include_rng_state") if "include_rng_state" in checkpoint_raw else True,
policy=_choice(_string(checkpoint_raw, "policy", default="full").lower(), {"full", "lightweight_scaling_probe"}, "checkpoint.policy"),
)
stability = StabilityConfig(
max_grad_norm=_optional_number(stability_raw, "max_grad_norm", minimum=0.0, exclusive_minimum=True),
dead_curve_patience_evals=_optional_integer(stability_raw, "dead_curve_patience_evals", minimum=1),
dead_curve_min_relative_improvement=_number(stability_raw, "dead_curve_min_relative_improvement", minimum=0.0, default=0.005),
dead_curve_warmup_steps=_integer(stability_raw, "dead_curve_warmup_steps", minimum=0, default=0),
)
precision = PrecisionConfig(
dtype=_choice(_string(precision_raw, "dtype", default="float32").lower(), {"float32", "bf16"}, "precision.dtype"),

15
src/airfrans_frontier/training/data.py
View file

@ -52,7 +52,7 @@ class SplitArrays:
features: FloatArray
targets: FloatArray
case_ids: tuple[str, ...]
point_counts: tuple[int, ...] = ()
@dataclass(frozen=True)
class DatasetBundle:
@ -177,7 +177,7 @@ def build_dataset_bundle(
train_cases: int,
val_cases: int,
test_cases: int,
points_per_case: int,
points_per_case: int | None,
seed: int,
) -> DatasetBundle:
split = create_case_split(
@ -194,7 +194,7 @@ def build_dataset_bundle_for_split(
samples: list[SimulationSample],
*,
split: CaseSplit,
points_per_case: int,
points_per_case: int | None,
seed: int,
) -> DatasetBundle:
validate_common_schema(samples)
@ -226,7 +226,7 @@ def build_dataset_bundle_for_split(
def build_split_arrays(
samples_by_id: dict[str, SimulationSample],
case_ids: tuple[str, ...],
points_per_case: int,
points_per_case: int | None,
*,
seed: int,
) -> SplitArrays:
@ -235,9 +235,11 @@ def build_split_arrays(
rng = np.random.default_rng(seed)
feature_parts: list[FloatArray] = []
target_parts: list[FloatArray] = []
point_counts: list[int] = []
for case_id in case_ids:
sample = samples_by_id[case_id]
indices = _sample_indices(sample.num_points, points_per_case, rng)
point_counts.append(int(indices.shape[0]))
feature_parts.append(np.ascontiguousarray(sample.features[indices], dtype=np.float32))
target_parts.append(np.ascontiguousarray(sample.targets[indices], dtype=np.float32))
@ -245,13 +247,14 @@ def build_split_arrays(
features=np.concatenate(feature_parts, axis=0),
targets=np.concatenate(target_parts, axis=0),
case_ids=case_ids,
point_counts=tuple(point_counts),
)
def _sample_indices(num_points: int, points_per_case: int, rng: np.random.Generator) -> NDArray[np.int64]:
def _sample_indices(num_points: int, points_per_case: int | None, rng: np.random.Generator) -> NDArray[np.int64]:
if num_points <= 0:
raise ValueError("Cannot sample from an empty simulation")
if points_per_case >= num_points:
if points_per_case is None or points_per_case >= num_points:
return np.arange(num_points, dtype=np.int64)
return np.sort(rng.choice(num_points, size=points_per_case, replace=False)).astype(np.int64)

165
src/airfrans_frontier/training/loop.py
View file

@ -16,6 +16,7 @@ from torch.nn import functional as F
from airfrans_frontier.models import (
DeepONetBranchTrunk,
EncodedPointwiseMLP,
FourierFiLMMLP,
LocalPointTransformer,
NeRFCFDMultiRes,
@ -69,6 +70,58 @@ class TrainingResult:
final_metrics: dict[str, Any]
class DeadCurveError(RuntimeError):
pass
@dataclass
class LossMovementMonitor:
patience_evals: int
min_relative_improvement: float
warmup_steps: int
baseline_loss: float | None = None
baseline_step: int | None = None
best_loss: float | None = None
evals_after_baseline: int = 0
@classmethod
def from_config(cls, config: TrainingConfig) -> "LossMovementMonitor | None":
if config.stability.dead_curve_patience_evals is None:
return None
return cls(
patience_evals=config.stability.dead_curve_patience_evals,
min_relative_improvement=config.stability.dead_curve_min_relative_improvement,
warmup_steps=config.stability.dead_curve_warmup_steps,
)
def observe(self, *, step: int, loss: float) -> dict[str, Any] | None:
if step < self.warmup_steps:
return None
if self.baseline_loss is None:
self.baseline_loss = float(loss)
self.baseline_step = int(step)
self.best_loss = float(loss)
self.evals_after_baseline = 0
return None
self.evals_after_baseline += 1
self.best_loss = min(float(self.best_loss), float(loss)) if self.best_loss is not None else float(loss)
denominator = max(abs(float(self.baseline_loss)), 1e-12)
relative_improvement = (float(self.baseline_loss) - float(self.best_loss)) / denominator
if self.evals_after_baseline >= self.patience_evals and relative_improvement < self.min_relative_improvement:
return {
"baseline_loss": float(self.baseline_loss),
"baseline_step": self.baseline_step,
"best_loss": float(self.best_loss),
"latest_loss": float(loss),
"relative_improvement": float(relative_improvement),
"min_relative_improvement": float(self.min_relative_improvement),
"patience_evals": int(self.patience_evals),
"evals_after_baseline": int(self.evals_after_baseline),
"warmup_steps": int(self.warmup_steps),
}
return None
def train_from_config_path(path: str | Path, resume_path: str | Path | None = None) -> TrainingResult:
config = load_training_config(path)
return train(config, resume_path=resume_path or os.environ.get("AIRFRANS_RESUME_CHECKPOINT"))
@ -204,8 +257,8 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
}
)
if config.data.source == "public_zip_streaming":
return _train_public_zip_streaming(
if config.data.source in {"public_zip_streaming", "huggingface_streaming"}:
return _train_streaming_data(
config=config,
resume=resume,
resume_info=resume_info,
@ -258,6 +311,9 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
"cache_dir": str(config.data.cache_dir) if config.data.cache_dir is not None else None,
"case_count": len(samples),
"total_points": sum(sample.num_points for sample in samples),
"selected_points": int(bundle.train.features.shape[0] + (bundle.val.features.shape[0] if bundle.val is not None else 0) + (bundle.test.features.shape[0] if bundle.test is not None else 0)),
"all_points_per_case": config.data.all_points_per_case,
"points_per_case": config.data.points_per_case,
"feature_names": list(bundle.feature_names),
"target_names": list(bundle.target_names),
"cases": [
@ -356,6 +412,7 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
start_step = 0
best_val_loss: float | None = None
initial_train_loss: float | None = None
loss_movement_monitor = LossMovementMonitor.from_config(config)
full_train_eval_cases = len(bundle.train.case_ids)
train_eval_cases = full_train_eval_cases
train_eval_features = train_features
@ -363,7 +420,12 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
train_eval_case_limit = config.data.streaming_normalization_cases
if train_eval_case_limit is not None and train_eval_case_limit < full_train_eval_cases:
train_eval_cases = max(1, train_eval_case_limit)
train_eval_rows = min(train_features.shape[0], train_eval_cases * config.data.points_per_case)
if bundle.train.point_counts:
train_eval_rows = min(train_features.shape[0], sum(bundle.train.point_counts[:train_eval_cases]))
elif config.data.points_per_case is not None:
train_eval_rows = min(train_features.shape[0], train_eval_cases * config.data.points_per_case)
else:
train_eval_rows = train_features.shape[0]
train_eval_features = train_features[:train_eval_rows]
train_eval_targets = train_targets[:train_eval_rows]
@ -678,6 +740,22 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
)
)
)
dead_curve = loss_movement_monitor.observe(step=step, loss=float(current_metric)) if loss_movement_monitor is not None else None
if dead_curve is not None:
_write_failure(
writer,
phase="training",
step=step,
error_type="DeadCurveError",
error_message="loss did not improve enough during the early measurement window",
latest_loss=float(current_metric),
latest_grad_norm=last_grad_norm,
latest_checkpoint=LATEST_CHECKPOINT if (writer.run_dir / LATEST_CHECKPOINT).is_file() else None,
config=config,
failure_category="plateau_or_dead_recipe",
dead_curve=dead_curve,
)
raise DeadCurveError("loss did not improve enough during the early measurement window")
last_log_at = time.perf_counter()
last_log_step = step
model.train()
@ -778,6 +856,10 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
"val_cases": len(bundle.split.val_ids),
"test_cases": len(bundle.split.test_ids),
"points_per_case": config.data.points_per_case,
"all_points_per_case": config.data.all_points_per_case,
"effective_point_updates": config.optim.steps * config.data.batch_size,
"checkpoint_policy": config.checkpoint.policy,
"checkpoint_include_optimizer_state": config.checkpoint.include_optimizer_state,
"steps": config.optim.steps,
"elapsed_seconds": elapsed,
"points_per_sec": last_points_per_sec,
@ -882,7 +964,7 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
observer.finish(exit_code=0)
return TrainingResult(run_dir=writer.run_dir, final_metrics=final_metrics)
def _train_public_zip_streaming(
def _train_streaming_data(
*,
config: TrainingConfig,
resume: Path | None,
@ -907,6 +989,7 @@ def _train_public_zip_streaming(
last_points_per_sec: float | None = None
first_streaming_metric_written = False
first_streaming_checkpoint_written = False
loss_movement_monitor = LossMovementMonitor.from_config(config)
def record_streaming_metrics(metrics: dict[str, Any]) -> None:
nonlocal first_streaming_metric_written
@ -950,7 +1033,7 @@ def _train_public_zip_streaming(
run_manifest.update(
{
"phase": "initialized",
"data_mode": "public_zip_streaming",
"data_mode": config.data.source,
"parameter_count": count_parameters(model),
**calibration_fields,
**protocol_fields,
@ -1284,6 +1367,23 @@ def _train_public_zip_streaming(
latest_checkpoint=LATEST_CHECKPOINT if should_checkpoint else None,
)
)
dead_curve = loss_movement_monitor.observe(step=step, loss=float(current_metric)) if loss_movement_monitor is not None else None
if dead_curve is not None:
_write_failure(
writer,
phase="streaming_training",
step=step,
error_type="DeadCurveError",
error_message="loss did not improve enough during the early measurement window",
latest_loss=float(current_metric),
latest_grad_norm=last_grad_norm,
latest_checkpoint=LATEST_CHECKPOINT if (writer.run_dir / LATEST_CHECKPOINT).is_file() else None,
streaming_summary=streaming.telemetry_summary(),
config=config,
failure_category="plateau_or_dead_recipe",
dead_curve=dead_curve,
)
raise DeadCurveError("loss did not improve enough during the early measurement window")
last_log_at = time.perf_counter()
last_log_step = step
model.train()
@ -1341,6 +1441,10 @@ def _train_public_zip_streaming(
"val_cases": len(bundle.split.val_ids),
"test_cases": len(bundle.split.test_ids),
"points_per_case": config.data.points_per_case,
"all_points_per_case": config.data.all_points_per_case,
"effective_point_updates": config.optim.steps * config.data.batch_size,
"checkpoint_policy": config.checkpoint.policy,
"checkpoint_include_optimizer_state": config.checkpoint.include_optimizer_state,
"steps": config.optim.steps,
"elapsed_seconds": elapsed,
"points_per_sec": last_points_per_sec,
@ -1348,7 +1452,7 @@ def _train_public_zip_streaming(
"validation_runtime_seconds": validation_runtime_seconds,
"checkpoint_interval_seconds": config.checkpoint.interval_seconds,
"data_source": config.data.source,
"data_mode": "public_zip_streaming",
"data_mode": config.data.source,
"data_public_source_url": config.data.public_source_url,
"data_hf_repo_id": config.data.hf_repo_id,
"data_hf_path_prefix": config.data.hf_path_prefix,
@ -1410,7 +1514,7 @@ def _train_public_zip_streaming(
"phase": "completed",
"finished_at": time.time(),
"exit_code": 0,
"data_mode": "public_zip_streaming",
"data_mode": config.data.source,
"final_metrics_path": str(writer.run_dir / "final_metrics.json"),
"checkpoint_latest_path": str(writer.run_dir / LATEST_CHECKPOINT),
"checkpoint_best_path": str(writer.run_dir / BEST_CHECKPOINT),
@ -1479,7 +1583,7 @@ def _train_public_zip_streaming(
config=config,
failure_category="streaming_training",
)
run_manifest.update({"phase": "failed", "finished_at": time.time(), "exit_code": 1, "data_mode": "public_zip_streaming", **uploader.finalize(training_success=False)})
run_manifest.update({"phase": "failed", "finished_at": time.time(), "exit_code": 1, "data_mode": config.data.source, **uploader.finalize(training_success=False)})
writer.write_json("run_manifest.json", run_manifest)
writer.write_artifact_manifest()
try:
@ -1545,6 +1649,16 @@ def _build_model(config: TrainingConfig, bundle: DatasetBundle, *, output_dim: i
depth=config.model.depth,
activation=config.model.activation,
)
if model_type == "mlp_encoded_baseline":
return EncodedPointwiseMLP(
feature_names=bundle.feature_names,
output_dim=output_dim,
coordinate_features=config.model.coordinate_features,
coordinate_encoding_spec=coord_spec,
hidden_width=config.model.hidden_width,
depth=config.model.depth,
activation=config.model.activation,
)
if model_type in {"film_fourier_mlp", "film_fourier_inr"}:
return FourierFiLMMLP(
feature_names=bundle.feature_names,
@ -1809,6 +1923,15 @@ def _job_manifest(*, config: TrainingConfig, run_id: str, run_dir: Path) -> dict
"proxy_for": config.model_metadata.proxy_for,
"proxy_notes": config.model_metadata.proxy_notes,
"model_metadata": asdict(config.model_metadata),
"data": {
"source": config.data.source,
"train_cases": config.data.train_cases,
"val_cases": config.data.val_cases,
"test_cases": config.data.test_cases,
"all_points_per_case": config.data.all_points_per_case,
"points_per_case": config.data.points_per_case,
"batch_size": config.data.batch_size,
},
"coordinate_encoding": {
"type": config.coordinate_encoding.type,
"features": list(config.coordinate_encoding.features),
@ -1820,10 +1943,19 @@ def _job_manifest(*, config: TrainingConfig, run_id: str, run_dir: Path) -> dict
},
"checkpoint_policy": {
"interval_seconds": config.checkpoint.interval_seconds,
"policy": config.checkpoint.policy,
"include_optimizer_state": config.checkpoint.include_optimizer_state,
"include_rng_state": config.checkpoint.include_rng_state,
"latest": LATEST_CHECKPOINT,
"best": BEST_CHECKPOINT,
"final": FINAL_CHECKPOINT,
},
"stability": {
"max_grad_norm": config.stability.max_grad_norm,
"dead_curve_patience_evals": config.stability.dead_curve_patience_evals,
"dead_curve_min_relative_improvement": config.stability.dead_curve_min_relative_improvement,
"dead_curve_warmup_steps": config.stability.dead_curve_warmup_steps,
},
"observability": {
"backend": config.observability.backend,
"project": config.observability.project,
@ -2111,6 +2243,8 @@ def _checkpoint_payload(
initial_train_loss: float | None,
final_metrics: dict[str, Any] | None,
) -> dict[str, Any]:
include_optimizer = config.checkpoint.include_optimizer_state
include_rng = config.checkpoint.include_rng_state
return {
"schema_version": CHECKPOINT_SCHEMA_VERSION,
"run_id": os.environ.get("AIRFRANS_REMOTE_RUN_ID", config.run.name),
@ -2123,18 +2257,21 @@ def _checkpoint_payload(
"output_dim": bundle.train.targets.shape[1],
"model_config": asdict(config.model),
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
"optimizer_state_dict": optimizer.state_dict() if include_optimizer else None,
"scheduler_state_dict": None,
"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,
"rng_state": random.getstate() if include_rng else None,
"numpy_rng_state": np.random.get_state() if include_rng else None,
"torch_rng_state": torch.get_rng_state() if include_rng else None,
"cuda_rng_state": torch.cuda.get_rng_state_all() if include_rng and torch.cuda.is_available() else None,
"batch_rng_state": rng.bit_generator.state if include_rng else None,
"checkpoint_policy": config.checkpoint.policy,
"include_optimizer_state": include_optimizer,
"include_rng_state": include_rng,
"final_metrics": final_metrics,
}

505
src/airfrans_frontier/training/streaming_data.py
View file

@ -1,9 +1,11 @@
from __future__ import annotations
import json
import queue
import os
import shutil
import time
import threading
from dataclasses import dataclass, field
from pathlib import Path, PurePosixPath
from typing import Any, Iterable, Iterator, Mapping
@ -335,6 +337,8 @@ class PublicZipStreamingCache:
"cache_high_water_bytes": self.data_config.streaming_cache_high_water_bytes,
"cache_low_water_bytes": self.data_config.streaming_cache_low_water_bytes,
"queue_max_cases": self.data_config.streaming_queue_max_cases,
"all_points_per_case": self.data_config.all_points_per_case,
"points_per_case": self.data_config.points_per_case,
"ranged_bytes_read": self.reader.bytes_read if self.reader is not None else None,
"cases": [
{
@ -500,11 +504,246 @@ class PublicZipStreamingCache:
return {"schema_version": 1, "cases": {}, "sampling": {"train": {}, "val": {}, "test": {}}}
class HuggingFaceStreamingCache:
def __init__(self, data_config: DataConfig, *, run_dir: Path, recorder: StreamingEventRecorder) -> None:
if not data_config.hf_repo_id:
raise ValueError("data.hf_repo_id is required when data.source = 'huggingface_streaming'")
self.data_config = data_config
self.run_dir = run_dir
self.recorder = recorder
self.repo_id = data_config.hf_repo_id
self.repo_type = data_config.hf_repo_type
self.path_prefix = data_config.hf_path_prefix.strip("/")
self.cache_dir = data_config.cache_dir or data_config.root
self.data_root = self.cache_dir / self.path_prefix if self.path_prefix else self.cache_dir
self.state_path = run_dir / STREAMING_STATE
self.indices_dir = run_dir / "streaming_sample_indices"
self.cache_dir.mkdir(parents=True, exist_ok=True)
self.data_root.mkdir(parents=True, exist_ok=True)
self.indices_dir.mkdir(parents=True, exist_ok=True)
self.case_files: dict[str, str] = {}
self.state: dict[str, Any] = _read_json(self.state_path) if self.state_path.is_file() else self._empty_state()
self._active_cases: set[str] = set()
self._last_access: dict[str, float] = {}
self._loaded: dict[str, SimulationSample] = {}
def enumerate_cases(self) -> tuple[str, ...]:
self.recorder.emit(
"dataset_enumeration_start",
phase="data",
source="huggingface_streaming",
hf_repo_id=self.repo_id,
hf_repo_type=self.repo_type,
hf_path_prefix=self.path_prefix,
)
token = _resolve_optional_secret("HF_TOKEN")
try:
from huggingface_hub import HfApi
except ModuleNotFoundError as exc:
raise RuntimeError("huggingface_hub is required when data.source = 'huggingface_streaming'") from exc
api = HfApi(token=token)
files = api.list_repo_files(repo_id=self.repo_id, repo_type=self.repo_type)
prefix = f"{self.path_prefix}/" if self.path_prefix else ""
case_files = {
Path(name).stem: name
for name in files
if name.startswith(prefix) and name.endswith(".npz") and Path(name).name
}
if not case_files:
raise FileNotFoundError(f"No .npz cases found in HF dataset {self.repo_id!r} under {self.path_prefix!r}")
self.case_files = dict(sorted(case_files.items()))
case_ids = tuple(self.case_files)
self.state["source"] = "huggingface_streaming"
self.state["hf_repo_id"] = self.repo_id
self.state["hf_repo_type"] = self.repo_type
self.state["hf_path_prefix"] = self.path_prefix
self.state["enumerated_case_count"] = len(case_ids)
self._write_state()
self.recorder.emit(
"dataset_enumeration_end",
phase="data",
source="huggingface_streaming",
hf_repo_id=self.repo_id,
hf_repo_type=self.repo_type,
hf_path_prefix=self.path_prefix,
case_count=len(case_ids),
)
return case_ids
def ensure_case(self, case_id: str) -> SimulationSample:
cached = self._load_valid_cached_case(case_id)
if cached is not None:
self._active_cases.add(case_id)
self._last_access[case_id] = time.time()
return cached
filename = self.case_files.get(case_id)
if filename is None:
raise KeyError(f"Unknown Hugging Face AirfRANS case: {case_id}")
token = _resolve_optional_secret("HF_TOKEN")
self._mark_case(case_id, status="downloading", path=str(self.data_root / f"{case_id}.npz"), download_started_at=time.time())
started = time.perf_counter()
self.recorder.emit("hf_case_download_start", phase="data", case_id=case_id, hf_filename=filename)
try:
from huggingface_hub import hf_hub_download
except ModuleNotFoundError as exc:
raise RuntimeError("huggingface_hub is required when data.source = 'huggingface_streaming'") from exc
try:
downloaded_path = Path(
hf_hub_download(
repo_id=self.repo_id,
filename=filename,
repo_type=self.repo_type,
local_dir=str(self.cache_dir),
token=token,
)
)
sample_path = self.data_root / f"{case_id}.npz"
if downloaded_path != sample_path and downloaded_path.is_file() and not sample_path.is_file():
sample_path.parent.mkdir(parents=True, exist_ok=True)
shutil.copy2(downloaded_path, sample_path)
sample = load_simulation_npz(sample_path)
except Exception as exc:
self._mark_case(case_id, status="failed", error_type=type(exc).__name__, error_message=str(exc), failed_at=time.time())
self.recorder.emit("hf_case_download_failure", phase="data", case_id=case_id, error_type=type(exc).__name__, error_message=str(exc))
raise
processed_bytes = sample.source_path.stat().st_size
self.recorder.add_downloaded_bytes(processed_bytes)
self._loaded[case_id] = sample
self._active_cases.add(case_id)
self._last_access[case_id] = time.time()
elapsed = time.perf_counter() - started
self._mark_case(
case_id,
status="processed",
path=str(sample.source_path),
processed_bytes=processed_bytes,
points=sample.num_points,
processed_at=time.time(),
processing_seconds=elapsed,
hf_filename=filename,
)
self.recorder.emit(
"hf_case_download_end",
phase="data",
case_id=case_id,
hf_filename=filename,
processed_bytes=processed_bytes,
points=sample.num_points,
download_seconds=elapsed,
)
self.recorder.emit(
"processing_end",
phase="data",
case_id=case_id,
processed_bytes=processed_bytes,
points=sample.num_points,
processing_seconds=elapsed,
)
self._observe_cache()
return sample
def release_case(self, case_id: str, *, consumed: bool = True) -> None:
self._active_cases.discard(case_id)
self._loaded.pop(case_id, None)
self._last_access[case_id] = time.time()
entry = self.state.setdefault("cases", {}).setdefault(case_id, {})
if consumed:
entry["consumed_count"] = int(entry.get("consumed_count", 0) or 0) + 1
entry["last_consumed_at"] = time.time()
self._write_state()
self._observe_cache()
def cache_size_bytes(self) -> int:
return _tree_size_bytes(self.data_root, suffix=".npz")
def data_manifest(self, *, split: CaseSplit, feature_names: tuple[str, ...], target_names: tuple[str, ...]) -> dict[str, Any]:
selected = set(split.train_ids) | set(split.val_ids) | set(split.test_ids)
cases_state = self.state.get("cases", {}) if isinstance(self.state.get("cases"), dict) else {}
return {
"source": "huggingface_streaming",
"hf_repo_id": self.repo_id,
"hf_repo_type": self.repo_type,
"hf_path_prefix": self.path_prefix,
"configured_root": str(self.data_config.root),
"cache_dir": str(self.cache_dir),
"data_root": str(self.data_root),
"case_count": self.state.get("enumerated_case_count"),
"selected_case_count": len(selected),
"train_cases": len(split.train_ids),
"val_cases": len(split.val_ids),
"test_cases": len(split.test_ids),
"feature_names": list(feature_names),
"target_names": list(target_names),
"queue_max_cases": self.data_config.streaming_queue_max_cases,
"all_points_per_case": self.data_config.all_points_per_case,
"points_per_case": self.data_config.points_per_case,
"cases": [
{
"case_id": case_id,
"split": _split_name_for_case(split, case_id),
"status": cases_state.get(case_id, {}).get("status", "pending") if isinstance(cases_state.get(case_id), dict) else "pending",
"points": cases_state.get(case_id, {}).get("points") if isinstance(cases_state.get(case_id), dict) else None,
"processed_bytes": cases_state.get(case_id, {}).get("processed_bytes") if isinstance(cases_state.get(case_id), dict) else None,
"hf_filename": self.case_files.get(case_id),
}
for case_id in sorted(selected)
],
}
def _load_valid_cached_case(self, case_id: str) -> SimulationSample | None:
if case_id in self._loaded:
return self._loaded[case_id]
path = self.data_root / f"{case_id}.npz"
if not path.is_file():
return None
try:
sample = load_simulation_npz(path)
except Exception:
path.unlink(missing_ok=True)
self.recorder.emit("partial_unit_discarded", phase="data", case_id=case_id, path=str(path))
self._mark_case(case_id, status="pending", discarded_invalid_at=time.time())
return None
self._loaded[case_id] = sample
self._last_access[case_id] = time.time()
entry = self.state.setdefault("cases", {}).setdefault(case_id, {})
if entry.get("status") == "processed":
self.recorder.emit("resume_validated_unit_reused", phase="data", case_id=case_id, path=str(path), processed_bytes=path.stat().st_size)
self._mark_case(case_id, status="processed", path=str(path), processed_bytes=path.stat().st_size, points=sample.num_points)
self._observe_cache()
return sample
def _observe_cache(self) -> None:
self.recorder.observe_cache(self.data_root, cache_bytes=self.cache_size_bytes())
def _mark_case(self, case_id: str, **fields: Any) -> None:
cases = self.state.setdefault("cases", {})
entry = cases.setdefault(case_id, {})
entry.update(fields)
entry["updated_at"] = time.time()
self._write_state()
def _write_state(self) -> None:
self.state["schema_version"] = 1
self.state["cache_dir"] = str(self.cache_dir)
self.state["data_root"] = str(self.data_root)
self.state["updated_at"] = time.time()
_atomic_write_json(self.state_path, self.state)
@staticmethod
def _empty_state() -> dict[str, Any]:
return {"schema_version": 1, "cases": {}, "sampling": {"train": {}, "val": {}, "test": {}}}
class StreamingTrainingData:
def __init__(self, config: TrainingConfig, *, run_dir: Path, recorder: StreamingEventRecorder) -> None:
self.config = config
self.recorder = recorder
self.cache = PublicZipStreamingCache(config.data, run_dir=run_dir, recorder=recorder)
if config.data.source == "public_zip_streaming":
self.cache = PublicZipStreamingCache(config.data, run_dir=run_dir, recorder=recorder)
elif config.data.source == "huggingface_streaming":
self.cache = HuggingFaceStreamingCache(config.data, run_dir=run_dir, recorder=recorder)
else:
raise ValueError(f"StreamingTrainingData does not support data.source = {config.data.source!r}")
self.split: CaseSplit | None = None
self.feature_names: tuple[str, ...] | None = None
self.target_names: tuple[str, ...] | None = None
@ -514,11 +753,18 @@ class StreamingTrainingData:
self._train_counts: list[int] = []
self._train_offsets: NDArray[np.int64] | None = None
self._upload_queue = ProcessedDataUploadQueue.from_config(config, run_dir=run_dir, recorder=recorder)
self._io_lock = threading.RLock()
self._batch_queue: queue.Queue[tuple[FloatArray, FloatArray] | BaseException] | None = None
self._batch_thread: threading.Thread | None = None
self._batch_stop = threading.Event()
self._acquisition_thread: threading.Thread | None = None
self._acquisition_stop = threading.Event()
self._acquisition_error: BaseException | None = None
@classmethod
def from_config(cls, config: TrainingConfig, *, run_dir: Path, recorder: StreamingEventRecorder) -> StreamingTrainingData:
if config.data.source != "public_zip_streaming":
raise ValueError(f"StreamingTrainingData requires data.source = 'public_zip_streaming', got {config.data.source!r}")
if config.data.source not in {"public_zip_streaming", "huggingface_streaming"}:
raise ValueError(f"StreamingTrainingData requires a streaming data source, got {config.data.source!r}")
return cls(config, run_dir=run_dir, recorder=recorder)
def prepare(self) -> None:
@ -543,11 +789,12 @@ class StreamingTrainingData:
seed=self.config.run.seed,
)
first_case = split.train_ids[0]
sample = self.cache.ensure_case(first_case)
self.feature_names = sample.feature_names
self.target_names = sample.target_names
self._upload_queue.enqueue(sample.source_path)
self.cache.release_case(first_case, consumed=False)
with self._io_lock:
sample = self.cache.ensure_case(first_case)
self.feature_names = sample.feature_names
self.target_names = sample.target_names
self._upload_queue.enqueue(sample.source_path)
self.cache.release_case(first_case, consumed=False)
def load_or_compute_normalization(self) -> NormalizationStats:
if self.split is None:
@ -558,6 +805,8 @@ class StreamingTrainingData:
self.stats = stats
self._restore_sampling_specs("train")
self._build_train_offsets()
if self._train_offsets is not None:
self._start_background_acquisition()
self.recorder.emit("normalization_end", phase="normalization", reused=True, normalization_runtime_seconds=0.0)
return stats
assert self.feature_names is not None
@ -576,15 +825,16 @@ class StreamingTrainingData:
accumulator = _StatsAccumulator(feature_names=self.feature_names, target_names=self.target_names)
rng = np.random.default_rng(self.config.run.seed + _SAMPLE_SEEDS["train"])
for case_id in normalization_case_ids:
sample = self.cache.ensure_case(case_id)
self._validate_schema(sample)
indices = self._sampling_spec_for_case("train", sample, rng)
features = _selected_rows(sample.features, indices)
targets = _selected_rows(sample.targets, indices)
accumulator.update(features, targets)
self._upload_queue.enqueue(sample.source_path)
self.cache.release_case(case_id)
self._upload_queue.drain()
with self._io_lock:
sample = self.cache.ensure_case(case_id)
self._validate_schema(sample)
indices = self._sampling_spec_for_case("train", sample, rng)
features = _selected_rows(sample.features, indices)
targets = _selected_rows(sample.targets, indices)
accumulator.update(features, targets)
self._upload_queue.enqueue(sample.source_path)
self.cache.release_case(case_id)
self._upload_queue.drain()
stats = accumulator.finish()
self.stats = stats
self._build_train_offsets()
@ -597,6 +847,8 @@ class StreamingTrainingData:
sample_count=accumulator.count,
normalization_runtime_seconds=runtime,
)
if self._train_offsets is not None:
self._start_background_acquisition()
return stats
def schema_bundle(self) -> DatasetBundle:
@ -622,31 +874,167 @@ class StreamingTrainingData:
self.recorder.set_upload_state(**self._upload_queue.telemetry_state())
return self.recorder.to_dict()
def _start_background_acquisition(self) -> None:
if self.config.data.source != "huggingface_streaming" or self._acquisition_thread is not None:
return
if self.split is None:
raise RuntimeError("split is missing")
targets = tuple(dict.fromkeys(self.split.train_ids + self.split.val_ids + self.split.test_ids))
self._acquisition_stop.clear()
self._acquisition_error = None
self._acquisition_thread = threading.Thread(
target=self._background_acquisition_loop,
args=(targets,),
name="airfrans-hf-dataset-acquisition",
daemon=True,
)
self._acquisition_thread.start()
self.recorder.emit("background_acquisition_start", phase="data", source="huggingface_streaming", case_count=len(targets))
def _background_acquisition_loop(self, case_ids: tuple[str, ...]) -> None:
acquired = 0
try:
for case_id in case_ids:
if self._acquisition_stop.is_set():
self.recorder.emit("background_acquisition_stopped", phase="data", source="huggingface_streaming", acquired_cases=acquired, target_cases=len(case_ids))
return
with self._io_lock:
sample = self.cache.ensure_case(case_id)
self._validate_schema(sample)
self.cache.release_case(case_id, consumed=False)
acquired += 1
self.recorder.emit("background_case_ready", phase="data", source="huggingface_streaming", case_id=case_id, acquired_cases=acquired, target_cases=len(case_ids))
self.recorder.emit("background_acquisition_complete", phase="data", source="huggingface_streaming", acquired_cases=acquired, target_cases=len(case_ids))
except BaseException as exc:
self._acquisition_error = exc
self.recorder.emit("background_acquisition_failure", phase="data", source="huggingface_streaming", error_type=type(exc).__name__, error_message=str(exc), acquired_cases=acquired, target_cases=len(case_ids))
def _check_background_acquisition(self) -> None:
if self._acquisition_error is not None:
raise self._acquisition_error
def stop_background_acquisition(self, *, wait: bool) -> None:
if self._acquisition_thread is None:
return
if not wait:
self._acquisition_stop.set()
self._acquisition_thread.join(timeout=None if wait else 2.0)
if self._acquisition_thread.is_alive():
return
self._acquisition_thread = None
self._check_background_acquisition()
def _ensure_train_batch_prefetch(self, *, batch_size: int) -> queue.Queue[tuple[FloatArray, FloatArray] | BaseException]:
if self._batch_queue is not None:
return self._batch_queue
max_batches = max(2, self.config.data.streaming_queue_max_cases * 4)
self._batch_queue = queue.Queue(maxsize=max_batches)
self._batch_stop.clear()
self._batch_thread = threading.Thread(
target=self._train_batch_prefetch_loop,
args=(batch_size,),
name="airfrans-streaming-batch-prefetch",
daemon=True,
)
self._batch_thread.start()
self.recorder.emit("train_batch_prefetch_start", phase="training", batch_size=batch_size, max_batches=max_batches)
return self._batch_queue
def _train_batch_prefetch_loop(self, batch_size: int) -> None:
assert self.stats is not None
assert self.split is not None
assert self._batch_queue is not None
rng = np.random.default_rng(self.config.run.seed + 505)
train_ids = list(self.split.train_ids)
batches_per_case = max(1, self.config.data.streaming_queue_max_cases)
try:
while not self._batch_stop.is_set():
for case_index in rng.permutation(len(train_ids)):
if self._batch_stop.is_set():
return
case_id = train_ids[int(case_index)]
batches: list[tuple[FloatArray, FloatArray]] = []
with self._io_lock:
sample = self.cache.ensure_case(case_id)
self._validate_schema(sample)
spec = self._sampling_spec_for_case("train", sample, rng)
max_batches_for_case = max(1, min(batches_per_case, (spec.count + batch_size - 1) // batch_size))
for _ in range(max_batches_for_case):
batches.append(self._sample_normalized_batch(sample, spec, rng, batch_size=batch_size))
self._upload_queue.enqueue(sample.source_path)
self.cache.release_case(case_id)
self._upload_queue.drain()
for batch in batches:
if not self._put_prefetched_batch(batch):
return
except BaseException as exc:
self.recorder.emit(
"train_batch_prefetch_failure",
phase="training",
error_type=type(exc).__name__,
error_message=str(exc),
)
self._put_prefetched_batch(exc)
def _sample_normalized_batch(
self,
sample: SimulationSample,
spec: SamplingSpec,
rng: np.random.Generator,
*,
batch_size: int,
) -> tuple[FloatArray, FloatArray]:
assert self.stats is not None
local_indices = rng.integers(0, spec.count, size=batch_size, dtype=np.int64)
source_indices = _source_indices_for_local(spec, local_indices)
selected_features = sample.features[source_indices]
selected_targets = sample.targets[source_indices]
features = ((selected_features - self.stats.feature_mean) / self.stats.feature_std).astype(np.float32, copy=False)
targets = ((selected_targets - self.stats.target_mean) / self.stats.target_std).astype(np.float32, copy=False)
return np.ascontiguousarray(features, dtype=np.float32), np.ascontiguousarray(targets, dtype=np.float32)
def _put_prefetched_batch(self, item: tuple[FloatArray, FloatArray] | BaseException) -> bool:
assert self._batch_queue is not None
while not self._batch_stop.is_set():
try:
self._batch_queue.put(item, timeout=0.1)
return True
except queue.Full:
continue
return False
def stop_train_batch_prefetch(self) -> None:
self._batch_stop.set()
if self._batch_thread is not None:
self._batch_thread.join(timeout=2.0)
self._batch_thread = None
self._batch_queue = None
def sample_train_batch(self, rng: np.random.Generator, *, batch_size: int, step: int) -> tuple[FloatArray, FloatArray]:
if self.stats is None or self.split is None:
raise RuntimeError("Streaming normalization must be computed before sampling")
self._check_background_acquisition()
started = time.perf_counter()
if self._train_offsets is None:
batch_queue = self._ensure_train_batch_prefetch(batch_size=batch_size)
item = batch_queue.get()
wait_seconds = time.perf_counter() - started
self.recorder.add_trainer_wait(wait_seconds, step=step)
self.recorder.mark_first_batch_ready()
if isinstance(item, BaseException):
raise item
self._start_background_acquisition()
return item
features = np.empty((batch_size, len(self.stats.feature_names)), dtype=np.float32)
targets = np.empty((batch_size, len(self.stats.target_names)), dtype=np.float32)
if self._train_offsets is None:
case_id = self.split.train_ids[int(rng.integers(0, len(self.split.train_ids)))]
sample = self.cache.ensure_case(case_id)
self._validate_schema(sample)
spec = self._sampling_spec_for_case("train", sample, rng)
local_indices = rng.integers(0, spec.count, size=batch_size, dtype=np.int64)
source_indices = _source_indices_for_local(spec, local_indices)
selected_features = sample.features[source_indices]
selected_targets = sample.targets[source_indices]
features[:] = ((selected_features - self.stats.feature_mean) / self.stats.feature_std).astype(np.float32, copy=False)
targets[:] = ((selected_targets - self.stats.target_mean) / self.stats.target_std).astype(np.float32, copy=False)
self._upload_queue.enqueue(sample.source_path)
self.cache.release_case(case_id)
else:
total = int(self._train_offsets[-1]) if self._train_offsets.size else 0
if total <= 0:
raise ValueError("Streaming train split has no sampled points")
global_indices = rng.integers(0, total, size=batch_size)
case_positions = np.searchsorted(self._train_offsets[1:], global_indices, side="right")
total = int(self._train_offsets[-1]) if self._train_offsets.size else 0
if total <= 0:
raise ValueError("Streaming train split has no sampled points")
global_indices = rng.integers(0, total, size=batch_size)
case_positions = np.searchsorted(self._train_offsets[1:], global_indices, side="right")
with self._io_lock:
for case_position in np.unique(case_positions):
mask = case_positions == case_position
case_id = self.split.train_ids[int(case_position)]
@ -661,10 +1049,10 @@ class StreamingTrainingData:
targets[mask] = ((selected_targets - self.stats.target_mean) / self.stats.target_std).astype(np.float32, copy=False)
self._upload_queue.enqueue(sample.source_path)
self.cache.release_case(case_id)
self._upload_queue.drain()
wait_seconds = time.perf_counter() - started
self.recorder.add_trainer_wait(wait_seconds, step=step)
self.recorder.mark_first_batch_ready()
self._upload_queue.drain()
return np.ascontiguousarray(features, dtype=np.float32), np.ascontiguousarray(targets, dtype=np.float32)
def iter_split_batches(self, split_name: str, *, batch_size: int) -> Iterator[tuple[FloatArray, FloatArray]]:
@ -673,22 +1061,26 @@ class StreamingTrainingData:
case_ids = self._split_case_ids.get(split_name)
if case_ids is None:
raise ValueError(f"Unknown split: {split_name}")
self.stop_train_batch_prefetch()
rng = np.random.default_rng(self.config.run.seed + _SAMPLE_SEEDS[split_name])
for case_id in case_ids:
sample = self.cache.ensure_case(case_id)
self._validate_schema(sample)
spec = self._sampling_spec_for_case(split_name, sample, rng)
for start in range(0, spec.count, batch_size):
stop = min(start + batch_size, spec.count)
source_indices = _source_indices_for_local(spec, np.arange(start, stop, dtype=np.int64))
features = ((sample.features[source_indices] - self.stats.feature_mean) / self.stats.feature_std).astype(np.float32, copy=False)
targets = ((sample.targets[source_indices] - self.stats.target_mean) / self.stats.target_std).astype(np.float32, copy=False)
yield np.ascontiguousarray(features, dtype=np.float32), np.ascontiguousarray(targets, dtype=np.float32)
self._upload_queue.enqueue(sample.source_path)
self.cache.release_case(case_id)
self._upload_queue.drain()
with self._io_lock:
sample = self.cache.ensure_case(case_id)
self._validate_schema(sample)
spec = self._sampling_spec_for_case(split_name, sample, rng)
for start in range(0, spec.count, batch_size):
stop = min(start + batch_size, spec.count)
source_indices = _source_indices_for_local(spec, np.arange(start, stop, dtype=np.int64))
features = ((sample.features[source_indices] - self.stats.feature_mean) / self.stats.feature_std).astype(np.float32, copy=False)
targets = ((sample.targets[source_indices] - self.stats.target_mean) / self.stats.target_std).astype(np.float32, copy=False)
yield np.ascontiguousarray(features, dtype=np.float32), np.ascontiguousarray(targets, dtype=np.float32)
self._upload_queue.enqueue(sample.source_path)
self.cache.release_case(case_id)
self._upload_queue.drain()
def finish(self, *, success: bool) -> None:
self.stop_train_batch_prefetch()
self.stop_background_acquisition(wait=success)
self._upload_queue.drain(force=True)
self._upload_queue.finalize(training_success=success)
self.cache.state["finished_at"] = time.time()
@ -712,7 +1104,7 @@ class StreamingTrainingData:
existing = self._sampling_specs[split_name].get(sample.case_id)
if existing is not None:
return existing
if self.config.data.points_per_case >= sample.num_points:
if self.config.data.all_points_per_case or self.config.data.points_per_case is None or self.config.data.points_per_case >= sample.num_points:
spec = SamplingSpec(case_id=sample.case_id, split_name=split_name, count=sample.num_points, mode="all")
else:
indices = np.sort(rng.choice(sample.num_points, size=self.config.data.points_per_case, replace=False)).astype(np.int64)
@ -1129,10 +1521,21 @@ def _read_json(path: Path) -> dict[str, Any]:
raise ValueError(f"Expected JSON object in {path}")
return data
def _resolve_optional_secret(name: str) -> str | None:
value = os.environ.get(name)
if value:
return value
path = Path(".env") / name
if path.is_file():
value = path.read_text().strip()
if value:
return value
return None
def _atomic_write_json(path: Path, payload: Mapping[str, Any]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
tmp = path.with_name(f"{path.name}.tmp")
tmp = path.with_name(f"{path.name}.{os.getpid()}.{threading.get_ident()}.tmp")
tmp.write_text(json.dumps(dict(payload), indent=2, sort_keys=True) + "\n")
tmp.replace(path)

1
tests/test_data_sources.py
View file

@ -37,6 +37,7 @@ class DataSourceTests(unittest.TestCase):
val_cases=0,
test_cases=0,
points_per_case=1,
all_points_per_case=False,
batch_size=1,
source="huggingface",
hf_repo_id="owner/airfrans-processed",

19
tests/test_mlp.py
View file

@ -9,7 +9,9 @@ remove_pythonpath_entries()
import torch
from airfrans_frontier.models import (
CoordinateEncodingSpec,
DeepONetBranchTrunk,
EncodedPointwiseMLP,
LocalPointTransformer,
NeRFCFDMultiRes,
PointContextPerceiver,
@ -28,6 +30,23 @@ class PointwiseMLPTests(unittest.TestCase):
self.assertEqual(tuple(output.shape), (7, 4))
def test_encoded_pointwise_mlp_uses_selected_coordinate_features(self) -> None:
feature_names = ("x", "y", "sdf", "u_inf", "aoa_deg")
model = EncodedPointwiseMLP(
feature_names=feature_names,
output_dim=4,
coordinate_features=("x", "y", "sdf"),
coordinate_encoding_spec=CoordinateEncodingSpec(type="fixed_fourier", scales=(1.0, 2.0)),
hidden_width=16,
depth=2,
activation="gelu",
)
batch = torch.randn(7, len(feature_names))
output = model(batch)
self.assertEqual(tuple(output.shape), (7, 4))
def test_frontier_models_return_batch_by_target_dim(self) -> None:
feature_names = ("x", "y", "sdf", "u_inf", "log_re", "aoa_deg")
batch = torch.randn(8, len(feature_names))

163
tests/test_streaming_data.py
View file

@ -2,9 +2,11 @@ from __future__ import annotations
import gzip
import json
import shutil
import os
import sys
import tempfile
import time
import types
import unittest
import zipfile
@ -124,6 +126,122 @@ interval_seconds = 0
+ "\n"
)
def write_huggingface_streaming_config(
path: Path,
*,
cache_dir: Path,
artifact_dir: Path,
train_cases: int = 4,
val_cases: int = 1,
test_cases: int = 1,
steps: int = 1,
normalization_cases: int = 1,
) -> None:
path.write_text(
f"""
[run]
name = "hf_streaming_test"
seed = 7
artifact_dir = "{artifact_dir}"
[data]
root = "{cache_dir / 'processed' / 'full'}"
source = "huggingface_streaming"
hf_repo_id = "owner/airfrans-processed"
hf_repo_type = "dataset"
hf_path_prefix = "processed/full"
cache_dir = "{cache_dir}"
train_cases = {train_cases}
val_cases = {val_cases}
test_cases = {test_cases}
all_points_per_case = true
batch_size = 2
streaming_queue_max_cases = 2
streaming_normalization_cases = {normalization_cases}
[model]
type = "mlp"
hidden_width = 16
depth = 2
activation = "gelu"
[optim]
lr = 0.01
weight_decay = 0.0
steps = {steps}
log_interval = 1
[device]
type = "cpu"
allow_cpu_fallback = false
benchmark_kernels = false
[loss]
type = "normalized_mse"
[checkpoint]
interval_seconds = 0
policy = "full"
include_optimizer_state = true
include_rng_state = true
""".strip()
+ "\n"
)
def write_processed_case(root: Path, relative_prefix: str, case_id: str, offset: float) -> None:
target = root / relative_prefix / f"{case_id}.npz"
target.parent.mkdir(parents=True, exist_ok=True)
features = np.asarray(
[
[offset + 0.0, 0.0, 1.0, 0.1],
[offset + 1.0, 1.0, 0.5, 0.2],
[offset + 2.0, 0.5, 0.25, 0.3],
[offset + 3.0, 0.25, 0.125, 0.4],
],
dtype=np.float32,
)
targets = np.asarray(
[
[offset + 0.0, 0.1, 0.2],
[offset + 0.2, 0.3, 0.4],
[offset + 0.4, 0.5, 0.6],
[offset + 0.6, 0.7, 0.8],
],
dtype=np.float32,
)
np.savez(
target,
features=features,
targets=targets,
feature_names=np.asarray(["x", "y", "sdf", "alpha"], dtype="U"),
target_names=np.asarray(["u", "v", "p"], dtype="U"),
)
def fake_huggingface_module(source_root: Path, calls: list[str]) -> types.ModuleType:
module = types.ModuleType("huggingface_hub")
class FakeHfApi:
def __init__(self, token: str | None = None) -> None:
self.token = token
def list_repo_files(self, *, repo_id: str, repo_type: str) -> list[str]:
return sorted(str(path.relative_to(source_root)) for path in source_root.rglob("*") if path.is_file())
def hf_hub_download(*, repo_id: str, filename: str, repo_type: str, local_dir: str, token: str | None = None) -> str:
calls.append(filename)
source = source_root / filename
destination = Path(local_dir) / filename
time.sleep(0.01)
destination.parent.mkdir(parents=True, exist_ok=True)
shutil.copy2(source, destination)
return str(destination)
module.HfApi = FakeHfApi
module.hf_hub_download = hf_hub_download
return module
def read_events(run_dir: Path) -> list[dict[str, object]]:
return [json.loads(line) for line in (run_dir / "streaming_events.jsonl").read_text().splitlines() if line.strip()]
@ -215,6 +333,51 @@ class FullDataBackpressureStreamingTests(unittest.TestCase):
normalization_end = next(event for event in events if event["event"] == "normalization_end")
self.assertEqual(normalization_end["normalization_cases"], 1)
def test_huggingface_streaming_fast_start_downloads_remaining_selected_cases(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
tmp_path = Path(tmp)
hf_source = tmp_path / "hf_source"
case_names = [f"case_{index:04d}" for index in range(6)]
for index, case_name in enumerate(case_names):
write_processed_case(hf_source, "processed/full", case_name, float(index))
config_path = tmp_path / "hf_streaming.toml"
cache_dir = tmp_path / "cache"
artifact_dir = tmp_path / "artifacts"
write_huggingface_streaming_config(
config_path,
cache_dir=cache_dir,
artifact_dir=artifact_dir,
train_cases=4,
val_cases=1,
test_cases=1,
steps=1,
normalization_cases=1,
)
calls: list[str] = []
fake_module = fake_huggingface_module(hf_source, calls)
with patch.dict(sys.modules, {"huggingface_hub": fake_module}):
result = train(load_training_config(config_path))
self.assertEqual(result.final_metrics["data_mode"], "huggingface_streaming")
events = read_events(result.run_dir)
event_names = {event["event"] for event in events}
self.assertIn("background_acquisition_start", event_names)
self.assertIn("hf_case_download_end", event_names)
self.assertIn("background_acquisition_complete", event_names)
first_gpu = next(index for index, event in enumerate(events) if event["event"] == "first_gpu_batch_consumed")
processed_before_gpu = {
str(event["case_id"])
for event in events[:first_gpu]
if event["event"] == "processing_end"
}
self.assertLess(len(processed_before_gpu), 4)
manifest = json.loads((result.run_dir / "data_manifest.json").read_text())
self.assertEqual(manifest["source"], "huggingface_streaming")
cached_cases = sorted(path.stem for path in (cache_dir / "processed" / "full").glob("*.npz"))
self.assertEqual(cached_cases, case_names)
self.assertEqual(sorted(set(calls)), [f"processed/full/{case_name}.npz" for case_name in case_names])
def test_backpressure_pauses_resumes_and_bounds_cache_with_inflight_slack(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
tmp_path = Path(tmp)

36
tests/test_sweep.py
View file

@ -13,6 +13,7 @@ from airfrans_frontier.sweep import (
collect_job_status,
complete_job_attempt,
generate_jobs,
generate_minimal_scaling_jobs,
read_jobs,
rescue_templates,
run_node,
@ -51,6 +52,41 @@ class SweepFoundationTests(unittest.TestCase):
self.assertEqual(config.model.hidden_width, 2048)
self.assertEqual(config.model.depth, 16)
def test_generate_minimal_scaling_jobs_are_all_points_and_non_raw_for_main_coordinate_models(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
root = Path(tmp)
jobs = generate_minimal_scaling_jobs(output_dir=root, include_proxies=False)
self.assertEqual(len(jobs), 23)
self.assertEqual([job["job_id"] for job in jobs[:3]], ["A_data_all_points_smoke", "A_checkpoint_700m_footprint_smoke", "A_plateau_interrupt_smoke"])
serious_raw = [
job["job_id"]
for job in jobs
if job["phase"] in {"B_family_viability", "C_scaling_ladder"}
and job["coordinate_encoding"] == "raw"
and job["model_family"] != "siren_conditioned_inr"
]
self.assertEqual(serious_raw, [])
self.assertTrue(all(job["all_points_per_case"] for job in jobs))
c10 = load_training_config(root / "configs" / "C_10m_film_fourier_inr_nerf_multires_scale.toml")
c100 = load_training_config(root / "configs" / "C_100m_film_fourier_inr_nerf_multires_scale.toml")
config = load_training_config(root / "configs" / "C_700m_film_fourier_inr_nerf_multires_scale.toml")
self.assertTrue(config.data.all_points_per_case)
self.assertEqual(config.data.source, "huggingface_streaming")
self.assertEqual(config.data.hf_repo_id, "zacheryasc/airfrans-processed")
self.assertEqual(config.data.hf_path_prefix, "processed/full")
self.assertEqual(config.data.train_cases, 900)
self.assertEqual(config.data.val_cases, 50)
self.assertEqual(config.data.test_cases, 50)
self.assertEqual(config.data.streaming_normalization_cases, 4)
self.assertEqual(config.data.streaming_queue_max_cases, 8)
self.assertEqual(c10.run.seed, c100.run.seed)
self.assertEqual(c100.run.seed, config.run.seed)
self.assertEqual(config.checkpoint.policy, "full")
self.assertTrue(config.checkpoint.include_optimizer_state)
self.assertEqual(config.stability.dead_curve_patience_evals, 8)
def test_generate_jobs_only_crosses_encoding_axis_for_compatible_families(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
root = Path(tmp)

123
tests/test_training_config.py
View file

@ -87,6 +87,129 @@ type = "normalized_mse"
self.assertTrue(parsed.model_metadata.is_proxy)
self.assertEqual(parsed.model_metadata.coordinate_encoding_compatibility, "raw_only")
def test_config_loader_accepts_all_points_dead_curve_and_lightweight_checkpoint(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
config_path = Path(tmp) / "all_points.toml"
config_path.write_text(
f"""
[run]
name = "all_points"
seed = 0
artifact_dir = "{Path(tmp) / "runs"}"
[data]
root = "{Path(tmp) / "data"}"
train_cases = 1
val_cases = 0
test_cases = 0
all_points_per_case = true
batch_size = 1
[coordinate_encoding]
type = "nerf_multires"
features = ["x", "y", "sdf"]
levels = 4
[model]
type = "mlp_encoded_baseline"
hidden_width = 8
depth = 1
activation = "gelu"
coordinate_features = ["x", "y", "sdf"]
[optim]
lr = 0.001
weight_decay = 0.0
steps = 1
[device]
type = "cpu"
allow_cpu_fallback = false
benchmark_kernels = false
[loss]
type = "normalized_mse"
[checkpoint]
policy = "lightweight_scaling_probe"
include_optimizer_state = false
include_rng_state = false
interval_seconds = 0
[stability]
dead_curve_patience_evals = 3
dead_curve_min_relative_improvement = 0.01
dead_curve_warmup_steps = 5
""".strip()
+ "\n"
)
parsed = load_training_config(config_path)
self.assertTrue(parsed.data.all_points_per_case)
self.assertIsNone(parsed.data.points_per_case)
self.assertEqual(parsed.model.type, "mlp_encoded_baseline")
self.assertEqual(parsed.checkpoint.policy, "lightweight_scaling_probe")
self.assertFalse(parsed.checkpoint.include_optimizer_state)
self.assertEqual(parsed.stability.dead_curve_patience_evals, 3)
def test_config_loader_accepts_huggingface_streaming_data_source(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
config_path = Path(tmp) / "hf_streaming.toml"
config_path.write_text(
f"""
[run]
name = "hf_streaming"
seed = 0
artifact_dir = "{Path(tmp) / "runs"}"
[data]
root = "{Path(tmp) / "cache" / "processed" / "full"}"
source = "huggingface_streaming"
hf_repo_id = "zacheryasc/airfrans-processed"
hf_repo_type = "dataset"
hf_path_prefix = "processed/full"
cache_dir = "{Path(tmp) / "cache"}"
train_cases = 1
val_cases = 0
test_cases = 0
all_points_per_case = true
batch_size = 1
streaming_queue_max_cases = 8
streaming_normalization_cases = 1
[model]
type = "mlp"
hidden_width = 8
depth = 1
activation = "gelu"
[optim]
lr = 0.001
weight_decay = 0.0
steps = 1
[device]
type = "cpu"
allow_cpu_fallback = false
benchmark_kernels = false
[loss]
type = "normalized_mse"
[checkpoint]
interval_seconds = 0
""".strip()
+ "\n"
)
parsed = load_training_config(config_path)
self.assertEqual(parsed.data.source, "huggingface_streaming")
self.assertEqual(parsed.data.hf_repo_id, "zacheryasc/airfrans-processed")
self.assertTrue(parsed.data.all_points_per_case)
def test_config_loader_rejects_missing_section(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
config_path = Path(tmp) / "bad.toml"

21
tests/test_training_data.py
View file

@ -10,7 +10,7 @@ remove_pythonpath_entries()
import numpy as np
from airfrans_frontier.training.data import create_case_split, load_processed_dataset, load_simulation_npz
from airfrans_frontier.training.data import build_dataset_bundle, create_case_split, load_processed_dataset, load_simulation_npz
from airfrans_frontier.training.normalize import compute_normalization_stats, normalize_targets
@ -74,6 +74,25 @@ class TrainingDataTests(unittest.TestCase):
self.assertEqual(len(first.val_ids), 2)
self.assertEqual(len(first.test_ids), 2)
def test_all_points_mode_keeps_every_row_per_selected_case(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
root = Path(tmp)
write_case(root / "case_a.npz", offset=0.0)
write_case(root / "case_b.npz", offset=10.0)
bundle = build_dataset_bundle(
load_processed_dataset(root),
train_cases=1,
val_cases=1,
test_cases=0,
points_per_case=None,
seed=0,
)
self.assertEqual(bundle.train.features.shape[0], 3)
self.assertEqual(bundle.val.features.shape[0], 3)
self.assertEqual(bundle.train.point_counts, (3,))
def test_normalization_uses_train_split_only(self) -> None:
train_features = np.array([[0.0], [2.0]], dtype=np.float32)
train_targets = np.array([[10.0], [14.0]], dtype=np.float32)

67
tests/test_training_loop.py
View file

@ -67,6 +67,9 @@ def write_training_config(
log_interval: int = 250,
checkpoint_interval_seconds: int = 1800,
hidden_width: int = 128,
checkpoint_include_optimizer_state: bool = True,
checkpoint_include_rng_state: bool = True,
dead_curve_patience_evals: int | None = None,
) -> None:
path.write_text(
f"""
@ -105,6 +108,14 @@ type = "normalized_mse"
[checkpoint]
interval_seconds = {checkpoint_interval_seconds}
include_optimizer_state = {str(checkpoint_include_optimizer_state).lower()}
include_rng_state = {str(checkpoint_include_rng_state).lower()}
{f'''
[stability]
dead_curve_patience_evals = {dead_curve_patience_evals}
dead_curve_min_relative_improvement = 0.01
dead_curve_warmup_steps = 0
''' if dead_curve_patience_evals is not None else ''}
""".strip()
+ "\n"
)
@ -421,6 +432,62 @@ class TrainingLoopTests(unittest.TestCase):
self.assertTrue((run_dir / "artifact_manifest.json").is_file())
self.assertTrue((run_dir / "checksums.txt").is_file())
def test_dead_curve_interruption_writes_plateau_failure(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
tmp_path = Path(tmp)
data_root = tmp_path / "data"
artifact_dir = tmp_path / "artifacts"
config_path = tmp_path / "config.toml"
write_toy_simulator_dataset(data_root, cases=4, points=64)
write_training_config(
config_path,
data_root=data_root,
artifact_dir=artifact_dir,
device_type="cpu",
steps=6,
log_interval=1,
checkpoint_interval_seconds=0,
hidden_width=32,
dead_curve_patience_evals=2,
)
config = load_training_config(config_path)
with patch("airfrans_frontier.training.loop.torch.optim.AdamW.step", lambda self: None):
with self.assertRaisesRegex(RuntimeError, "loss did not improve"):
train(config)
run_dir = next(path for path in artifact_dir.iterdir() if path.is_dir())
report = json.loads((run_dir / "failure_report.json").read_text())
self.assertEqual(report["failure_category"], "plateau_or_dead_recipe")
self.assertEqual(report["error_type"], "DeadCurveError")
self.assertIn("dead_curve", report)
def test_lightweight_checkpoint_omits_optimizer_payload(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
tmp_path = Path(tmp)
data_root = tmp_path / "data"
artifact_dir = tmp_path / "artifacts"
config_path = tmp_path / "config.toml"
write_toy_simulator_dataset(data_root, cases=4, points=64)
write_training_config(
config_path,
data_root=data_root,
artifact_dir=artifact_dir,
device_type="cpu",
steps=2,
log_interval=1,
checkpoint_interval_seconds=0,
hidden_width=32,
checkpoint_include_optimizer_state=False,
checkpoint_include_rng_state=False,
)
result = train(load_training_config(config_path))
checkpoint = torch.load(result.run_dir / "checkpoint_final.pt", map_location="cpu", weights_only=False)
self.assertIsNone(checkpoint["optimizer_state_dict"])
self.assertIsNone(checkpoint["rng_state"])
self.assertFalse(result.final_metrics["checkpoint_include_optimizer_state"])
def test_nonfinite_gradient_writes_failure_report(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
tmp_path = Path(tmp)