stash: pre-preflight fixes

.gitignore

@@ -6,6 +6,7 @@ artifacts/
 .env
 WANDB_API_KEY
 HF_TOKEN
+sky_logs/
 
 # Python
 .venv/

.skyignore

@@ -9,3 +9,4 @@ __pycache__
 HF_TOKEN
 WANDB_API_KEY
 .env
+sky_logs/

configs/archive/2026-07-27_pre_deployment_cleanup/aggressive_smoke_local.toml

@@ -44,3 +44,10 @@ interval_seconds = 0
 
 [stability]
 max_grad_norm = 1.0
+
+[observability]
+backend = "wandb"
+entity = "zacheryasc-personal"
+project = "airfRANS-model-sweep"
+group = "local_smoke"
+tags = ["airfrans", "local", "aggressive-smoke"]

configs/mlp_tiny.toml

@@ -30,3 +30,10 @@ benchmark_kernels = true
 
 [loss]
 type = "normalized_mse"
+
+[observability]
+backend = "wandb"
+entity = "zacheryasc-personal"
+project = "airfRANS-model-sweep"
+group = "local_smoke"
+tags = ["airfrans", "local", "tiny", "mlp"]

notebooks/explore_raw_subset.ipynb

@@ -1646,7 +1646,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.12"
+   "version": "3.12.12"
   }
  },
  "nbformat": 4,

src/airfrans_frontier/cli.py

@@ -68,6 +68,35 @@ def build_parser() -> argparse.ArgumentParser:
     sanity.add_argument("--families", nargs="*", help="model families to check; defaults to every frontier family")
     sanity.set_defaults(command="model-sanity")
 
+    sweep_generate = subparsers.add_parser("sweep-generate", help="generate aggressive OOM sweep pool.toml, jobs.jsonl, and configs")
+    sweep_generate.add_argument("--output-dir", default="artifacts/aggressive_oom_sweep")
+    sweep_generate.add_argument("--data-root", default="artifacts/data_cache/airfrans_processed/processed/full")
+    sweep_generate.add_argument("--artifact-dir", default="artifacts/runs")
+    sweep_generate.add_argument("--hf-repo-id", default="zacheryasc/airfrans-frontier-checkpoints")
+    sweep_generate.add_argument("--group", default="aggressive_oom_sweep_01")
+    sweep_generate.add_argument("--bands", nargs="*", default=["100m", "700m"])
+    sweep_generate.add_argument("--families", nargs="*")
+    sweep_generate.add_argument("--encodings", nargs="*")
+    sweep_generate.set_defaults(command="sweep-generate")
+
+    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")
+
+    sweep_gate = subparsers.add_parser("sweep-gate", help="evaluate GPU pool expansion gate")
+    sweep_gate.add_argument("--pool", default="artifacts/aggressive_oom_sweep/pool.toml")
+    sweep_gate.add_argument("--utilization", required=True)
+    sweep_gate.add_argument("--backlog", type=int, required=True)
+    sweep_gate.add_argument("--recent-failures", type=int, default=0)
+    sweep_gate.add_argument("--spent-usd", type=float, default=0.0)
+    sweep_gate.add_argument("--reserved-usd", type=float, default=0.0)
+    sweep_gate.add_argument("--next-pool-cost-usd", type=float, required=True)
+    sweep_gate.set_defaults(command="sweep-gate")
+
+    sweep_sample_util = subparsers.add_parser("sweep-sample-util", help="append one nvidia-smi utilization sample as JSONL")
+    sweep_sample_util.add_argument("--output", default="artifacts/aggressive_oom_sweep/utilization.jsonl")
+    sweep_sample_util.set_defaults(command="sweep-sample-util")
+
     return parser
 
 
@@ -214,6 +243,73 @@ def main(argv: list[str] | None = None) -> int:
         print(f"families: {len(result['families'])}")
         return 0
 
+    if args.command == "sweep-generate":
+        from airfrans_frontier.sweep import DEFAULT_ENCODINGS, DEFAULT_FAMILIES, generate_jobs
+
+        try:
+            jobs = generate_jobs(
+                output_dir=resolve_path(args.output_dir),
+                data_root=args.data_root,
+                artifact_dir=args.artifact_dir,
+                hf_repo_id=args.hf_repo_id,
+                group=args.group,
+                bands=tuple(args.bands),
+                families=tuple(args.families) if args.families else DEFAULT_FAMILIES,
+                encodings=tuple(args.encodings) if args.encodings else DEFAULT_ENCODINGS,
+            )
+        except (FileNotFoundError, NotADirectoryError, ValueError, RuntimeError) as exc:
+            print(f"error: {exc}", file=sys.stderr)
+            return 1
+        print(f"output_dir: {resolve_path(args.output_dir)}")
+        print(f"jobs: {len(jobs)}")
+        print(f"jobs_jsonl: {resolve_path(args.output_dir) / 'jobs.jsonl'}")
+        return 0
+
+    if args.command == "sweep-collect":
+        from airfrans_frontier.sweep import collect_job_status
+
+        try:
+            status = collect_job_status(jobs_path=resolve_path(args.jobs))
+        except (FileNotFoundError, NotADirectoryError, ValueError, RuntimeError) as exc:
+            print(f"error: {exc}", file=sys.stderr)
+            return 1
+        print(json.dumps(status, indent=2, sort_keys=True))
+        return 0
+
+    if args.command == "sweep-gate":
+        from airfrans_frontier.sweep import BudgetLedger, can_expand_pool, load_pool_config, utilization_summary
+
+        if args.backlog < 0:
+            print("error: --backlog must be non-negative", file=sys.stderr)
+            return 1
+        if args.recent_failures < 0:
+            print("error: --recent-failures must be non-negative", file=sys.stderr)
+            return 1
+        try:
+            pool = load_pool_config(resolve_path(args.pool))
+            utilization = utilization_summary(resolve_path(args.utilization))
+            ledger = BudgetLedger(budget_usd=pool.budget_usd, spent_usd=args.spent_usd, reserved_usd=args.reserved_usd)
+            allowed, reasons = can_expand_pool(
+                pool=pool,
+                utilization=utilization,
+                backlog=args.backlog,
+                recent_failures=args.recent_failures,
+                ledger=ledger,
+                next_pool_cost_usd=args.next_pool_cost_usd,
+            )
+        except (FileNotFoundError, NotADirectoryError, ValueError, RuntimeError) as exc:
+            print(f"error: {exc}", file=sys.stderr)
+            return 1
+        print(json.dumps({"allowed": allowed, "reasons": reasons, "utilization": utilization}, indent=2, sort_keys=True))
+        return 0 if allowed else 1
+
+    if args.command == "sweep-sample-util":
+        from airfrans_frontier.sweep import append_utilization_sample
+
+        sample = append_utilization_sample(resolve_path(args.output))
+        print(json.dumps(sample, indent=2, sort_keys=True))
+        return 0
+
     parser.error(f"unknown command: {args.command}")
     return 2
 

src/airfrans_frontier/models/__init__.py

@@ -1,5 +1,6 @@
 """Baseline model definitions."""
 
+from airfrans_frontier.models.coordinate_encoding import CoordinateEncoder, CoordinateEncodingSpec
 from airfrans_frontier.models.film import FourierFiLMMLP
 from airfrans_frontier.models.frontier import (
     DeepONetBranchTrunk,
@@ -12,6 +13,8 @@ from airfrans_frontier.models.frontier import (
 from airfrans_frontier.models.mlp import PointwiseMLP
 
 __all__ = [
+    "CoordinateEncoder",
+    "CoordinateEncodingSpec",
     "DeepONetBranchTrunk",
     "FourierFiLMMLP",
     "LocalPointTransformer",

src/airfrans_frontier/models/coordinate_encoding.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from collections.abc import Sequence
+
+import torch
+from torch import nn
+
+
+@dataclass(frozen=True)
+class CoordinateEncodingSpec:
+    type: str = "fixed_fourier"
+    features: tuple[str, ...] = ("x", "y", "sdf")
+    scales: tuple[float, ...] = (1.0, 2.0, 4.0, 8.0, 16.0)
+    levels: int = 8
+    num_features: int = 256
+    sigma: float = 16.0
+    seed: int = 0
+
+
+class CoordinateEncoder(nn.Module):
+    def __init__(self, spec: CoordinateEncodingSpec, *, coordinate_dim: int) -> None:
+        super().__init__()
+        if coordinate_dim <= 0:
+            raise ValueError("coordinate_dim must be positive")
+        self.spec = spec
+        self.coordinate_dim = coordinate_dim
+        if spec.type == "raw":
+            self.output_dim = coordinate_dim
+        elif spec.type == "fixed_fourier":
+            self.register_buffer("scales", torch.tensor(spec.scales, dtype=torch.float32), persistent=False)
+            self.output_dim = coordinate_dim * (1 + 2 * len(spec.scales))
+        elif spec.type == "nerf_multires":
+            if spec.levels < 0:
+                raise ValueError("levels must be non-negative")
+            self.output_dim = coordinate_dim * (1 + 2 * spec.levels)
+        elif spec.type == "random_fourier":
+            if spec.num_features <= 0:
+                raise ValueError("num_features must be positive")
+            generator = torch.Generator(device="cpu")
+            generator.manual_seed(spec.seed)
+            matrix = torch.randn(coordinate_dim, spec.num_features, generator=generator, dtype=torch.float32) * float(spec.sigma)
+            self.register_buffer("random_matrix", matrix, persistent=True)
+            self.output_dim = coordinate_dim + 2 * spec.num_features
+        else:
+            raise ValueError(f"Unsupported coordinate encoding: {spec.type}")
+
+    def forward(self, coordinates: torch.Tensor) -> torch.Tensor:
+        if self.spec.type == "raw":
+            return coordinates
+        if self.spec.type == "fixed_fourier":
+            if self.scales.numel() == 0:
+                return coordinates
+            phases = coordinates.unsqueeze(-1) * self.scales.to(device=coordinates.device, dtype=coordinates.dtype) * math.pi
+            return torch.cat((coordinates, torch.sin(phases).flatten(1), torch.cos(phases).flatten(1)), dim=1)
+        if self.spec.type == "nerf_multires":
+            if self.spec.levels <= 0:
+                return coordinates
+            scales = torch.pow(
+                coordinates.new_tensor(2.0),
+                torch.arange(self.spec.levels, device=coordinates.device, dtype=coordinates.dtype),
+            )
+            phases = coordinates.unsqueeze(-1) * scales * math.pi
+            return torch.cat((coordinates, torch.sin(phases).flatten(1), torch.cos(phases).flatten(1)), dim=1)
+        if self.spec.type == "random_fourier":
+            projection = coordinates @ self.random_matrix.to(device=coordinates.device, dtype=coordinates.dtype)
+            phases = projection * (2.0 * math.pi)
+            return torch.cat((coordinates, torch.sin(phases), torch.cos(phases)), dim=1)
+        raise RuntimeError(f"Unsupported coordinate encoding: {self.spec.type}")
+
+
+def coordinate_encoding_spec(
+    *,
+    encoding_type: str,
+    features: Sequence[str],
+    scales: Sequence[float] = (),
+    levels: int = 0,
+    num_features: int = 256,
+    sigma: float = 16.0,
+    seed: int = 0,
+) -> CoordinateEncodingSpec:
+    return CoordinateEncodingSpec(
+        type=encoding_type,
+        features=tuple(features),
+        scales=tuple(float(scale) for scale in scales),
+        levels=int(levels),
+        num_features=int(num_features),
+        sigma=float(sigma),
+        seed=int(seed),
+    )

src/airfrans_frontier/models/film.py

@@ -6,6 +6,8 @@ from collections.abc import Sequence
 import torch
 from torch import nn
 
+from airfrans_frontier.models.coordinate_encoding import CoordinateEncoder, CoordinateEncodingSpec
+
 
 class FourierFiLMMLP(nn.Module):
     """Coordinate trunk modulated by per-simulation condition features.
@@ -23,6 +25,7 @@ class FourierFiLMMLP(nn.Module):
         output_dim: int,
         coordinate_features: Sequence[str] = ("x", "y", "sdf"),
         fourier_scales: Sequence[float] = (1.0, 2.0, 4.0, 8.0, 16.0),
+        coordinate_encoding_spec: CoordinateEncodingSpec | None = None,
         trunk_width: int = 1024,
         trunk_depth: int = 8,
         condition_width: int = 512,
@@ -60,11 +63,14 @@ class FourierFiLMMLP(nn.Module):
         self.condition_names = condition_names
         self.register_buffer("coordinate_indices", torch.tensor([names.index(name) for name in coordinate_names], dtype=torch.long), persistent=False)
         self.register_buffer("condition_indices", torch.tensor([names.index(name) for name in condition_names], dtype=torch.long), persistent=False)
-        self.register_buffer("fourier_scales", torch.tensor(tuple(float(scale) for scale in fourier_scales), dtype=torch.float32), persistent=False)
-
-        coordinate_dim = len(coordinate_names)
-        fourier_dim = coordinate_dim * (1 + 2 * len(tuple(fourier_scales)))
-        self.input = nn.Linear(fourier_dim, trunk_width)
+        spec = coordinate_encoding_spec or CoordinateEncodingSpec(
+            type="fixed_fourier",
+            features=coordinate_names,
+            scales=tuple(float(scale) for scale in fourier_scales),
+            levels=0,
+        )
+        self.coordinate_encoder = CoordinateEncoder(spec, coordinate_dim=len(coordinate_names))
+        self.input = nn.Linear(self.coordinate_encoder.output_dim, trunk_width)
         self.condition_encoder = _mlp(
             input_dim=len(condition_names),
             hidden_width=condition_width,
@@ -83,7 +89,7 @@ class FourierFiLMMLP(nn.Module):
     def forward(self, features: torch.Tensor) -> torch.Tensor:
         coordinates = features.index_select(dim=1, index=self.coordinate_indices)
         conditions = features.index_select(dim=1, index=self.condition_indices)
-        coordinate_embedding = _fourier_features(coordinates, self.fourier_scales)
+        coordinate_embedding = self.coordinate_encoder(coordinates)
         hidden = self.input(coordinate_embedding)
 
         unique_conditions, inverse = torch.unique(conditions, dim=0, return_inverse=True)

src/airfrans_frontier/models/frontier.py

@@ -7,6 +7,8 @@ import torch
 from torch import nn
 from torch.nn import functional as F
 
+from airfrans_frontier.models.coordinate_encoding import CoordinateEncoder, CoordinateEncodingSpec
+
 
 class NeRFCFDMultiRes(nn.Module):
     def __init__(
@@ -21,14 +23,21 @@ class NeRFCFDMultiRes(nn.Module):
         condition_width: int,
         condition_depth: int,
         activation: str,
+        coordinate_encoding_spec: CoordinateEncodingSpec | None = None,
     ) -> None:
         super().__init__()
         coordinate_indices = _indices(feature_names, coordinate_features)
         condition_indices = _complement_indices(len(feature_names), coordinate_indices)
         self.register_buffer("coordinate_indices", torch.tensor(coordinate_indices, dtype=torch.long), persistent=False)
         self.register_buffer("condition_indices", torch.tensor(condition_indices, dtype=torch.long), persistent=False)
-        self.encoding_levels = int(encoding_levels)
-        encoded_dim = len(coordinate_indices) * (1 + 2 * self.encoding_levels)
+        spec = coordinate_encoding_spec or CoordinateEncodingSpec(
+            type="nerf_multires",
+            features=tuple(coordinate_features),
+            levels=int(encoding_levels),
+            scales=(),
+        )
+        self.coordinate_encoder = CoordinateEncoder(spec, coordinate_dim=len(coordinate_indices))
+        encoded_dim = self.coordinate_encoder.output_dim
         condition_input_dim = len(condition_indices) if condition_indices else 1
         self.condition_encoder = _mlp(
             input_dim=condition_input_dim,
@@ -49,7 +58,7 @@ class NeRFCFDMultiRes(nn.Module):
     def forward(self, features: torch.Tensor) -> torch.Tensor:
         coordinates = features.index_select(dim=1, index=self.coordinate_indices)
         condition = _gather_or_zeros(features, self.condition_indices)
-        encoded = _multires_encode(coordinates, self.encoding_levels)
+        encoded = self.coordinate_encoder(coordinates)
         condition_embedding = self.condition_encoder(condition)
         return self.decoder(torch.cat((encoded, condition_embedding), dim=1))
 
@@ -67,14 +76,20 @@ class DeepONetBranchTrunk(nn.Module):
         condition_width: int,
         condition_depth: int,
         activation: str,
+        coordinate_encoding_spec: CoordinateEncodingSpec | None = None,
     ) -> None:
         super().__init__()
         coordinate_indices = _indices(feature_names, coordinate_features)
         condition_indices = _complement_indices(len(feature_names), coordinate_indices)
         self.register_buffer("coordinate_indices", torch.tensor(coordinate_indices, dtype=torch.long), persistent=False)
         self.register_buffer("condition_indices", torch.tensor(condition_indices, dtype=torch.long), persistent=False)
-        self.register_buffer("fourier_scales", torch.tensor(tuple(float(scale) for scale in fourier_scales), dtype=torch.float32), persistent=False)
-        trunk_input_dim = len(coordinate_indices) * (1 + 2 * len(fourier_scales))
+        spec = coordinate_encoding_spec or CoordinateEncodingSpec(
+            type="fixed_fourier",
+            features=tuple(coordinate_features),
+            scales=tuple(float(scale) for scale in fourier_scales),
+            levels=0,
+        )
+        self.coordinate_encoder = CoordinateEncoder(spec, coordinate_dim=len(coordinate_indices))
         condition_input_dim = len(condition_indices) if condition_indices else 1
         self.branch = _mlp(
             input_dim=condition_input_dim,
@@ -84,7 +99,7 @@ class DeepONetBranchTrunk(nn.Module):
             activation=activation,
         )
         self.trunk = _mlp(
-            input_dim=trunk_input_dim,
+            input_dim=self.coordinate_encoder.output_dim,
             hidden_width=hidden_width,
             output_dim=hidden_width,
             depth=depth,
@@ -95,7 +110,7 @@ class DeepONetBranchTrunk(nn.Module):
     def forward(self, features: torch.Tensor) -> torch.Tensor:
         coordinates = features.index_select(dim=1, index=self.coordinate_indices)
         condition = _gather_or_zeros(features, self.condition_indices)
-        trunk = self.trunk(_fourier_features(coordinates, self.fourier_scales))
+        trunk = self.trunk(self.coordinate_encoder(coordinates))
         branch = self.branch(condition)
         return self.head(trunk * branch)
 

src/airfrans_frontier/sweep.py

@@ -0,0 +1,447 @@
+from __future__ import annotations
+
+import json
+import math
+import os
+import statistics
+import subprocess
+import sys
+import time
+import tomllib
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Iterable, Mapping
+
+EXPANSION_LADDER = (1, 2, 4, 8)
+REQUIRED_JOB_ARTIFACTS = (
+    "config.toml",
+    "job_manifest.json",
+    "metrics.jsonl",
+    "latest_metrics.json",
+    "run_manifest.json",
+    "environment_manifest.json",
+    "utilization.jsonl",
+)
+SUCCESS_ARTIFACTS = REQUIRED_JOB_ARTIFACTS + ("final_metrics.json",)
+FAILURE_ARTIFACTS = REQUIRED_JOB_ARTIFACTS + ("failure_report.json",)
+
+OOM_BANDS: dict[str, dict[str, int]] = {
+    "10m": {"hidden_width": 1024, "depth": 8, "steps": 300, "batch_size": 4096},
+    "100m": {"hidden_width": 2048, "depth": 16, "steps": 1500, "batch_size": 4096},
+    "700m": {"hidden_width": 4096, "depth": 40, "steps": 500, "batch_size": 2048},
+}
+FILM_DEPTH_OVERRIDES = {"700m": 32}
+DEFAULT_FAMILIES = (
+    "mlp",
+    "nerf_cfd_multires",
+    "deeponet_branch_trunk",
+    "film_fourier_inr",
+    "siren_conditioned_inr",
+    "raster_fno_unet",
+    "point_context_perceiver",
+    "meshgraphnet_or_point_transformer_local",
+)
+DEFAULT_ENCODINGS = (
+    "raw",
+    "fixed_fourier",
+    "nerf_multires",
+    "random_fourier",
+)
+ENCODING_COMPATIBLE_FAMILIES = {
+    "nerf_cfd_multires",
+    "deeponet_branch_trunk",
+    "film_fourier_mlp",
+    "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",
+}
+
+
+@dataclass(frozen=True)
+class PoolNode:
+    name: str
+    gpus: int = 1
+    enabled: bool = True
+    max_parallel_jobs: int = 1
+
+
+@dataclass(frozen=True)
+class PoolConfig:
+    nodes: tuple[PoolNode, ...]
+    budget_usd: float = 25.0
+    expansion_ladder: tuple[int, ...] = EXPANSION_LADDER
+    min_gpu_util_median: float = 85.0
+    max_gpu_idle_fraction: float = 0.10
+
+
+@dataclass(frozen=True)
+class BudgetLedger:
+    budget_usd: float
+    spent_usd: float
+    reserved_usd: float = 0.0
+
+    @property
+    def remaining_usd(self) -> float:
+        return max(0.0, self.budget_usd - self.spent_usd - self.reserved_usd)
+
+
+def load_pool_config(path: str | Path) -> PoolConfig:
+    raw = tomllib.loads(Path(path).read_text())
+    pool_raw = raw.get("pool", {})
+    nodes_raw = raw.get("nodes", [])
+    if not isinstance(nodes_raw, list) or not nodes_raw:
+        raise ValueError("pool.toml must contain at least one [[nodes]] entry")
+    nodes = tuple(
+        PoolNode(
+            name=str(node["name"]),
+            gpus=int(node.get("gpus", 1)),
+            enabled=bool(node.get("enabled", True)),
+            max_parallel_jobs=int(node.get("max_parallel_jobs", 1)),
+        )
+        for node in nodes_raw
+    )
+    ladder = tuple(int(value) for value in pool_raw.get("expansion_ladder", EXPANSION_LADDER))
+    return PoolConfig(
+        nodes=nodes,
+        budget_usd=float(pool_raw.get("budget_usd", 25.0)),
+        expansion_ladder=ladder,
+        min_gpu_util_median=float(pool_raw.get("min_gpu_util_median", 85.0)),
+        max_gpu_idle_fraction=float(pool_raw.get("max_gpu_idle_fraction", 0.10)),
+    )
+
+
+def write_default_pool(path: str | Path) -> Path:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    target.write_text(
+        "[pool]\n"
+        "budget_usd = 25.0\n"
+        "expansion_ladder = [1, 2, 4, 8]\n"
+        "min_gpu_util_median = 85.0\n"
+        "max_gpu_idle_fraction = 0.10\n\n"
+        "[[nodes]]\n"
+        "name = \"node-0\"\n"
+        "gpus = 1\n"
+        "enabled = true\n"
+        "max_parallel_jobs = 1\n"
+    )
+    return target
+
+
+def generate_jobs(
+    *,
+    output_dir: str | Path,
+    data_root: str,
+    artifact_dir: str = "artifacts/runs",
+    hf_repo_id: str = "zacheryasc/airfrans-frontier-checkpoints",
+    group: str = "aggressive_oom_sweep_01",
+    bands: Iterable[str] = ("100m", "700m"),
+    families: Iterable[str] = DEFAULT_FAMILIES,
+    encodings: Iterable[str] = DEFAULT_ENCODINGS,
+    seed: int = 20260727,
+) -> list[dict[str, Any]]:
+    root = Path(output_dir)
+    config_dir = root / "configs"
+    config_dir.mkdir(parents=True, exist_ok=True)
+    jobs: list[dict[str, Any]] = []
+    for band in bands:
+        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}")
+                job_id = f"{band}_{family}_{encoding}"
+                run_name = job_id
+                config_path = config_dir / f"{job_id}.toml"
+                config_text = training_config_text(
+                    run_name=run_name,
+                    seed=seed + len(jobs),
+                    data_root=data_root,
+                    artifact_dir=artifact_dir,
+                    model_family=family,
+                    band=band,
+                    coordinate_encoding=encoding,
+                    hf_repo_id=hf_repo_id,
+                    group=group,
+                )
+                config_path.write_text(config_text)
+                jobs.append(
+                    {
+                        "job_id": job_id,
+                        "status": "pending",
+                        "band": band,
+                        "model_family": family,
+                        "reported_family": PROXY_FAMILY_LABELS.get(family, family),
+                        "coordinate_encoding": encoding,
+                        "config_path": str(config_path),
+                        "attempts": 0,
+                    }
+                )
+    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")
+    return jobs
+
+
+def training_config_text(
+    *,
+    run_name: str,
+    seed: int,
+    data_root: str,
+    artifact_dir: str,
+    model_family: str,
+    band: str,
+    coordinate_encoding: str,
+    hf_repo_id: str,
+    group: str,
+) -> str:
+    shape = dict(OOM_BANDS[band])
+    if model_family == "film_fourier_inr" and band in FILM_DEPTH_OVERRIDES:
+        shape["depth"] = FILM_DEPTH_OVERRIDES[band]
+    condition_width = max(512, shape["hidden_width"] // 2)
+    condition_depth = 4 if band in {"100m", "700m"} else 2
+    condition_dim = min(shape["hidden_width"], 1024)
+    return "\n".join(
+        line for line in (
+            "[run]",
+            f"name = {json.dumps(run_name)}",
+            f"seed = {seed}",
+            f"artifact_dir = {json.dumps(artifact_dir)}",
+            "",
+            "[data]",
+            f"root = {json.dumps(data_root)}",
+            "source = \"huggingface\"",
+            "hf_repo_id = \"zacheryasc/airfrans-processed\"",
+            "hf_repo_type = \"dataset\"",
+            "hf_path_prefix = \"processed/full\"",
+            "cache_dir = \"artifacts/data_cache/airfrans_processed\"",
+            "train_cases = 900",
+            "val_cases = 50",
+            "test_cases = 50",
+            "points_per_case = 8192",
+            f"batch_size = {shape['batch_size']}",
+            "",
+            "[coordinate_encoding]",
+            _coordinate_encoding_toml(coordinate_encoding),
+            "",
+            "[model]",
+            f"type = {json.dumps(model_family)}",
+            f"hidden_width = {shape['hidden_width']}",
+            f"depth = {shape['depth']}",
+            "activation = \"gelu\"",
+            "coordinate_features = [\"x\", \"y\", \"sdf\"]",
+            f"condition_width = {condition_width}",
+            f"condition_depth = {condition_depth}",
+            f"condition_dim = {condition_dim}",
+            "context_points = 1024",
+            "latent_width = 1024",
+            "attention_depth = 4",
+            "neighbors = 16",
+            "grid_resolution = 96",
+            "siren_omega0 = 30.0",
+            "",
+            "[optim]",
+            "lr = 0.0001",
+            "weight_decay = 0.0001",
+            f"steps = {shape['steps']}",
+            "log_interval = 25",
+            "",
+            "[device]",
+            "type = \"cuda\"",
+            "allow_cpu_fallback = false",
+            "benchmark_kernels = true",
+            "",
+            "[loss]",
+            "type = \"normalized_mse\"",
+            "",
+            "[precision]",
+            "dtype = \"bf16\"",
+            "",
+            "[checkpoint]",
+            "interval_seconds = 900",
+            "",
+            "[stability]",
+            "max_grad_norm = 1.0",
+            "",
+            "[observability]",
+            "backend = \"wandb\"",
+            "entity = \"zacheryasc-personal\"",
+            "project = \"airfRANS-model-sweep\"",
+            f"group = {json.dumps(group)}",
+            f"tags = [\"airfrans\", \"aggressive-oom-sweep\", {json.dumps(band)}, {json.dumps(model_family)}, {json.dumps(coordinate_encoding)}]",
+            "",
+            "[huggingface]",
+            "enabled = true",
+            f"repo_id = {json.dumps(hf_repo_id)}",
+            "repo_type = \"model\"",
+            f"path_prefix = {json.dumps(group)}",
+            "private = false",
+        )
+    ) + "\n"
+
+
+def _coordinate_encoding_toml(encoding: str) -> str:
+    if encoding == "raw":
+        return "type = \"raw\"\nfeatures = [\"x\", \"y\", \"sdf\"]"
+    if encoding == "fixed_fourier":
+        return "type = \"fixed_fourier\"\nfeatures = [\"x\", \"y\", \"sdf\"]\nscales = [1, 2, 4, 8, 16, 32]"
+    if encoding == "nerf_multires":
+        return "type = \"nerf_multires\"\nfeatures = [\"x\", \"y\", \"sdf\"]\nlevels = 16"
+    if encoding == "random_fourier":
+        return "type = \"random_fourier\"\nfeatures = [\"x\", \"y\", \"sdf\"]\nnum_features = 256\nsigma = 16.0\nseed = 0"
+    raise ValueError(f"unsupported coordinate encoding: {encoding}")
+
+
+def read_jobs(path: str | Path) -> list[dict[str, Any]]:
+    jobs: list[dict[str, Any]] = []
+    for line in Path(path).read_text().splitlines():
+        if line.strip():
+            jobs.append(json.loads(line))
+    return jobs
+
+
+def write_jobs(path: str | Path, jobs: Iterable[Mapping[str, Any]]) -> None:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    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)
+
+    env = os.environ.copy()
+    if extra_env:
+        env.update(extra_env)
+    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
+
+
+def collect_job_status(*, jobs_path: str | Path) -> dict[str, Any]:
+    jobs = read_jobs(jobs_path)
+    records = []
+    counts: dict[str, int] = {}
+    for job in jobs:
+        status = reconstruct_status(job)
+        counts[status["status"]] = counts.get(status["status"], 0) + 1
+        records.append(status)
+    return {"jobs": records, "counts": counts, "total": len(records)}
+
+
+def reconstruct_status(job: Mapping[str, Any]) -> dict[str, Any]:
+    config_path = Path(str(job["config_path"]))
+    run_name = _read_run_name(config_path)
+    artifact_root = _read_artifact_dir(config_path)
+    candidates = sorted(artifact_root.glob(f"*_{_safe_name(run_name)}*")) if artifact_root.exists() else []
+    run_dir = candidates[-1] if candidates else None
+    if run_dir is None:
+        return {"job_id": job["job_id"], "status": str(job.get("status", "pending")), "run_dir": None, "missing_artifacts": list(REQUIRED_JOB_ARTIFACTS)}
+    success_missing = [name for name in SUCCESS_ARTIFACTS if not (run_dir / name).is_file()]
+    failure_missing = [name for name in FAILURE_ARTIFACTS if not (run_dir / name).is_file()]
+    if not success_missing:
+        status = "succeeded"
+        missing = []
+    elif not failure_missing:
+        status = "failed"
+        missing = []
+    else:
+        status = "incomplete"
+        missing = success_missing
+    return {"job_id": job["job_id"], "status": status, "run_dir": str(run_dir), "missing_artifacts": missing}
+
+
+def utilization_summary(path: str | Path) -> dict[str, Any]:
+    samples = []
+    for line in Path(path).read_text().splitlines():
+        if line.strip():
+            samples.append(json.loads(line))
+    gpu_utils = [float(sample["gpu_util_percent"]) for sample in samples if sample.get("gpu_util_percent") is not None]
+    memory_values = [float(sample["memory_used_mb"]) for sample in samples if sample.get("memory_used_mb") is not None]
+    if not samples or not gpu_utils:
+        return {"samples": len(samples), "gpu_util_median": None, "gpu_idle_fraction": None, "memory_used_peak": max(memory_values) if memory_values else None}
+    return {
+        "samples": len(samples),
+        "gpu_util_median": statistics.median(gpu_utils),
+        "gpu_idle_fraction": sum(1 for value in gpu_utils if value <= 5.0) / len(gpu_utils),
+        "memory_used_peak": max(memory_values) if memory_values else None,
+    }
+
+
+def can_expand_pool(*, pool: PoolConfig, utilization: Mapping[str, Any], backlog: int, recent_failures: int, ledger: BudgetLedger, next_pool_cost_usd: float) -> tuple[bool, tuple[str, ...]]:
+    reasons: list[str] = []
+    if utilization.get("gpu_util_median") is None or float(utilization["gpu_util_median"]) < pool.min_gpu_util_median:
+        reasons.append("median GPU utilization below gate")
+    if utilization.get("gpu_idle_fraction") is None or float(utilization["gpu_idle_fraction"]) > pool.max_gpu_idle_fraction:
+        reasons.append("GPU idle fraction above gate")
+    if backlog <= 0:
+        reasons.append("no job backlog")
+    if recent_failures > 0:
+        reasons.append("recent failures are not isolated")
+    if ledger.remaining_usd < next_pool_cost_usd:
+        reasons.append("remaining budget does not support larger pool")
+    return (not reasons, tuple(reasons))
+
+
+def sample_gpu_utilization() -> dict[str, Any]:
+    command = ["nvidia-smi", "--query-gpu=utilization.gpu,memory.used", "--format=csv,noheader,nounits"]
+    try:
+        completed = subprocess.run(command, text=True, capture_output=True, check=True)
+    except (FileNotFoundError, subprocess.CalledProcessError) as exc:
+        return {"timestamp": time.time(), "gpu_util_percent": None, "memory_used_mb": None, "error": str(exc)}
+    rows = [row.strip().split(",") for row in completed.stdout.splitlines() if row.strip()]
+    utils = [float(row[0].strip()) for row in rows]
+    memory = [float(row[1].strip()) for row in rows]
+    return {
+        "timestamp": time.time(),
+        "gpu_util_percent": statistics.mean(utils) if utils else None,
+        "memory_used_mb": max(memory) if memory else None,
+    }
+
+
+def append_utilization_sample(path: str | Path) -> dict[str, Any]:
+    sample = sample_gpu_utilization()
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    with target.open("a") as file:
+        file.write(json.dumps(sample, sort_keys=True) + "\n")
+    return sample
+
+
+def _read_run_name(config_path: Path) -> str:
+    raw = tomllib.loads(config_path.read_text())
+    return str(raw["run"]["name"])
+
+
+def _read_artifact_dir(config_path: Path) -> Path:
+    raw = tomllib.loads(config_path.read_text())
+    artifact_dir = Path(str(raw["run"]["artifact_dir"])).expanduser()
+    return artifact_dir if artifact_dir.is_absolute() else Path.cwd() / artifact_dir
+
+
+def _safe_name(value: str) -> str:
+    safe = "".join(character if character.isalnum() or character in "-_" else "_" for character in value)
+    return safe or "run"

src/airfrans_frontier/training/artifacts.py

@@ -62,6 +62,9 @@ class ArtifactWriter:
     def write_json(self, name: str, data: Mapping[str, Any]) -> None:
         self._write_json_file(name, dict(data))
 
+    def append_jsonl(self, name: str, data: Mapping[str, Any]) -> None:
+        self._append_text(name, json.dumps(dict(data), sort_keys=True) + "\n")
+
     def write_failure_report(self, report: Mapping[str, Any]) -> None:
         self.write_json("failure_report.json", dict(report))
 

src/airfrans_frontier/training/config.py

@@ -70,6 +70,17 @@ class ModelConfig:
     siren_omega0: float
 
 
+@dataclass(frozen=True)
+class CoordinateEncodingConfig:
+    type: str
+    features: tuple[str, ...]
+    scales: tuple[float, ...]
+    levels: int
+    num_features: int
+    sigma: float
+    seed: int
+
+
 @dataclass(frozen=True)
 class OptimConfig:
     lr: float
@@ -139,6 +150,7 @@ class TrainingConfig:
     precision: PrecisionConfig
     observability: ObservabilityConfig
     huggingface: HuggingFaceConfig
+    coordinate_encoding: CoordinateEncodingConfig
 
 _REQUIRED_SECTIONS = ("run", "data", "model", "optim", "device", "loss")
 
@@ -191,6 +203,39 @@ def load_training_config(path: str | Path) -> TrainingConfig:
         huggingface_raw = {}
     if not isinstance(huggingface_raw, dict):
         raise ValueError("Training config [huggingface] section must be a table")
+    coordinate_encoding_raw = raw.get("coordinate_encoding", {})
+    if coordinate_encoding_raw is None:
+        coordinate_encoding_raw = {}
+    if not isinstance(coordinate_encoding_raw, dict):
+        raise ValueError("Training config [coordinate_encoding] section must be a table")
+
+
+    coordinate_encoding = CoordinateEncodingConfig(
+        type=_choice(
+            _string(coordinate_encoding_raw, "type", default="fixed_fourier").lower(),
+            {"raw", "fixed_fourier", "nerf_multires", "random_fourier"},
+            "coordinate_encoding.type",
+        ),
+        features=_string_tuple(
+            coordinate_encoding_raw,
+            "features",
+            default=_string_tuple(model_raw, "coordinate_features", default=("x", "y", "sdf")),
+        ),
+        scales=_number_tuple(
+            coordinate_encoding_raw,
+            "scales",
+            default=_number_tuple(model_raw, "fourier_scales", default=(1.0, 2.0, 4.0, 8.0, 16.0)),
+        ),
+        levels=_integer(coordinate_encoding_raw, "levels", minimum=0, default=_integer(model_raw, "encoding_levels", minimum=0, default=8)),
+        num_features=_integer(coordinate_encoding_raw, "num_features", minimum=1, default=_integer(model_raw, "features_per_level", minimum=1, default=2) * 128),
+        sigma=_number(coordinate_encoding_raw, "sigma", minimum=0.0, exclusive_minimum=True, default=16.0),
+        seed=_integer(coordinate_encoding_raw, "seed", minimum=0, default=0),
+    )
+
+
+    model_coordinate_features = coordinate_encoding.features
+    model_fourier_scales = coordinate_encoding.scales if coordinate_encoding.type == "fixed_fourier" else ()
+    model_encoding_levels = coordinate_encoding.levels if coordinate_encoding.type == "nerf_multires" else 0
 
 
 
@@ -244,13 +289,13 @@ def load_training_config(path: str | Path) -> TrainingConfig:
         hidden_width=_integer(model_raw, "hidden_width", minimum=1),
         depth=_integer(model_raw, "depth", minimum=1),
         activation=_choice(_string(model_raw, "activation").lower(), {"gelu", "relu", "silu", "tanh"}, "model.activation"),
-        coordinate_features=_string_tuple(model_raw, "coordinate_features", default=("x", "y", "sdf")),
-        fourier_scales=_number_tuple(model_raw, "fourier_scales", default=(1.0, 2.0, 4.0, 8.0, 16.0)),
+        coordinate_features=model_coordinate_features,
+        fourier_scales=model_fourier_scales,
         condition_width=_integer(model_raw, "condition_width", minimum=1, default=_integer(model_raw, "hidden_width", minimum=1)),
         condition_depth=_integer(model_raw, "condition_depth", minimum=1, default=2),
         condition_dim=_integer(model_raw, "condition_dim", minimum=1, default=_integer(model_raw, "hidden_width", minimum=1)),
-        encoding_levels=_integer(model_raw, "encoding_levels", minimum=0, default=8),
-        features_per_level=_integer(model_raw, "features_per_level", minimum=1, default=2),
+        encoding_levels=model_encoding_levels,
+        features_per_level=max(1, coordinate_encoding.num_features // 128),
         context_points=_integer(model_raw, "context_points", minimum=1, default=512),
         latent_width=_integer(model_raw, "latent_width", minimum=1, default=_integer(model_raw, "hidden_width", minimum=1)),
         attention_depth=_integer(model_raw, "attention_depth", minimum=1, default=2),
@@ -314,6 +359,7 @@ def load_training_config(path: str | Path) -> TrainingConfig:
         precision=precision,
         observability=observability,
         huggingface=huggingface,
+        coordinate_encoding=coordinate_encoding,
     )
 
 

src/airfrans_frontier/training/loop.py

@@ -4,6 +4,7 @@ import hashlib
 import json
 import os
 import random
+import subprocess
 import time
 from dataclasses import asdict, dataclass
 from pathlib import Path
@@ -23,6 +24,7 @@ from airfrans_frontier.models import (
     RasterFNOUNet,
     SirenConditionedINR,
 )
+from airfrans_frontier.models.coordinate_encoding import coordinate_encoding_spec
 from airfrans_frontier.remote.artifacts import verify_artifacts
 from airfrans_frontier.training.artifacts import ArtifactWriter
 from airfrans_frontier.training.calibration import checkpoint_size_bytes, static_calibration_fields
@@ -88,6 +90,17 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
         hf_path_in_repo=uploader.path_in_repo,
     )
     writer.write_json("run_manifest.json", run_manifest)
+    writer.write_json("job_manifest.json", _job_manifest(config=config, run_id=run_id, run_dir=writer.run_dir))
+    writer.append_jsonl(
+        "utilization.jsonl",
+        {
+            "timestamp": time.time(),
+            "phase": "starting",
+            "gpu_util_percent": None,
+            "gpu_idle": None,
+            **_memory_metrics(device),
+        },
+    )
     observer.update_config(
         {
             "run_id": run_id,
@@ -120,6 +133,20 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
     def record_metrics(metrics: dict[str, Any]) -> None:
         nonlocal first_metric_timeline_written
         writer.append_metrics(metrics)
+        writer.append_jsonl(
+            "utilization.jsonl",
+            {
+                "timestamp": time.time(),
+                "phase": metrics.get("phase"),
+                "step": metrics.get("step"),
+                "gpu_util_percent": metrics.get("gpu_util_percent"),
+                "gpu_idle": metrics.get("gpu_idle"),
+                "memory_used_mb": metrics.get("gpu_memory_used_mb"),
+                "gpu_memory_allocated_mb": metrics.get("gpu_memory_allocated_mb"),
+                "gpu_memory_reserved_mb": metrics.get("gpu_memory_reserved_mb"),
+                "gpu_memory_peak_allocated_mb": metrics.get("gpu_memory_peak_allocated_mb"),
+            },
+        )
         observer.log(metrics)
         if not first_metric_timeline_written:
             record_timeline("training", "first_metric", step=metrics.get("step"), metric_event=metrics.get("event"))
@@ -127,6 +154,13 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
 
     def publish_artifacts(names: tuple[str, ...], *, event: str, step: int) -> None:
         nonlocal first_checkpoint_upload_timeline_written
+        observer.log_artifact_files(
+            writer.run_dir,
+            names,
+            event=event,
+            step=step,
+            aliases=(event, f"step-{step}", "latest"),
+        )
         if not config.huggingface.enabled:
             return
         try:
@@ -253,6 +287,8 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
     publish_artifacts(
         (
             "config.toml",
+            "job_manifest.json",
+            "utilization.jsonl",
             "environment_manifest.json",
             "split_manifest.json",
             "data_manifest.json",
@@ -651,6 +687,7 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
         **resume_info,
         **calibration_fields,
         **protocol_fields,
+        **_nvidia_smi_metrics(device),
         **device_metrics(device),
     }
     if observer.url is not None:
@@ -1201,6 +1238,7 @@ def _train_public_zip_streaming(
             **resume_info,
             **calibration_fields,
             **protocol_fields,
+            **_nvidia_smi_metrics(device),
             **device_metrics(device),
             **_streaming_metric_fields(streaming_summary),
         }
@@ -1365,6 +1403,15 @@ def _autocast_context(config: TrainingConfig, device: torch.device):
 def _build_model(config: TrainingConfig, bundle: DatasetBundle, *, output_dim: int) -> torch.nn.Module:
     input_dim = bundle.train.features.shape[1]
     model_type = config.model.type
+    coord_spec = coordinate_encoding_spec(
+        encoding_type=config.coordinate_encoding.type,
+        features=config.coordinate_encoding.features,
+        scales=config.coordinate_encoding.scales,
+        levels=config.coordinate_encoding.levels,
+        num_features=config.coordinate_encoding.num_features,
+        sigma=config.coordinate_encoding.sigma,
+        seed=config.coordinate_encoding.seed,
+    )
     if model_type == "mlp":
         return PointwiseMLP(
             input_dim=input_dim,
@@ -1379,6 +1426,7 @@ def _build_model(config: TrainingConfig, bundle: DatasetBundle, *, output_dim: i
             output_dim=output_dim,
             coordinate_features=config.model.coordinate_features,
             fourier_scales=config.model.fourier_scales,
+            coordinate_encoding_spec=coord_spec,
             trunk_width=config.model.hidden_width,
             trunk_depth=config.model.depth,
             condition_width=config.model.condition_width,
@@ -1392,6 +1440,7 @@ def _build_model(config: TrainingConfig, bundle: DatasetBundle, *, output_dim: i
             output_dim=output_dim,
             coordinate_features=config.model.coordinate_features,
             encoding_levels=config.model.encoding_levels,
+            coordinate_encoding_spec=coord_spec,
             hidden_width=config.model.hidden_width,
             depth=config.model.depth,
             condition_width=config.model.condition_width,
@@ -1404,6 +1453,7 @@ def _build_model(config: TrainingConfig, bundle: DatasetBundle, *, output_dim: i
             output_dim=output_dim,
             coordinate_features=config.model.coordinate_features,
             fourier_scales=config.model.fourier_scales,
+            coordinate_encoding_spec=coord_spec,
             hidden_width=config.model.hidden_width,
             depth=config.model.depth,
             condition_width=config.model.condition_width,
@@ -1583,10 +1633,49 @@ def _log_metrics(
         "elapsed_seconds": elapsed_seconds,
         "points_per_sec": points_per_sec,
         "latest_checkpoint": latest_checkpoint,
+        **_nvidia_smi_metrics(device),
         **_memory_metrics(device),
     }
 
 
+def _job_manifest(*, config: TrainingConfig, run_id: str, run_dir: Path) -> dict[str, Any]:
+    return {
+        "run_id": run_id,
+        "run_name": config.run.name,
+        "run_dir": str(run_dir),
+        "config_path": str(config.path),
+        "model_family": config.model.type,
+        "coordinate_encoding": {
+            "type": config.coordinate_encoding.type,
+            "features": list(config.coordinate_encoding.features),
+            "scales": list(config.coordinate_encoding.scales),
+            "levels": config.coordinate_encoding.levels,
+            "num_features": config.coordinate_encoding.num_features,
+            "sigma": config.coordinate_encoding.sigma,
+            "seed": config.coordinate_encoding.seed,
+        },
+        "checkpoint_policy": {
+            "interval_seconds": config.checkpoint.interval_seconds,
+            "latest": LATEST_CHECKPOINT,
+            "best": BEST_CHECKPOINT,
+            "final": FINAL_CHECKPOINT,
+        },
+        "observability": {
+            "backend": config.observability.backend,
+            "project": config.observability.project,
+            "group": config.observability.group,
+            "tags": list(config.observability.tags),
+        },
+        "huggingface": {
+            "enabled": config.huggingface.enabled,
+            "repo_id": config.huggingface.repo_id,
+            "repo_type": config.huggingface.repo_type,
+            "path_prefix": config.huggingface.path_prefix,
+            "private": config.huggingface.private,
+        },
+    }
+
+
 def _run_manifest(
     *,
     config: TrainingConfig,
@@ -1603,6 +1692,8 @@ def _run_manifest(
         "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),
         "phase": phase,
         "started_at": started_at,
         "artifact_dir": str(run_dir),
@@ -1658,6 +1749,8 @@ def _evaluation_protocol(model_type: str) -> dict[str, Any]:
 def _verification_required(*, success: bool) -> tuple[str, ...]:
     required = [
         "config.toml",
+        "job_manifest.json",
+        "utilization.jsonl",
         "metrics.jsonl",
         "latest_metrics.json",
         "heartbeat.json",
@@ -1684,6 +1777,8 @@ def _verification_required(*, success: bool) -> tuple[str, ...]:
 def _final_upload_names() -> tuple[str, ...]:
     return (
         "config.toml",
+        "job_manifest.json",
+        "utilization.jsonl",
         "metrics.jsonl",
         "latest_metrics.json",
         "heartbeat.json",
@@ -1761,6 +1856,44 @@ def _memory_metrics(device: torch.device) -> dict[str, int | None]:
     }
 
 
+def _nvidia_smi_metrics(device: torch.device) -> dict[str, Any]:
+    if device.type != "cuda":
+        return {
+            "gpu_util_percent": None,
+            "gpu_idle": None,
+            "gpu_memory_used_mb": None,
+        }
+    try:
+        completed = subprocess.run(
+            ["nvidia-smi", "--query-gpu=utilization.gpu,memory.used", "--format=csv,noheader,nounits"],
+            text=True,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            check=True,
+        )
+    except (FileNotFoundError, subprocess.CalledProcessError):
+        return {
+            "gpu_util_percent": None,
+            "gpu_idle": None,
+            "gpu_memory_used_mb": None,
+        }
+    rows = [row.strip().split(",") for row in completed.stdout.splitlines() if row.strip()]
+    if not rows:
+        return {
+            "gpu_util_percent": None,
+            "gpu_idle": None,
+            "gpu_memory_used_mb": None,
+        }
+    gpu_utils = [float(row[0].strip()) for row in rows]
+    memory_used = [float(row[1].strip()) for row in rows]
+    gpu_util_percent = sum(gpu_utils) / len(gpu_utils)
+    return {
+        "gpu_util_percent": gpu_util_percent,
+        "gpu_idle": gpu_util_percent <= 5.0,
+        "gpu_memory_used_mb": max(memory_used),
+    }
+
+
 def _learning_rate(optimizer: torch.optim.Optimizer) -> float:
     return float(optimizer.param_groups[0]["lr"])
 

src/airfrans_frontier/training/observability.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from dataclasses import asdict, is_dataclass
 from pathlib import Path
-from typing import Any, Mapping
+from typing import Any, Iterable, Mapping
 
 from airfrans_frontier.training.config import TrainingConfig
 
@@ -20,6 +20,16 @@ class TrainingObserver:
 
     def update_config(self, values: Mapping[str, Any]) -> None:
         return None
+    def log_artifact_files(
+        self,
+        run_dir: Path,
+        names: Iterable[str],
+        *,
+        event: str,
+        step: int,
+        aliases: Iterable[str] = (),
+    ) -> None:
+        return None
 
     def finish(self, *, exit_code: int = 0) -> None:
         return None
@@ -52,6 +62,31 @@ class WandbObserver(TrainingObserver):
 
     def update_config(self, values: Mapping[str, Any]) -> None:
         self._run.config.update(_json_safe(dict(values)), allow_val_change=True)
+    def log_artifact_files(
+        self,
+        run_dir: Path,
+        names: Iterable[str],
+        *,
+        event: str,
+        step: int,
+        aliases: Iterable[str] = (),
+    ) -> None:
+        existing = [name for name in names if (run_dir / name).is_file()]
+        if not existing:
+            return
+        artifact_name = _artifact_name(getattr(self._run, "name", "airfrans-run"), event, step)
+        artifact = self._wandb.Artifact(
+            artifact_name,
+            type="airfrans-run-artifacts",
+            metadata={
+                "event": event,
+                "step": step,
+                "file_count": len(existing),
+            },
+        )
+        for name in existing:
+            artifact.add_file(str(run_dir / name), name=name)
+        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)
@@ -85,6 +120,12 @@ def start_observer(config: TrainingConfig, *, run_dir: Path) -> TrainingObserver
     return WandbObserver(run, wandb)
 
 
+def _artifact_name(run_name: Any, event: str, step: int) -> str:
+    raw = f"{run_name}-{event}-step-{step}"
+    safe = "".join(character if character.isalnum() or character in "-_." else "-" for character in str(raw))
+    return safe.strip("-") or "airfrans-run-artifacts"
+
+
 def _json_safe(value: Any) -> Any:
     if is_dataclass(value):
         return _json_safe(asdict(value))

tests/test_sweep.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+import json
+import tempfile
+import unittest
+from pathlib import Path
+
+from airfrans_frontier.sweep import (
+    BudgetLedger,
+    PoolConfig,
+    can_expand_pool,
+    collect_job_status,
+    generate_jobs,
+    read_jobs,
+    utilization_summary,
+)
+from airfrans_frontier.training.config import load_training_config
+
+
+class SweepFoundationTests(unittest.TestCase):
+    def test_generate_jobs_writes_deterministic_manifests_and_coordinate_encoding_configs(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            jobs = generate_jobs(
+                output_dir=root,
+                data_root="artifacts/data_cache/airfrans_processed/processed/full",
+                bands=("100m",),
+                families=("film_fourier_inr",),
+                encodings=("nerf_multires", "raw"),
+                group="test_sweep",
+            )
+
+            persisted = read_jobs(root / "jobs.jsonl")
+            self.assertEqual([job["job_id"] for job in persisted], [job["job_id"] for job in jobs])
+            self.assertTrue((root / "pool.toml").is_file())
+            self.assertTrue((root / "budget_ledger.json").is_file())
+
+            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.encoding_levels, 16)
+            self.assertEqual(config.model.fourier_scales, ())
+            self.assertEqual(config.model.hidden_width, 2048)
+            self.assertEqual(config.model.depth, 16)
+
+    def test_generate_jobs_only_crosses_encoding_axis_for_compatible_families(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            generate_jobs(
+                output_dir=root,
+                data_root="data/full",
+                bands=("100m",),
+                families=("film_fourier_inr", "siren_conditioned_inr"),
+                encodings=("raw", "random_fourier"),
+            )
+
+            job_ids = [job["job_id"] for job in read_jobs(root / "jobs.jsonl")]
+
+            self.assertIn("100m_film_fourier_inr_random_fourier", job_ids)
+            self.assertIn("100m_siren_conditioned_inr_raw", job_ids)
+            self.assertNotIn("100m_siren_conditioned_inr_random_fourier", job_ids)
+
+    def test_collect_reconstructs_success_from_required_job_artifacts(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            run_root = root / "runs"
+            jobs = generate_jobs(
+                output_dir=root / "sweep",
+                data_root="data/full",
+                artifact_dir=str(run_root),
+                bands=("100m",),
+                families=("mlp",),
+                encodings=("raw",),
+            )
+            run_dir = run_root / "20260727T000000Z_100m_mlp_raw"
+            run_dir.mkdir(parents=True)
+            for name in (
+                "config.toml",
+                "job_manifest.json",
+                "metrics.jsonl",
+                "latest_metrics.json",
+                "final_metrics.json",
+                "run_manifest.json",
+                "environment_manifest.json",
+                "utilization.jsonl",
+            ):
+                (run_dir / name).write_text("{}\n")
+
+            status = collect_job_status(jobs_path=root / "sweep" / "jobs.jsonl")
+
+            self.assertEqual(status["counts"], {"succeeded": 1})
+            self.assertEqual(status["jobs"][0]["job_id"], jobs[0]["job_id"])
+            self.assertEqual(status["jobs"][0]["missing_artifacts"], [])
+
+    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}
+
+        allowed, reasons = can_expand_pool(
+            pool=pool,
+            utilization=good_utilization,
+            backlog=3,
+            recent_failures=0,
+            ledger=BudgetLedger(budget_usd=25.0, spent_usd=5.0),
+            next_pool_cost_usd=10.0,
+        )
+        self.assertTrue(allowed)
+        self.assertEqual(reasons, ())
+
+        blocked, reasons = can_expand_pool(
+            pool=pool,
+            utilization={"gpu_util_median": 70.0, "gpu_idle_fraction": 0.20},
+            backlog=0,
+            recent_failures=1,
+            ledger=BudgetLedger(budget_usd=25.0, spent_usd=24.0),
+            next_pool_cost_usd=2.0,
+        )
+        self.assertFalse(blocked)
+        self.assertIn("median GPU utilization below gate", reasons)
+        self.assertIn("GPU idle fraction above gate", reasons)
+        self.assertIn("no job backlog", reasons)
+        self.assertIn("recent failures are not isolated", reasons)
+        self.assertIn("remaining budget does not support larger pool", reasons)
+
+    def test_utilization_summary_reports_gate_metrics(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            path = Path(tmp) / "utilization.jsonl"
+            path.write_text(
+                "".join(
+                    json.dumps(sample) + "\n"
+                    for sample in (
+                        {"gpu_util_percent": 90, "memory_used_mb": 1000},
+                        {"gpu_util_percent": 95, "memory_used_mb": 1500},
+                        {"gpu_util_percent": 0, "memory_used_mb": 1200},
+                    )
+                )
+            )
+
+            summary = utilization_summary(path)
+
+            self.assertEqual(summary["gpu_util_median"], 90.0)
+            self.assertAlmostEqual(summary["gpu_idle_fraction"], 1 / 3)
+            self.assertEqual(summary["memory_used_peak"], 1500.0)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_training_loop.py

@@ -19,6 +19,7 @@ from unittest.mock import Mock
 
 from airfrans_frontier.training.config import load_training_config
 from airfrans_frontier.training.hf_upload import HfArtifactUploader
+from airfrans_frontier.training.observability import WandbObserver
 from airfrans_frontier.training.loop import train, select_device
 
 
@@ -141,6 +142,42 @@ class HfArtifactUploaderTests(unittest.TestCase):
             self.assertEqual(len(manifest["commits"]), 1)
             self.assertEqual(set(manifest["uploaded_paths"]), set(result["uploaded"]))
 
+class WandbObserverTests(unittest.TestCase):
+    def test_logs_existing_run_artifacts_with_event_and_step_aliases(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            run_dir = Path(tmp)
+            (run_dir / "metrics.jsonl").write_text("{}\n")
+            (run_dir / "checkpoint_latest.pt").write_text("checkpoint")
+            fake_artifact = Mock()
+            fake_wandb = Mock()
+            fake_wandb.Artifact.return_value = fake_artifact
+            fake_run = Mock()
+            fake_run.name = "run with spaces"
+            observer = WandbObserver(fake_run, fake_wandb)
+
+            observer.log_artifact_files(
+                run_dir,
+                ("metrics.jsonl", "checkpoint_latest.pt", "missing.json"),
+                event="latest_checkpoint",
+                step=5,
+                aliases=("latest_checkpoint", "step-5", "latest"),
+            )
+
+            fake_wandb.Artifact.assert_called_once()
+            artifact_kwargs = fake_wandb.Artifact.call_args.kwargs
+            self.assertEqual(artifact_kwargs["type"], "airfrans-run-artifacts")
+            self.assertEqual(artifact_kwargs["metadata"]["event"], "latest_checkpoint")
+            self.assertEqual(artifact_kwargs["metadata"]["step"], 5)
+            self.assertEqual(fake_artifact.add_file.call_count, 2)
+            self.assertEqual(
+                [call.kwargs["name"] for call in fake_artifact.add_file.call_args_list],
+                ["metrics.jsonl", "checkpoint_latest.pt"],
+            )
+            fake_run.log_artifact.assert_called_once_with(
+                fake_artifact,
+                aliases=["latest_checkpoint", "step-5", "latest"],
+            )
+
 
 class TrainingLoopTests(unittest.TestCase):
     def test_cuda_config_fails_clearly_when_cuda_unavailable(self) -> None:
@@ -224,6 +261,9 @@ class TrainingLoopTests(unittest.TestCase):
             self.assertTrue((run_dir / "evaluation_protocol.json").is_file())
             self.assertTrue((run_dir / "run_manifest.json").is_file())
             self.assertTrue((run_dir / "metrics.jsonl").exists())
+            utilization_lines = [json.loads(line) for line in (run_dir / "utilization.jsonl").read_text().splitlines()]
+            self.assertGreaterEqual(len(utilization_lines), 2)
+            self.assertIn("gpu_util_percent", utilization_lines[-1])
             heartbeat = json.loads((live_dir / "heartbeat.json").read_text())
             self.assertEqual(heartbeat["run_id"], "test-run")
             self.assertEqual(heartbeat["phase"], "completed")