Clean up baseline persistence and comparison tests

baseline/rg_baselines/comparison.py

@@ -132,10 +132,16 @@ class BaselineComparisonResult:
     final_epoch_summary: pd.DataFrame
     convergence_by_seed: pd.DataFrame
     convergence_summary: pd.DataFrame
-    paired_final_differences: pd.DataFrame
+    paired_terminal_differences: pd.DataFrame
     plot_paths: tuple[Path, ...]
     expected_outputs: tuple[Path, ...]
 
+    @property
+    def paired_final_differences(self) -> pd.DataFrame:
+        """Backward-compatible alias for historical notebooks."""
+
+        return self.paired_terminal_differences
+
 
 def _required_seed_paths(seed_dir: Path, epochs: int) -> list[Path]:
     paths = [
@@ -546,7 +552,7 @@ def run_baseline_comparison(
         final_epoch_summary=final_summary,
         convergence_by_seed=convergence,
         convergence_summary=convergence_summary,
-        paired_final_differences=paired,
+        paired_terminal_differences=paired,
         plot_paths=plot_paths,
         expected_outputs=tuple(expected),
     )

baseline/rg_baselines/config.py

@@ -101,6 +101,10 @@ def validate(self) -> None:
             raise ValueError("recipe_version must be positive")
         if self.initialization != MNIST_REFERENCE_INITIALIZATION:
             raise ValueError("unsupported MLP3 initialization contract")
+        if self.seed < 0:
+            raise ValueError("seed must be non-negative")
+        if self.split_seed < 0:
+            raise ValueError("split_seed must be non-negative")
         if self.epochs < 2:
             raise ValueError(
                 "epochs must be at least two for warm-up/cosine schedules"
@@ -113,6 +117,13 @@ def validate(self) -> None:
             raise ValueError("num_workers must be non-negative")
         if self.grad_clip_norm <= 0.0:
             raise ValueError("grad_clip_norm must be positive")
+        if (
+            self.train_eval_max_batches is not None
+            and self.train_eval_max_batches < 1
+        ):
+            raise ValueError(
+                "train_eval_max_batches must be positive or None"
+            )
         if self.checkpoint_every_epochs < 1:
             raise ValueError("checkpoint_every_epochs must be positive")
         if not self.test_monitoring_only:
@@ -162,15 +173,20 @@ def validate(self) -> None:
         }.items():
             if not 0.0 <= value < 1.0:
                 raise ValueError(f"{name} must lie in [0, 1)")
-        if self.sgd_dampening < 0.0:
-            raise ValueError("dampening must be non-negative")
+        if not 0.0 <= self.sgd_dampening < 1.0:
+            raise ValueError("dampening must lie in [0, 1)")
         if self.sgd_nesterov and (
             self.sgd_momentum <= 0.0 or self.sgd_dampening != 0.0
         ):
             raise ValueError(
                 "Nesterov SGD requires positive momentum and zero dampening"
             )
 
+        if self.muon_nesterov and self.muon_momentum <= 0.0:
+            raise ValueError(
+                "Nesterov Muon requires positive momentum"
+            )
+
         for name, value in {
             "adamw_beta1": self.adamw_beta1,
             "adamw_beta2": self.adamw_beta2,
@@ -181,10 +197,32 @@ def validate(self) -> None:
                 raise ValueError(f"{name} must lie in [0, 1)")
         if self.muon_newton_schulz_steps < 1:
             raise ValueError("muon_newton_schulz_steps must be positive")
+        if not self.muon_parameter_names:
+            raise ValueError("muon_parameter_names must not be empty")
+        if len(set(self.muon_parameter_names)) != len(
+            self.muon_parameter_names
+        ):
+            raise ValueError("muon_parameter_names must be unique")
+        if any(
+            not isinstance(name, str) or not name.strip()
+            for name in self.muon_parameter_names
+        ):
+            raise ValueError(
+                "muon_parameter_names must contain non-empty strings"
+            )
         if min(self.muon_eps, self.adamw_eps, self.muon_aux_eps) <= 0.0:
             raise ValueError("optimizer eps values must be positive")
         if self.ww_min_evals < 2:
             raise ValueError("ww_min_evals must be at least two")
+        if (
+            self.ww_max_evals is not None
+            and self.ww_max_evals < self.ww_min_evals
+        ):
+            raise ValueError(
+                "ww_max_evals must be at least ww_min_evals or None"
+            )
+        if not str(self.ww_svd_method).strip():
+            raise ValueError("ww_svd_method must not be empty")
         if not self.ww_randomize:
             raise ValueError(
                 "ww_randomize must be True because WeightWatcher "

baseline/rg_baselines/io_utils.py

@@ -0,0 +1,41 @@
+"""Atomic persistence helpers for baseline progress artifacts."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+
+def atomic_csv(frame: pd.DataFrame, path: str | Path) -> Path:
+    """Replace a CSV only after the temporary file is complete."""
+
+    destination = Path(path)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    temporary = destination.with_suffix(destination.suffix + ".tmp")
+    try:
+        frame.to_csv(temporary, index=False)
+        temporary.replace(destination)
+    finally:
+        temporary.unlink(missing_ok=True)
+    return destination
+
+
+def atomic_npz(
+    arrays: Mapping[str, np.ndarray],
+    path: str | Path,
+) -> Path:
+    """Replace a compressed NumPy archive atomically."""
+
+    destination = Path(path)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    temporary = destination.with_suffix(destination.suffix + ".tmp")
+    try:
+        with temporary.open("wb") as handle:
+            np.savez_compressed(handle, **arrays)
+        temporary.replace(destination)
+    finally:
+        temporary.unlink(missing_ok=True)
+    return destination

baseline/rg_baselines/runner.py

@@ -27,6 +27,7 @@
     set_seed,
     train_one_epoch,
 )
+from .io_utils import atomic_csv, atomic_npz
 from .model import MLP3
 from .optimizers import (
     build_optimizer,
@@ -240,13 +241,17 @@ def _write_progress(
     esds: dict[str, np.ndarray],
 ) -> None:
     run_dir.mkdir(parents=True, exist_ok=True)
-    performance.to_csv(run_dir / "performance_by_epoch.csv", index=False)
-    spectral.to_csv(
-        run_dir / "spectral_metrics_by_epoch_and_layer.csv", index=False
+    atomic_csv(performance, run_dir / "performance_by_epoch.csv")
+    atomic_csv(
+        spectral,
+        run_dir / "spectral_metrics_by_epoch_and_layer.csv",
+    )
+    atomic_csv(
+        details,
+        run_dir / "weightwatcher_details_by_epoch.csv",
     )
-    details.to_csv(run_dir / "weightwatcher_details_by_epoch.csv", index=False)
-    groups.to_csv(run_dir / "optimizer_groups_by_epoch.csv", index=False)
-    np.savez_compressed(run_dir / "esd_history.npz", **esds)
+    atomic_csv(groups, run_dir / "optimizer_groups_by_epoch.csv")
+    atomic_npz(esds, run_dir / "esd_history.npz")
 
 
 def _load_completed_result(
@@ -428,12 +433,6 @@ def run_baseline(
             expected_fingerprint=fingerprint,
         )
         model.to(device)
-        for frame_name in ("performance", "spectral", "details", "groups"):
-            frame = locals()[frame_name]
-            if not frame.empty and "epoch" in frame:
-                locals()[frame_name] = frame[
-                    frame["epoch"].astype(int) <= start_epoch
-                ].copy()
         performance = performance[
             performance["epoch"].astype(int) <= start_epoch
         ].copy() if not performance.empty else performance

baseline/tests/test_mnist_cleanup.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import tempfile
+import unittest
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+from rg_baselines.config import BaselineConfig
+from rg_baselines.io_utils import atomic_csv, atomic_npz
+
+
+class BaselineConfigCleanupTests(unittest.TestCase):
+    def test_optional_limits_are_validated(self) -> None:
+        invalid = (
+            BaselineConfig(
+                optimizer="adamw",
+                train_eval_max_batches=0,
+            ),
+            BaselineConfig(
+                optimizer="adamw",
+                ww_min_evals=8,
+                ww_max_evals=7,
+            ),
+            BaselineConfig(
+                optimizer="adamw",
+                ww_svd_method="",
+            ),
+            BaselineConfig(
+                optimizer="sgd_momentum_muon",
+                muon_parameter_names=(),
+            ),
+            BaselineConfig(
+                optimizer="sgd_momentum_muon",
+                muon_parameter_names=(
+                    "fc1.weight",
+                    "fc1.weight",
+                ),
+            ),
+            BaselineConfig(
+                optimizer="sgd_momentum_muon",
+                muon_momentum=0.0,
+                muon_nesterov=True,
+            ),
+        )
+        for config in invalid:
+            with self.subTest(config=config):
+                with self.assertRaises(ValueError):
+                    config.validate()
+
+    def test_seed_and_dampening_ranges_are_validated(self) -> None:
+        for config in (
+            BaselineConfig(optimizer="adamw", seed=-1),
+            BaselineConfig(optimizer="adamw", split_seed=-1),
+            BaselineConfig(
+                optimizer="sgd_momentum",
+                sgd_dampening=1.0,
+                sgd_nesterov=False,
+            ),
+        ):
+            with self.subTest(config=config):
+                with self.assertRaises(ValueError):
+                    config.validate()
+
+
+class AtomicPersistenceTests(unittest.TestCase):
+    def test_csv_and_npz_replace_existing_files(self) -> None:
+        with tempfile.TemporaryDirectory() as temporary:
+            root = Path(temporary)
+            csv_path = root / "progress.csv"
+            csv_path.write_text("old\n", encoding="utf-8")
+            frame = pd.DataFrame(
+                [
+                    {"epoch": 0, "loss": 1.0},
+                    {"epoch": 1, "loss": 0.5},
+                ]
+            )
+            atomic_csv(frame, csv_path)
+            pd.testing.assert_frame_equal(pd.read_csv(csv_path), frame)
+            self.assertFalse(
+                csv_path.with_suffix(".csv.tmp").exists()
+            )
+
+            npz_path = root / "history.npz"
+            npz_path.write_bytes(b"old")
+            arrays = {
+                "epoch_000": np.asarray([1.0, 2.0]),
+                "epoch_001": np.asarray([3.0, 4.0]),
+            }
+            atomic_npz(arrays, npz_path)
+            with np.load(npz_path) as archive:
+                self.assertEqual(set(archive.files), set(arrays))
+                for name, expected in arrays.items():
+                    np.testing.assert_array_equal(
+                        archive[name], expected
+                    )
+            self.assertFalse(
+                npz_path.with_suffix(".npz.tmp").exists()
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()