| import pytest |
|
|
| torch = pytest.importorskip("torch") |
|
|
| from sgjm.research.cards import ExperimentCard, SweepResult |
| from sgjm.research.runner import _apply_override, _make_variant_config, run_sweep |
| from sgjm.research.sweep import ( |
| Sweep, |
| SweepEntry, |
| ablation_sweep, |
| available_sweeps, |
| get_sweep, |
| ) |
| from sgjm.training.config import TrainingConfig |
|
|
|
|
| def test_sweep_registry_lists_known_sweeps(): |
| assert "ablation" in available_sweeps() |
| assert "loss_weight" in available_sweeps() |
| sweep = get_sweep("ablation") |
| assert len(sweep) >= 3 |
| for entry in sweep: |
| assert isinstance(entry.card, ExperimentCard) |
| assert entry.card.hypothesis |
|
|
|
|
| def test_apply_override_nested_loss_weights(): |
| base = TrainingConfig.smoke() |
| new_cfg, eval_overrides = _apply_override(base, "loss.jepa", 0.0) |
| assert eval_overrides == {} |
| assert new_cfg.loss.jepa == 0.0 |
| assert base.loss.jepa != 0.0 |
|
|
|
|
| def test_apply_override_model_block_size(): |
| base = TrainingConfig.smoke() |
| new_cfg, _ = _apply_override(base, "model.block_size", 8) |
| assert new_cfg.model.block_size == 8 |
|
|
|
|
| def test_apply_override_eval_prefix_routes_to_eval_overrides(): |
| base = TrainingConfig.smoke() |
| new_cfg, eval_overrides = _apply_override(base, "_eval.merge_radius_bits", 12) |
| assert eval_overrides == {"merge_radius_bits": 12} |
| assert new_cfg is base or new_cfg.model.block_size == base.model.block_size |
|
|
|
|
| def test_make_variant_config_applies_all_overrides(tmp_path): |
| base = TrainingConfig.smoke() |
| card = ExperimentCard( |
| name="t", |
| hypothesis="h", |
| overrides={"loss.jepa": 0.0, "loss.drafter": 0.0, "_eval.merge_radius_bits": 2}, |
| ) |
| cfg, eo = _make_variant_config(base, card, tmp_path) |
| assert cfg.loss.jepa == 0.0 |
| assert cfg.loss.drafter == 0.0 |
| assert cfg.checkpoint_dir.endswith("/t") |
| assert eo["merge_radius_bits"] == 2 |
|
|
|
|
| def test_sweep_result_primary_score_handles_error(): |
| card = ExperimentCard(name="x", hypothesis="h", overrides={}) |
| err_result = SweepResult(card=card, elapsed_sec=0, sgjm_metrics=None, |
| baseline_metrics=None, comparison=None, error="boom") |
| assert err_result.primary_score == float("-inf") |
|
|
|
|
| def test_run_smoke_ablation_sweep_end_to_end(tmp_path): |
| base_cfg = TrainingConfig.smoke() |
| base_cfg.optim.max_steps = 2 |
| |
| sweep = ablation_sweep() |
| sweep.entries = sweep.entries[:2] |
| results = run_sweep(sweep, base_cfg, backend="cpu", out_dir=tmp_path, eval_batches=2) |
| assert len(results) == 2 |
| assert all(r.error is None for r in results) |
| assert (tmp_path / "summary.json").exists() |
| for r in results: |
| assert r.sgjm_metrics is not None |
| assert r.baseline_metrics is not None |
| assert r.comparison is not None |
|
|
