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

stash: post failed large run

Browse source
This commit is contained in:
Zachery Aaron Shores-Chmielewski 2026-07-28 22:04:50 +04:00
parent 028f811f78
commit 2071670a4d
31 changed files with 2919 additions and 326 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
.skyignore
View file

@ -1,12 +1,27 @@
/artifacts
/artifacts/remote_runs
/artifacts/runs
/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
/data/airfrans_processed_full_50cases.tar.gz
/.venv
/.git
/.airfrans_hf_resume
/sky_logs
/notebooks
__pycache__
*.pyc
HF_TOKEN
WANDB_API_KEY
.env
sky_logs/

103
configs/remote_aggressive_oom_node.toml Normal file
View file

@ -0,0 +1,103 @@
[run]
name = "aggressive_oom_sweep_node"
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]
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/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 -c "from airfrans_frontier.sweep import read_jobs; jobs=read_jobs('artifacts/aggressive_oom_sweep/jobs.jsonl'); pending=sum(1 for j in jobs if j['status'] == 'pending'); assert pending > 0; print('pending_jobs=' + str(pending))"
"""
[job]
command = """
uv run --no-dev python scripts/aggressive_oom_node_wrapper.py --jobs artifacts/aggressive_oom_sweep/jobs.jsonl --node "$AIRFRANS_REMOTE_RUN_ID" --artifact-dir artifacts/current_run --utilization artifacts/aggressive_oom_sweep/utilization.jsonl --sample-interval-seconds 30 --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"

103
configs/remote_aggressive_oom_node1_700m.toml Normal file
View file

@ -0,0 +1,103 @@
[run]
name = "aggressive_oom_sweep_node1_700m"
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, 135723]
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/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 -c "from airfrans_frontier.sweep import read_jobs; jobs=read_jobs('artifacts/aggressive_oom_sweep/jobs_node1_700m.jsonl'); pending=sum(1 for j in jobs if j['status'] == 'pending'); assert pending > 0; print('pending_jobs=' + str(pending))"
"""
[job]
command = """
uv run --no-dev python scripts/aggressive_oom_node_wrapper.py --jobs artifacts/aggressive_oom_sweep/jobs_node1_700m.jsonl --node node-1 --artifact-dir artifacts/current_run --utilization artifacts/aggressive_oom_sweep/utilization_node1.jsonl --sample-interval-seconds 30 --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"

103
configs/remote_aggressive_oom_node2_100m.toml Normal file
View file

@ -0,0 +1,103 @@
[run]
name = "aggressive_oom_sweep_node2_100m"
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, 135723, 213498, 4535, 566749, 307216]
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/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 -c "from airfrans_frontier.sweep import read_jobs; jobs=read_jobs('artifacts/aggressive_oom_sweep/jobs_node2_100m.jsonl'); pending=sum(1 for j in jobs if j['status'] == 'pending'); assert pending > 0; print('pending_jobs=' + str(pending))"
"""
[job]
command = """
uv run --no-dev python scripts/aggressive_oom_node_wrapper.py --jobs artifacts/aggressive_oom_sweep/jobs_node2_100m.jsonl --node node-2 --artifact-dir artifacts/current_run --utilization artifacts/aggressive_oom_sweep/utilization_node2.jsonl --sample-interval-seconds 30 --stale-after-seconds 21600 --max-attempts 2
"""
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"

103
configs/remote_aggressive_oom_node3_700m.toml Normal file
View file

@ -0,0 +1,103 @@
[run]
name = "aggressive_oom_sweep_node3_700m"
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, 135723, 213498, 4535, 566749, 307216]
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/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 -c "from airfrans_frontier.sweep import read_jobs; jobs=read_jobs('artifacts/aggressive_oom_sweep/jobs_node3_700m.jsonl'); pending=sum(1 for j in jobs if j['status'] == 'pending'); assert pending > 0; print('pending_jobs=' + str(pending))"
"""
[job]
command = """
uv run --no-dev python scripts/aggressive_oom_node_wrapper.py --jobs artifacts/aggressive_oom_sweep/jobs_node3_700m.jsonl --node node-3 --artifact-dir artifacts/current_run --utilization artifacts/aggressive_oom_sweep/utilization_node3.jsonl --sample-interval-seconds 30 --stale-after-seconds 21600 --max-attempts 2
"""
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"

104
configs/remote_aggressive_oom_probe.toml Normal file
View file

@ -0,0 +1,104 @@
[run]
name = "aggressive_oom_sweep_probe_node"
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]
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/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 -c "from airfrans_frontier.sweep import read_jobs; jobs=read_jobs('artifacts/aggressive_oom_sweep/jobs.jsonl'); assert len(jobs) == 34; assert any(j['status'] == 'pending' for j in jobs); print('pending_jobs=' + str(sum(1 for j in jobs if j['status'] == 'pending')))"
uv run --no-dev python -c "from airfrans_frontier.training.config import load_training_config; c=load_training_config('artifacts/aggressive_oom_sweep/configs/100m_mlp_raw.toml'); assert c.data.source == 'huggingface'; assert c.huggingface.enabled; print('probe_config=' + c.run.name + ' source=' + c.data.source)"
"""
[job]
command = """
uv run --no-dev python scripts/aggressive_oom_node_wrapper.py --jobs artifacts/aggressive_oom_sweep/jobs.jsonl --node "$AIRFRANS_REMOTE_RUN_ID" --artifact-dir artifacts/current_run --utilization artifacts/aggressive_oom_sweep/utilization.jsonl --max-jobs 1 --sample-interval-seconds 30 --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"

2
configs/remote_hf_smoke.toml
View file

@ -58,6 +58,7 @@ metrics_file = "artifacts/current_run/metrics.jsonl"
mode = "object_store_upload"
required = [
"config.toml",
"job_manifest.json",
"metrics.jsonl",
"latest_metrics.json",
"heartbeat.json",
@ -70,6 +71,7 @@ required = [
"hf_upload_manifest.json",
"artifact_manifest.json",
"checksums.txt",
"utilization.jsonl",
]
[cleanup]

2
configs/remote_smoke.toml
View file

@ -60,6 +60,7 @@ metrics_file = "artifacts/current_run/metrics.jsonl"
mode = "object_store_upload"
required = [
"config.toml",
"job_manifest.json",
"metrics.jsonl",
"latest_metrics.json",
"heartbeat.json",
@ -78,6 +79,7 @@ required = [
"artifact_manifest.json",
"checksums.txt",
"verification_report.json",
"utilization.jsonl",
]
[cleanup]

2
configs/remote_wandb_smoke.toml
View file

@ -59,6 +59,7 @@ metrics_file = "artifacts/current_run/metrics.jsonl"
mode = "rsync"
required = [
"config.toml",
"job_manifest.json",
"metrics.jsonl",
"latest_metrics.json",
"heartbeat.json",
@ -71,6 +72,7 @@ required = [
"wandb_smoke_manifest.json",
"artifact_manifest.json",
"checksums.txt",
"utilization.jsonl",
]
[cleanup]

337
scripts/aggressive_oom_node_wrapper.py Executable file
View file

@ -0,0 +1,337 @@
#!/usr/bin/env python3
from __future__ import annotations
import argparse
import json
import shutil
import threading
import time
import traceback
from pathlib import Path
from typing import Any
from airfrans_frontier.remote.artifacts import verify_artifacts
from airfrans_frontier.sweep import collect_job_status, run_node, sample_gpu_utilization
_BASE_ARTIFACTS = (
"config.toml",
"job_manifest.json",
"metrics.jsonl",
"latest_metrics.json",
"heartbeat.json",
"run_manifest.json",
"environment_manifest.json",
"utilization.jsonl",
)
_TERMINAL_ARTIFACTS = (
"final_metrics.json",
"failure_report.json",
"checkpoint_latest.pt",
"checkpoint_best.pt",
"checkpoint_final.pt",
"split_manifest.json",
"data_manifest.json",
"normalization.json",
"calibration_manifest.json",
"evaluation_protocol.json",
"hf_upload_manifest.json",
"verification_report.json",
"artifact_manifest.json",
"checksums.txt",
)
_PRESERVE_IN_CURRENT_RUN = {"startup_timeline.jsonl", "nvidia_smi.txt", "disk_telemetry.json"}
def main() -> int:
parser = argparse.ArgumentParser(description="Run one aggressive OOM sweep node and flatten the latest terminal job artifact for remote-run collection.")
parser.add_argument("--jobs", default="artifacts/aggressive_oom_sweep/jobs.jsonl")
parser.add_argument("--node", required=True)
parser.add_argument("--artifact-dir", default="artifacts/current_run")
parser.add_argument("--utilization", default="artifacts/aggressive_oom_sweep/utilization.jsonl")
parser.add_argument("--max-jobs", type=int)
parser.add_argument("--sample-interval-seconds", type=float, default=30.0)
parser.add_argument("--stale-after-seconds", type=float, default=21600.0)
parser.add_argument("--max-attempts", type=int, default=2)
parser.add_argument("--require-success", action="store_true")
args = parser.parse_args()
jobs_path = Path(args.jobs)
artifact_dir = Path(args.artifact_dir)
utilization_path = Path(args.utilization)
artifact_dir.mkdir(parents=True, exist_ok=True)
utilization_path.parent.mkdir(parents=True, exist_ok=True)
started_at = time.time()
stop_sampling = threading.Event()
sampler = threading.Thread(
target=_sample_until_stopped,
args=(stop_sampling, utilization_path, artifact_dir, args.sample_interval_seconds, started_at, args.node),
daemon=True,
)
_write_node_heartbeat(artifact_dir, node=args.node, phase="starting", started_at=started_at)
_append_jsonl(artifact_dir / "metrics.jsonl", _node_metric(args.node, phase="starting", started_at=started_at))
sampler.start()
exit_code = 0
summary: dict[str, Any] | None = None
status: dict[str, Any] | None = None
try:
summary = run_node(
jobs_path=jobs_path,
node_name=args.node,
max_jobs=args.max_jobs,
stale_after_seconds=args.stale_after_seconds,
max_attempts=args.max_attempts,
)
status = collect_job_status(jobs_path=jobs_path)
if args.require_success and int(summary.get("succeeded", 0)) <= 0:
exit_code = 1
except Exception as exc: # noqa: BLE001 - terminal failure artifact is the contract here.
exit_code = 1
status = _safe_collect(jobs_path)
summary = {
"node": args.node,
"claimed": 0,
"succeeded": 0,
"failed": 1,
"error_type": type(exc).__name__,
"error_message": str(exc),
"traceback_tail": traceback.format_exc()[-4000:],
}
finally:
stop_sampling.set()
sampler.join(timeout=max(1.0, min(10.0, args.sample_interval_seconds)))
finished_at = time.time()
status = status or _safe_collect(jobs_path)
summary = summary or {"node": args.node, "claimed": 0, "succeeded": 0, "failed": 0}
summary = dict(summary)
summary.update(
{
"started_at": started_at,
"finished_at": finished_at,
"duration_seconds": finished_at - started_at,
"jobs_path": str(jobs_path),
"utilization_path": str(utilization_path),
"status_counts": (status or {}).get("counts", {}),
}
)
selected = _select_terminal_record(status or {})
if selected is not None and selected.get("run_dir"):
_flatten_job_artifacts(Path(str(selected["run_dir"])), artifact_dir)
summary["flattened_job_id"] = selected.get("job_id")
summary["flattened_run_dir"] = selected.get("run_dir")
summary["flattened_status"] = selected.get("status")
else:
exit_code = 1
_write_minimal_failure_artifacts(
artifact_dir,
node=args.node,
started_at=started_at,
finished_at=finished_at,
summary=summary,
message="no terminal job artifacts were available to flatten",
)
_write_sweep_sidecars(
artifact_dir=artifact_dir,
jobs_path=jobs_path,
utilization_path=utilization_path,
summary=summary,
status=status or {},
node=args.node,
started_at=started_at,
finished_at=finished_at,
)
try:
verify_artifacts(artifact_dir)
except Exception as exc: # noqa: BLE001
exit_code = 1
if not (artifact_dir / "failure_report.json").is_file():
for terminal_name in ("final_metrics.json", "checkpoint_latest.pt", "checkpoint_best.pt", "checkpoint_final.pt"):
path = artifact_dir / terminal_name
if path.exists():
path.unlink()
_write_minimal_failure_artifacts(
artifact_dir,
node=args.node,
started_at=started_at,
finished_at=time.time(),
summary=summary,
message=f"flattened artifact verification failed: {exc}",
)
verify_artifacts(artifact_dir)
print(json.dumps({"exit_code": exit_code, "summary": summary, "status_counts": (status or {}).get("counts", {})}, indent=2, sort_keys=True))
return exit_code
def _sample_until_stopped(stop: threading.Event, utilization_path: Path, artifact_dir: Path, interval: float, started_at: float, node: str) -> None:
interval = max(1.0, interval)
while not stop.is_set():
sample = sample_gpu_utilization()
sample.update({"node": node, "elapsed_seconds": time.time() - started_at})
_append_jsonl(utilization_path, sample)
_append_jsonl(artifact_dir / "node_utilization.jsonl", sample)
_write_node_heartbeat(artifact_dir, node=node, phase="running", started_at=started_at, latest_utilization=sample)
_append_jsonl(artifact_dir / "metrics.jsonl", _node_metric(node, phase="running", started_at=started_at, latest_utilization=sample))
stop.wait(interval)
def _node_metric(node: str, *, phase: str, started_at: float, latest_utilization: dict[str, Any] | None = None) -> dict[str, Any]:
payload: dict[str, Any] = {
"timestamp": time.time(),
"node": node,
"phase": phase,
"elapsed_seconds": time.time() - started_at,
}
if latest_utilization:
payload.update(
{
"gpu_util_percent": latest_utilization.get("gpu_util_percent"),
"memory_used_mb": latest_utilization.get("memory_used_mb"),
}
)
return payload
def _write_node_heartbeat(artifact_dir: Path, *, node: str, phase: str, started_at: float, latest_utilization: dict[str, Any] | None = None) -> None:
_write_json(
artifact_dir / "heartbeat.json",
{
"node": node,
"phase": phase,
"timestamp": time.time(),
"started_at": started_at,
"elapsed_seconds": time.time() - started_at,
"latest_utilization": latest_utilization,
},
)
def _safe_collect(jobs_path: Path) -> dict[str, Any]:
try:
return collect_job_status(jobs_path=jobs_path)
except Exception as exc: # noqa: BLE001
return {"jobs": [], "counts": {"collect_error": 1}, "total": 0, "error": str(exc)}
def _select_terminal_record(status: dict[str, Any]) -> dict[str, Any] | None:
jobs = [dict(item) for item in status.get("jobs", []) if isinstance(item, dict)]
succeeded = [item for item in jobs if item.get("status") == "succeeded" and item.get("run_dir")]
if succeeded:
return max(succeeded, key=lambda item: _mtime(Path(str(item["run_dir"]))))
failed = [item for item in jobs if item.get("status") == "failed" and item.get("run_dir")]
if failed:
return max(failed, key=lambda item: _mtime(Path(str(item["run_dir"]))))
incomplete = [item for item in jobs if item.get("run_dir")]
if incomplete:
return max(incomplete, key=lambda item: _mtime(Path(str(item["run_dir"]))))
return None
def _mtime(path: Path) -> float:
try:
return path.stat().st_mtime
except FileNotFoundError:
return 0.0
def _flatten_job_artifacts(run_dir: Path, artifact_dir: Path) -> None:
if not run_dir.is_dir():
raise FileNotFoundError(f"terminal run directory not found: {run_dir}")
for path in artifact_dir.iterdir():
if path.is_file() and path.name not in _PRESERVE_IN_CURRENT_RUN:
path.unlink()
for name in (*_BASE_ARTIFACTS, *_TERMINAL_ARTIFACTS):
source = run_dir / name
if source.is_file():
shutil.copy2(source, artifact_dir / name)
def _write_minimal_failure_artifacts(artifact_dir: Path, *, node: str, started_at: float, finished_at: float, summary: dict[str, Any], message: str) -> None:
for name in ("final_metrics.json", "checkpoint_latest.pt", "checkpoint_best.pt", "checkpoint_final.pt"):
path = artifact_dir / name
if path.exists():
path.unlink()
base = {
"run_id": node,
"run_name": node,
"node": node,
"started_at": started_at,
"finished_at": finished_at,
"artifact_dir": str(artifact_dir),
"summary": summary,
}
_write_json(artifact_dir / "config.toml", {"note": "placeholder"}) if False else None
if not (artifact_dir / "config.toml").is_file():
(artifact_dir / "config.toml").write_text("[run]\nname = \"sweep_node_failure\"\n")
_write_json(artifact_dir / "job_manifest.json", {**base, "command": "aggressive_oom_node_wrapper"})
_write_json(artifact_dir / "run_manifest.json", base)
_write_json(artifact_dir / "environment_manifest.json", {"node": node, "recorded_at": time.time()})
if not (artifact_dir / "metrics.jsonl").is_file():
_append_jsonl(artifact_dir / "metrics.jsonl", {"timestamp": finished_at, "phase": "failed", "node": node})
_write_json(artifact_dir / "latest_metrics.json", {"timestamp": finished_at, "phase": "failed", "node": node})
_write_node_heartbeat(artifact_dir, node=node, phase="failed", started_at=started_at)
if not (artifact_dir / "utilization.jsonl").is_file():
_append_jsonl(artifact_dir / "utilization.jsonl", {"timestamp": finished_at, "gpu_util_percent": None, "memory_used_mb": None})
_write_json(
artifact_dir / "failure_report.json",
{
"run_id": node,
"phase": "sweep_node",
"failure_phase": "sweep_node",
"failure_category": "sweep_node_artifact_collection",
"error_type": "RuntimeError",
"error_message": message,
"summary": summary,
"timestamp": finished_at,
},
)
def _write_sweep_sidecars(*, artifact_dir: Path, jobs_path: Path, utilization_path: Path, summary: dict[str, Any], status: dict[str, Any], node: str, started_at: float, finished_at: float) -> None:
_write_json(artifact_dir / "node_summary.json", summary)
_write_json(artifact_dir / "sweep_collect.json", status)
if jobs_path.is_file():
shutil.copy2(jobs_path, artifact_dir / "sweep_jobs.jsonl")
else:
(artifact_dir / "sweep_jobs.jsonl").write_text("")
attempts_path = jobs_path.parent / "attempts.jsonl"
if attempts_path.is_file():
shutil.copy2(attempts_path, artifact_dir / "sweep_attempts.jsonl")
else:
(artifact_dir / "sweep_attempts.jsonl").write_text("")
if utilization_path.is_file():
shutil.copy2(utilization_path, artifact_dir / "sweep_utilization.jsonl")
shutil.copy2(utilization_path, artifact_dir / "utilization.jsonl")
elif not (artifact_dir / "utilization.jsonl").is_file():
_append_jsonl(artifact_dir / "utilization.jsonl", {"timestamp": finished_at, "gpu_util_percent": None, "memory_used_mb": None})
_append_jsonl(
artifact_dir / "node_metrics.jsonl",
{
"timestamp": finished_at,
"node": node,
"phase": "finished",
"elapsed_seconds": finished_at - started_at,
"summary": summary,
"status_counts": status.get("counts", {}),
},
)
def _write_json(path: Path, payload: Any) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
def _append_jsonl(path: Path, payload: Any) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("a", encoding="utf-8") as handle:
handle.write(json.dumps(payload, sort_keys=True) + "\n")
if __name__ == "__main__":
raise SystemExit(main())

90
scripts/rsync_sweep_node_artifacts.py Executable file
View file

@ -0,0 +1,90 @@
#!/usr/bin/env python3
from __future__ import annotations
import argparse
import signal
import subprocess
import time
from pathlib import Path
_stop = False
def _handle_stop(signum: int, frame: object) -> None:
global _stop
_stop = True
def _run(command: list[str]) -> int:
completed = subprocess.run(command, check=False)
return int(completed.returncode)
def _rsync(host: str, port: int, remote: str, local: Path, *, include_metrics_only: bool = False) -> int:
local.mkdir(parents=True, exist_ok=True)
command = [
'rsync',
'-az',
'--delete',
'--exclude',
'*.pt',
'--exclude',
'.wandb/',
]
if include_metrics_only:
command.extend([
'--include',
'*/',
'--include',
'*.json',
'--include',
'*.jsonl',
'--include',
'*.toml',
'--include',
'*.txt',
'--exclude',
'*',
])
command.extend([
'-e',
f'ssh -p {port} -o StrictHostKeyChecking=no',
f'root@{host}:{remote}',
str(local),
])
return _run(command)
def main() -> int:
parser = argparse.ArgumentParser(description='Continuously rsync lightweight sweep artifacts from one node.')
parser.add_argument('--host', required=True)
parser.add_argument('--port', type=int, required=True)
parser.add_argument('--remote-root', default='/root/sky_workdir')
parser.add_argument('--local-root', default='artifacts/remote_runs/live_node0')
parser.add_argument('--interval-seconds', type=float, default=60.0)
args = parser.parse_args()
signal.signal(signal.SIGINT, _handle_stop)
signal.signal(signal.SIGTERM, _handle_stop)
local_root = Path(args.local_root)
remote_root = args.remote_root.rstrip('/')
interval = max(5.0, args.interval_seconds)
cycle = 0
while not _stop:
started = time.time()
cycle += 1
rcodes = [
_rsync(args.host, args.port, f'{remote_root}/artifacts/aggressive_oom_sweep/', local_root / 'aggressive_oom_sweep/'),
_rsync(args.host, args.port, f'{remote_root}/artifacts/current_run/', local_root / 'current_run/'),
_rsync(args.host, args.port, f'{remote_root}/artifacts/runs/', local_root / 'runs/', include_metrics_only=True),
]
print(f'collector_cycle={cycle} rcodes={rcodes} elapsed_seconds={time.time() - started:.1f}', flush=True)
deadline = time.time() + interval
while not _stop and time.time() < deadline:
time.sleep(min(1.0, deadline - time.time()))
print('collector_stopped=true', flush=True)
return 0
if __name__ == '__main__':
raise SystemExit(main())

63
scripts/run_quick_probe_on_cluster.py Executable file
View file

@ -0,0 +1,63 @@
#!/usr/bin/env python3
from __future__ import annotations
import os
import subprocess
import sys
from pathlib import Path
def _secret(name: str) -> str:
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
raise RuntimeError(f'{name} is required')
def main(argv: list[str]) -> int:
if len(argv) != 2:
print('usage: run_quick_probe_on_cluster.py CLUSTER', file=sys.stderr)
return 2
cluster = argv[1]
env = os.environ.copy()
env.pop('PYTHONPATH', None)
env['HF_TOKEN'] = _secret('HF_TOKEN')
env['WANDB_API_KEY'] = _secret('WANDB_API_KEY')
remote_cmd = (
'uv run --no-dev python scripts/aggressive_oom_node_wrapper.py '
'--jobs artifacts/aggressive_oom_sweep_quick_probe/jobs.jsonl '
'--node quick_probe_node '
'--artifact-dir artifacts/current_run_quick_probe '
'--utilization artifacts/aggressive_oom_sweep_quick_probe/utilization.jsonl '
'--max-jobs 1 '
'--sample-interval-seconds 5 '
'--require-success'
)
command = [
'uv',
'run',
'sky',
'exec',
'--workdir',
'.',
'--gpus',
'RTX4090:1',
'--secret',
'HF_TOKEN',
'--secret',
'WANDB_API_KEY',
cluster,
remote_cmd,
]
print(f'quick_probe_cluster={cluster}', flush=True)
print('quick_probe_command=sky exec --workdir . --gpus RTX4090:1 --secret HF_TOKEN --secret WANDB_API_KEY <cluster> <wrapper>', flush=True)
return subprocess.run(command, env=env, check=False).returncode
if __name__ == '__main__':
raise SystemExit(main(sys.argv))

83
scripts/run_sweep_node_direct_ssh.py Executable file
View file

@ -0,0 +1,83 @@
#!/usr/bin/env python3
from __future__ import annotations
import argparse
import os
import shlex
import subprocess
from pathlib import Path
def _secret(name: str) -> str:
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
raise RuntimeError(f'{name} is required')
def main() -> int:
parser = argparse.ArgumentParser(description='Run an aggressive sweep node directly over SSH, outside Sky/Ray job workers.')
parser.add_argument('--host', required=True)
parser.add_argument('--port', required=True)
parser.add_argument('--node', required=True)
parser.add_argument('--jobs', required=True)
parser.add_argument('--artifact-dir', default='artifacts/current_run')
parser.add_argument('--utilization', required=True)
parser.add_argument('--remote-root', default='/root/sky_workdir')
parser.add_argument('--max-jobs', type=int)
parser.add_argument('--sample-interval-seconds', type=float, default=30.0)
parser.add_argument('--stale-after-seconds', type=float, default=21600.0)
parser.add_argument('--max-attempts', type=int, default=2)
args = parser.parse_args()
hf_token = _secret('HF_TOKEN')
wandb_api_key = _secret('WANDB_API_KEY')
parts = [
'uv run --no-dev python scripts/aggressive_oom_node_wrapper.py',
f'--jobs {shlex.quote(args.jobs)}',
f'--node {shlex.quote(args.node)}',
f'--artifact-dir {shlex.quote(args.artifact_dir)}',
f'--utilization {shlex.quote(args.utilization)}',
f'--sample-interval-seconds {args.sample_interval_seconds:g}',
f'--stale-after-seconds {args.stale_after_seconds:g}',
f'--max-attempts {args.max_attempts}',
]
if args.max_jobs is not None:
parts.append(f'--max-jobs {args.max_jobs}')
remote_cmd = ' '.join(parts)
remote_script = '\n'.join(
[
'set -euo pipefail',
f'cd {shlex.quote(args.remote_root)}',
f'export HF_TOKEN={shlex.quote(hf_token)}',
f'export WANDB_API_KEY={shlex.quote(wandb_api_key)}',
'export PATH=\"$HOME/.local/bin:/root/.local/bin:$PATH\"',
'export PYTHONUNBUFFERED=1',
remote_cmd,
'',
]
)
command = [
'ssh',
'-o',
'StrictHostKeyChecking=no',
'-p',
str(args.port),
f'root@{args.host}',
'bash',
'-s',
]
print(f'sweep_direct_host={args.host}', flush=True)
print(f'sweep_direct_port={args.port}', flush=True)
print(f'sweep_direct_node={args.node}', flush=True)
print('sweep_direct_command=<redacted secrets> uv run --no-dev python scripts/aggressive_oom_node_wrapper.py', flush=True)
return subprocess.run(command, input=remote_script, text=True, check=False).returncode
if __name__ == '__main__':
raise SystemExit(main())

75
scripts/run_sweep_node_on_cluster.py Executable file
View file

@ -0,0 +1,75 @@
#!/usr/bin/env python3
from __future__ import annotations
import argparse
import os
import subprocess
from pathlib import Path
def _secret(name: str) -> str:
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
raise RuntimeError(f'{name} is required')
def main() -> int:
parser = argparse.ArgumentParser(description='Run an aggressive sweep node on an existing Sky cluster.')
parser.add_argument('cluster')
parser.add_argument('--node', required=True)
parser.add_argument('--jobs', default='artifacts/aggressive_oom_sweep/jobs.jsonl')
parser.add_argument('--artifact-dir', default='artifacts/current_run')
parser.add_argument('--utilization', default='artifacts/aggressive_oom_sweep/utilization.jsonl')
parser.add_argument('--max-jobs', type=int)
parser.add_argument('--sample-interval-seconds', type=float, default=30.0)
parser.add_argument('--stale-after-seconds', type=float, default=21600.0)
parser.add_argument('--max-attempts', type=int, default=2)
args = parser.parse_args()
env = os.environ.copy()
env.pop('PYTHONPATH', None)
env['HF_TOKEN'] = _secret('HF_TOKEN')
env['WANDB_API_KEY'] = _secret('WANDB_API_KEY')
parts = [
'uv run --no-dev python scripts/aggressive_oom_node_wrapper.py',
f'--jobs {args.jobs}',
f'--node {args.node}',
f'--artifact-dir {args.artifact_dir}',
f'--utilization {args.utilization}',
f'--sample-interval-seconds {args.sample_interval_seconds:g}',
f'--stale-after-seconds {args.stale_after_seconds:g}',
f'--max-attempts {args.max_attempts}',
]
if args.max_jobs is not None:
parts.append(f'--max-jobs {args.max_jobs}')
remote_cmd = ' '.join(parts)
command = [
'uv',
'run',
'sky',
'exec',
'--workdir',
'.',
'--gpus',
'RTX4090:1',
'--secret',
'HF_TOKEN',
'--secret',
'WANDB_API_KEY',
args.cluster,
remote_cmd,
]
print(f'sweep_cluster={args.cluster}', flush=True)
print(f'sweep_node={args.node}', flush=True)
print('sweep_command=sky exec --workdir . --gpus RTX4090:1 --secret HF_TOKEN --secret WANDB_API_KEY <cluster> <wrapper>', flush=True)
return subprocess.run(command, env=env, check=False).returncode
if __name__ == '__main__':
raise SystemExit(main())

34
src/airfrans_frontier/cli.py
View file

@ -79,6 +79,14 @@ def build_parser() -> argparse.ArgumentParser:
sweep_generate.add_argument("--encodings", nargs="*")
sweep_generate.set_defaults(command="sweep-generate")
sweep_run_node = subparsers.add_parser("sweep-run-node", help="drain sweep jobs sequentially for one named node")
sweep_run_node.add_argument("--jobs", default="artifacts/aggressive_oom_sweep/jobs.jsonl")
sweep_run_node.add_argument("--node", required=True)
sweep_run_node.add_argument("--max-jobs", type=int)
sweep_run_node.add_argument("--stale-after-seconds", type=float, default=21600.0)
sweep_run_node.add_argument("--max-attempts", type=int, default=2)
sweep_run_node.set_defaults(command="sweep-run-node")
sweep_collect = subparsers.add_parser("sweep-collect", help="reconstruct sweep job status from artifact files")
sweep_collect.add_argument("--jobs", default="artifacts/aggressive_oom_sweep/jobs.jsonl")
sweep_collect.set_defaults(command="sweep-collect")
@ -265,6 +273,32 @@ def main(argv: list[str] | None = None) -> int:
print(f"jobs_jsonl: {resolve_path(args.output_dir) / 'jobs.jsonl'}")
return 0
if args.command == "sweep-run-node":
from airfrans_frontier.sweep import run_node
if args.max_jobs is not None and args.max_jobs < 0:
print("error: --max-jobs must be non-negative", file=sys.stderr)
return 1
if args.stale_after_seconds < 0:
print("error: --stale-after-seconds must be non-negative", file=sys.stderr)
return 1
if args.max_attempts <= 0:
print("error: --max-attempts must be positive", file=sys.stderr)
return 1
try:
summary = run_node(
jobs_path=resolve_path(args.jobs),
node_name=args.node,
max_jobs=args.max_jobs,
stale_after_seconds=args.stale_after_seconds,
max_attempts=args.max_attempts,
)
except (FileNotFoundError, NotADirectoryError, ValueError, RuntimeError) as exc:
print(f"error: {exc}", file=sys.stderr)
return 1
print(json.dumps(summary, indent=2, sort_keys=True))
return 0
if args.command == "sweep-collect":
from airfrans_frontier.sweep import collect_job_status

34
src/airfrans_frontier/model_metadata.py Normal file
View file

@ -0,0 +1,34 @@
from __future__ import annotations
from typing import Any
PROXY_FAMILY_LABELS: dict[str, str] = {
"raster_fno_unet": "raster_fno_unet_proxy",
"meshgraphnet_or_point_transformer_local": "local_point_transformer_proxy",
"point_context_perceiver": "point_context_perceiver_query_proxy",
}
PROXY_FAMILY_TARGETS: dict[str, str] = {
"raster_fno_unet": "fno_or_raster_field_model",
"meshgraphnet_or_point_transformer_local": "meshgraphnet_or_point_transformer",
"point_context_perceiver": "query_context_perceiver",
}
PROXY_FAMILY_NOTES: dict[str, str] = {
"raster_fno_unet": "Raster-grid sampling proxy; not a true FNO/rasterized field pipeline.",
"meshgraphnet_or_point_transformer_local": "KNN point-local proxy; not a true MeshGraphNet without mesh adjacency and edge features.",
"point_context_perceiver": "Query/context proxy; target coordinates are explicit and target values are not used as context.",
}
def model_metadata_for_family(family: str, *, coordinate_encoding_compatibility: str | None = None) -> dict[str, Any]:
reported = PROXY_FAMILY_LABELS.get(family, family)
return {
"requested_family": family,
"implementation_family": family,
"reported_family": reported,
"is_proxy": reported != family,
"proxy_for": PROXY_FAMILY_TARGETS.get(family),
"proxy_notes": PROXY_FAMILY_NOTES.get(family),
"coordinate_encoding_compatibility": coordinate_encoding_compatibility,
}

22
src/airfrans_frontier/remote/artifacts.py
View file

@ -8,13 +8,15 @@ from typing import Any, Iterable
BASE_REQUIRED = (
"config.toml",
"job_manifest.json",
"metrics.jsonl",
"latest_metrics.json",
"heartbeat.json",
"checkpoint_latest.pt",
"checkpoint_best.pt",
"run_manifest.json",
"environment_manifest.json",
"utilization.jsonl",
)
SUCCESS_REQUIRED = ("final_metrics.json", "checkpoint_final.pt")
SUCCESS_REQUIRED = ("final_metrics.json", "checkpoint_latest.pt", "checkpoint_best.pt", "checkpoint_final.pt")
FAILURE_REQUIRED = ("failure_report.json",)
DEFAULT_REQUIRED = BASE_REQUIRED
@ -39,6 +41,8 @@ def verify_artifacts(
has_final = (root / "final_metrics.json").is_file()
has_failure = (root / "failure_report.json").is_file()
if has_final and has_failure:
raise ValueError("Artifact directory has both final_metrics.json and failure_report.json")
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:
@ -46,6 +50,7 @@ def verify_artifacts(
if missing_success:
raise ValueError(f"Successful artifact directory missing files: {', '.join(missing_success)}")
if has_failure:
_validate_failure_report(root / "failure_report.json")
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)}")
@ -151,6 +156,17 @@ def _validate_json(path: Path) -> None:
except json.JSONDecodeError as exc:
raise ValueError(f"Invalid JSON artifact {path}: {exc}") from exc
def _validate_failure_report(path: Path) -> None:
_validate_json(path)
data = json.loads(path.read_text())
if not isinstance(data, dict):
raise ValueError(f"failure_report.json must be a mapping: {path}")
for field in ("error_type", "error_message", "phase"):
if not isinstance(data.get(field), str) or not data[field]:
raise ValueError(f"failure_report.json missing non-empty {field}: {path}")
if "failure_category" not in data:
raise ValueError(f"failure_report.json missing failure_category: {path}")
def _validate_jsonl(path: Path) -> None:
with path.open() as file:

2
src/airfrans_frontier/remote/cli.py
View file

@ -545,6 +545,7 @@ _TERMINAL_ARTIFACT_NAMES = (
"disk_telemetry.json",
"environment_manifest.json",
"evaluation_protocol.json",
"job_manifest.json",
"failure_report.json",
"final_metrics.json",
"heartbeat.json",
@ -554,6 +555,7 @@ _TERMINAL_ARTIFACT_NAMES = (
"normalization.json",
"run_manifest.json",
"split_manifest.json",
"utilization.jsonl",
"startup_timeline.jsonl",
"wandb_smoke_manifest.json",
"verification_report.json",

19
src/airfrans_frontier/remote/config.py
View file

@ -176,7 +176,24 @@ def load_remote_run_config(path: str | Path) -> RemoteRunConfig:
)
artifacts = ArtifactConfig(
mode=_choice(_string(artifacts_raw, "mode", default="rsync"), {"rsync", "object_store_upload"}, "artifacts.mode"),
required=_string_tuple(artifacts_raw, "required", default=("final_metrics.json", "metrics.jsonl", "checkpoint.pt", "run_manifest.json")),
required=_string_tuple(
artifacts_raw,
"required",
default=(
"config.toml",
"job_manifest.json",
"metrics.jsonl",
"latest_metrics.json",
"heartbeat.json",
"run_manifest.json",
"environment_manifest.json",
"utilization.jsonl",
"final_metrics.json",
"checkpoint_latest.pt",
"checkpoint_best.pt",
"checkpoint_final.pt",
),
),
)
cleanup = CleanupConfig(
on_success=_choice(_string(cleanup_raw, "on_success", default="sky_down"), {"sky_down", "keep"}, "cleanup.on_success"),

194
src/airfrans_frontier/remote/smoke.py
View file

@ -12,6 +12,7 @@ from pathlib import Path
from typing import Any
from airfrans_frontier.remote.artifacts import verify_artifacts
from airfrans_frontier.training.hf_upload import HfArtifactUploader
from airfrans_frontier.runtime import remove_pythonpath_entries
@ -96,6 +97,8 @@ def run_smoke_training(
"timestamp": time.time(),
},
)
if error is not None:
_ensure_smoke_failure_contract(output_dir, config_path=config_path, run_id=run_id, error=error)
latest_metrics = _read_json(output_dir / "latest_metrics.json")
run_manifest: dict[str, Any] = {
@ -138,6 +141,7 @@ def run_smoke_training(
def _copy_training_artifacts(training_dir: Path, output_dir: Path) -> None:
names = (
"config.toml",
"job_manifest.json",
"metrics.jsonl",
"latest_metrics.json",
"heartbeat.json",
@ -156,6 +160,7 @@ def _copy_training_artifacts(training_dir: Path, output_dir: Path) -> None:
"artifact_manifest.json",
"checksums.txt",
"verification_report.json",
"utilization.jsonl",
"streaming_events.jsonl",
"streaming_state.json",
"streaming_summary.json",
@ -167,6 +172,55 @@ def _copy_training_artifacts(training_dir: Path, output_dir: Path) -> None:
shutil.copy2(source, output_dir / name)
def _ensure_smoke_failure_contract(output_dir: Path, *, config_path: str | Path, run_id: str, error: Exception) -> None:
if not (output_dir / "config.toml").is_file():
source_config = Path(config_path)
if source_config.is_file():
shutil.copy2(source_config, output_dir / "config.toml")
else:
(output_dir / "config.toml").write_text(
"[run]\n"
f"name = {json.dumps(run_id)}\n"
"kind = \"smoke_train_failure\"\n"
)
if not (output_dir / "job_manifest.json").is_file():
_write_json(
output_dir / "job_manifest.json",
{
"run_id": run_id,
"run_name": run_id,
"command": f"remote-run smoke-train {config_path}",
"artifact_dir": str(output_dir),
},
)
if not (output_dir / "metrics.jsonl").is_file():
metric = {
"run_id": run_id,
"event": "failed",
"phase": "failed",
"step": 0,
"error_type": type(error).__name__,
"timestamp": time.time(),
}
(output_dir / "metrics.jsonl").write_text(json.dumps(metric, sort_keys=True) + "\n")
_write_json(output_dir / "latest_metrics.json", metric)
elif not (output_dir / "latest_metrics.json").is_file():
_write_json(output_dir / "latest_metrics.json", _read_jsonl_last(output_dir / "metrics.jsonl"))
if not (output_dir / "utilization.jsonl").is_file():
(output_dir / "utilization.jsonl").write_text(
json.dumps({"timestamp": time.time(), "phase": "failed", "gpu_util_percent": None, "memory_used_mb": None}, sort_keys=True) + "\n"
)
report = _read_json(output_dir / "failure_report.json")
report.setdefault("run_id", run_id)
report.setdefault("phase", "training")
report.setdefault("failure_phase", report["phase"])
report.setdefault("failure_category", "training")
report.setdefault("error_type", type(error).__name__)
report.setdefault("error_message", str(error))
report.setdefault("timestamp", time.time())
_write_json(output_dir / "failure_report.json", report)
def _latest_training_run_dir(config_path: str | Path) -> Path | None:
try:
from airfrans_frontier.training.config import load_training_config
@ -186,14 +240,20 @@ def _latest_training_run_dir(config_path: str | Path) -> Path | None:
def _smoke_required(*, success: bool) -> tuple[str, ...]:
if not success:
return (
"config.toml",
"job_manifest.json",
"metrics.jsonl",
"latest_metrics.json",
"heartbeat.json",
"environment_manifest.json",
"run_manifest.json",
"utilization.jsonl",
"failure_report.json",
)
return (
"config.toml",
"metrics.jsonl",
"job_manifest.json",
"latest_metrics.json",
"heartbeat.json",
"checkpoint_latest.pt",
@ -208,6 +268,7 @@ def _smoke_required(*, success: bool) -> tuple[str, ...]:
"run_manifest.json",
"artifact_manifest.json",
"checksums.txt",
"utilization.jsonl",
"checkpoint_final.pt",
"final_metrics.json",
)
@ -265,25 +326,33 @@ def run_hf_upload_smoke(
(output_dir / "metrics.jsonl").write_text(json.dumps(metric, sort_keys=True) + "\n")
_write_json(output_dir / "latest_metrics.json", metric)
_write_json(output_dir / "final_metrics.json", {"hf_smoke": True, "loss": 0.0, "step": 0})
payload_dir = output_dir / "hf_payload"
final_dir = payload_dir / "final"
final_dir.mkdir(parents=True, exist_ok=True)
model_bytes = f"airfrans HF upload smoke\nrun_id={run_id}\n".encode()
model_path = final_dir / "smoke_model.bin"
model_path.write_bytes(model_bytes)
payload_manifest = {
"run_id": run_id,
"created_at": time.time(),
"files": [
_write_json(
output_dir / "job_manifest.json",
{
"run_id": run_id,
"run_name": run_id,
"command": "remote-run hf-smoke",
"artifact_dir": str(output_dir),
"checkpoint_policy": {
"latest": "checkpoint_latest.pt",
"best": "checkpoint_best.pt",
"final": "checkpoint_final.pt",
},
},
)
(output_dir / "utilization.jsonl").write_text(
json.dumps(
{
"path": "final/smoke_model.bin",
"bytes": len(model_bytes),
"sha256": hashlib.sha256(model_bytes).hexdigest(),
}
],
}
_write_json(payload_dir / "hf_smoke_manifest.json", payload_manifest)
"timestamp": time.time(),
"phase": "completed",
"step": 0,
"gpu_util_percent": None,
"memory_used_mb": None,
},
sort_keys=True,
)
+ "\n"
)
token = _resolve_hf_token()
try:
@ -293,7 +362,6 @@ def run_hf_upload_smoke(
api = HfApi(token=token)
resolved_repo_id = _resolve_hf_repo_id(api, repo_id or os.environ.get("AIRFRANS_HF_REPO_ID") or "airfrans-hf-smoke")
api.create_repo(repo_id=resolved_repo_id, repo_type="model", private=False, exist_ok=True)
path_in_repo = f"smoke/{run_id}"
_write_json(
heartbeat_path,
@ -305,34 +373,35 @@ def run_hf_upload_smoke(
"timestamp": time.time(),
},
)
commit = api.upload_folder(
uploader = HfArtifactUploader(
enabled=True,
run_dir=output_dir,
repo_id=resolved_repo_id,
repo_type="model",
folder_path=str(payload_dir),
path_in_repo=path_in_repo,
commit_message=f"Add AirfRANS HF smoke artifact {run_id}",
private=False,
)
upload_result = uploader.upload_files(
(
"checkpoint_latest.pt",
"checkpoint_best.pt",
"checkpoint_final.pt",
"metrics.jsonl",
"latest_metrics.json",
"final_metrics.json",
),
commit_message=f"Add AirfRANS checkpoint upload smoke {run_id}",
)
repo_files = set(api.list_repo_files(repo_id=resolved_repo_id, repo_type="model"))
expected_paths = [
f"{path_in_repo}/final/smoke_model.bin",
f"{path_in_repo}/hf_smoke_manifest.json",
]
expected_paths = list(upload_result["uploaded"])
missing = [path for path in expected_paths if path not in repo_files]
if missing:
raise RuntimeError(f"HF upload completed but repo listing is missing: {', '.join(missing)}")
finished = time.time()
hf_manifest = {
"repo_id": resolved_repo_id,
"repo_url": f"https://huggingface.co/{resolved_repo_id}",
"path_in_repo": path_in_repo,
"uploaded_paths": expected_paths,
"commit": _commit_payload(commit),
"started_at": started,
"finished_at": finished,
"elapsed_seconds": finished - started,
}
_write_json(output_dir / "hf_upload_manifest.json", hf_manifest)
hf_report = uploader.finalize(training_success=True)
final_metrics = {"hf_smoke": True, "loss": 0.0, "step": 0, **hf_report}
_write_json(output_dir / "final_metrics.json", final_metrics)
_write_json(
output_dir / "run_manifest.json",
{
@ -343,8 +412,9 @@ def run_hf_upload_smoke(
"elapsed_seconds": finished - started,
"exit_code": 0,
"artifact_dir": str(output_dir),
"hf_repo_url": hf_manifest["repo_url"],
"hf_repo_url": uploader.repo_url,
"hf_path_in_repo": path_in_repo,
**hf_report,
},
)
_write_json(
@ -353,8 +423,9 @@ def run_hf_upload_smoke(
"run_id": run_id,
"phase": "completed",
"latest_metrics": metric,
"hf_repo_url": hf_manifest["repo_url"],
"hf_repo_url": uploader.repo_url,
"hf_path_in_repo": path_in_repo,
**hf_report,
"started_at": started,
"finished_at": finished,
"updated_at": time.time(),
@ -436,7 +507,7 @@ def run_wandb_smoke(
}
wandb.log(latest_metric, step=step)
metrics_lines.append(json.dumps(latest_metric, sort_keys=True))
if hf_repo_url:
if hf_repo_url and run is not None:
run.summary["hf_repo_url"] = hf_repo_url
wandb_run_url = run.get_url() if run is not None else None
wandb.finish(exit_code=0)
@ -455,12 +526,36 @@ def run_wandb_smoke(
(output_dir / "metrics.jsonl").write_text("\n".join(metrics_lines) + "\n")
_write_json(output_dir / "latest_metrics.json", latest_metric)
_write_json(
output_dir / "job_manifest.json",
{
"run_id": run_id,
"run_name": run_id,
"command": "remote-run wandb-smoke",
"artifact_dir": str(output_dir),
},
)
(output_dir / "utilization.jsonl").write_text(
json.dumps(
{
"timestamp": time.time(),
"phase": "completed",
"step": latest_metric["step"],
"gpu_util_percent": None,
"memory_used_mb": None,
},
sort_keys=True,
)
+ "\n"
)
final_metrics = {
"wandb_smoke": True,
"loss": latest_metric["loss"],
"step": latest_metric["step"],
"wandb_run_url": wandb_run_url,
"hf_repo_url": hf_repo_url,
"status": "completed",
"wandb_run_url_present": bool(wandb_run_url),
}
_write_json(output_dir / "final_metrics.json", final_metrics)
finished = time.time()
@ -471,6 +566,9 @@ def run_wandb_smoke(
"project": project,
"run_id": run_id,
"run_url": wandb_run_url,
"status": "completed",
"run_url_present": bool(wandb_run_url),
"hf_repo_url": hf_repo_url,
"logged_steps": 20,
"started_at": started,
"finished_at": finished,
@ -488,6 +586,8 @@ def run_wandb_smoke(
"exit_code": 0,
"artifact_dir": str(output_dir),
"wandb_run_url": wandb_run_url,
"status": "completed",
"hf_repo_url": hf_repo_url,
},
)
_write_json(
@ -497,6 +597,8 @@ def run_wandb_smoke(
"phase": "completed",
"latest_metrics": latest_metric,
"wandb_run_url": wandb_run_url,
"status": "completed",
"hf_repo_url": hf_repo_url,
"started_at": started,
"finished_at": finished,
"updated_at": time.time(),
@ -555,6 +657,18 @@ def _read_json(path: Path) -> dict[str, Any]:
return data if isinstance(data, dict) else {}
def _read_jsonl_last(path: Path) -> dict[str, Any]:
last: dict[str, Any] = {}
if not path.is_file():
return last
for line in path.read_text().splitlines():
if line.strip():
data = json.loads(line)
if isinstance(data, dict):
last = data
return last
def _write_json(path: Path, data: Any) -> None:
path.write_text(json.dumps(data, indent=2, sort_keys=True) + "\n")

478
src/airfrans_frontier/sweep.py
View file

@ -1,4 +1,5 @@
from __future__ import annotations
import fcntl
import json
import math
@ -8,9 +9,12 @@ import subprocess
import sys
import time
import tomllib
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable, Mapping
from typing import Any, Iterable, Iterator, Mapping
from airfrans_frontier.model_metadata import PROXY_FAMILY_LABELS, model_metadata_for_family
EXPANSION_LADDER = (1, 2, 4, 8)
REQUIRED_JOB_ARTIFACTS = (
@ -54,12 +58,19 @@ ENCODING_COMPATIBLE_FAMILIES = {
"film_fourier_inr",
}
PROXY_FAMILY_LABELS = {
"raster_fno_unet": "raster_fno_unet_proxy",
"meshgraphnet_or_point_transformer_local": "local_point_transformer_proxy",
"point_context_perceiver": "point_context_perceiver_query_proxy",
COORDINATE_ENCODING_EXCLUSION_REASONS = {
"mlp": "architecture_uses_full_raw_point_features",
"siren_conditioned_inr": "siren_uses_raw_coordinates_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",
}
DEFAULT_STALE_AFTER_SECONDS = 6 * 60 * 60
DEFAULT_MAX_ATTEMPTS = 2
ATTEMPTS_JSONL = "attempts.jsonl"
RESCUE_TEMPLATES_JSON = "rescue_templates.json"
@dataclass(frozen=True)
class PoolNode:
@ -131,6 +142,34 @@ 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:
raise ValueError(f"unsupported coordinate encoding: {encoding}")
if family in ENCODING_COMPATIBLE_FAMILIES:
return {
"family": family,
"encoding": encoding,
"compatible": True,
"compatibility": "shared_coordinate_encoder",
"reason": "family accepts the shared coordinate encoding config",
}
if encoding == "raw":
return {
"family": family,
"encoding": encoding,
"compatible": True,
"compatibility": "raw_only",
"reason": COORDINATE_ENCODING_EXCLUSION_REASONS.get(family, "raw encoding is the compatible fallback"),
}
return {
"family": family,
"encoding": encoding,
"compatible": False,
"compatibility": "excluded",
"reason": COORDINATE_ENCODING_EXCLUSION_REASONS.get(family, "family does not use the shared coordinate encoder"),
}
def generate_jobs(
*,
@ -152,13 +191,17 @@ def generate_jobs(
if band not in OOM_BANDS:
raise ValueError(f"unsupported OOM band: {band}")
for family in families:
family_encodings = tuple(encodings) if family in ENCODING_COMPATIBLE_FAMILIES else ("raw",)
for encoding in family_encodings:
if encoding not in DEFAULT_ENCODINGS:
raise ValueError(f"unsupported coordinate encoding: {encoding}")
for encoding in encodings:
decision = coordinate_encoding_decision(family, encoding)
if not decision["compatible"]:
continue
job_id = f"{band}_{family}_{encoding}"
run_name = job_id
config_path = config_dir / f"{job_id}.toml"
metadata = model_metadata_for_family(
family,
coordinate_encoding_compatibility=str(decision["compatibility"]),
)
config_text = training_config_text(
run_name=run_name,
seed=seed + len(jobs),
@ -177,8 +220,15 @@ def generate_jobs(
"status": "pending",
"band": band,
"model_family": family,
"reported_family": PROXY_FAMILY_LABELS.get(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"],
"config_path": str(config_path),
"attempts": 0,
}
@ -186,6 +236,7 @@ def generate_jobs(
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
@ -207,6 +258,11 @@ def training_config_text(
condition_width = max(512, shape["hidden_width"] // 2)
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)
metadata = model_metadata_for_family(
model_family,
coordinate_encoding_compatibility=str(decision["compatibility"]),
)
return "\n".join(
line for line in (
"[run]",
@ -226,6 +282,7 @@ def training_config_text(
"test_cases = 50",
"points_per_case = 8192",
f"batch_size = {shape['batch_size']}",
"streaming_normalization_cases = 16" if band == "700m" else None,
"",
"[coordinate_encoding]",
_coordinate_encoding_toml(coordinate_encoding),
@ -246,6 +303,15 @@ def training_config_text(
"grid_resolution = 96",
"siren_omega0 = 30.0",
"",
"[model_metadata]",
f"requested_family = {json.dumps(metadata['requested_family'])}",
f"implementation_family = {json.dumps(metadata['implementation_family'])}",
f"reported_family = {json.dumps(metadata['reported_family'])}",
f"is_proxy = {str(bool(metadata['is_proxy'])).lower()}",
f"proxy_for = {json.dumps(metadata['proxy_for'])}" if metadata["proxy_for"] is not None else None,
f"proxy_notes = {json.dumps(metadata['proxy_notes'])}" if metadata["proxy_notes"] is not None else None,
f"coordinate_encoding_compatibility = {json.dumps(metadata['coordinate_encoding_compatibility'])}",
"",
"[optim]",
"lr = 0.0001",
"weight_decay = 0.0001",
@ -283,6 +349,7 @@ def training_config_text(
f"path_prefix = {json.dumps(group)}",
"private = false",
)
if line is not None
) + "\n"
@ -297,6 +364,96 @@ def _coordinate_encoding_toml(encoding: str) -> str:
return "type = \"random_fourier\"\nfeatures = [\"x\", \"y\", \"sdf\"]\nnum_features = 256\nsigma = 16.0\nseed = 0"
raise ValueError(f"unsupported coordinate encoding: {encoding}")
def rescue_templates(families: Iterable[str] = DEFAULT_FAMILIES) -> list[dict[str, Any]]:
templates: list[dict[str, Any]] = [
{
"template_id": "optimizer_lr_3e-4",
"scope": "all",
"axis": "optim.lr",
"overrides": {"optim": {"lr": 3.0e-4}},
},
{
"template_id": "optimizer_lr_1e-4",
"scope": "all",
"axis": "optim.lr",
"overrides": {"optim": {"lr": 1.0e-4}},
},
{
"template_id": "optimizer_lr_3e-5",
"scope": "all",
"axis": "optim.lr",
"overrides": {"optim": {"lr": 3.0e-5}},
},
{
"template_id": "stability_grad_clip_0_5",
"scope": "all",
"axis": "stability.max_grad_norm",
"overrides": {"stability": {"max_grad_norm": 0.5}},
},
{
"template_id": "stability_grad_clip_2_0",
"scope": "all",
"axis": "stability.max_grad_norm",
"overrides": {"stability": {"max_grad_norm": 2.0}},
},
]
family_templates: dict[str, list[dict[str, Any]]] = {
"siren_conditioned_inr": [
{"template_id": "siren_omega0_10", "axis": "model.siren_omega0", "overrides": {"model": {"siren_omega0": 10.0}}},
{"template_id": "siren_omega0_30", "axis": "model.siren_omega0", "overrides": {"model": {"siren_omega0": 30.0}}},
{"template_id": "siren_omega0_60", "axis": "model.siren_omega0", "overrides": {"model": {"siren_omega0": 60.0}}},
{"template_id": "siren_lr_1e-5", "axis": "optim.lr", "overrides": {"optim": {"lr": 1.0e-5}}},
{"template_id": "siren_lr_3e-5", "axis": "optim.lr", "overrides": {"optim": {"lr": 3.0e-5}}},
{"template_id": "siren_lr_1e-4", "axis": "optim.lr", "overrides": {"optim": {"lr": 1.0e-4}}},
],
"film_fourier_inr": [
{"template_id": "film_condition_dim_512", "axis": "model.condition_dim", "overrides": {"model": {"condition_dim": 512}}},
{"template_id": "film_condition_dim_1024", "axis": "model.condition_dim", "overrides": {"model": {"condition_dim": 1024}}},
{"template_id": "film_condition_width_1024", "axis": "model.condition_width", "overrides": {"model": {"condition_width": 1024}}},
{"template_id": "film_condition_width_2048", "axis": "model.condition_width", "overrides": {"model": {"condition_width": 2048}}},
],
"deeponet_branch_trunk": [
{"template_id": "deeponet_fourier_max_scale_16", "axis": "coordinate_encoding.scales", "overrides": {"coordinate_encoding": {"scales": [1, 2, 4, 8, 16]}}},
{"template_id": "deeponet_fourier_max_scale_32", "axis": "coordinate_encoding.scales", "overrides": {"coordinate_encoding": {"scales": [1, 2, 4, 8, 16, 32]}}},
{"template_id": "deeponet_fourier_max_scale_64", "axis": "coordinate_encoding.scales", "overrides": {"coordinate_encoding": {"scales": [1, 2, 4, 8, 16, 32, 64]}}},
{"template_id": "deeponet_condition_depth_3", "axis": "model.condition_depth", "overrides": {"model": {"condition_depth": 3}}},
{"template_id": "deeponet_condition_depth_4", "axis": "model.condition_depth", "overrides": {"model": {"condition_depth": 4}}},
],
"raster_fno_unet": [
{"template_id": "raster_grid_64", "axis": "model.grid_resolution", "overrides": {"model": {"grid_resolution": 64}}},
{"template_id": "raster_grid_96", "axis": "model.grid_resolution", "overrides": {"model": {"grid_resolution": 96}}},
{"template_id": "raster_grid_128", "axis": "model.grid_resolution", "overrides": {"model": {"grid_resolution": 128}}},
],
"point_context_perceiver": [
{"template_id": "perceiver_context_points_512", "axis": "model.context_points", "overrides": {"model": {"context_points": 512}}},
{"template_id": "perceiver_context_points_1024", "axis": "model.context_points", "overrides": {"model": {"context_points": 1024}}},
{"template_id": "perceiver_latent_width_512", "axis": "model.latent_width", "overrides": {"model": {"latent_width": 512}}},
{"template_id": "perceiver_latent_width_1024", "axis": "model.latent_width", "overrides": {"model": {"latent_width": 1024}}},
{"template_id": "perceiver_attention_depth_2", "axis": "model.attention_depth", "overrides": {"model": {"attention_depth": 2}}},
{"template_id": "perceiver_attention_depth_4", "axis": "model.attention_depth", "overrides": {"model": {"attention_depth": 4}}},
],
"meshgraphnet_or_point_transformer_local": [
{"template_id": "local_neighbors_8", "axis": "model.neighbors", "overrides": {"model": {"neighbors": 8}}},
{"template_id": "local_neighbors_16", "axis": "model.neighbors", "overrides": {"model": {"neighbors": 16}}},
{"template_id": "local_neighbors_32", "axis": "model.neighbors", "overrides": {"model": {"neighbors": 32}}},
{"template_id": "local_batch_limit_1024", "axis": "data.batch_size", "overrides": {"data": {"batch_size": 1024}}},
{"template_id": "local_batch_limit_2048", "axis": "data.batch_size", "overrides": {"data": {"batch_size": 2048}}},
],
}
for family in families:
for template in family_templates.get(family, ()):
item = dict(template)
item["scope"] = family
templates.append(item)
return sorted(templates, key=lambda item: (str(item["scope"]), str(item["template_id"])))
def write_rescue_templates(path: str | Path, families: Iterable[str] = DEFAULT_FAMILIES) -> Path:
target = Path(path)
target.parent.mkdir(parents=True, exist_ok=True)
target.write_text(json.dumps(rescue_templates(families), indent=2, sort_keys=True) + "\n")
return target
def read_jobs(path: str | Path) -> list[dict[str, Any]]:
jobs: list[dict[str, Any]] = []
@ -312,33 +469,292 @@ def write_jobs(path: str | Path, jobs: Iterable[Mapping[str, Any]]) -> None:
target.write_text("".join(json.dumps(dict(job), sort_keys=True) + "\n" for job in jobs))
def run_next_job(*, jobs_path: str | Path, node_name: str, extra_env: Mapping[str, str] | None = None) -> dict[str, Any] | None:
path = Path(jobs_path)
jobs = read_jobs(path)
selected_index = next((index for index, job in enumerate(jobs) if job.get("status") == "pending"), None)
if selected_index is None:
return None
job = dict(jobs[selected_index])
job["status"] = "running"
job["node"] = node_name
job["started_at"] = time.time()
job["attempts"] = int(job.get("attempts", 0)) + 1
jobs[selected_index] = job
write_jobs(path, jobs)
@contextmanager
def _locked_jobs_file(jobs_path: Path) -> Iterator[None]:
jobs_path.parent.mkdir(parents=True, exist_ok=True)
lock_path = jobs_path.with_suffix(jobs_path.suffix + ".lock")
with lock_path.open("a") as lock_file:
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
try:
yield
finally:
fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
def recover_stale_jobs(
*,
jobs_path: str | Path,
stale_after_seconds: float = DEFAULT_STALE_AFTER_SECONDS,
now: float | None = None,
) -> list[dict[str, Any]]:
path = Path(jobs_path)
current_time = time.time() if now is None else float(now)
with _locked_jobs_file(path):
jobs = read_jobs(path)
recovered = _recover_stale_jobs_locked(path, jobs, stale_after_seconds=stale_after_seconds, now=current_time)
if recovered:
write_jobs(path, jobs)
return recovered
def claim_next_job(
*,
jobs_path: str | Path,
node_name: str,
stale_after_seconds: float = DEFAULT_STALE_AFTER_SECONDS,
max_attempts: int = DEFAULT_MAX_ATTEMPTS,
now: float | None = None,
) -> dict[str, Any] | None:
if max_attempts <= 0:
raise ValueError("max_attempts must be positive")
path = Path(jobs_path)
current_time = time.time() if now is None else float(now)
with _locked_jobs_file(path):
jobs = read_jobs(path)
_recover_stale_jobs_locked(path, jobs, stale_after_seconds=stale_after_seconds, now=current_time)
selected_index: int | None = None
for index, job in enumerate(jobs):
if job.get("status") != "pending":
continue
if int(job.get("attempts", 0)) >= max_attempts:
exhausted = dict(job)
exhausted.update({"status": "exhausted", "exhausted_at": current_time, "last_error": "max attempts exhausted"})
jobs[index] = exhausted
_append_attempt_record(
path,
{
"event": "exhausted",
"job_id": exhausted.get("job_id"),
"attempts": exhausted.get("attempts", 0),
"timestamp": current_time,
},
)
continue
selected_index = index
break
if selected_index is None:
write_jobs(path, jobs)
return None
job = dict(jobs[selected_index])
attempt_number = int(job.get("attempts", 0)) + 1
attempt_id = f"{job['job_id']}-attempt-{attempt_number}"
job.update(
{
"status": "running",
"node": node_name,
"claimed_by": node_name,
"started_at": current_time,
"lease_updated_at": current_time,
"attempts": attempt_number,
"last_attempt_id": attempt_id,
}
)
jobs[selected_index] = job
write_jobs(path, jobs)
_append_attempt_record(
path,
{
"event": "started",
"attempt_id": attempt_id,
"attempt": attempt_number,
"job_id": job["job_id"],
"node": node_name,
"config_path": job.get("config_path"),
"timestamp": current_time,
},
)
return job
def complete_job_attempt(
*,
jobs_path: str | Path,
attempt_id: str,
returncode: int,
stdout: str = "",
stderr: str = "",
now: float | None = None,
) -> dict[str, Any]:
path = Path(jobs_path)
current_time = time.time() if now is None else float(now)
terminal_status = "succeeded" if returncode == 0 else "failed"
with _locked_jobs_file(path):
jobs = read_jobs(path)
for index, raw_job in enumerate(jobs):
if raw_job.get("last_attempt_id") != attempt_id:
continue
job = dict(raw_job)
job.update(
{
"status": terminal_status,
"finished_at": current_time,
"returncode": int(returncode),
"stdout": stdout[-4000:],
"stderr": stderr[-4000:],
"last_error": None if returncode == 0 else (stderr or stdout)[-1000:],
}
)
jobs[index] = job
write_jobs(path, jobs)
_append_attempt_record(
path,
{
"event": terminal_status,
"attempt_id": attempt_id,
"job_id": job.get("job_id"),
"timestamp": current_time,
"returncode": int(returncode),
"stdout_tail": stdout[-4000:],
"stderr_tail": stderr[-4000:],
},
)
return job
raise ValueError(f"No running job attempt found for attempt_id={attempt_id}")
def run_next_job(*, jobs_path: str | Path, node_name: str, extra_env: Mapping[str, str] | None = None) -> dict[str, Any] | None:
job = claim_next_job(jobs_path=jobs_path, node_name=node_name)
if job is None:
return None
env = os.environ.copy()
if extra_env:
env.update(extra_env)
env.update(
{
"AIRFRANS_SWEEP_JOB_ID": str(job["job_id"]),
"AIRFRANS_SWEEP_ATTEMPT_ID": str(job["last_attempt_id"]),
}
)
command = [sys.executable, "-m", "airfrans_frontier.cli", "train", str(job["config_path"])]
completed = subprocess.run(command, text=True, capture_output=True, env=env, check=False)
job["finished_at"] = time.time()
job["returncode"] = completed.returncode
job["stdout"] = completed.stdout[-4000:]
job["stderr"] = completed.stderr[-4000:]
job["status"] = "succeeded" if completed.returncode == 0 else "failed"
jobs[selected_index] = job
write_jobs(path, jobs)
return job
return complete_job_attempt(
jobs_path=jobs_path,
attempt_id=str(job["last_attempt_id"]),
returncode=completed.returncode,
stdout=completed.stdout,
stderr=completed.stderr,
)
def run_node(
*,
jobs_path: str | Path,
node_name: str,
max_jobs: int | None = None,
stale_after_seconds: float = DEFAULT_STALE_AFTER_SECONDS,
max_attempts: int = DEFAULT_MAX_ATTEMPTS,
extra_env: Mapping[str, str] | None = None,
) -> dict[str, Any]:
if max_jobs is not None and max_jobs < 0:
raise ValueError("max_jobs must be non-negative")
recovered = recover_stale_jobs(jobs_path=jobs_path, stale_after_seconds=stale_after_seconds)
summary = {
"node": node_name,
"claimed": 0,
"succeeded": 0,
"failed": 0,
"stale_recovered": len(recovered),
"attempts_path": str(Path(jobs_path).parent / ATTEMPTS_JSONL),
}
while max_jobs is None or summary["claimed"] < max_jobs:
job = claim_next_job(
jobs_path=jobs_path,
node_name=node_name,
stale_after_seconds=stale_after_seconds,
max_attempts=max_attempts,
)
if job is None:
break
summary["claimed"] += 1
env = os.environ.copy()
if extra_env:
env.update(extra_env)
env.update(
{
"AIRFRANS_SWEEP_JOB_ID": str(job["job_id"]),
"AIRFRANS_SWEEP_ATTEMPT_ID": str(job["last_attempt_id"]),
}
)
completed = subprocess.run(
[sys.executable, "-m", "airfrans_frontier.cli", "train", str(job["config_path"])],
text=True,
capture_output=True,
env=env,
check=False,
)
terminal = complete_job_attempt(
jobs_path=jobs_path,
attempt_id=str(job["last_attempt_id"]),
returncode=completed.returncode,
stdout=completed.stdout,
stderr=completed.stderr,
)
if terminal["status"] == "succeeded":
summary["succeeded"] += 1
else:
summary["failed"] += 1
return summary
def _recover_stale_jobs_locked(
jobs_path: Path,
jobs: list[dict[str, Any]],
*,
stale_after_seconds: float,
now: float,
) -> list[dict[str, Any]]:
if stale_after_seconds < 0:
return []
recovered: list[dict[str, Any]] = []
for index, raw_job in enumerate(jobs):
if raw_job.get("status") != "running":
continue
started_at = raw_job.get("lease_updated_at", raw_job.get("started_at"))
if started_at is None or now - float(started_at) <= stale_after_seconds:
continue
job = dict(raw_job)
reconstructed = reconstruct_status(job)
if reconstructed["status"] in {"succeeded", "failed"}:
job.update(
{
"status": reconstructed["status"],
"finished_at": now,
"terminal_run_dir": reconstructed["run_dir"],
"last_error": None,
}
)
event = "stale_terminal_reconciled"
else:
job.update(
{
"status": "pending",
"stale_recoveries": int(job.get("stale_recoveries", 0)) + 1,
"last_error": "stale running attempt recovered to pending",
"recovered_at": now,
}
)
event = "stale_recovered"
jobs[index] = job
record = {
"event": event,
"attempt_id": job.get("last_attempt_id"),
"job_id": job.get("job_id"),
"node": job.get("node"),
"timestamp": now,
"previous_status": "running",
"new_status": job["status"],
"run_dir": reconstructed.get("run_dir"),
}
_append_attempt_record(jobs_path, record)
recovered.append(record)
return recovered
def _append_attempt_record(jobs_path: Path, record: Mapping[str, Any]) -> None:
attempts_path = jobs_path.parent / ATTEMPTS_JSONL
attempts_path.parent.mkdir(parents=True, exist_ok=True)
with attempts_path.open("a") as file:
file.write(json.dumps(dict(record), sort_keys=True) + "\n")
def collect_job_status(*, jobs_path: str | Path) -> dict[str, Any]:

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

@ -4,6 +4,7 @@ import tomllib
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from airfrans_frontier.model_metadata import model_metadata_for_family
_MODEL_TYPES = {
@ -70,6 +71,17 @@ class ModelConfig:
siren_omega0: float
@dataclass(frozen=True)
class ModelMetadataConfig:
requested_family: str
reported_family: str
implementation_family: str
is_proxy: bool
proxy_for: str | None
proxy_notes: str | None
coordinate_encoding_compatibility: str | None
@dataclass(frozen=True)
class CoordinateEncodingConfig:
type: str
@ -142,6 +154,7 @@ class TrainingConfig:
run: RunConfig
data: DataConfig
model: ModelConfig
model_metadata: ModelMetadataConfig
optim: OptimConfig
device: DeviceConfig
loss: LossConfig
@ -208,6 +221,11 @@ def load_training_config(path: str | Path) -> TrainingConfig:
coordinate_encoding_raw = {}
if not isinstance(coordinate_encoding_raw, dict):
raise ValueError("Training config [coordinate_encoding] section must be a table")
model_metadata_raw = raw.get("model_metadata", {})
if model_metadata_raw is None:
model_metadata_raw = {}
if not isinstance(model_metadata_raw, dict):
raise ValueError("Training config [model_metadata] section must be a table")
coordinate_encoding = CoordinateEncodingConfig(
@ -303,6 +321,25 @@ def load_training_config(path: str | Path) -> TrainingConfig:
grid_resolution=_integer(model_raw, "grid_resolution", minimum=2, default=32),
siren_omega0=_number(model_raw, "siren_omega0", minimum=0.0, exclusive_minimum=True, default=30.0),
)
metadata_defaults = model_metadata_for_family(model.type)
model_metadata = ModelMetadataConfig(
requested_family=_string(model_metadata_raw, "requested_family", default=str(metadata_defaults["requested_family"])),
reported_family=_string(model_metadata_raw, "reported_family", default=str(metadata_defaults["reported_family"])),
implementation_family=_string(model_metadata_raw, "implementation_family", default=str(metadata_defaults["implementation_family"])),
is_proxy=_boolean(model_metadata_raw, "is_proxy") if "is_proxy" in model_metadata_raw else bool(metadata_defaults["is_proxy"]),
proxy_for=_optional_string(model_metadata_raw, "proxy_for") if "proxy_for" in model_metadata_raw else metadata_defaults["proxy_for"],
proxy_notes=_optional_string(model_metadata_raw, "proxy_notes") if "proxy_notes" in model_metadata_raw else metadata_defaults["proxy_notes"],
coordinate_encoding_compatibility=(
_optional_string(model_metadata_raw, "coordinate_encoding_compatibility")
if "coordinate_encoding_compatibility" in model_metadata_raw
else metadata_defaults["coordinate_encoding_compatibility"]
),
)
if model_metadata.is_proxy and not (model_metadata.proxy_for or model_metadata.reported_family != model_metadata.implementation_family):
raise ValueError("model_metadata.is_proxy requires proxy_for or a reported_family distinct from implementation_family")
if not model_metadata.is_proxy and model_metadata.reported_family != model_metadata.implementation_family:
raise ValueError("model_metadata reported_family cannot differ from implementation_family unless is_proxy = true")
optim = OptimConfig(
lr=_number(optim_raw, "lr", minimum=0.0, exclusive_minimum=True),
weight_decay=_number(optim_raw, "weight_decay", minimum=0.0),
@ -351,6 +388,7 @@ def load_training_config(path: str | Path) -> TrainingConfig:
run=run,
data=data,
model=model,
model_metadata=model_metadata,
optim=optim,
device=device,
loss=loss,

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

@ -64,16 +64,32 @@ class DatasetBundle:
target_names: tuple[str, ...]
def load_processed_dataset(root: str | Path) -> list[SimulationSample]:
def processed_npz_paths(root: str | Path) -> tuple[Path, ...]:
root_path = Path(root).expanduser()
if not root_path.exists():
raise FileNotFoundError(f"Processed data directory not found: {root_path}")
if not root_path.is_dir():
raise NotADirectoryError(f"Processed data path is not a directory: {root_path}")
paths = sorted(root_path.glob("*.npz"))
paths = tuple(sorted(root_path.glob("*.npz")))
if not paths:
raise ValueError(f"No .npz simulation files found under: {root_path}")
return paths
def processed_case_ids(root: str | Path) -> tuple[str, ...]:
return tuple(path.stem for path in processed_npz_paths(root))
def load_processed_dataset(root: str | Path, case_ids: tuple[str, ...] | None = None) -> list[SimulationSample]:
paths = processed_npz_paths(root)
if case_ids is not None:
wanted = set(case_ids)
paths_by_id = {path.stem: path for path in paths}
missing = sorted(wanted.difference(paths_by_id))
if missing:
raise FileNotFoundError(f"Processed data is missing selected cases: {missing[:10]}")
paths = tuple(paths_by_id[case_id] for case_id in case_ids)
samples = [load_simulation_npz(path) for path in paths]
validate_common_schema(samples)
@ -164,15 +180,28 @@ def build_dataset_bundle(
points_per_case: int,
seed: int,
) -> DatasetBundle:
validate_common_schema(samples)
samples_by_id = {sample.case_id: sample for sample in samples}
split = create_case_split(
tuple(samples_by_id),
tuple(sample.case_id for sample in samples),
train_cases=train_cases,
val_cases=val_cases,
test_cases=test_cases,
seed=seed,
)
return build_dataset_bundle_for_split(samples, split=split, points_per_case=points_per_case, seed=seed)
def build_dataset_bundle_for_split(
samples: list[SimulationSample],
*,
split: CaseSplit,
points_per_case: int,
seed: int,
) -> DatasetBundle:
validate_common_schema(samples)
samples_by_id = {sample.case_id: sample for sample in samples}
missing = sorted(set(split.train_ids + split.val_ids + split.test_ids).difference(samples_by_id))
if missing:
raise ValueError(f"Selected split references unloaded cases: {missing[:10]}")
train = build_split_arrays(samples_by_id, split.train_ids, points_per_case, seed=seed + 101)
val = (
build_split_arrays(samples_by_id, split.val_ids, points_per_case, seed=seed + 202)

245
src/airfrans_frontier/training/hf_upload.py
View file

@ -11,6 +11,9 @@ from typing import Any, Iterable
from airfrans_frontier.training.config import TrainingConfig
DEFAULT_MAX_HF_FILE_BYTES = 512 * 1024**2
@dataclass
class UploadRecord:
@ -32,6 +35,8 @@ class UploadManifest:
uploaded_files: list[dict[str, Any]] = field(default_factory=list)
commits: list[dict[str, Any]] = field(default_factory=list)
suppressed_uploads: list[dict[str, Any]] = field(default_factory=list)
upload_attempts: list[dict[str, Any]] = field(default_factory=list)
pending_uploads: list[dict[str, Any]] = field(default_factory=list)
last_error: str | None = None
rate_limit_until: float | None = None
rate_limit_retry_after_seconds: float | None = None
@ -39,6 +44,12 @@ class UploadManifest:
publication_complete: bool = False
publication_status: str = "disabled"
finalized_at: float | None = None
uploaded_bytes_total: int = 0
uploaded_file_count: int = 0
commit_count: int = 0
pending_bytes_total: int = 0
storage_warning: dict[str, Any] | None = None
storage_pricing_url: str = "https://huggingface.co/pricing#storage"
class HfArtifactUploader:
def __init__(
@ -60,12 +71,14 @@ class HfArtifactUploader:
self.private = private
self.max_rate_limit_sleep_seconds = max_rate_limit_sleep_seconds
self._api: Any | None = None
self._manifest = UploadManifest(
enabled=enabled,
repo_id=repo_id,
repo_type=repo_type,
repo_url=f"https://huggingface.co/{repo_id}" if repo_id else None,
path_in_repo=self.path_in_repo,
self._manifest = self._load_existing_manifest(
UploadManifest(
enabled=enabled,
repo_id=repo_id,
repo_type=repo_type,
repo_url=f"https://huggingface.co/{repo_id}" if repo_id else None,
path_in_repo=self.path_in_repo,
)
)
self.write_manifest()
@ -105,6 +118,7 @@ class HfArtifactUploader:
return self.final_report()
def final_report(self) -> dict[str, Any]:
self._refresh_totals()
self._refresh_publication_status()
return {
"hf_publication_status": self._manifest.publication_status,
@ -112,7 +126,14 @@ class HfArtifactUploader:
"hf_training_success": self._manifest.training_success,
"hf_rate_limit_until": self._manifest.rate_limit_until,
"hf_last_error": self._manifest.last_error,
"hf_uploaded_bytes_total": self._manifest.uploaded_bytes_total,
"hf_uploaded_file_count": self._manifest.uploaded_file_count,
"hf_commit_count": self._manifest.commit_count,
"hf_pending_bytes_total": self._manifest.pending_bytes_total,
"hf_storage_warning": self._manifest.storage_warning,
"hf_storage_pricing_url": self._manifest.storage_pricing_url,
}
def upload_files(self, names: Iterable[str], *, commit_message: str) -> dict[str, Any]:
names = tuple(dict.fromkeys(names))
if not self.enabled:
@ -123,18 +144,63 @@ class HfArtifactUploader:
suppressed = self._suppress_if_rate_limited(names, commit_message=commit_message)
if suppressed is not None:
return suppressed
api = self._ensure_api()
paths: list[tuple[Path, str]] = []
skipped: list[str] = []
suppressed_large_files: list[str] = []
max_file_bytes = _max_hf_file_bytes()
for name in names:
local_path = self.run_dir / name
repo_path = f"{self.path_in_repo}/{name}" if self.path_in_repo else name
paths.append((local_path, repo_path))
size = local_path.stat().st_size
if max_file_bytes is not None and size > max_file_bytes:
suppressed_large_files.append(repo_path)
self._manifest.suppressed_uploads.append(
{
"event": "file_too_large",
"local_path": str(local_path),
"repo_path": repo_path,
"bytes": size,
"max_file_bytes": max_file_bytes,
"commit_message": commit_message,
"timestamp": time.time(),
}
)
continue
if self._already_uploaded(local_path, repo_path):
skipped.append(repo_path)
else:
paths.append((local_path, repo_path))
if not paths:
self.write_manifest()
return {
"enabled": True,
"uploaded": [],
"skipped": skipped,
"suppressed_large_files": suppressed_large_files,
"missing": [],
"rate_limited": False,
**self.final_report(),
}
api = self._ensure_api()
uploaded = [repo_path for _, repo_path in paths]
attempts = 0
max_attempts = 3
while True:
try:
from huggingface_hub import CommitOperationAdd
self._manifest.upload_attempts.append(
{
"event": "attempt_started",
"attempt": attempts + 1,
"names": list(names),
"repo_paths": uploaded,
"commit_message": commit_message,
"timestamp": time.time(),
}
)
self._set_pending_uploads(paths)
self.write_manifest()
operations = [
CommitOperationAdd(path_in_repo=repo_path, path_or_fileobj=str(local_path))
for local_path, repo_path in paths
@ -146,6 +212,7 @@ class HfArtifactUploader:
commit_message=commit_message,
)
uploaded_at = time.time()
self._drop_pending_uploads(uploaded)
for local_path, repo_path in paths:
record = UploadRecord(
local_path=str(local_path),
@ -155,22 +222,37 @@ class HfArtifactUploader:
uploaded_at=uploaded_at,
)
self._manifest.uploaded_paths.append(repo_path)
self._manifest.uploaded_files = [
item for item in self._manifest.uploaded_files if item.get("repo_path") != repo_path
]
self._manifest.uploaded_files.append(record.__dict__)
self._manifest.commits.append(_commit_payload(commit))
self._manifest.uploaded_paths = sorted(set(self._manifest.uploaded_paths))
self._manifest.last_error = None
self._manifest.rate_limit_until = None
self._manifest.rate_limit_retry_after_seconds = None
self._manifest.upload_attempts.append(
{
"event": "attempt_succeeded",
"attempt": attempts + 1,
"repo_paths": uploaded,
"timestamp": uploaded_at,
}
)
self.write_manifest()
return {"enabled": True, "uploaded": uploaded, "missing": [], "rate_limited": False}
return {"enabled": True, "uploaded": uploaded, "skipped": skipped, "suppressed_large_files": suppressed_large_files, "missing": [], "rate_limited": False, **self.final_report()}
except Exception as exc:
retry_after = _retry_after_seconds(exc)
if retry_after is not None:
self._record_rate_limit(exc, retry_after, names=names, commit_message=commit_message)
if attempts == 0 and retry_after <= self.max_rate_limit_sleep_seconds:
attempts += 1
time.sleep(max(0.0, retry_after))
continue
else:
retry_after = _transient_retry_seconds(exc, attempts)
self._record_upload_failure(exc, paths=paths, attempt=attempts + 1, retry_after_seconds=retry_after)
if retry_after is not None and attempts + 1 < max_attempts and retry_after <= self.max_rate_limit_sleep_seconds:
attempts += 1
time.sleep(max(0.0, retry_after))
continue
self._set_pending_uploads(paths)
self._manifest.last_error = str(exc)
self.write_manifest()
raise
@ -208,10 +290,104 @@ class HfArtifactUploader:
)
self.write_manifest()
def _load_existing_manifest(self, default: UploadManifest) -> UploadManifest:
path = self.run_dir / "hf_upload_manifest.json"
if not path.is_file():
return default
try:
data = json.loads(path.read_text())
except (OSError, json.JSONDecodeError):
return default
if not isinstance(data, dict):
return default
allowed = set(default.__dict__)
merged = default.__dict__ | {key: value for key, value in data.items() if key in allowed}
merged.update(
{
"enabled": self.enabled,
"repo_id": self.repo_id,
"repo_type": self.repo_type,
"repo_url": f"https://huggingface.co/{self.repo_id}" if self.repo_id else None,
"path_in_repo": self.path_in_repo,
}
)
return UploadManifest(**merged)
def _already_uploaded(self, local_path: Path, repo_path: str) -> bool:
size = local_path.stat().st_size
digest = _sha256_file(local_path)
for record in self._manifest.uploaded_files:
if record.get("repo_path") == repo_path and record.get("bytes") == size and record.get("sha256") == digest:
return True
return False
def _set_pending_uploads(self, paths: list[tuple[Path, str]]) -> None:
records = []
for local_path, repo_path in paths:
records.append(
{
"local_path": str(local_path),
"repo_path": repo_path,
"bytes": local_path.stat().st_size,
"sha256": _sha256_file(local_path),
"updated_at": time.time(),
}
)
existing = [item for item in self._manifest.pending_uploads if item.get("repo_path") not in {record["repo_path"] for record in records}]
self._manifest.pending_uploads = existing + records
def _drop_pending_uploads(self, repo_paths: list[str]) -> None:
pending = set(repo_paths)
self._manifest.pending_uploads = [item for item in self._manifest.pending_uploads if item.get("repo_path") not in pending]
def _record_upload_failure(
self,
exc: Exception,
*,
paths: list[tuple[Path, str]],
attempt: int,
retry_after_seconds: float | None,
) -> None:
self._manifest.last_error = str(exc)
self._manifest.upload_attempts.append(
{
"event": "attempt_failed",
"attempt": attempt,
"error_type": type(exc).__name__,
"error_message": str(exc),
"transient": retry_after_seconds is not None,
"retry_after_seconds": retry_after_seconds,
"repo_paths": [repo_path for _, repo_path in paths],
"timestamp": time.time(),
}
)
self._set_pending_uploads(paths)
def _refresh_totals(self) -> None:
self._manifest.uploaded_bytes_total = int(sum(int(record.get("bytes", 0)) for record in self._manifest.uploaded_files))
self._manifest.uploaded_file_count = len(self._manifest.uploaded_paths)
self._manifest.commit_count = len(self._manifest.commits)
self._manifest.pending_bytes_total = int(sum(int(record.get("bytes", 0)) for record in self._manifest.pending_uploads))
projected = self._manifest.uploaded_bytes_total + self._manifest.pending_bytes_total
threshold = 500 * 1024**3 if self.private else 5 * 1024**4
self._manifest.storage_warning = (
{
"projected_bytes": projected,
"threshold_bytes": threshold,
"scope": "private" if self.private else "public",
"pricing_url": self._manifest.storage_pricing_url,
}
if projected >= threshold
else None
)
def write_manifest(self) -> Path:
self._refresh_totals()
self._refresh_publication_status()
path = self.run_dir / "hf_upload_manifest.json"
path.write_text(json.dumps(self._manifest.__dict__, indent=2, sort_keys=True) + "\n")
tmp_path = path.with_suffix(path.suffix + ".tmp")
tmp_path.write_text(json.dumps(self._manifest.__dict__, indent=2, sort_keys=True) + "\n")
tmp_path.replace(path)
return path
def _refresh_publication_status(self) -> None:
@ -223,7 +399,7 @@ class HfArtifactUploader:
self._manifest.publication_status = "training_failed"
self._manifest.publication_complete = False
return
incomplete = self._manifest.last_error is not None or self._manifest.rate_limit_until is not None
incomplete = self._manifest.last_error is not None or self._manifest.rate_limit_until is not None or bool(self._manifest.pending_uploads)
if incomplete:
self._manifest.publication_status = (
"training_succeeded_hf_incomplete" if self._manifest.training_success is True else "hf_publication_incomplete"
@ -253,6 +429,16 @@ class HfArtifactUploader:
return api
def _max_hf_file_bytes() -> int | None:
raw = os.environ.get("AIRFRANS_HF_MAX_FILE_BYTES")
if raw is None:
return DEFAULT_MAX_HF_FILE_BYTES
raw = raw.strip().lower()
if raw in {"", "0", "none", "off", "false"}:
return None
return max(0, int(raw))
def resolve_resume_checkpoint(resume_path: str | Path | None) -> tuple[Path | None, dict[str, Any]]:
if resume_path is None:
return None, {"resume_source": None, "resume_downloaded": False, "resume_downloaded_path": None}
@ -279,6 +465,35 @@ def resolve_resume_checkpoint(resume_path: str | Path | None) -> tuple[Path | No
return downloaded, {"resume_source": raw, "resume_downloaded": True, "resume_downloaded_path": str(downloaded)}
def _transient_retry_seconds(exc: Exception, attempt: int) -> float | None:
response = getattr(exc, "response", None)
status = getattr(response, "status_code", None) or getattr(response, "status", None)
try:
status_int = int(status) if status is not None else None
except (TypeError, ValueError):
status_int = None
if status_int == 408 or (status_int is not None and 500 <= status_int <= 599):
return min(0.25 * (2**attempt), 5.0)
lowered = str(exc).lower()
transient_markers = (
"connection reset",
"connection aborted",
"timed out",
"timeout",
"temporarily unavailable",
"incomplete read",
"incomplete upload",
"server error",
"bad gateway",
"service unavailable",
"gateway timeout",
)
if isinstance(exc, (ConnectionResetError, TimeoutError)) or type(exc).__name__ in {"IncompleteRead", "ReadTimeout", "ConnectTimeout"}:
return min(0.25 * (2**attempt), 5.0)
if any(marker in lowered for marker in transient_markers):
return min(0.25 * (2**attempt), 5.0)
return None
def _parse_hf_checkpoint_uri(uri: str) -> tuple[str, str]:
rest = uri.removeprefix("hf://")
parts = rest.split("/")

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

@ -33,7 +33,14 @@ from airfrans_frontier.training.data_sources import resolve_training_data_root
from airfrans_frontier.training.environment import environment_manifest
from airfrans_frontier.training.hf_upload import HfArtifactUploader, resolve_resume_checkpoint
from airfrans_frontier.training.observability import start_observer
from airfrans_frontier.training.data import DatasetBundle, build_dataset_bundle, load_processed_dataset
from airfrans_frontier.training.data import (
DatasetBundle,
build_dataset_bundle,
build_dataset_bundle_for_split,
create_case_split,
load_processed_dataset,
processed_case_ids,
)
from airfrans_frontier.training.metrics import count_parameters, device_metrics, overall_mse, per_channel_mse
from airfrans_frontier.training.streaming_data import (
PROCESSED_UPLOAD_MANIFEST,
@ -104,7 +111,7 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
observer.update_config(
{
"run_id": run_id,
"model_family": config.model.type,
**_analysis_metadata(config),
"hf_repo_url": uploader.repo_url,
"hf_path_in_repo": uploader.path_in_repo,
}
@ -132,6 +139,7 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
def record_metrics(metrics: dict[str, Any]) -> None:
nonlocal first_metric_timeline_written
metrics.update(_analysis_metadata(config))
writer.append_metrics(metrics)
writer.append_jsonl(
"utilization.jsonl",
@ -212,102 +220,158 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
device=device,
)
data_root = resolve_training_data_root(config.data)
samples = load_processed_dataset(data_root)
bundle = build_dataset_bundle(
samples,
train_cases=config.data.train_cases,
val_cases=config.data.val_cases,
test_cases=config.data.test_cases,
points_per_case=config.data.points_per_case,
seed=config.run.seed,
)
stats = compute_normalization_stats(
bundle.train.features,
bundle.train.targets,
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(data_root),
"configured_root": str(config.data.root),
"source": config.data.source,
"hf_repo_id": config.data.hf_repo_id,
"hf_repo_type": config.data.hf_repo_type,
"hf_path_prefix": config.data.hf_path_prefix,
"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),
"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
val_targets = normalize_targets(bundle.val.targets, stats) if bundle.val is not None else None
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 = _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,
)
calibration_fields = static_calibration_fields(config, model)
protocol_fields = _evaluation_protocol(config.model.type)
writer.write_json("calibration_manifest.json", calibration_fields)
writer.write_json("evaluation_protocol.json", protocol_fields)
observer.update_config({**calibration_fields, **protocol_fields})
run_manifest.update(
{
"phase": "initialized",
"parameter_count": count_parameters(model),
**calibration_fields,
**protocol_fields,
}
)
writer.write_json("run_manifest.json", run_manifest)
writer.write_artifact_manifest()
publish_artifacts(
(
"config.toml",
"job_manifest.json",
"utilization.jsonl",
"environment_manifest.json",
"split_manifest.json",
started = time.perf_counter()
try:
data_root = resolve_training_data_root(config.data)
available_case_ids = processed_case_ids(data_root)
split = create_case_split(
available_case_ids,
train_cases=config.data.train_cases,
val_cases=config.data.val_cases,
test_cases=config.data.test_cases,
seed=config.run.seed,
)
selected_case_ids = split.train_ids + split.val_ids + split.test_ids
samples = load_processed_dataset(data_root, case_ids=selected_case_ids)
bundle = build_dataset_bundle_for_split(
samples,
split=split,
points_per_case=config.data.points_per_case,
seed=config.run.seed,
)
stats = compute_normalization_stats(
bundle.train.features,
bundle.train.targets,
feature_names=bundle.feature_names,
target_names=bundle.target_names,
)
writer.write_split_manifest(bundle.split.to_dict())
writer.write_json(
"data_manifest.json",
"normalization.json",
"calibration_manifest.json",
"evaluation_protocol.json",
"run_manifest.json",
"artifact_manifest.json",
"checksums.txt",
),
event="initialized",
step=0,
)
{
"root": str(data_root),
"configured_root": str(config.data.root),
"source": config.data.source,
"hf_repo_id": config.data.hf_repo_id,
"hf_repo_type": config.data.hf_repo_type,
"hf_path_prefix": config.data.hf_path_prefix,
"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),
"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
val_targets = normalize_targets(bundle.val.targets, stats) if bundle.val is not None else None
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
except Exception as exc:
_write_terminal_failure_bundle(
writer,
config=config,
run_manifest=run_manifest,
uploader=uploader,
observer=observer,
phase="data",
step=0,
exc=exc,
device=device,
started_perf=started,
)
observer.finish(exit_code=1)
raise
try:
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,
)
calibration_fields = static_calibration_fields(config, model)
protocol_fields = _evaluation_protocol(config.model.type)
writer.write_json("calibration_manifest.json", calibration_fields)
writer.write_json("evaluation_protocol.json", protocol_fields)
observer.update_config({**calibration_fields, **protocol_fields})
run_manifest.update(
{
"phase": "initialized",
"parameter_count": count_parameters(model),
**calibration_fields,
**protocol_fields,
}
)
writer.write_json("run_manifest.json", run_manifest)
writer.write_artifact_manifest()
publish_artifacts(
(
"config.toml",
"job_manifest.json",
"utilization.jsonl",
"environment_manifest.json",
"split_manifest.json",
"data_manifest.json",
"normalization.json",
"calibration_manifest.json",
"evaluation_protocol.json",
"run_manifest.json",
"artifact_manifest.json",
"checksums.txt",
),
event="initialized",
step=0,
)
except Exception as exc:
_write_terminal_failure_bundle(
writer,
config=config,
run_manifest=run_manifest,
uploader=uploader,
observer=observer,
phase="model",
step=0,
exc=exc,
device=device,
started_perf=started,
)
observer.finish(exit_code=1)
raise
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
full_train_eval_cases = len(bundle.train.case_ids)
train_eval_cases = full_train_eval_cases
train_eval_features = train_features
train_eval_targets = train_targets
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)
train_eval_features = train_features[:train_eval_rows]
train_eval_targets = train_targets[:train_eval_rows]
def with_train_eval_scope(metrics: dict[str, Any]) -> dict[str, Any]:
if train_eval_cases != full_train_eval_cases:
metrics["train_eval_cases"] = train_eval_cases
metrics["train_eval_full_cases"] = full_train_eval_cases
return metrics
if resume is not None:
try:
@ -327,6 +391,8 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
error_type=type(exc).__name__,
error_message=str(exc),
latest_checkpoint=str(resume),
config=config,
failure_category="resume",
)
raise
record_metrics(
@ -344,88 +410,122 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
)
)
initial_train = evaluate_arrays(
model,
train_features,
train_targets,
batch_size=config.data.batch_size,
device=device,
target_names=bundle.target_names,
)
initial_val = (
evaluate_arrays(
try:
initial_train = evaluate_arrays(
model,
val_features,
val_targets,
train_eval_features,
train_eval_targets,
batch_size=config.data.batch_size,
device=device,
target_names=bundle.target_names,
)
if val_features is not None and val_targets is not None
else None
)
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"]
record_metrics(
_log_metrics(
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,
initial_val = (
evaluate_arrays(
model,
val_features,
val_targets,
batch_size=config.data.batch_size,
device=device,
target_names=bundle.target_names,
)
if val_features is not None and val_targets is not None
else None
)
)
_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,
)
if not first_checkpoint_timeline_written:
record_timeline("training", "first_checkpoint_written", step=start_step)
first_checkpoint_timeline_written = True
writer.write_artifact_manifest()
publish_artifacts(
(
"metrics.jsonl",
"latest_metrics.json",
"heartbeat.json",
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"]
record_metrics(
with_train_eval_scope(
_log_metrics(
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,
)
)
)
except Exception as exc:
_write_terminal_failure_bundle(
writer,
config=config,
run_manifest=run_manifest,
uploader=uploader,
observer=observer,
phase="eval",
step=start_step,
exc=exc,
device=device,
started_perf=started,
)
observer.finish(exit_code=1)
raise
try:
_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,
"artifact_manifest.json",
"checksums.txt",
),
event="initial_checkpoint",
step=start_step,
)
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,
)
if not first_checkpoint_timeline_written:
record_timeline("training", "first_checkpoint_written", step=start_step)
first_checkpoint_timeline_written = True
writer.write_artifact_manifest()
publish_artifacts(
(
"metrics.jsonl",
"latest_metrics.json",
"heartbeat.json",
LATEST_CHECKPOINT,
BEST_CHECKPOINT,
"artifact_manifest.json",
"checksums.txt",
),
event="initial_checkpoint",
step=start_step,
)
except Exception as exc:
_write_terminal_failure_bundle(
writer,
config=config,
run_manifest=run_manifest,
uploader=uploader,
observer=observer,
phase="checkpoint",
step=start_step,
exc=exc,
device=device,
started_perf=started,
)
observer.finish(exit_code=1)
raise
log_interval = config.optim.log_interval or max(1, config.optim.steps // 10)
last_checkpoint_at = time.monotonic()
@ -459,6 +559,8 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
latest_loss=float(loss.detach().cpu().item()),
latest_grad_norm=last_grad_norm,
latest_checkpoint=LATEST_CHECKPOINT,
config=config,
failure_category="nonfinite",
)
raise RuntimeError("nonfinite loss")
loss.backward()
@ -478,6 +580,8 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
latest_loss=float(loss.detach().cpu().item()),
latest_grad_norm=last_grad_norm,
latest_checkpoint=LATEST_CHECKPOINT,
config=config,
failure_category="nonfinite",
)
raise RuntimeError("nonfinite gradients") from exc
last_grad_norm = float(grad_norm_tensor.detach().cpu().item())
@ -514,8 +618,8 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
if step % log_interval == 0 or step == config.optim.steps:
train_eval = evaluate_arrays(
model,
train_features,
train_targets,
train_eval_features,
train_eval_targets,
batch_size=config.data.batch_size,
device=device,
target_names=bundle.target_names,
@ -559,17 +663,19 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
points_per_sec = (step - last_log_step) * config.data.batch_size / interval_elapsed
last_points_per_sec = points_per_sec
record_metrics(
_log_metrics(
event="train_eval",
step=step,
train_loss=train_eval["loss"],
val_loss=val_eval["loss"] if val_eval is not None else None,
elapsed_seconds=elapsed,
lr=_learning_rate(optimizer),
grad_norm=last_grad_norm,
points_per_sec=points_per_sec,
device=device,
latest_checkpoint=LATEST_CHECKPOINT if should_checkpoint else None,
with_train_eval_scope(
_log_metrics(
event="train_eval",
step=step,
train_loss=train_eval["loss"],
val_loss=val_eval["loss"] if val_eval is not None else None,
elapsed_seconds=elapsed,
lr=_learning_rate(optimizer),
grad_norm=last_grad_norm,
points_per_sec=points_per_sec,
device=device,
latest_checkpoint=LATEST_CHECKPOINT if should_checkpoint else None,
)
)
)
last_log_at = time.perf_counter()
@ -586,6 +692,8 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
error_message=str(exc),
latest_grad_norm=last_grad_norm,
latest_checkpoint=LATEST_CHECKPOINT,
config=config,
failure_category="training",
)
run_manifest.update({"phase": "failed", "finished_at": time.time(), "exit_code": 1, **uploader.finalize(training_success=False)})
writer.write_json("run_manifest.json", run_manifest)
@ -622,8 +730,8 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
final_train = evaluate_arrays(
model,
train_features,
train_targets,
train_eval_features,
train_eval_targets,
batch_size=config.data.batch_size,
device=device,
target_names=bundle.target_names,
@ -664,10 +772,7 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
"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,
"model_family": config.model.type,
**_analysis_metadata(config),
"precision": config.precision.dtype,
"train_cases": len(bundle.split.train_ids),
"val_cases": len(bundle.split.val_ids),
@ -695,6 +800,9 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
if uploader.repo_url is not None:
final_metrics["hf_repo_url"] = uploader.repo_url
final_metrics["hf_path_in_repo"] = uploader.path_in_repo
if train_eval_cases != full_train_eval_cases:
final_metrics["train_eval_cases"] = train_eval_cases
final_metrics["train_eval_full_cases"] = full_train_eval_cases
writer.write_final_metrics(final_metrics)
_save_training_checkpoint(
writer,
@ -720,18 +828,20 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
writer.write_final_metrics(final_metrics)
observer.update_summary(final_metrics)
record_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,
with_train_eval_scope(
_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,
)
)
)
run_manifest.update(
@ -761,7 +871,11 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
event="verification",
step=config.optim.steps,
)
run_manifest.update(uploader.finalize(training_success=True))
hf_report = uploader.finalize(training_success=True)
final_metrics.update(hf_report)
writer.write_final_metrics(final_metrics)
observer.update_summary(final_metrics)
run_manifest.update(hf_report)
writer.write_json("run_manifest.json", run_manifest)
writer.write_artifact_manifest()
verify_artifacts(writer.run_dir, required=_verification_required(success=True))
@ -888,6 +1002,8 @@ def _train_public_zip_streaming(
error_type=type(exc).__name__,
error_message=str(exc),
latest_checkpoint=str(resume),
config=config,
failure_category="resume",
)
raise
record_streaming_metrics(
@ -1038,6 +1154,8 @@ def _train_public_zip_streaming(
latest_loss=float(loss.detach().cpu().item()),
latest_grad_norm=last_grad_norm,
latest_checkpoint=LATEST_CHECKPOINT,
config=config,
failure_category="nonfinite",
)
raise RuntimeError("nonfinite loss")
latest_loss_value = float(loss.detach().cpu().item())
@ -1058,6 +1176,8 @@ def _train_public_zip_streaming(
latest_loss=float(loss.detach().cpu().item()),
latest_grad_norm=last_grad_norm,
latest_checkpoint=LATEST_CHECKPOINT,
config=config,
failure_category="nonfinite",
)
raise RuntimeError("nonfinite gradients") from exc
last_grad_norm = float(grad_norm_tensor.detach().cpu().item())
@ -1215,8 +1335,7 @@ def _train_public_zip_streaming(
"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,
"model_family": config.model.type,
**_analysis_metadata(config),
"precision": config.precision.dtype,
"train_cases": len(bundle.split.train_ids),
"val_cases": len(bundle.split.val_ids),
@ -1315,7 +1434,11 @@ def _train_public_zip_streaming(
event="verification",
step=config.optim.steps,
)
run_manifest.update(uploader.finalize(training_success=True))
hf_report = uploader.finalize(training_success=True)
final_metrics.update(hf_report)
writer.write_final_metrics(final_metrics)
observer.update_summary(final_metrics)
run_manifest.update(hf_report)
writer.write_json("run_manifest.json", run_manifest)
writer.write_artifact_manifest()
verify_artifacts(writer.run_dir, required=_streaming_verification_required(success=True))
@ -1353,6 +1476,8 @@ def _train_public_zip_streaming(
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="streaming_training",
)
run_manifest.update({"phase": "failed", "finished_at": time.time(), "exit_code": 1, "data_mode": "public_zip_streaming", **uploader.finalize(training_success=False)})
writer.write_json("run_manifest.json", run_manifest)
@ -1637,6 +1762,37 @@ def _log_metrics(
**_memory_metrics(device),
}
def _coordinate_encoding_metadata(config: TrainingConfig) -> dict[str, Any]:
return {
"coordinate_encoding_type": config.coordinate_encoding.type,
"coordinate_encoding_features": list(config.coordinate_encoding.features),
"coordinate_encoding_scales": list(config.coordinate_encoding.scales),
"coordinate_encoding_levels": config.coordinate_encoding.levels,
"coordinate_encoding_num_features": config.coordinate_encoding.num_features,
"coordinate_encoding_sigma": config.coordinate_encoding.sigma,
"coordinate_encoding_seed": config.coordinate_encoding.seed,
}
def _model_metadata(config: TrainingConfig) -> dict[str, Any]:
metadata = asdict(config.model_metadata)
return {
"model_type": config.model.type,
"model_family": metadata["reported_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_compatibility": metadata["coordinate_encoding_compatibility"],
}
def _analysis_metadata(config: TrainingConfig) -> dict[str, Any]:
return {**_model_metadata(config), **_coordinate_encoding_metadata(config)}
def _job_manifest(*, config: TrainingConfig, run_id: str, run_dir: Path) -> dict[str, Any]:
return {
@ -1644,7 +1800,15 @@ def _job_manifest(*, config: TrainingConfig, run_id: str, run_dir: Path) -> dict
"run_name": config.run.name,
"run_dir": str(run_dir),
"config_path": str(config.path),
"model_family": config.model.type,
"model_type": config.model.type,
"model_family": config.model_metadata.reported_family,
"requested_family": config.model_metadata.requested_family,
"implementation_family": config.model_metadata.implementation_family,
"reported_family": config.model_metadata.reported_family,
"is_proxy": config.model_metadata.is_proxy,
"proxy_for": config.model_metadata.proxy_for,
"proxy_notes": config.model_metadata.proxy_notes,
"model_metadata": asdict(config.model_metadata),
"coordinate_encoding": {
"type": config.coordinate_encoding.type,
"features": list(config.coordinate_encoding.features),
@ -1691,9 +1855,8 @@ def _run_manifest(
return {
"run_id": run_id,
"run_name": config.run.name,
"model_family": config.model.type,
"coordinate_encoding_type": config.coordinate_encoding.type,
"coordinate_encoding_features": list(config.coordinate_encoding.features),
**_analysis_metadata(config),
"model_metadata": asdict(config.model_metadata),
"phase": phase,
"started_at": started_at,
"artifact_dir": str(run_dir),
@ -1754,21 +1917,26 @@ def _verification_required(*, success: bool) -> tuple[str, ...]:
"metrics.jsonl",
"latest_metrics.json",
"heartbeat.json",
LATEST_CHECKPOINT,
BEST_CHECKPOINT,
"split_manifest.json",
"data_manifest.json",
"normalization.json",
"environment_manifest.json",
"calibration_manifest.json",
"evaluation_protocol.json",
"run_manifest.json",
"hf_upload_manifest.json",
"artifact_manifest.json",
"checksums.txt",
]
if success:
required.extend(("final_metrics.json", FINAL_CHECKPOINT))
required.extend(
(
LATEST_CHECKPOINT,
BEST_CHECKPOINT,
"split_manifest.json",
"data_manifest.json",
"normalization.json",
"calibration_manifest.json",
"evaluation_protocol.json",
"final_metrics.json",
FINAL_CHECKPOINT,
)
)
else:
required.append("failure_report.json")
return tuple(required)
@ -2043,21 +2211,111 @@ def _write_failure(
latest_loss: float | None = None,
latest_grad_norm: float | None = None,
latest_checkpoint: str | None = None,
config: TrainingConfig | None = None,
**extra: Any,
) -> 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(),
payload = {
"run_id": os.environ.get("AIRFRANS_REMOTE_RUN_ID"),
"phase": phase,
"failure_phase": phase,
"failure_category": str(extra.pop("failure_category", 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(),
**extra,
}
if config is not None:
payload.update(_analysis_metadata(config))
writer.write_failure_report(payload)
writer.write_artifact_manifest()
def _write_terminal_failure_bundle(
writer: ArtifactWriter,
*,
config: TrainingConfig,
run_manifest: dict[str, Any],
uploader: HfArtifactUploader,
observer: Any,
phase: str,
step: int,
exc: Exception,
device: torch.device,
started_perf: float,
latest_checkpoint: str | None = None,
latest_grad_norm: float | None = None,
latest_loss: float | None = None,
**extra: Any,
) -> None:
if not (writer.run_dir / "metrics.jsonl").is_file():
metrics = _log_metrics(
event="failed",
phase="failed",
step=step,
train_loss=None,
val_loss=None,
elapsed_seconds=max(0.0, time.perf_counter() - started_perf),
lr=0.0,
grad_norm=latest_grad_norm,
points_per_sec=None,
device=device,
latest_checkpoint=latest_checkpoint,
)
metrics.update(_analysis_metadata(config))
writer.append_metrics(metrics)
if not (writer.run_dir / "failure_report.json").is_file():
_write_failure(
writer,
phase=phase,
step=step,
error_type=type(exc).__name__,
error_message=str(exc),
latest_loss=latest_loss,
latest_grad_norm=latest_grad_norm,
latest_checkpoint=latest_checkpoint,
config=config,
failure_category=_failure_category(phase, exc),
**extra,
)
hf_report = uploader.finalize(training_success=False)
run_manifest.update(
{
"phase": "failed",
"failure_phase": phase,
"failure_category": _failure_category(phase, exc),
"finished_at": time.time(),
"exit_code": 1,
**hf_report,
}
)
writer.write_json("run_manifest.json", run_manifest)
writer.write_artifact_manifest()
try:
verify_artifacts(writer.run_dir, required=_verification_required(success=False))
except Exception as verification_exc:
writer.write_json(
"verification_report.json",
{
"ok": False,
"error_type": type(verification_exc).__name__,
"error_message": str(verification_exc),
"checked_at": time.time(),
},
)
observer.update_summary(hf_report)
def _failure_category(phase: str, exc: Exception) -> str:
message = str(exc).lower()
if "cuda requested" in message or ("cuda" in message and phase in {"setup", "device"}):
return "device"
if phase in {"data", "normalization", "model", "eval", "checkpoint", "resume", "streaming_training"}:
return phase
if isinstance(exc, (FileNotFoundError, NotADirectoryError)):
return "data"
return phase

29
src/airfrans_frontier/training/observability.py
View file

@ -1,5 +1,7 @@
from __future__ import annotations
import os
import threading
from dataclasses import asdict, is_dataclass
from pathlib import Path
from typing import Any, Iterable, Mapping
@ -89,7 +91,32 @@ class WandbObserver(TrainingObserver):
self._run.log_artifact(artifact, aliases=list(aliases) or [event, f"step-{step}"])
def finish(self, *, exit_code: int = 0) -> None:
self._wandb.finish(exit_code=exit_code)
timeout = _wandb_finish_timeout_seconds()
if timeout is None:
self._wandb.finish(exit_code=exit_code)
return
error: list[BaseException] = []
def finish() -> None:
try:
self._wandb.finish(exit_code=exit_code)
except BaseException as exc: # noqa: BLE001 - preserve observer failure when not timed out.
error.append(exc)
thread = threading.Thread(target=finish, name="airfrans-wandb-finish", daemon=True)
thread.start()
thread.join(timeout=timeout)
if thread.is_alive():
return
if error:
raise error[0]
def _wandb_finish_timeout_seconds() -> float | None:
raw = os.environ.get("AIRFRANS_WANDB_FINISH_TIMEOUT_SECONDS", "60").strip().lower()
if raw in {"", "none", "off", "false"}:
return None
return max(0.0, float(raw))
def start_observer(config: TrainingConfig, *, run_dir: Path) -> TrainingObserver:

6
tests/test_remote_run.py
View file

@ -241,7 +241,7 @@ class ArtifactVerificationTests(unittest.TestCase):
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")
(root / "failure_report.json").write_text(json.dumps({"phase": "training", "failure_category": "nonfinite", "error_type": "NonFiniteLoss", "error_message": "bad"}) + "\n")
manifest = verify_artifacts(root)
@ -274,6 +274,10 @@ 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 / "job_manifest.json").write_text(json.dumps({"run_id": "test"}) + "\n")
(root / "run_manifest.json").write_text(json.dumps({"phase": "completed" if success else "failed"}) + "\n")
(root / "environment_manifest.json").write_text(json.dumps({"python": "test"}) + "\n")
(root / "utilization.jsonl").write_text(json.dumps({"gpu_util_percent": None}) + "\n")
(root / "heartbeat.json").write_text(json.dumps({"phase": "training"}) + "\n")
checkpoint = {
"schema_version": 1,

3
tests/test_streaming_data.py
View file

@ -400,7 +400,8 @@ class FullDataBackpressureStreamingTests(unittest.TestCase):
events = {event["event"] for event in read_events(run_dir)}
self.assertIn("processing_failure", events)
verification = json.loads((run_dir / "verification_report.json").read_text())
self.assertFalse(verification["ok"])
self.assertTrue(verification["ok"])
self.assertEqual(verification["checks"]["terminal_artifact"], "failure_report.json")
if __name__ == "__main__":

87
tests/test_sweep.py
View file

@ -4,15 +4,20 @@ import json
import tempfile
import unittest
from pathlib import Path
from unittest.mock import Mock, patch
from airfrans_frontier.sweep import (
BudgetLedger,
PoolConfig,
can_expand_pool,
claim_next_job,
collect_job_status,
complete_job_attempt,
generate_jobs,
read_jobs,
rescue_templates,
run_node,
utilization_summary,
can_expand_pool,
)
from airfrans_frontier.training.config import load_training_config
@ -38,6 +43,9 @@ class SweepFoundationTests(unittest.TestCase):
config = load_training_config(root / "configs" / "100m_film_fourier_inr_nerf_multires.toml")
self.assertEqual(config.coordinate_encoding.type, "nerf_multires")
self.assertEqual(config.coordinate_encoding.features, ("x", "y", "sdf"))
self.assertEqual(config.model_metadata.reported_family, "film_fourier_inr")
self.assertFalse(config.model_metadata.is_proxy)
self.assertTrue((root / "rescue_templates.json").is_file())
self.assertEqual(config.model.encoding_levels, 16)
self.assertEqual(config.model.fourier_scales, ())
self.assertEqual(config.model.hidden_width, 2048)
@ -55,6 +63,9 @@ class SweepFoundationTests(unittest.TestCase):
)
job_ids = [job["job_id"] for job in read_jobs(root / "jobs.jsonl")]
proxy_config = load_training_config(root / "configs" / "100m_raster_fno_unet_raw.toml") if (root / "configs" / "100m_raster_fno_unet_raw.toml").is_file() else None
if proxy_config is not None:
self.assertEqual(proxy_config.model_metadata.reported_family, "raster_fno_unet_proxy")
self.assertIn("100m_film_fourier_inr_random_fourier", job_ids)
self.assertIn("100m_siren_conditioned_inr_raw", job_ids)
@ -92,6 +103,80 @@ class SweepFoundationTests(unittest.TestCase):
self.assertEqual(status["jobs"][0]["job_id"], jobs[0]["job_id"])
self.assertEqual(status["jobs"][0]["missing_artifacts"], [])
def test_claim_complete_and_stale_recovery_are_durable(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
root = Path(tmp)
jobs = generate_jobs(
output_dir=root / "sweep",
data_root="data/full",
artifact_dir=str(root / "runs"),
bands=("100m",),
families=("mlp",),
encodings=("raw",),
)
jobs_path = root / "sweep" / "jobs.jsonl"
claimed = claim_next_job(jobs_path=jobs_path, node_name="node-a", now=100.0)
assert claimed is not None
self.assertEqual(claimed["status"], "running")
self.assertEqual(claimed["attempts"], 1)
self.assertEqual(claimed["node"], "node-a")
self.assertIsNone(claim_next_job(jobs_path=jobs_path, node_name="node-b", now=101.0))
completed = complete_job_attempt(
jobs_path=jobs_path,
attempt_id=claimed["last_attempt_id"],
returncode=7,
stdout="out",
stderr="err",
now=102.0,
)
self.assertEqual(completed["status"], "failed")
attempts = [json.loads(line) for line in (root / "sweep" / "attempts.jsonl").read_text().splitlines()]
self.assertEqual([item["event"] for item in attempts], ["started", "failed"])
persisted = read_jobs(jobs_path)
persisted[0].update({"status": "running", "lease_updated_at": 10.0, "last_attempt_id": "stale"})
from airfrans_frontier.sweep import write_jobs
write_jobs(jobs_path, persisted)
reclaimed = claim_next_job(jobs_path=jobs_path, node_name="node-c", stale_after_seconds=5.0, now=20.0)
assert reclaimed is not None
self.assertEqual(reclaimed["job_id"], jobs[0]["job_id"])
self.assertEqual(reclaimed["node"], "node-c")
events = [json.loads(line)["event"] for line in (root / "sweep" / "attempts.jsonl").read_text().splitlines()]
self.assertIn("stale_recovered", events)
def test_run_node_drains_jobs_with_patched_training_process(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
root = Path(tmp)
generate_jobs(
output_dir=root / "sweep",
data_root="data/full",
artifact_dir=str(root / "runs"),
bands=("100m",),
families=("mlp",),
encodings=("raw",),
)
completed = Mock(returncode=0, stdout="ok", stderr="")
with patch("airfrans_frontier.sweep.subprocess.run", Mock(return_value=completed)):
summary = run_node(jobs_path=root / "sweep" / "jobs.jsonl", node_name="node-a", max_jobs=1)
self.assertEqual(summary["claimed"], 1)
self.assertEqual(summary["succeeded"], 1)
self.assertEqual(read_jobs(root / "sweep" / "jobs.jsonl")[0]["status"], "succeeded")
def test_rescue_templates_cover_spec_axes_deterministically(self) -> None:
templates = rescue_templates(("siren_conditioned_inr", "film_fourier_inr", "meshgraphnet_or_point_transformer_local"))
keys = [(str(template["scope"]), str(template["template_id"])) for template in templates]
ids = [template["template_id"] for template in templates]
self.assertEqual(keys, sorted(keys))
self.assertIn("optimizer_lr_3e-4", ids)
self.assertIn("siren_omega0_10", ids)
self.assertIn("film_condition_width_2048", ids)
self.assertIn("local_neighbors_32", ids)
def test_expansion_gate_requires_utilization_backlog_stability_and_budget(self) -> None:
pool = PoolConfig(nodes=(), budget_usd=25.0)
good_utilization = {"gpu_util_median": 91.0, "gpu_idle_fraction": 0.04}

59
tests/test_training_config.py
View file

@ -28,6 +28,65 @@ class TrainingConfigTests(unittest.TestCase):
self.assertEqual(config.data.hf_path_prefix, "processed/full")
self.assertTrue(config.data.cache_dir is not None)
def test_config_loader_parses_model_metadata_and_defaults_old_configs(self) -> None:
config = load_training_config("configs/mlp_tiny.toml")
self.assertEqual(config.model_metadata.reported_family, "mlp")
self.assertFalse(config.model_metadata.is_proxy)
with tempfile.TemporaryDirectory() as tmp:
config_path = Path(tmp) / "proxy.toml"
config_path.write_text(
f"""
[run]
name = "proxy"
seed = 0
artifact_dir = "{Path(tmp) / "runs"}"
[data]
root = "{Path(tmp) / "data"}"
train_cases = 1
val_cases = 0
test_cases = 0
points_per_case = 1
batch_size = 1
[model]
type = "raster_fno_unet"
hidden_width = 8
depth = 1
activation = "gelu"
[model_metadata]
requested_family = "raster_fno_unet"
implementation_family = "raster_fno_unet"
reported_family = "raster_fno_unet_proxy"
is_proxy = true
proxy_for = "fno_or_raster_field_model"
proxy_notes = "proxy"
coordinate_encoding_compatibility = "raw_only"
[optim]
lr = 0.001
weight_decay = 0.0
steps = 1
[device]
type = "cpu"
allow_cpu_fallback = false
benchmark_kernels = false
[loss]
type = "normalized_mse"
""".strip()
+ "\n"
)
parsed = load_training_config(config_path)
self.assertEqual(parsed.model_metadata.reported_family, "raster_fno_unet_proxy")
self.assertTrue(parsed.model_metadata.is_proxy)
self.assertEqual(parsed.model_metadata.coordinate_encoding_compatibility, "raw_only")
def test_config_loader_rejects_missing_section(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
config_path = Path(tmp) / "bad.toml"

59
tests/test_training_loop.py
View file

@ -5,6 +5,7 @@ import os
import subprocess
import sys
import tempfile
import time
import unittest
from pathlib import Path
from unittest.mock import patch
@ -142,6 +143,39 @@ class HfArtifactUploaderTests(unittest.TestCase):
self.assertEqual(len(manifest["commits"]), 1)
self.assertEqual(set(manifest["uploaded_paths"]), set(result["uploaded"]))
def test_upload_files_suppresses_large_artifacts_and_uploads_metadata(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
run_dir = Path(tmp)
(run_dir / "metrics.jsonl").write_text("{}\n")
(run_dir / "checkpoint_large.pt").write_text("checkpoint")
uploader = HfArtifactUploader(
enabled=True,
run_dir=run_dir,
repo_id="owner/repo",
repo_type="model",
path_in_repo="runs/model",
)
fake_api = Mock()
fake_api.create_commit.return_value = Mock(oid="abc123", commit_url="https://hf/commit/abc123", pr_url=None)
uploader._api = fake_api
with patch.dict(os.environ, {"AIRFRANS_HF_MAX_FILE_BYTES": "4"}):
result = uploader.upload_files(
("metrics.jsonl", "checkpoint_large.pt"),
commit_message="skip oversized checkpoints",
)
self.assertEqual(result["uploaded"], ["runs/model/metrics.jsonl"])
self.assertEqual(result["suppressed_large_files"], ["runs/model/checkpoint_large.pt"])
fake_api.create_commit.assert_called_once()
self.assertEqual(
[operation.path_in_repo for operation in fake_api.create_commit.call_args.kwargs["operations"]],
["runs/model/metrics.jsonl"],
)
manifest = json.loads((run_dir / "hf_upload_manifest.json").read_text())
self.assertEqual(manifest["suppressed_uploads"][-1]["event"], "file_too_large")
self.assertEqual(manifest["suppressed_uploads"][-1]["repo_path"], "runs/model/checkpoint_large.pt")
class WandbObserverTests(unittest.TestCase):
def test_logs_existing_run_artifacts_with_event_and_step_aliases(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
@ -178,6 +212,21 @@ class WandbObserverTests(unittest.TestCase):
aliases=["latest_checkpoint", "step-5", "latest"],
)
def test_finish_timeout_does_not_block_completed_training(self) -> None:
fake_wandb = Mock()
def slow_finish(*, exit_code: int) -> None:
time.sleep(1.0)
fake_wandb.finish.side_effect = slow_finish
observer = WandbObserver(Mock(), fake_wandb)
started = time.perf_counter()
with patch.dict(os.environ, {"AIRFRANS_WANDB_FINISH_TIMEOUT_SECONDS": "0.01"}):
observer.finish(exit_code=0)
self.assertLess(time.perf_counter() - started, 0.5)
class TrainingLoopTests(unittest.TestCase):
def test_cuda_config_fails_clearly_when_cuda_unavailable(self) -> None:
@ -274,6 +323,13 @@ class TrainingLoopTests(unittest.TestCase):
self.assertIn("estimated_train_flops", final_metrics)
self.assertIn("checkpoint_final_bytes", final_metrics)
self.assertFalse(final_metrics["context_target_values_allowed"])
self.assertEqual(final_metrics["reported_family"], "mlp")
self.assertFalse(final_metrics["is_proxy"])
self.assertEqual(final_metrics["coordinate_encoding_type"], "fixed_fourier")
run_manifest = json.loads((run_dir / "run_manifest.json").read_text())
self.assertEqual(run_manifest["reported_family"], "mlp")
job_manifest = json.loads((run_dir / "job_manifest.json").read_text())
self.assertEqual(job_manifest["model_metadata"]["reported_family"], "mlp")
if torch.cuda.is_available():
self.assertIn("T550", final_metrics["gpu_name"])
@ -356,6 +412,8 @@ class TrainingLoopTests(unittest.TestCase):
run_dir = run_dirs[0]
report = json.loads((run_dir / "failure_report.json").read_text())
self.assertEqual(report["error_type"], "NonFiniteLoss")
self.assertEqual(report["failure_category"], "nonfinite")
self.assertEqual(report["reported_family"], "mlp")
self.assertEqual(report["step"], 1)
checkpoint = torch.load(run_dir / "checkpoint_latest.pt", map_location="cpu", weights_only=False)
self.assertEqual(checkpoint["step"], 0)
@ -388,6 +446,7 @@ class TrainingLoopTests(unittest.TestCase):
run_dir = next(path for path in artifact_dir.iterdir() if path.is_dir())
report = json.loads((run_dir / "failure_report.json").read_text())
self.assertEqual(report["error_type"], "NonFiniteGradient")
self.assertEqual(report["failure_category"], "nonfinite")
self.assertEqual(report["step"], 1)
self.assertFalse((run_dir / "checkpoint_final.pt").exists())
self.assertTrue((run_dir / "artifact_manifest.json").is_file())