"""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 # YAML null -> omitted key 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 # --------------------------------------------------------------------------- # G0: clipping on by default, armed tripwires, control-plane bounds # --------------------------------------------------------------------------- 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 # explicit opt-out still possible 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 # struct mode: a typo'd name raises instead of silently disabling the bound 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"]) # round trip through dump rt = _write(tmp_workdir, dumps_toml(to_dict(cfg)), "rt.toml") assert to_dict(load_config(RunConfig, rt)) == to_dict(cfg)