ci: apply and test baseline minor cleanup #1
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: Apply baseline minor cleanup | |
| on: | |
| push: | |
| branches: | |
| - "agent/baseline-minor-cleanup" | |
| paths: | |
| - ".github/workflows/baseline-minor-cleanup-apply.yml" | |
| permissions: | |
| contents: write | |
| jobs: | |
| cleanup-and-test: | |
| runs-on: ubuntu-latest | |
| timeout-minutes: 45 | |
| steps: | |
| - uses: actions/checkout@v4 | |
| with: | |
| ref: agent/baseline-minor-cleanup | |
| fetch-depth: 0 | |
| - uses: actions/setup-python@v5 | |
| with: | |
| python-version: "3.11" | |
| cache: pip | |
| cache-dependency-path: | | |
| baseline/pyproject.toml | |
| baseline/requirements.txt | |
| - name: Apply narrowly scoped cleanup | |
| shell: bash | |
| run: | | |
| python - <<'PY' | |
| from __future__ import annotations | |
| from pathlib import Path | |
| from textwrap import dedent | |
| def replace_once(path: Path, old: str, new: str) -> None: | |
| text = path.read_text(encoding="utf-8") | |
| count = text.count(old) | |
| if count != 1: | |
| raise RuntimeError( | |
| f"{path}: expected one replacement target, found {count}" | |
| ) | |
| path.write_text(text.replace(old, new, 1), encoding="utf-8") | |
| config_path = Path("baseline/rg_baselines/config.py") | |
| replace_once( | |
| config_path, | |
| " if self.epochs < 2:\n", | |
| ( | |
| " if self.seed < 0:\n" | |
| " raise ValueError(\"seed must be non-negative\")\n" | |
| " if self.split_seed < 0:\n" | |
| " raise ValueError(\"split_seed must be non-negative\")\n" | |
| " if self.epochs < 2:\n" | |
| ), | |
| ) | |
| replace_once( | |
| config_path, | |
| " if self.checkpoint_every_epochs < 1:\n", | |
| ( | |
| " if (\n" | |
| " self.train_eval_max_batches is not None\n" | |
| " and self.train_eval_max_batches < 1\n" | |
| " ):\n" | |
| " raise ValueError(\n" | |
| " \"train_eval_max_batches must be positive or None\"\n" | |
| " )\n" | |
| " if self.checkpoint_every_epochs < 1:\n" | |
| ), | |
| ) | |
| replace_once( | |
| config_path, | |
| ( | |
| " if self.sgd_dampening < 0.0:\n" | |
| " raise ValueError(\"dampening must be non-negative\")\n" | |
| ), | |
| ( | |
| " if not 0.0 <= self.sgd_dampening < 1.0:\n" | |
| " raise ValueError(\"dampening must lie in [0, 1)\")\n" | |
| ), | |
| ) | |
| replace_once( | |
| config_path, | |
| ( | |
| " for name, value in {\n" | |
| " \"adamw_beta1\": self.adamw_beta1,\n" | |
| ), | |
| ( | |
| " if self.muon_nesterov and self.muon_momentum <= 0.0:\n" | |
| " raise ValueError(\n" | |
| " \"Nesterov Muon requires positive momentum\"\n" | |
| " )\n\n" | |
| " for name, value in {\n" | |
| " \"adamw_beta1\": self.adamw_beta1,\n" | |
| ), | |
| ) | |
| replace_once( | |
| config_path, | |
| " if min(self.muon_eps, self.adamw_eps, self.muon_aux_eps) <= 0.0:\n", | |
| ( | |
| " if not self.muon_parameter_names:\n" | |
| " raise ValueError(\"muon_parameter_names must not be empty\")\n" | |
| " if len(set(self.muon_parameter_names)) != len(\n" | |
| " self.muon_parameter_names\n" | |
| " ):\n" | |
| " raise ValueError(\"muon_parameter_names must be unique\")\n" | |
| " if any(\n" | |
| " not isinstance(name, str) or not name.strip()\n" | |
| " for name in self.muon_parameter_names\n" | |
| " ):\n" | |
| " raise ValueError(\n" | |
| " \"muon_parameter_names must contain non-empty strings\"\n" | |
| " )\n" | |
| " if min(self.muon_eps, self.adamw_eps, self.muon_aux_eps) <= 0.0:\n" | |
| ), | |
| ) | |
| replace_once( | |
| config_path, | |
| " if not self.ww_randomize:\n", | |
| ( | |
| " if (\n" | |
| " self.ww_max_evals is not None\n" | |
| " and self.ww_max_evals < self.ww_min_evals\n" | |
| " ):\n" | |
| " raise ValueError(\n" | |
| " \"ww_max_evals must be at least ww_min_evals or None\"\n" | |
| " )\n" | |
| " if not str(self.ww_svd_method).strip():\n" | |
| " raise ValueError(\"ww_svd_method must not be empty\")\n" | |
| " if not self.ww_randomize:\n" | |
| ), | |
| ) | |
| io_utils = dedent( | |
| '''\ | |
| """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 | |
| ''' | |
| ) | |
| Path("baseline/rg_baselines/io_utils.py").write_text( | |
| io_utils, encoding="utf-8" | |
| ) | |
| runner_path = Path("baseline/rg_baselines/runner.py") | |
| replace_once( | |
| runner_path, | |
| "from .model import MLP3\n", | |
| "from .io_utils import atomic_csv, atomic_npz\nfrom .model import MLP3\n", | |
| ) | |
| replace_once( | |
| runner_path, | |
| ( | |
| " performance.to_csv(run_dir / \"performance_by_epoch.csv\", index=False)\n" | |
| " spectral.to_csv(\n" | |
| " run_dir / \"spectral_metrics_by_epoch_and_layer.csv\", index=False\n" | |
| " )\n" | |
| " details.to_csv(run_dir / \"weightwatcher_details_by_epoch.csv\", index=False)\n" | |
| " groups.to_csv(run_dir / \"optimizer_groups_by_epoch.csv\", index=False)\n" | |
| " np.savez_compressed(run_dir / \"esd_history.npz\", **esds)\n" | |
| ), | |
| ( | |
| " atomic_csv(performance, run_dir / \"performance_by_epoch.csv\")\n" | |
| " atomic_csv(\n" | |
| " spectral,\n" | |
| " run_dir / \"spectral_metrics_by_epoch_and_layer.csv\",\n" | |
| " )\n" | |
| " atomic_csv(\n" | |
| " details,\n" | |
| " run_dir / \"weightwatcher_details_by_epoch.csv\",\n" | |
| " )\n" | |
| " atomic_csv(groups, run_dir / \"optimizer_groups_by_epoch.csv\")\n" | |
| " atomic_npz(esds, run_dir / \"esd_history.npz\")\n" | |
| ), | |
| ) | |
| replace_once( | |
| runner_path, | |
| ( | |
| " for frame_name in (\"performance\", \"spectral\", \"details\", \"groups\"):\n" | |
| " frame = locals()[frame_name]\n" | |
| " if not frame.empty and \"epoch\" in frame:\n" | |
| " locals()[frame_name] = frame[\n" | |
| " frame[\"epoch\"].astype(int) <= start_epoch\n" | |
| " ].copy()\n" | |
| ), | |
| "", | |
| ) | |
| comparison_path = Path("baseline/rg_baselines/comparison.py") | |
| replace_once( | |
| comparison_path, | |
| " paired_final_differences: pd.DataFrame\n", | |
| " paired_terminal_differences: pd.DataFrame\n", | |
| ) | |
| replace_once( | |
| comparison_path, | |
| ( | |
| " plot_paths: tuple[Path, ...]\n" | |
| " expected_outputs: tuple[Path, ...]\n\n\n" | |
| "def _required_seed_paths" | |
| ), | |
| ( | |
| " plot_paths: tuple[Path, ...]\n" | |
| " expected_outputs: tuple[Path, ...]\n\n" | |
| " @property\n" | |
| " def paired_final_differences(self) -> pd.DataFrame:\n" | |
| " \"\"\"Backward-compatible alias for historical notebooks.\"\"\"\n\n" | |
| " return self.paired_terminal_differences\n\n\n" | |
| "def _required_seed_paths" | |
| ), | |
| ) | |
| replace_once( | |
| comparison_path, | |
| " paired_final_differences=paired,\n", | |
| " paired_terminal_differences=paired,\n", | |
| ) | |
| cleanup_tests = dedent( | |
| '''\ | |
| 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() | |
| ''' | |
| ) | |
| Path("baseline/tests/test_mnist_cleanup.py").write_text( | |
| cleanup_tests, encoding="utf-8" | |
| ) | |
| comparison_tests = dedent( | |
| '''\ | |
| from __future__ import annotations | |
| from dataclasses import asdict | |
| import json | |
| import tempfile | |
| import unittest | |
| from pathlib import Path | |
| from unittest import mock | |
| import pandas as pd | |
| from rg_baselines.comparison import ( | |
| LAYER_ORDER, | |
| OPTIMIZER_LABELS, | |
| OPTIMIZER_ORDER, | |
| _load_and_validate, | |
| run_baseline_comparison, | |
| ) | |
| from rg_baselines.config import BaselineConfig | |
| SEEDS = (1337, 2027, 31415) | |
| EPOCHS = 3 | |
| OFFSETS = { | |
| "sgd_momentum": 0.00, | |
| "adamw": 0.01, | |
| "sgd_momentum_muon": 0.02, | |
| } | |
| def _touch(path: Path, content: bytes = b"placeholder") -> None: | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| path.write_bytes(content) | |
| def _write_suite(root: Path) -> None: | |
| for optimizer in OPTIMIZER_ORDER: | |
| directory = root / optimizer | |
| directory.mkdir(parents=True, exist_ok=True) | |
| config = BaselineConfig( | |
| optimizer=optimizer, | |
| epochs=EPOCHS, | |
| sgd_warmup_epochs=1, | |
| muon_warmup_epochs=1, | |
| ) | |
| config.validate() | |
| manifest = { | |
| "optimizer": optimizer, | |
| "optimizer_label": OPTIMIZER_LABELS[optimizer], | |
| "seeds": list(SEEDS), | |
| "replicate_count": len(SEEDS), | |
| "confidence": 0.95, | |
| "config_template": asdict(config), | |
| } | |
| (directory / "replicate_manifest.json").write_text( | |
| json.dumps(manifest), encoding="utf-8" | |
| ) | |
| performance_rows = [] | |
| spectral_rows = [] | |
| offset = OFFSETS[optimizer] | |
| for seed_index, seed in enumerate(SEEDS): | |
| seed_effect = seed_index * 0.001 | |
| for epoch in range(EPOCHS + 1): | |
| validation_loss = ( | |
| (1.0, 0.7, 0.5, 0.6)[epoch] - offset | |
| ) | |
| train_accuracy = ( | |
| 0.70 + 0.05 * epoch + offset + seed_effect | |
| ) | |
| validation_accuracy = train_accuracy - 0.02 | |
| test_accuracy = train_accuracy - 0.03 | |
| train_loss = 1.2 - 0.2 * epoch - offset | |
| test_loss = train_loss + 0.08 | |
| performance_rows.append( | |
| { | |
| "run": OPTIMIZER_LABELS[optimizer], | |
| "optimizer": optimizer, | |
| "optimizer_label": OPTIMIZER_LABELS[optimizer], | |
| "seed": seed, | |
| "epoch": epoch, | |
| "global_step": 10 * epoch, | |
| "train_loss": train_loss, | |
| "validation_loss": validation_loss, | |
| "test_loss": test_loss, | |
| "train_accuracy": train_accuracy, | |
| "validation_accuracy": validation_accuracy, | |
| "test_accuracy": test_accuracy, | |
| "validation_accuracy_gap": ( | |
| train_accuracy - validation_accuracy | |
| ), | |
| "test_accuracy_gap": ( | |
| train_accuracy - test_accuracy | |
| ), | |
| "validation_loss_gap": ( | |
| validation_loss - train_loss | |
| ), | |
| "test_loss_gap": test_loss - train_loss, | |
| "primary_lr": 0.01, | |
| "test_monitoring_only": 1, | |
| } | |
| ) | |
| for layer_index, layer in enumerate(LAYER_ORDER): | |
| spectral_rows.append( | |
| { | |
| "run": OPTIMIZER_LABELS[optimizer], | |
| "optimizer": optimizer, | |
| "optimizer_label": OPTIMIZER_LABELS[optimizer], | |
| "seed": seed, | |
| "epoch": epoch, | |
| "global_step": 10 * epoch, | |
| "layer_id": layer_index + 1, | |
| "layer": layer, | |
| "status": "ok", | |
| "alpha": 2.0 + 0.1 * layer_index, | |
| "num_traps": layer_index, | |
| "detX_num": 4, | |
| "num_pl_spikes": 2, | |
| "ERG_gap": 2, | |
| "m_midpoint": 3, | |
| "trace_log_midpoint_per_eval": 0.0, | |
| "trace_log_midpoint_total": 0.0, | |
| } | |
| ) | |
| seed_dir = directory / "seeds" / f"seed_{seed}" | |
| required = ( | |
| "final_state.pt", | |
| "checkpoint_latest.pt", | |
| "checkpoint_best.pt", | |
| "test_results.json", | |
| "manifest.json", | |
| "config.json", | |
| "performance_by_epoch.csv", | |
| "spectral_metrics_by_epoch_and_layer.csv", | |
| "esd_history.npz", | |
| ) | |
| for filename in required: | |
| _touch(seed_dir / filename) | |
| for epoch in range(1, EPOCHS + 1): | |
| _touch( | |
| seed_dir | |
| / "checkpoints" | |
| / f"epoch_{epoch:03d}.pt" | |
| ) | |
| (seed_dir / "run_complete.json").write_text( | |
| json.dumps( | |
| { | |
| "completed": True, | |
| "optimizer": optimizer, | |
| "seed": seed, | |
| "epochs": EPOCHS, | |
| "best_validation_epoch": 2, | |
| } | |
| ), | |
| encoding="utf-8", | |
| ) | |
| pd.DataFrame(performance_rows).to_csv( | |
| directory / "performance_by_epoch_and_seed.csv", | |
| index=False, | |
| ) | |
| pd.DataFrame(spectral_rows).to_csv( | |
| directory | |
| / "spectral_metrics_by_epoch_layer_and_seed.csv", | |
| index=False, | |
| ) | |
| (directory / "performance_summary_95ci.csv").write_text( | |
| "placeholder\\n", encoding="utf-8" | |
| ) | |
| (directory / "spectral_summary_95ci.csv").write_text( | |
| "placeholder\\n", encoding="utf-8" | |
| ) | |
| class MnistComparisonTests(unittest.TestCase): | |
| def test_synthetic_complete_suite_and_paired_alias(self) -> None: | |
| with tempfile.TemporaryDirectory() as temporary: | |
| root = Path(temporary) / "runs" | |
| output = Path(temporary) / "comparison" | |
| _write_suite(root) | |
| with mock.patch( | |
| "rg_baselines.comparison.plot_all_comparisons", | |
| return_value=(), | |
| ): | |
| result = run_baseline_comparison( | |
| root, | |
| output_dir=output, | |
| show_plots=False, | |
| ) | |
| self.assertEqual(result.seeds, SEEDS) | |
| self.assertEqual(result.epochs, EPOCHS) | |
| pd.testing.assert_frame_equal( | |
| result.paired_terminal_differences, | |
| result.paired_final_differences, | |
| ) | |
| row = result.paired_terminal_differences.loc[ | |
| result.paired_terminal_differences[ | |
| "checkpoint" | |
| ].eq("final") | |
| & result.paired_terminal_differences[ | |
| "optimizer_a" | |
| ].eq("adamw") | |
| & result.paired_terminal_differences[ | |
| "optimizer_b" | |
| ].eq("sgd_momentum") | |
| & result.paired_terminal_differences[ | |
| "metric" | |
| ].eq("test_accuracy") | |
| ].iloc[0] | |
| self.assertEqual(int(row["n"]), 3) | |
| self.assertAlmostEqual( | |
| float(row["mean_difference"]), 0.01 | |
| ) | |
| self.assertTrue( | |
| ( | |
| output | |
| / "paired_terminal_differences_95ci.csv" | |
| ).is_file() | |
| ) | |
| def test_duplicate_performance_grid_is_rejected(self) -> None: | |
| with tempfile.TemporaryDirectory() as temporary: | |
| root = Path(temporary) / "runs" | |
| _write_suite(root) | |
| path = ( | |
| root | |
| / "adamw" | |
| / "performance_by_epoch_and_seed.csv" | |
| ) | |
| frame = pd.read_csv(path) | |
| frame = pd.concat( | |
| [frame, frame.iloc[[0]]], | |
| ignore_index=True, | |
| ) | |
| frame.to_csv(path, index=False) | |
| with self.assertRaisesRegex( | |
| RuntimeError, | |
| "incomplete or duplicate performance grid", | |
| ): | |
| _load_and_validate(root) | |
| if __name__ == "__main__": | |
| unittest.main() | |
| ''' | |
| ) | |
| Path("baseline/tests/test_mnist_comparison.py").write_text( | |
| comparison_tests, encoding="utf-8" | |
| ) | |
| PY | |
| - name: Install baseline test environment | |
| run: | | |
| python -m pip install --upgrade pip | |
| python -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu | |
| python -m pip install -e './baseline[experiment]' | |
| python -m pip install nbformat pytest | |
| python -m pip check | |
| - name: Validate cleanup and complete baseline suite | |
| env: | |
| PYTHONPATH: baseline | |
| MPLBACKEND: Agg | |
| run: | | |
| git diff --check | |
| python -m compileall -q baseline/rg_baselines baseline/tests | |
| python -m unittest discover -s baseline/tests -v | |
| - name: Commit and push tested cleanup | |
| shell: bash | |
| run: | | |
| set -euo pipefail | |
| git config user.name "github-actions[bot]" | |
| git config user.email "41898282+github-actions[bot]@users.noreply.github.com" | |
| git add \ | |
| baseline/rg_baselines/config.py \ | |
| baseline/rg_baselines/io_utils.py \ | |
| baseline/rg_baselines/runner.py \ | |
| baseline/rg_baselines/comparison.py \ | |
| baseline/tests/test_mnist_cleanup.py \ | |
| baseline/tests/test_mnist_comparison.py | |
| git commit -m "Clean up baseline persistence and comparison tests" | |
| git push origin HEAD:agent/baseline-minor-cleanup |