Skip to content

ci: apply and test baseline minor cleanup #1

ci: apply and test baseline minor cleanup

ci: apply and test baseline minor cleanup #1

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