From 2071670a4d30680613f229a7129d808cc89b911f Mon Sep 17 00:00:00 2001 From: Zachery Aaron Shores-Chmielewski Date: Tue, 28 Jul 2026 22:04:50 +0400 Subject: [PATCH] stash: post failed large run --- .skyignore | 21 +- configs/remote_aggressive_oom_node.toml | 103 +++ configs/remote_aggressive_oom_node1_700m.toml | 103 +++ configs/remote_aggressive_oom_node2_100m.toml | 103 +++ configs/remote_aggressive_oom_node3_700m.toml | 103 +++ configs/remote_aggressive_oom_probe.toml | 104 +++ configs/remote_hf_smoke.toml | 2 + configs/remote_smoke.toml | 2 + configs/remote_wandb_smoke.toml | 2 + scripts/aggressive_oom_node_wrapper.py | 337 +++++++++ scripts/rsync_sweep_node_artifacts.py | 90 +++ scripts/run_quick_probe_on_cluster.py | 63 ++ scripts/run_sweep_node_direct_ssh.py | 83 ++ scripts/run_sweep_node_on_cluster.py | 75 ++ src/airfrans_frontier/cli.py | 34 + src/airfrans_frontier/model_metadata.py | 34 + src/airfrans_frontier/remote/artifacts.py | 22 +- src/airfrans_frontier/remote/cli.py | 2 + src/airfrans_frontier/remote/config.py | 19 +- src/airfrans_frontier/remote/smoke.py | 194 ++++- src/airfrans_frontier/sweep.py | 478 +++++++++++- src/airfrans_frontier/training/config.py | 38 + src/airfrans_frontier/training/data.py | 39 +- src/airfrans_frontier/training/hf_upload.py | 245 +++++- src/airfrans_frontier/training/loop.py | 706 ++++++++++++------ .../training/observability.py | 29 +- tests/test_remote_run.py | 6 +- tests/test_streaming_data.py | 3 +- tests/test_sweep.py | 87 ++- tests/test_training_config.py | 59 ++ tests/test_training_loop.py | 59 ++ 31 files changed, 2919 insertions(+), 326 deletions(-) create mode 100644 configs/remote_aggressive_oom_node.toml create mode 100644 configs/remote_aggressive_oom_node1_700m.toml create mode 100644 configs/remote_aggressive_oom_node2_100m.toml create mode 100644 configs/remote_aggressive_oom_node3_700m.toml create mode 100644 configs/remote_aggressive_oom_probe.toml create mode 100755 scripts/aggressive_oom_node_wrapper.py create mode 100755 scripts/rsync_sweep_node_artifacts.py create mode 100755 scripts/run_quick_probe_on_cluster.py create mode 100755 scripts/run_sweep_node_direct_ssh.py create mode 100755 scripts/run_sweep_node_on_cluster.py create mode 100644 src/airfrans_frontier/model_metadata.py diff --git a/.skyignore b/.skyignore index d0a5577..b5b3709 100644 --- a/.skyignore +++ b/.skyignore @@ -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/ diff --git a/configs/remote_aggressive_oom_node.toml b/configs/remote_aggressive_oom_node.toml new file mode 100644 index 0000000..6ef939c --- /dev/null +++ b/configs/remote_aggressive_oom_node.toml @@ -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" diff --git a/configs/remote_aggressive_oom_node1_700m.toml b/configs/remote_aggressive_oom_node1_700m.toml new file mode 100644 index 0000000..17d146e --- /dev/null +++ b/configs/remote_aggressive_oom_node1_700m.toml @@ -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" diff --git a/configs/remote_aggressive_oom_node2_100m.toml b/configs/remote_aggressive_oom_node2_100m.toml new file mode 100644 index 0000000..970c5c7 --- /dev/null +++ b/configs/remote_aggressive_oom_node2_100m.toml @@ -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" diff --git a/configs/remote_aggressive_oom_node3_700m.toml b/configs/remote_aggressive_oom_node3_700m.toml new file mode 100644 index 0000000..9b42f09 --- /dev/null +++ b/configs/remote_aggressive_oom_node3_700m.toml @@ -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" diff --git a/configs/remote_aggressive_oom_probe.toml b/configs/remote_aggressive_oom_probe.toml new file mode 100644 index 0000000..801a46e --- /dev/null +++ b/configs/remote_aggressive_oom_probe.toml @@ -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" diff --git a/configs/remote_hf_smoke.toml b/configs/remote_hf_smoke.toml index dfb7070..d829a8d 100644 --- a/configs/remote_hf_smoke.toml +++ b/configs/remote_hf_smoke.toml @@ -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] diff --git a/configs/remote_smoke.toml b/configs/remote_smoke.toml index 40c9373..2d8680e 100644 --- a/configs/remote_smoke.toml +++ b/configs/remote_smoke.toml @@ -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] diff --git a/configs/remote_wandb_smoke.toml b/configs/remote_wandb_smoke.toml index fa4fcfc..fd2ee7a 100644 --- a/configs/remote_wandb_smoke.toml +++ b/configs/remote_wandb_smoke.toml @@ -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] diff --git a/scripts/aggressive_oom_node_wrapper.py b/scripts/aggressive_oom_node_wrapper.py new file mode 100755 index 0000000..04edccf --- /dev/null +++ b/scripts/aggressive_oom_node_wrapper.py @@ -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()) diff --git a/scripts/rsync_sweep_node_artifacts.py b/scripts/rsync_sweep_node_artifacts.py new file mode 100755 index 0000000..d588976 --- /dev/null +++ b/scripts/rsync_sweep_node_artifacts.py @@ -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()) diff --git a/scripts/run_quick_probe_on_cluster.py b/scripts/run_quick_probe_on_cluster.py new file mode 100755 index 0000000..49dabf8 --- /dev/null +++ b/scripts/run_quick_probe_on_cluster.py @@ -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 ', flush=True) + return subprocess.run(command, env=env, check=False).returncode + + +if __name__ == '__main__': + raise SystemExit(main(sys.argv)) diff --git a/scripts/run_sweep_node_direct_ssh.py b/scripts/run_sweep_node_direct_ssh.py new file mode 100755 index 0000000..55aeb5e --- /dev/null +++ b/scripts/run_sweep_node_direct_ssh.py @@ -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= 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()) diff --git a/scripts/run_sweep_node_on_cluster.py b/scripts/run_sweep_node_on_cluster.py new file mode 100755 index 0000000..7da3798 --- /dev/null +++ b/scripts/run_sweep_node_on_cluster.py @@ -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 ', flush=True) + return subprocess.run(command, env=env, check=False).returncode + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/src/airfrans_frontier/cli.py b/src/airfrans_frontier/cli.py index 57dfd6b..20bad88 100644 --- a/src/airfrans_frontier/cli.py +++ b/src/airfrans_frontier/cli.py @@ -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 diff --git a/src/airfrans_frontier/model_metadata.py b/src/airfrans_frontier/model_metadata.py new file mode 100644 index 0000000..1efe632 --- /dev/null +++ b/src/airfrans_frontier/model_metadata.py @@ -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, + } diff --git a/src/airfrans_frontier/remote/artifacts.py b/src/airfrans_frontier/remote/artifacts.py index 48f4ed3..2e108b6 100644 --- a/src/airfrans_frontier/remote/artifacts.py +++ b/src/airfrans_frontier/remote/artifacts.py @@ -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: diff --git a/src/airfrans_frontier/remote/cli.py b/src/airfrans_frontier/remote/cli.py index 0bf94c8..5d0334c 100644 --- a/src/airfrans_frontier/remote/cli.py +++ b/src/airfrans_frontier/remote/cli.py @@ -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", diff --git a/src/airfrans_frontier/remote/config.py b/src/airfrans_frontier/remote/config.py index dfcfd63..0a522bc 100644 --- a/src/airfrans_frontier/remote/config.py +++ b/src/airfrans_frontier/remote/config.py @@ -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"), diff --git a/src/airfrans_frontier/remote/smoke.py b/src/airfrans_frontier/remote/smoke.py index ad9e1d8..4bdb30d 100644 --- a/src/airfrans_frontier/remote/smoke.py +++ b/src/airfrans_frontier/remote/smoke.py @@ -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") diff --git a/src/airfrans_frontier/sweep.py b/src/airfrans_frontier/sweep.py index 557fb6a..9657c13 100644 --- a/src/airfrans_frontier/sweep.py +++ b/src/airfrans_frontier/sweep.py @@ -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]: diff --git a/src/airfrans_frontier/training/config.py b/src/airfrans_frontier/training/config.py index 3c66bd1..064a7fb 100644 --- a/src/airfrans_frontier/training/config.py +++ b/src/airfrans_frontier/training/config.py @@ -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, diff --git a/src/airfrans_frontier/training/data.py b/src/airfrans_frontier/training/data.py index 5a0d4f4..0301fb1 100644 --- a/src/airfrans_frontier/training/data.py +++ b/src/airfrans_frontier/training/data.py @@ -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) diff --git a/src/airfrans_frontier/training/hf_upload.py b/src/airfrans_frontier/training/hf_upload.py index 8690d9f..efac7cb 100644 --- a/src/airfrans_frontier/training/hf_upload.py +++ b/src/airfrans_frontier/training/hf_upload.py @@ -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("/") diff --git a/src/airfrans_frontier/training/loop.py b/src/airfrans_frontier/training/loop.py index db0912d..9ca5cfb 100644 --- a/src/airfrans_frontier/training/loop.py +++ b/src/airfrans_frontier/training/loop.py @@ -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 diff --git a/src/airfrans_frontier/training/observability.py b/src/airfrans_frontier/training/observability.py index 1eeff46..34d5df7 100644 --- a/src/airfrans_frontier/training/observability.py +++ b/src/airfrans_frontier/training/observability.py @@ -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: diff --git a/tests/test_remote_run.py b/tests/test_remote_run.py index 5bd4c86..11bd3d0 100644 --- a/tests/test_remote_run.py +++ b/tests/test_remote_run.py @@ -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, diff --git a/tests/test_streaming_data.py b/tests/test_streaming_data.py index 4404457..89485b0 100644 --- a/tests/test_streaming_data.py +++ b/tests/test_streaming_data.py @@ -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__": diff --git a/tests/test_sweep.py b/tests/test_sweep.py index 0a07a70..4a8aeaa 100644 --- a/tests/test_sweep.py +++ b/tests/test_sweep.py @@ -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} diff --git a/tests/test_training_config.py b/tests/test_training_config.py index 7b8ecc2..c610449 100644 --- a/tests/test_training_config.py +++ b/tests/test_training_config.py @@ -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" diff --git a/tests/test_training_loop.py b/tests/test_training_loop.py index e915de3..350edc9 100644 --- a/tests/test_training_loop.py +++ b/tests/test_training_loop.py @@ -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())