From 7d1e24403af60f43509b2bac69bf349ea4fe3b65 Mon Sep 17 00:00:00 2001 From: Alpamys Date: Sun, 5 Jul 2026 17:46:51 +0500 Subject: [PATCH] test(prm): tdd-review coverage (input cap, multi-message render, env correctness, recipe resolve, trust delegation, multi-shard) (v0.71.30) --- tests/test_v07130.py | 251 +++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 242 insertions(+), 9 deletions(-) diff --git a/tests/test_v07130.py b/tests/test_v07130.py index 446c3fd..d19f8d9 100644 --- a/tests/test_v07130.py +++ b/tests/test_v07130.py @@ -186,9 +186,17 @@ class TestPrmSchema: def test_rejects_oversize(self): from soup_cli.config.schema import TrainingConfig - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="512|chars"): TrainingConfig(prm_reward="x" * 5000) + def test_prm_aggregate_invalid_literal_rejected_by_schema(self): + from pydantic import ValidationError + + from soup_cli.config.schema import TrainingConfig + + with pytest.raises(ValidationError, match="prm_aggregate"): + TrainingConfig(prm_aggregate="mean") + # --------------------------------------------------------------------------- # Task 3 — PRMScorer + build_prm_reward_fn @@ -316,6 +324,94 @@ class TestPRMScorer: # prefix len 3 -> boundaries at 4 and 7 -> min 4.0 assert out == pytest.approx([4.0]) + def test_prompt_multi_message_joined(self): + # Two non-empty prompt messages join with "\n" -> 1 + 2 = 3 prefix tokens. + s = _make_fake_scorer("min") + out = s( + [[{"role": "assistant", "content": "a"}]], + prompts=[ + [ + {"role": "system", "content": "one"}, + {"role": "user", "content": "two three"}, + ] + ], + ) + # prefix 3 tokens, single step "a" (1 tok) -> boundary at position 3 + assert out == pytest.approx([3.0]) + + def test_string_prompt_context(self): + s = _make_fake_scorer("min") + out = s(["a"], prompts=["one two"]) + # prefix 2 tokens + step "a" -> boundary at position 2 + assert out == pytest.approx([2.0]) + + def test_bare_dict_completion(self): + # _completion_text dict branch (completion not wrapped in a list). + s = _make_fake_scorer("min") + out = s([{"role": "assistant", "content": "a b\nc d e"}]) + assert out == pytest.approx([1.0]) + + def test_tuple_completion(self): + # _completion_text tuple branch -> "".join -> single step. + s = _make_fake_scorer("last") + out = s([("a", "b")]) + assert len(out) == 1 + assert isinstance(out[0], float) + + +def _make_capped_scorer(max_pos, aggregate="min"): + """PRMScorer whose fake model advertises a tiny max_position_embeddings.""" + from types import SimpleNamespace + + import torch + from torch import nn + + from soup_cli.utils.prm_reward import PRMScorer + + hidden = 8 + + class _FakeTok: + pad_token = "" + eos_token = "" + + def __call__(self, text, add_special_tokens=False): + n = max(1, len(text.split())) if text else 0 + return {"input_ids": [1] * n} + + head = nn.Linear(hidden, 1, bias=True) + with torch.no_grad(): + head.weight.copy_(torch.ones(1, hidden) / hidden) + head.bias.zero_() + + class _FakeModelCapped: + reward_head = head + config = SimpleNamespace(max_position_embeddings=max_pos) + + def __call__(self, input_ids, output_hidden_states=False): + seq_len = input_ids.shape[1] + hs = torch.arange(seq_len).float().reshape(1, seq_len, 1).repeat(1, 1, hidden) + return SimpleNamespace(hidden_states=[hs]) + + scorer = PRMScorer("./prm", aggregate=aggregate, device="cpu") + scorer._model = _FakeModelCapped() + scorer._tokenizer = _FakeTok() + return scorer + + +class TestPRMScorerInputCap: + def test_truncates_to_max_position_embeddings(self): + # steps "a b"(2 tok)->boundary 1, "c d"(2 tok)->boundary 3; cap=3 keeps + # only positions < 3, dropping boundary 3. + s = _make_capped_scorer(max_pos=3, aggregate="min") + out = s([[{"role": "assistant", "content": "a b\nc d"}]]) + assert out == pytest.approx([1.0]) + + def test_cap_below_first_boundary_returns_zero(self): + # cap=0 -> no boundary survives -> 0.0 + s = _make_capped_scorer(max_pos=0, aggregate="min") + out = s([[{"role": "assistant", "content": "a b\nc d"}]]) + assert out == [0.0] + class TestBuildPrmRewardFn: def test_returns_named_scorer(self, monkeypatch): @@ -354,11 +450,69 @@ class TestBuildPrmRewardFn: with pytest.raises(ValueError, match="prm_reward=None"): mod.build_prm_reward_fn(_T(), device="cpu", trust_remote_code=False) + def test_local_existing_dir_under_cwd_accepted(self, tmp_path, monkeypatch): + import os + + import soup_cli.utils.prm_reward as mod + + monkeypatch.setattr(mod, "_resolve_trust", lambda *a, **k: False) + monkeypatch.chdir(tmp_path) + (tmp_path / "my_prm").mkdir() + + class _T: + prm_reward = "./my_prm" + prm_aggregate = "min" + + fn = mod.build_prm_reward_fn(_T(), device="cpu", trust_remote_code=False) + assert fn.prm_path == os.path.realpath(str(tmp_path / "my_prm")) + + +class TestResolveTrust: + def test_delegates_with_requires_flag(self, monkeypatch): + import soup_cli.utils.trust_remote as trust_mod + from soup_cli.utils.prm_reward import _resolve_trust + + captured = {} + monkeypatch.setattr( + trust_mod, "model_requires_trust_remote_code", lambda base: True + ) + + def _fake_resolve(base, requested, console, requires_remote_code): + captured["requires"] = requires_remote_code + captured["requested"] = requested + return True + + monkeypatch.setattr(trust_mod, "resolve_trust_remote_code", _fake_resolve) + assert _resolve_trust("some/base", False, console=object()) is True + assert captured["requires"] is True + assert captured["requested"] is False + + +class TestLoadRewardHeadMultiShard: + def test_collects_head_across_shards(self, tmp_path): + import torch + from safetensors.torch import save_file + + from soup_cli.utils.prm_reward import load_reward_head_weights + + # Head split across two shards (weight in one, bias in the other). + save_file( + {"model.a": torch.zeros(2, 2), "reward_head.weight": torch.ones(1, 8)}, + str(tmp_path / "model-00001-of-00002.safetensors"), + ) + save_file( + {"model.b": torch.zeros(2, 2), "reward_head.bias": torch.zeros(1)}, + str(tmp_path / "model-00002-of-00002.safetensors"), + ) + head = load_reward_head_weights(str(tmp_path)) + assert set(head.keys()) == {"weight", "bias"} + # --------------------------------------------------------------------------- # Task 5 — bundled rollout envs # --------------------------------------------------------------------------- _ENV_MODULES = ["calculator", "retrieval_qa", "guess_number"] +_NEW_RECIPES = ["grpo-env-calculator", "grpo-env-retrieval-qa", "grpo-env-guess-number"] class TestEnvs: @@ -404,18 +558,78 @@ class TestEnvs: expected = {"+": a + b, "-": a - b, "*": a * b}[op] assert row["answer"] == str(expected) - def test_guess_number_answer_is_int(self): + def test_guess_number_answer_matches_product(self): + import re + from soup_cli.envs.guess_number import rollout for row in rollout([]): - assert row["answer"].lstrip("-").isdigit() + m = re.search(r"equals (\d+) times (\d+)", row["prompt"]) + assert m is not None, row["prompt"] + a, b = int(m.group(1)), int(m.group(2)) + assert row["answer"] == str(a * b) + + def test_retrieval_qa_answer_matches_asked_fact(self): + import re - def test_retrieval_qa_answer_in_prompt(self): from soup_cli.envs.retrieval_qa import rollout for row in rollout([]): - # The answer span must appear in the document/prompt. - assert row["answer"] in row["prompt"] + # Tie the answer to the SPECIFIC asked fact, not just "any embedded + # value" — a wrong-index pairing bug must fail here. + m = re.search(r"completes '(.+?) ___'", row["prompt"]) + assert m is not None, row["prompt"] + entity_attr = m.group(1) + assert f"{entity_attr} {row['answer']}." in row["prompt"], row + + @pytest.mark.parametrize("modname", _ENV_MODULES) + def test_row_count_is_default(self, modname): + import importlib + + from soup_cli.envs._common import DEFAULT_ROWS + + mod = importlib.import_module(f"soup_cli.envs.{modname}") + assert len(mod.rollout([])) == DEFAULT_ROWS + + def test_envs_produce_distinct_output(self): + # Guard against an accidental seed / module copy-paste across envs. + from soup_cli.envs import calculator, guess_number, retrieval_qa + + outs = [ + tuple(r["prompt"] for r in calculator.rollout([])), + tuple(r["prompt"] for r in guess_number.rollout([])), + tuple(r["prompt"] for r in retrieval_qa.rollout([])), + ] + assert len(set(outs)) == 3 + + +class TestRecipeRolloutFunc: + _EXPECTED_REWARD = { + "grpo-env-calculator": ("verifiable", "math"), + "grpo-env-retrieval-qa": ("accuracy", None), + "grpo-env-guess-number": ("verifiable", "math"), + } + + @pytest.mark.parametrize("name", _NEW_RECIPES) + def test_rollout_func_resolves_and_runs(self, name): + from soup_cli.config.loader import load_config_from_string + from soup_cli.recipes.catalog import get_recipe + from soup_cli.utils.agent_rollout import resolve_rollout_func + + cfg = load_config_from_string(get_recipe(name).yaml_str) + fn = resolve_rollout_func(cfg.training.rollout_func) + rows = fn([]) + assert rows, "resolved rollout_func must produce rows" + + @pytest.mark.parametrize("name", _NEW_RECIPES) + def test_recipe_reward_matches_env(self, name): + from soup_cli.config.loader import load_config_from_string + from soup_cli.recipes.catalog import get_recipe + + cfg = load_config_from_string(get_recipe(name).yaml_str) + expected_fn, expected_domain = self._EXPECTED_REWARD[name] + assert cfg.training.reward_fn == expected_fn + assert cfg.training.verifiable_domain == expected_domain # --------------------------------------------------------------------------- @@ -461,13 +675,32 @@ class TestGrpoPrmWiring: assert out == "REWARD_FN" assert captured["spec"] == "accuracy" + def test_verifiable_domain_threaded(self, monkeypatch): + import soup_cli.trainer.grpo as grpo + + captured = {} + + def _fake_load(spec, verifiable_domain=None): + captured["spec"] = spec + captured["domain"] = verifiable_domain + return "REWARD_FN" + + monkeypatch.setattr("soup_cli.trainer.rewards.load_reward_fn", _fake_load) + + class _T: + prm_reward = None + prm_aggregate = "min" + reward_fn = "verifiable" + verifiable_domain = "math" + + grpo._select_reward_fn(_T(), "cpu", False) + assert captured["spec"] == "verifiable" + assert captured["domain"] == "math" + # --------------------------------------------------------------------------- # Task 6 — recipes # --------------------------------------------------------------------------- -_NEW_RECIPES = ["grpo-env-calculator", "grpo-env-retrieval-qa", "grpo-env-guess-number"] - - class TestRecipes: @pytest.mark.parametrize("name", _NEW_RECIPES) def test_recipe_resolves(self, name):