Blame for code/tests/test_config_toml.py

"""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)