baseline-minor-cleanup-apply.yml

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