soup/tests/test_v0410_part_b.py

"""v0.41.0 Part B — lr_groups tests."""

from __future__ import annotations

import pytest
from pydantic import ValidationError

from soup_cli.config.schema import TrainingConfig
from soup_cli.utils.lr_groups import (
    MAX_LR_GROUPS,
    LrGroup,
    build_optimizer_param_groups,
    lr_groups_from_schema,
    parse_lr_groups,
)


class TestParseLrGroups:
    def test_none_returns_none(self):
        assert parse_lr_groups(None) is None

    def test_empty_list_returns_none(self):
        assert parse_lr_groups([]) is None

    def test_empty_dict_returns_none(self):
        assert parse_lr_groups({}) is None

    def test_dict_form(self):
        out = parse_lr_groups({"q_proj": 1e-4, "v_proj": 5e-5})
        assert len(out) == 2
        assert out[0] == LrGroup(pattern="q_proj", lr=1e-4)
        assert out[1] == LrGroup(pattern="v_proj", lr=5e-5)

    def test_pair_list_form(self):
        out = parse_lr_groups([("q_proj", 1e-4), ("v_proj", 5e-5)])
        assert len(out) == 2

    def test_dict_entries_form(self):
        out = parse_lr_groups([
            {"pattern": "q_proj", "lr": 1e-4},
            {"pattern": "v_proj", "lr": 5e-5},
        ])
        assert len(out) == 2

    def test_dict_entry_extra_keys_rejected(self):
        with pytest.raises(ValueError, match="exactly"):
            parse_lr_groups([{"pattern": "q", "lr": 1e-4, "extra": 1}])

    def test_pair_wrong_arity_rejected(self):
        with pytest.raises(ValueError, match=r"\(pattern, lr\)"):
            parse_lr_groups([("q",)])

    def test_non_string_pattern_rejected(self):
        with pytest.raises(ValueError, match="must be a string"):
            parse_lr_groups([(123, 1e-4)])

    def test_empty_pattern_rejected(self):
        with pytest.raises(ValueError, match="non-empty"):
            parse_lr_groups([("", 1e-4)])

    def test_null_byte_pattern_rejected(self):
        with pytest.raises(ValueError, match="null bytes"):
            parse_lr_groups([("q\x00", 1e-4)])

    def test_oversize_pattern_rejected(self):
        with pytest.raises(ValueError, match="exceeds"):
            parse_lr_groups([("a" * 257, 1e-4)])

    def test_invalid_regex_rejected(self):
        with pytest.raises(ValueError, match="not a valid regex"):
            parse_lr_groups([("(unclosed", 1e-4)])

    def test_too_many_groups_rejected(self):
        too_many = {f"p{i}": 1e-4 for i in range(MAX_LR_GROUPS + 1)}
        with pytest.raises(ValueError, match="exceeds cap"):
            parse_lr_groups(too_many)

    def test_duplicate_pattern_rejected(self):
        with pytest.raises(ValueError, match="duplicate"):
            parse_lr_groups([("q_proj", 1e-4), ("q_proj", 5e-5)])

    def test_lr_bool_rejected(self):
        with pytest.raises(ValueError, match="must be a number"):
            parse_lr_groups([("q", True)])

    def test_lr_zero_rejected(self):
        with pytest.raises(ValueError, match="must be in"):
            parse_lr_groups([("q", 0)])

    def test_lr_negative_rejected(self):
        with pytest.raises(ValueError, match="must be in"):
            parse_lr_groups([("q", -1e-4)])

    def test_lr_above_one_rejected(self):
        with pytest.raises(ValueError, match="must be in"):
            parse_lr_groups([("q", 1.5)])

    def test_lr_nan_rejected(self):
        with pytest.raises(ValueError, match="finite"):
            parse_lr_groups([("q", float("nan"))])

    def test_int_lr_coerced(self):
        # Integer 1 is on the boundary (lr_upper_inclusive=1.0).
        out = parse_lr_groups([("q", 1)])
        assert out[0].lr == 1.0

    def test_str_form_float_accepted(self):
        # PyYAML parses ``1e-4`` (no dot) as a string; coerce to float
        # so YAML round-trip is friendly.
        out = parse_lr_groups([("q", "1e-4")])
        assert out[0].lr == 1e-4

    def test_str_non_numeric_rejected(self):
        with pytest.raises(ValueError, match="got string"):
            parse_lr_groups([("q", "not-a-number")])

    def test_non_dict_non_list_rejected(self):
        with pytest.raises(ValueError, match="must be a list"):
            parse_lr_groups("q_proj=1e-4")

    def test_list_entry_wrong_type_rejected(self):
        with pytest.raises(ValueError, match="dicts or"):
            parse_lr_groups(["q_proj"])


class TestBuildOptimizerParamGroups:
    @staticmethod
    def _named_params():
        # Stand-in for model.named_parameters() — Tensors not required for
        # routing logic; we just check the bucket assignments.
        return [
            ("model.layers.0.self_attn.q_proj.weight", "T_q0"),
            ("model.layers.0.self_attn.v_proj.weight", "T_v0"),
            ("model.layers.1.mlp.gate_proj.weight", "T_g1"),
            ("lm_head.weight", "T_head"),
        ]

    def test_no_groups_single_base_bucket(self):
        out = build_optimizer_param_groups(self._named_params(), 2e-5, None)
        assert len(out) == 1
        assert out[0]["lr"] == 2e-5
        assert out[0]["name"] == "base"
        assert len(out[0]["params"]) == 4

    def test_single_pattern_routes(self):
        groups = parse_lr_groups([("q_proj", 1e-4)])
        out = build_optimizer_param_groups(self._named_params(), 2e-5, groups)
        # Two buckets: q_proj override + base
        assert len(out) == 2
        q_bucket = next(g for g in out if g["name"] == "lr_group:q_proj")
        assert q_bucket["lr"] == 1e-4
        assert "T_q0" in q_bucket["params"]
        base_bucket = next(g for g in out if g["name"] == "base")
        assert "T_v0" in base_bucket["params"]

    def test_first_match_wins(self):
        groups = parse_lr_groups([
            ("self_attn", 1e-4),
            ("q_proj", 9e-9),  # would also match q_proj; must NOT win
        ])
        out = build_optimizer_param_groups(self._named_params(), 2e-5, groups)
        attn_bucket = next(g for g in out if g["name"] == "lr_group:self_attn")
        assert "T_q0" in attn_bucket["params"]
        # q_proj bucket should not exist (no params left to claim)
        names = [g["name"] for g in out]
        assert "lr_group:q_proj" not in names

    def test_empty_buckets_omitted(self):
        groups = parse_lr_groups([("nonexistent", 1e-4)])
        out = build_optimizer_param_groups(self._named_params(), 2e-5, groups)
        assert len(out) == 1
        assert out[0]["name"] == "base"

    def test_base_lr_bool_rejected(self):
        with pytest.raises(ValueError, match="must be a number"):
            build_optimizer_param_groups([], True, None)  # type: ignore[arg-type]

    def test_base_lr_non_positive_rejected(self):
        with pytest.raises(ValueError, match="must be > 0"):
            build_optimizer_param_groups([], 0.0, None)


class TestSchemaIntegration:
    def test_lr_groups_default_none(self):
        cfg = TrainingConfig()
        assert cfg.lr_groups is None

    def test_lr_groups_dict_form(self):
        cfg = TrainingConfig(lr_groups={"q_proj": 1e-4, "v_proj": 5e-5})
        assert cfg.lr_groups is not None
        patterns = [entry["pattern"] for entry in cfg.lr_groups]
        assert "q_proj" in patterns

    def test_lr_groups_list_form(self):
        cfg = TrainingConfig(lr_groups=[
            {"pattern": "q_proj", "lr": 1e-4},
            {"pattern": "v_proj", "lr": 5e-5},
        ])
        assert len(cfg.lr_groups) == 2

    def test_lr_groups_invalid_lr(self):
        with pytest.raises(ValidationError, match="must be in"):
            TrainingConfig(lr_groups={"q": 2.0})

    def test_lr_groups_too_many(self):
        with pytest.raises(ValidationError, match="exceeds cap"):
            TrainingConfig(
                lr_groups={f"p{i}": 1e-4 for i in range(MAX_LR_GROUPS + 1)}
            )


class TestLrGroupFrozen:
    def test_lr_group_frozen(self):
        from dataclasses import FrozenInstanceError

        g = LrGroup(pattern="q_proj", lr=1e-4)
        with pytest.raises(FrozenInstanceError):
            g.lr = 5e-5  # type: ignore[misc]


class TestLrInfNan:
    def test_inf_lr_rejected(self):
        with pytest.raises(ValueError, match="must be finite"):
            parse_lr_groups([("q", float("inf"))])

    def test_neg_inf_lr_rejected(self):
        with pytest.raises(ValueError, match="must be finite"):
            parse_lr_groups([("q", float("-inf"))])


class TestBuildOptimizerEdgeCases:
    def test_empty_lr_groups_list(self):
        out = build_optimizer_param_groups(
            [("p", "T")], 2e-5, []
        )
        # Empty groups → single base bucket.
        assert len(out) == 1
        assert out[0]["name"] == "base"

    def test_base_lr_inf_rejected(self):
        # +inf passes the "> 0" guard but is non-finite — current
        # implementation accepts it; we assert the *current* behaviour
        # so a future stricter check shows up here.
        out = build_optimizer_param_groups([], 1.0, None)
        assert out[0]["lr"] == 1.0


class TestLrGroupsFromSchema:
    def test_none_input(self):
        assert lr_groups_from_schema(None) is None

    def test_empty_input(self):
        assert lr_groups_from_schema([]) is None

    def test_roundtrip(self):
        cfg = TrainingConfig(lr_groups={"q_proj": 1e-4, "v_proj": 5e-5})
        runtime = lr_groups_from_schema(cfg.lr_groups)
        assert runtime is not None
        assert all(isinstance(g, LrGroup) for g in runtime)
        assert runtime[0].pattern == "q_proj"
        assert runtime[0].lr == 1e-4

    def test_runtime_consumer_accepts_converted(self):
        cfg = TrainingConfig(lr_groups={"q_proj": 1e-4})
        runtime = lr_groups_from_schema(cfg.lr_groups)
        out = build_optimizer_param_groups(
            [("model.q_proj.weight", "T_q"), ("model.other.weight", "T_o")],
            2e-5, runtime,
        )
        assert any(g["name"] == "lr_group:q_proj" for g in out)