| """TOML config loader. Run: uv run pytest tests/test_config_toml.py -v -m "not gpu" |
| |
| The loader replaced OmegaConf; these tests pin the three behaviours OmegaConf |
| gave us that the stage scripts silently depend on: |
| * struct mode at EVERY depth (a typo'd key raises, naming its dotted path); |
| * CLI override coercion from the field annotation (`lr=1e-5` is a float); |
| * every shipped configs/*.toml loads, and survives a dump/reload round trip. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import glob |
| import os |
| import tomllib |
|
|
| import pytest |
|
|
| from align.config import ( |
| ConfigError, |
| RunConfig, |
| TrainConfig, |
| dumps_toml, |
| load_config, |
| load_toml, |
| to_dict, |
| ) |
| from tests.conftest import REPO_ROOT |
|
|
| CONFIG_DIR = os.path.join(REPO_ROOT, "configs") |
| SHIPPED = sorted(glob.glob(os.path.join(CONFIG_DIR, "*.toml"))) |
|
|
|
|
| def _write(tmp_workdir, text, name="c.toml"): |
| p = os.path.join(tmp_workdir, name) |
| with open(p, "w") as f: |
| f.write(text) |
| return p |
|
|
|
|
| def test_every_legacy_yaml_has_a_toml(): |
| """Conversion completeness: 7 stage configs shipped as YAML, 7 as TOML.""" |
| names = {os.path.basename(p) for p in SHIPPED} |
| assert { |
| "sft.toml", "pretrain.toml", "trace.toml", "eval.toml", |
| "dpo_argilla.toml", "dpo_steppairs.toml", "grpo_stageb.toml", |
| } <= names |
|
|
|
|
| @pytest.mark.parametrize("path", SHIPPED, ids=os.path.basename) |
| def test_shipped_toml_parses_and_has_no_yaml_residue(path): |
| with open(path, "rb") as f: |
| tomllib.load(f) |
| text = open(path).read() |
| assert "OmegaConf" not in text and ".yaml" not in text |
|
|
|
|
| @pytest.mark.parametrize("name", ["sft", "pretrain", "trace"]) |
| def test_runconfig_tomls_round_trip(name, tmp_workdir): |
| """load -> to_dict -> dumps_toml -> load is the identity. None-valued |
| keys are omitted by the dump and come back as the dataclass default.""" |
| cfg = load_config(os.path.join(CONFIG_DIR, f"{name}.toml")) |
| again = load_config(RunConfig, _write(tmp_workdir, dumps_toml(cfg))) |
| assert to_dict(again) == to_dict(cfg) |
| assert isinstance(again.optimizer.betas, tuple) |
|
|
|
|
| def test_sft_values_match_the_legacy_yaml(): |
| cfg = load_config(os.path.join(CONFIG_DIR, "sft.toml")) |
| assert cfg.run_name == "sft-stage-11" |
| assert cfg.optimizer.name == "adamw8bit" and cfg.optimizer.grad_accum == 2 |
| assert cfg.optimizer.lr == 1e-5 and cfg.optimizer.betas == (0.9, 0.95) |
| assert cfg.train.global_batch_rows == 24 and cfg.train.micro_batch_size == 3 |
| assert cfg.model.attn_implementation is None |
| assert cfg.gates.epoch_budget_s == 14400.0 |
|
|
|
|
| @pytest.mark.parametrize("name", ["dpo_argilla", "dpo_steppairs"]) |
| def test_dpo_tomls_load_through_dpo_stage(name): |
| from align.train.dpo_stage import load_stage_config |
|
|
| cfg = load_stage_config(os.path.join(CONFIG_DIR, f"{name}.toml")) |
| assert "logps_chosen" in cfg.armed_tripwires |
| assert "length_drift" in cfg.armed_tripwires |
|
|
|
|
| def test_grpo_stageb_toml_invariants(): |
| """The regression-tested GRPO pins survive the format change.""" |
| from align.rl.stage_config import load_stage_config |
|
|
| cfg = load_stage_config(os.path.join(CONFIG_DIR, "grpo_stageb.toml")) |
| assert cfg.grpo.beta == 0.0 and cfg.grpo.load_ref_policy is False |
| assert (cfg.grpo.clip_lower, cfg.grpo.clip_higher) == (0.2, 0.272) |
| assert cfg.grpo.num_samples_per_prompt_rollout >= 8 |
| assert isinstance(cfg.grpo.loss_denominator, float) |
| assert cfg.grpo.use_rho_correction is True |
|
|
|
|
| def test_eval_toml_is_a_plain_dict_with_overrides(): |
| cfg = load_toml(os.path.join(CONFIG_DIR, "eval.toml"), ["eval.limit=8", "eval.temperature=0.7"]) |
| assert cfg["eval"]["limit"] == 8 |
| assert cfg["eval"]["temperature"] == 0.7 |
| assert cfg["eval"]["verifier"]["outer_timeout_seconds"] >= 4 * cfg["eval"]["verifier"]["timeout_seconds"] |
|
|
|
|
| def test_unknown_key_names_full_dotted_path(tmp_workdir): |
| p = _write(tmp_workdir, "[optimizer]\nlr = 1e-6\nnot_a_field = 3\n") |
| with pytest.raises(ConfigError, match=r"optimizer\.not_a_field"): |
| load_config(RunConfig, p) |
| p = _write(tmp_workdir, "[optimzer]\nlr = 1e-6\n") |
| with pytest.raises(ConfigError, match="optimzer"): |
| load_config(RunConfig, p) |
|
|
|
|
| def test_unknown_override_raises_at_any_depth(tmp_workdir): |
| p = _write(tmp_workdir, "run_name = 'x'\n") |
| with pytest.raises(ConfigError, match=r"train\.nope"): |
| load_config(RunConfig, p, ["train.nope=1"]) |
| with pytest.raises(ConfigError, match="top_level_typo"): |
| load_config(RunConfig, p, ["--top_level_typo=1"]) |
|
|
|
|
| def test_override_coercion_from_annotation(tmp_workdir): |
| p = _write(tmp_workdir, "run_name = 'x'\n") |
| cfg = load_config( |
| RunConfig, |
| p, |
| [ |
| "--optimizer.lr=1e-5", |
| "--train.epochs=2", |
| "--train.full_ft=false", |
| "--train.max_micro_steps=40", |
| "--optimizer.betas=[0.8, 0.9]", |
| "--tripwires.armed=logps_chosen,length_drift", |
| "--trace.tier1_layers=[1, 2, 3]", |
| "--model.attn_implementation=sdpa", |
| "--run_name=hello", |
| ], |
| ) |
| assert isinstance(cfg.optimizer.lr, float) and cfg.optimizer.lr == 1e-5 |
| assert isinstance(cfg.train.epochs, int) and cfg.train.epochs == 2 |
| assert cfg.train.full_ft is False |
| assert cfg.train.max_micro_steps == 40 |
| assert cfg.optimizer.betas == (0.8, 0.9) |
| assert cfg.tripwires.armed == ("logps_chosen", "length_drift") |
| assert cfg.trace.tier1_layers == [1, 2, 3] |
| assert cfg.model.attn_implementation == "sdpa" |
| assert cfg.run_name == "hello" |
|
|
|
|
| def test_optional_null_override_and_bad_values(tmp_workdir): |
| p = _write(tmp_workdir, "[train]\nmax_micro_steps = 5\n") |
| assert load_config(RunConfig, p).train.max_micro_steps == 5 |
| assert load_config(RunConfig, p, ["train.max_micro_steps=null"]).train.max_micro_steps is None |
| with pytest.raises(ConfigError, match=r"train\.epochs"): |
| load_config(RunConfig, p, ["train.epochs=abc"]) |
| with pytest.raises(ConfigError, match=r"train\.epochs"): |
| load_config(RunConfig, p, ["train.epochs=2.5"]) |
| with pytest.raises(ConfigError, match=r"train\.full_ft"): |
| load_config(RunConfig, p, ["train.full_ft=maybe"]) |
| with pytest.raises(ConfigError, match=r"train\.epochs"): |
| load_config(RunConfig, _write(tmp_workdir, "[train]\nepochs = true\n", "b.toml")) |
|
|
|
|
| def test_int_widens_to_float_but_not_the_reverse(tmp_workdir): |
| p = _write(tmp_workdir, "[optimizer]\nweight_decay = 1\n[train]\nseed = 3.0\n") |
| cfg = load_config(RunConfig, p) |
| assert isinstance(cfg.optimizer.weight_decay, float) |
| assert cfg.train.seed == 3 and isinstance(cfg.train.seed, int) |
|
|
|
|
| def test_two_call_shapes_agree(tmp_workdir): |
| p = _write(tmp_workdir, "[train]\nepochs = 3\n") |
| a = load_config(p, ["train.seed=9"]) |
| b = load_config(RunConfig, p, ["train.seed=9"]) |
| assert to_dict(a) == to_dict(b) |
| assert isinstance(a.train, TrainConfig) |
|
|
|
|
| def test_shipped_ppo_stage_toml_loads_through_the_struct_mode_loader(): |
| """The PPO stage config ships as a DRAFT: it must still parse through the |
| struct-mode loader (an unknown key or a missing launch-refusal field is a |
| typo that would otherwise only surface at launch), and it must keep its |
| pre-registration TBD markers so nobody mistakes a legacy default for a |
| frozen decision.""" |
| from align.ppo.stage_config import load_stage_config |
|
|
| path = os.path.join(CONFIG_DIR, "ppo_stage.toml") |
| cfg = load_stage_config(path) |
| assert cfg.arm == "prm" and cfg.ppo.temperature == cfg.rollout.temperature |
| assert cfg.rows_per_iteration == 64 |
| text = open(path).read() |
| assert text.count("[TBD pre-reg freeze]") >= 4 and "NOT LAUNCHABLE AS SHIPPED" in text |
|
|
|
|
| |
| |
| |
|
|
|
|
| def test_grad_clipping_is_on_by_default_and_every_runconfig_toml_states_it(): |
| """Default max_grad_norm is 1.0 (OLMoE/OLMo 2). The pre-G0 default None |
| left every SFT/pretrain run unclipped. Each RunConfig-shaped TOML with an |
| [optimizer] table must NAME the value rather than inherit it.""" |
| from align.config import OptimizerConfig |
|
|
| assert OptimizerConfig().max_grad_norm == 1.0 |
| for name in ("sft", "pretrain"): |
| raw = tomllib.load(open(os.path.join(CONFIG_DIR, f"{name}.toml"), "rb")) |
| assert raw["optimizer"]["max_grad_norm"] == 1.0, name |
| assert load_config(os.path.join(CONFIG_DIR, f"{name}.toml")).optimizer.max_grad_norm == 1.0 |
| |
| assert load_config(RunConfig, os.path.join(CONFIG_DIR, "sft.toml"), |
| ["optimizer.max_grad_norm=null"]).optimizer.max_grad_norm is None |
|
|
|
|
| @pytest.mark.parametrize("name", ["sft", "pretrain"]) |
| def test_sft_and_pretrain_arm_exactly_the_grad_norm_tripwire(name): |
| """grad_norm is the only tripwire these stages can FEED; arming |
| logps_chosen/length_drift there would halt at step one (no inputs). The |
| armed names must also be ones the Trainer can build.""" |
| from align.train.trainer import TRIPWIRE_BUILDERS |
|
|
| cfg = load_config(os.path.join(CONFIG_DIR, f"{name}.toml")) |
| assert cfg.tripwires.armed == ("grad_norm",) |
| assert set(cfg.tripwires.armed) <= set(TRIPWIRE_BUILDERS) |
|
|
|
|
| def test_control_plane_train_fields_names_defaults_and_toml(tmp_workdir): |
| """The control-plane agent reads these EXACT names (handoff/guards-configs.md).""" |
| t = TrainConfig() |
| assert t.max_wall_seconds is None |
| assert t.max_consecutive_skipped_steps == 3 and isinstance(t.max_consecutive_skipped_steps, int) |
| assert t.max_skipped_fraction == 0.05 |
| p = _write(tmp_workdir, |
| "[train]\nmax_wall_seconds = 3600\nmax_consecutive_skipped_steps = 5\n" |
| "max_skipped_fraction = 0.1\n") |
| cfg = load_config(RunConfig, p) |
| assert cfg.train.max_wall_seconds == 3600.0 and isinstance(cfg.train.max_wall_seconds, float) |
| assert cfg.train.max_consecutive_skipped_steps == 5 and cfg.train.max_skipped_fraction == 0.1 |
| assert load_config(RunConfig, p, ["train.max_wall_seconds=null"]).train.max_wall_seconds is None |
| |
| with pytest.raises(ConfigError, match="max_wall_second"): |
| load_config(RunConfig, _write(tmp_workdir, "[train]\nmax_wall_second = 1\n", "t.toml")) |
| with pytest.raises(ConfigError, match="max_consecutive_skipped_steps"): |
| load_config(RunConfig, p, ["train.max_consecutive_skipped_steps=2.5"]) |
| |
| rt = _write(tmp_workdir, dumps_toml(to_dict(cfg)), "rt.toml") |
| assert to_dict(load_config(RunConfig, rt)) == to_dict(cfg) |
|
|
