"""v0.71.21 "Precision & rollout lift" tests. Closes #141 (live fp8_attention + NVFP4), #124 (vLLM sleep mode), #125 (multi-turn agent rollout launchers), #228 (apple-adapter converter), #97 (delinearize-llama4 runtime). Hardware notes: fp8_attention / nvfp4 / vllm-sleep are BETA hw-gated — the gates + dispatch are exercised here via monkeypatched capability probes and fake torchao / vllm modules (sys.modules swap via ``monkeypatch.setitem``, the established test_v0531_142 pattern). The apple-adapter conversion and the delinearize runtime are fully CPU-validatable and run on real arrays. """ from __future__ import annotations import itertools import json import sys import types from pathlib import Path from types import SimpleNamespace import numpy as np import pytest from typer.testing import CliRunner _SRC = Path(__file__).resolve().parent.parent / "src" / "soup_cli" runner = CliRunner() def _torch_or_skip(): return pytest.importorskip("torch") # --------------------------------------------------------------------------- # #141 — attention-projection detection # --------------------------------------------------------------------------- class TestIsAttentionProjection: def test_canonical_projections(self): from soup_cli.utils.advanced_precision import is_attention_projection for name in ("q_proj", "k_proj", "v_proj", "o_proj"): assert is_attention_projection(f"model.layers.0.self_attn.{name}") def test_fused_variants(self): from soup_cli.utils.advanced_precision import is_attention_projection assert is_attention_projection("transformer.h.0.attn.c_attn") assert is_attention_projection("model.layers.3.self_attn.qkv_proj") assert is_attention_projection( "transformer.layers.1.attention.query_key_value" ) def test_mlp_projections_rejected(self): from soup_cli.utils.advanced_precision import is_attention_projection assert not is_attention_projection("model.layers.0.mlp.gate_proj") assert not is_attention_projection("model.layers.0.mlp.down_proj") assert not is_attention_projection("lm_head") def test_defensive_inputs(self): from soup_cli.utils.advanced_precision import is_attention_projection assert not is_attention_projection("") assert not is_attention_projection(None) # type: ignore[arg-type] assert not is_attention_projection(123) # type: ignore[arg-type] assert not is_attention_projection("a\x00q_proj") def test_substring_not_enough(self): """Last-component match only — 'my_q_proj_extra' is not q_proj.""" from soup_cli.utils.advanced_precision import is_attention_projection assert not is_attention_projection("model.layers.0.my_q_proj_extra") class TestIsBlackwellGpu: def test_false_without_cuda(self, monkeypatch): torch = _torch_or_skip() from soup_cli.utils.advanced_precision import is_blackwell_gpu monkeypatch.setattr(torch.cuda, "is_available", lambda: False) assert is_blackwell_gpu() is False @pytest.mark.parametrize( ("capability", "expected"), [((8, 6), False), ((9, 0), False), ((10, 0), True), ((12, 0), True)], ) def test_capability_matrix(self, monkeypatch, capability, expected): torch = _torch_or_skip() from soup_cli.utils.advanced_precision import is_blackwell_gpu monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr( torch.cuda, "get_device_capability", lambda _i=0: capability ) assert is_blackwell_gpu() is expected # --------------------------------------------------------------------------- # #141 — apply_fp8_attention (live) # --------------------------------------------------------------------------- def _tiny_attn_model(with_attention: bool = True): """A tiny real-torch model with q/k/v/o projections + an MLP linear.""" torch = _torch_or_skip() nn = torch.nn class _Attn(nn.Module): def __init__(self): super().__init__() self.q_proj = nn.Linear(8, 8) self.k_proj = nn.Linear(8, 8) self.v_proj = nn.Linear(8, 8) self.o_proj = nn.Linear(8, 8) class _Mlp(nn.Module): def __init__(self): super().__init__() self.gate_proj = nn.Linear(8, 8) class _Block(nn.Module): def __init__(self): super().__init__() if with_attention: self.self_attn = _Attn() self.mlp = _Mlp() class _Model(nn.Module): def __init__(self): super().__init__() self.layers = nn.ModuleList([_Block()]) return _Model() def _install_fake_torchao_float8(monkeypatch, record): """Inject a fake ``torchao.float8`` that performs a real module swap.""" torch = _torch_or_skip() nn = torch.nn class _F8(nn.Linear): pass _F8.__name__ = "Float8Linear" class _Cfg: @classmethod def from_recipe_name(cls, name): record["recipe"] = name return cls() def _convert(model, config=None, module_filter_fn=None): record["config"] = config swapped = [] modules = dict(model.named_modules()) for fqn, mod in list(modules.items()): if not isinstance(mod, nn.Linear): continue if type(mod).__name__ == "Float8Linear": continue if module_filter_fn is not None and not module_filter_fn(mod, fqn): continue parent_fqn, _, child = fqn.rpartition(".") parent = modules[parent_fqn] if parent_fqn else model new = _F8(mod.in_features, mod.out_features, bias=mod.bias is not None) setattr(parent, child, new) swapped.append(fqn) record["swapped"] = swapped return model fake_cfg_mod = types.ModuleType("torchao.float8.config") fake_cfg_mod.Float8LinearConfig = _Cfg fake_f8 = types.ModuleType("torchao.float8") fake_f8.convert_to_float8_training = _convert fake_f8.config = fake_cfg_mod fake_root = types.ModuleType("torchao") fake_root.float8 = fake_f8 monkeypatch.setitem(sys.modules, "torchao", fake_root) monkeypatch.setitem(sys.modules, "torchao.float8", fake_f8) monkeypatch.setitem(sys.modules, "torchao.float8.config", fake_cfg_mod) return _F8 def _enable_fp8_gates(monkeypatch): # The torchao probe checks sys.modules first (v0.27.0 None-stub # idiom), so a bare fake root is enough for the gate; conversion # tests install the full fake via _install_fake_torchao_float8. if "torchao" not in sys.modules: monkeypatch.setitem(sys.modules, "torchao", types.ModuleType("torchao")) monkeypatch.setattr("soup_cli.utils.fp8.is_fp8_gpu_supported", lambda: True) class TestApplyFp8Attention: def test_none_model_rejected(self): from soup_cli.utils.advanced_precision import apply_fp8_attention with pytest.raises(TypeError, match="model"): apply_fp8_attention(None) def test_no_torchao_friendly_gate(self, monkeypatch): from soup_cli.utils.advanced_precision import apply_fp8_attention # v0.27.0 None-stub idiom — the probe must treat this as absent. monkeypatch.setitem(sys.modules, "torchao", None) with pytest.raises(RuntimeError, match="torchao"): apply_fp8_attention(_tiny_attn_model()) def test_non_hopper_friendly_gate(self, monkeypatch): from soup_cli.utils.advanced_precision import apply_fp8_attention monkeypatch.setitem( sys.modules, "torchao", types.ModuleType("torchao") ) monkeypatch.setattr( "soup_cli.utils.fp8.is_fp8_gpu_supported", lambda: False ) with pytest.raises(RuntimeError, match="Hopper"): apply_fp8_attention(_tiny_attn_model()) @pytest.mark.parametrize("recipe", [True, "", None, 123, "row\x00wise"]) def test_bad_recipe_rejected(self, monkeypatch, recipe): from soup_cli.utils.advanced_precision import apply_fp8_attention _enable_fp8_gates(monkeypatch) with pytest.raises((TypeError, ValueError), match="recipe"): apply_fp8_attention(_tiny_attn_model(), recipe=recipe) def test_converts_only_attention(self, monkeypatch): from soup_cli.utils.advanced_precision import apply_fp8_attention record: dict = {} _install_fake_torchao_float8(monkeypatch, record) _enable_fp8_gates(monkeypatch) model = _tiny_attn_model() count = apply_fp8_attention(model, recipe="rowwise") assert count == 4 assert record["recipe"] == "rowwise" block = model.layers[0] assert type(block.self_attn.q_proj).__name__ == "Float8Linear" assert type(block.self_attn.o_proj).__name__ == "Float8Linear" # MLP linears untouched. assert type(block.mlp.gate_proj).__name__ == "Linear" def test_already_converted_counted_not_reswapped(self, monkeypatch): from soup_cli.utils.advanced_precision import apply_fp8_attention record: dict = {} f8_cls = _install_fake_torchao_float8(monkeypatch, record) _enable_fp8_gates(monkeypatch) model = _tiny_attn_model() # Pre-convert q_proj manually. model.layers[0].self_attn.q_proj = f8_cls(8, 8) count = apply_fp8_attention(model) assert count == 4 # all four attention projections are now Float8 assert "self_attn.q_proj" not in " ".join(record["swapped"]) assert len(record["swapped"]) == 3 def test_no_attention_projections_value_error(self, monkeypatch): from soup_cli.utils.advanced_precision import apply_fp8_attention record: dict = {} _install_fake_torchao_float8(monkeypatch, record) _enable_fp8_gates(monkeypatch) with pytest.raises(ValueError, match="no attention projections"): apply_fp8_attention(_tiny_attn_model(with_attention=False)) # --------------------------------------------------------------------------- # #141 — apply_nvfp4 (live, Blackwell-gated) # --------------------------------------------------------------------------- def _install_fake_torchao_quant(monkeypatch, record, *, with_nvfp4=True): fake_q = types.ModuleType("torchao.quantization") class _NVFP4Config: pass def _quantize(model, config): record["model"] = model record["config"] = config fake_q.quantize_ = _quantize if with_nvfp4: fake_q.NVFP4Config = _NVFP4Config fake_root = types.ModuleType("torchao") fake_root.quantization = fake_q monkeypatch.setitem(sys.modules, "torchao", fake_root) monkeypatch.setitem(sys.modules, "torchao.quantization", fake_q) class TestApplyNvfp4: def test_none_model_rejected(self): from soup_cli.utils.advanced_precision import apply_nvfp4 with pytest.raises(TypeError, match="model"): apply_nvfp4(None) def test_non_blackwell_friendly_gate(self, monkeypatch): from soup_cli.utils.advanced_precision import apply_nvfp4 monkeypatch.setattr( "soup_cli.utils.advanced_precision.is_blackwell_gpu", lambda: False ) with pytest.raises(RuntimeError, match="Blackwell"): apply_nvfp4(_tiny_attn_model()) def test_missing_torchao_friendly(self, monkeypatch): from soup_cli.utils.advanced_precision import apply_nvfp4 monkeypatch.setattr( "soup_cli.utils.advanced_precision.is_blackwell_gpu", lambda: True ) monkeypatch.setitem(sys.modules, "torchao", None) with pytest.raises(RuntimeError, match="torchao"): apply_nvfp4(_tiny_attn_model()) def test_old_torchao_missing_nvfp4config(self, monkeypatch): from soup_cli.utils.advanced_precision import apply_nvfp4 record: dict = {} _install_fake_torchao_quant(monkeypatch, record, with_nvfp4=False) monkeypatch.setattr( "soup_cli.utils.advanced_precision.is_blackwell_gpu", lambda: True ) with pytest.raises(RuntimeError, match="NVFP4Config"): apply_nvfp4(_tiny_attn_model()) def test_happy_path_quantizes(self, monkeypatch): from soup_cli.utils.advanced_precision import apply_nvfp4 record: dict = {} _install_fake_torchao_quant(monkeypatch, record) monkeypatch.setattr( "soup_cli.utils.advanced_precision.is_blackwell_gpu", lambda: True ) model = _tiny_attn_model() count = apply_nvfp4(model) assert count == 5 # 4 attention + 1 mlp linear assert record["model"] is model assert type(record["config"]).__name__ == "_NVFP4Config" # --------------------------------------------------------------------------- # #141 — apply_v028_speed_memory wiring # --------------------------------------------------------------------------- class TestV028PrecisionWiring: def _tcfg(self, **kwargs): base = { "use_cut_ce": False, "quantization_aware": False, "kernel_auto_compose": False, } base.update(kwargs) return SimpleNamespace(**base) def test_no_features_dict_unchanged(self): """Back-compat regression: 3-key exact dict on the no-features path.""" from soup_cli.utils.v028_features import apply_v028_speed_memory result = apply_v028_speed_memory( model=object(), tcfg=self._tcfg(), base_model="x/y", ) assert result == { "cut_ce": False, "fp8": False, "kernel_auto_compose": False, } def test_fp8_attention_gate_failure_degrades(self, monkeypatch): from soup_cli.utils.v028_features import apply_v028_speed_memory def _gate(model, **kwargs): raise RuntimeError("no Hopper") # Hermetic: force the gate (review fix — the bare call passed only # because the host lacks torchao/Hopper). monkeypatch.setattr( "soup_cli.utils.advanced_precision.apply_fp8_attention", _gate ) result = apply_v028_speed_memory( model=object(), tcfg=self._tcfg(quantization_aware="fp8", fp8_attention=True), base_model="x/y", ) assert result["fp8_attention"] is False def test_fp8_attention_applied(self, monkeypatch): from soup_cli.utils.v028_features import apply_v028_speed_memory monkeypatch.setattr( "soup_cli.utils.advanced_precision.apply_fp8_attention", lambda model, recipe="tensorwise": 4, ) result = apply_v028_speed_memory( model=object(), tcfg=self._tcfg(quantization_aware="fp8", fp8_attention=True), base_model="x/y", ) assert result["fp8_attention"] is True def test_nvfp4_gate_failure_degrades(self, monkeypatch): from soup_cli.utils.v028_features import apply_v028_speed_memory def _gate(model, **kwargs): raise RuntimeError("no Blackwell") monkeypatch.setattr( "soup_cli.utils.advanced_precision.apply_nvfp4", _gate ) result = apply_v028_speed_memory( model=object(), tcfg=self._tcfg(nvfp4=True), base_model="x/y", ) assert result["nvfp4"] is False def test_nvfp4_applied(self, monkeypatch): from soup_cli.utils.v028_features import apply_v028_speed_memory monkeypatch.setattr( "soup_cli.utils.advanced_precision.apply_nvfp4", lambda model: 2, ) result = apply_v028_speed_memory( model=object(), tcfg=self._tcfg(nvfp4=True), base_model="x/y", ) assert result["nvfp4"] is True # --------------------------------------------------------------------------- # #124 — vLLM sleep mode (live) # --------------------------------------------------------------------------- def _install_fake_vllm(monkeypatch, version="0.8.1"): fake = types.ModuleType("vllm") fake.__version__ = version monkeypatch.setitem(sys.modules, "vllm", fake) return fake class TestParseVersionTuple: @pytest.mark.parametrize( ("raw", "expected"), [ ("0.8.1", (0, 8, 1)), ("0.7.0.dev0", (0, 7, 0)), ("0.6", (0, 6)), ("garbage", ()), ("", ()), ], ) def test_matrix(self, raw, expected): from soup_cli.utils.grpo_long_context import _parse_version_tuple assert _parse_version_tuple(raw) == expected class TestVllmSupportsSleepMode: def test_false_when_vllm_missing(self, monkeypatch): from soup_cli.utils.grpo_long_context import vllm_supports_sleep_mode monkeypatch.setitem(sys.modules, "vllm", None) assert vllm_supports_sleep_mode() is False def test_true_on_modern_vllm(self, monkeypatch): from soup_cli.utils.grpo_long_context import vllm_supports_sleep_mode _install_fake_vllm(monkeypatch, "0.8.1") assert vllm_supports_sleep_mode() is True def test_false_on_old_vllm(self, monkeypatch): from soup_cli.utils.grpo_long_context import vllm_supports_sleep_mode _install_fake_vllm(monkeypatch, "0.6.2") assert vllm_supports_sleep_mode() is False class TestApplyVllmSleepMode: def test_none_rejected(self): from soup_cli.utils.grpo_long_context import apply_vllm_sleep_mode with pytest.raises(TypeError, match="engine_args"): apply_vllm_sleep_mode(None) def test_missing_vllm_friendly(self, monkeypatch): from soup_cli.utils.grpo_long_context import apply_vllm_sleep_mode monkeypatch.setitem(sys.modules, "vllm", None) with pytest.raises(RuntimeError, match=r"vLLM >= 0\.7"): apply_vllm_sleep_mode(SimpleNamespace()) def test_old_vllm_friendly(self, monkeypatch): from soup_cli.utils.grpo_long_context import apply_vllm_sleep_mode _install_fake_vllm(monkeypatch, "0.6.2") with pytest.raises(RuntimeError, match=r"0\.6\.2"): apply_vllm_sleep_mode(SimpleNamespace()) def test_happy_sets_enable_sleep_mode(self, monkeypatch): from soup_cli.utils.grpo_long_context import apply_vllm_sleep_mode _install_fake_vllm(monkeypatch, "0.8.1") args = SimpleNamespace() out = apply_vllm_sleep_mode(args) assert out is args assert args.enable_sleep_mode is True class TestVllmSleepCycle: def test_sleep_then_wake(self): from soup_cli.utils.grpo_long_context import vllm_sleep_cycle calls: list = [] class _Engine: def sleep(self, level=1): calls.append(("sleep", level)) def wake_up(self): calls.append(("wake_up",)) with vllm_sleep_cycle(_Engine()): calls.append(("body",)) assert calls == [("sleep", 1), ("body",), ("wake_up",)] def test_wakes_even_on_body_exception(self): from soup_cli.utils.grpo_long_context import vllm_sleep_cycle calls: list = [] class _Engine: def sleep(self, level=1): calls.append("sleep") def wake_up(self): calls.append("wake_up") with pytest.raises(RuntimeError, match="boom"): with vllm_sleep_cycle(_Engine()): raise RuntimeError("boom") assert calls == ["sleep", "wake_up"] def test_engine_without_sleep_warns_not_crashes(self, caplog): import logging from soup_cli.utils.grpo_long_context import vllm_sleep_cycle with caplog.at_level(logging.WARNING, logger="soup_cli.utils.grpo_long_context"): with vllm_sleep_cycle(object()): pass assert any("sleep" in rec.message for rec in caplog.records) class TestMaybeEnableTrlSleepMode: def test_sets_kwarg_when_trl_exposes_it(self): from soup_cli.utils.grpo_long_context import maybe_enable_trl_sleep_mode kwargs: dict = {} ok = maybe_enable_trl_sleep_mode( kwargs, ("output_dir", "vllm_enable_sleep_mode"), None, ) assert ok is True assert kwargs["vllm_enable_sleep_mode"] is True def test_advisory_when_trl_lacks_hook(self): import io from rich.console import Console from soup_cli.utils.grpo_long_context import maybe_enable_trl_sleep_mode buf = io.StringIO() kwargs: dict = {} ok = maybe_enable_trl_sleep_mode( kwargs, ("output_dir",), Console(file=buf), ) assert ok is False assert kwargs == {} assert "sleep" in buf.getvalue().lower() class TestSleepModeSourceWiring: def test_create_vllm_engine_accepts_sleep_mode(self): source = (_SRC / "utils" / "vllm.py").read_text(encoding="utf-8") assert "sleep_mode" in source assert "apply_vllm_sleep_mode" in source def test_grpo_trainer_wires_sleep_mode(self): source = (_SRC / "trainer" / "grpo.py").read_text(encoding="utf-8") assert "maybe_enable_trl_sleep_mode" in source # --------------------------------------------------------------------------- # #125 — rollout launchers # --------------------------------------------------------------------------- _MODULE_COUNTER = itertools.count() def _write_rollout_module(tmp_path, monkeypatch, body: str) -> str: """Write an importable rollout module; return the 'module:fn' spec.""" name = f"soup_test_rollout_{next(_MODULE_COUNTER)}" (tmp_path / f"{name}.py").write_text(body, encoding="utf-8") monkeypatch.syspath_prepend(str(tmp_path)) return f"{name}:rollout" class TestValidateRolloutFunc: def test_none_passthrough(self): from soup_cli.utils.agent_rollout import validate_rollout_func assert validate_rollout_func(None) is None def test_happy(self): from soup_cli.utils.agent_rollout import validate_rollout_func assert validate_rollout_func("my_mod.sub:my_fn") == "my_mod.sub:my_fn" @pytest.mark.parametrize( "bad", ["", "no-colon", "mod:", ":fn", "mod:fn:extra", "a b:c", "mod\x00:fn", 7], ) def test_rejection_matrix(self, bad): from soup_cli.utils.agent_rollout import validate_rollout_func with pytest.raises(ValueError, match="rollout_func"): validate_rollout_func(bad) def test_oversize_rejected(self): from soup_cli.utils.agent_rollout import validate_rollout_func with pytest.raises(ValueError, match="rollout_func"): validate_rollout_func("a" * 300 + ":fn") class TestResolveRolloutFunc: def test_resolves_callable(self, tmp_path, monkeypatch): from soup_cli.utils.agent_rollout import resolve_rollout_func spec = _write_rollout_module( tmp_path, monkeypatch, "def rollout(prompts):\n return []\n" ) fn = resolve_rollout_func(spec) assert callable(fn) def test_missing_module_friendly(self): from soup_cli.utils.agent_rollout import resolve_rollout_func with pytest.raises(ValueError, match="could not be imported"): resolve_rollout_func("definitely_not_a_module_xyz:fn") def test_missing_attr_friendly(self, tmp_path, monkeypatch): from soup_cli.utils.agent_rollout import resolve_rollout_func spec = _write_rollout_module( tmp_path, monkeypatch, "def rollout(prompts):\n return []\n" ) module_name = spec.split(":")[0] with pytest.raises(ValueError, match="no attribute"): resolve_rollout_func(f"{module_name}:nope") def test_non_callable_rejected(self, tmp_path, monkeypatch): from soup_cli.utils.agent_rollout import resolve_rollout_func spec = _write_rollout_module(tmp_path, monkeypatch, "rollout = 42\n") with pytest.raises(ValueError, match="not callable"): resolve_rollout_func(spec) class TestRolloutResult: def test_frozen(self): import dataclasses from soup_cli.utils.agent_rollout import RolloutResult result = RolloutResult(backend="openenv", rows=({"prompt": "x"},)) with pytest.raises(dataclasses.FrozenInstanceError): result.backend = "art" # type: ignore[misc] def test_rows_must_be_tuple(self): from soup_cli.utils.agent_rollout import RolloutResult with pytest.raises(TypeError, match="tuple"): RolloutResult(backend="openenv", rows=[{"prompt": "x"}]) # type: ignore[arg-type] def test_unknown_backend_rejected(self): from soup_cli.utils.agent_rollout import RolloutResult with pytest.raises(ValueError, match="not supported"): RolloutResult(backend="trlx", rows=({"prompt": "x"},)) def test_row_without_prompt_rejected(self): from soup_cli.utils.agent_rollout import RolloutResult with pytest.raises(ValueError, match="prompt"): RolloutResult(backend="openenv", rows=({"answer": "y"},)) def test_row_cap(self): from soup_cli.utils.agent_rollout import ( _MAX_ROLLOUT_ROWS, RolloutResult, ) rows = tuple({"prompt": "x"} for _ in range(_MAX_ROLLOUT_ROWS + 1)) with pytest.raises(ValueError, match="rows"): RolloutResult(backend="openenv", rows=rows) class TestLaunchRolloutOpenenv: def test_happy_path(self, tmp_path, monkeypatch): from soup_cli.utils.agent_rollout import launch_rollout spec = _write_rollout_module( tmp_path, monkeypatch, "def rollout(prompts):\n" " return [\n" " {'prompt': p, 'answer': 'a'} for p in prompts\n" " ] + [{'prompt': 'extra'}]\n", ) result = launch_rollout( "openenv", prompts=["hello", "world"], rollout_func=spec, ) assert result.backend == "openenv" assert len(result.rows) == 3 assert result.rows[0]["prompt"] == "hello" assert result.rows[0]["answer"] == "a" assert "answer" not in result.rows[2] def test_message_list_prompts_preserved(self, tmp_path, monkeypatch): from soup_cli.utils.agent_rollout import launch_rollout spec = _write_rollout_module( tmp_path, monkeypatch, "def rollout(prompts):\n" " return [{'prompt': [{'role': 'user', 'content': 'hi'}]}]\n", ) result = launch_rollout("openenv", prompts=[], rollout_func=spec) assert result.rows[0]["prompt"] == [{"role": "user", "content": "hi"}] def test_openenv_requires_rollout_func(self): from soup_cli.utils.agent_rollout import launch_rollout with pytest.raises(ValueError, match="rollout_func"): launch_rollout("openenv", prompts=["x"]) def test_empty_rows_rejected(self, tmp_path, monkeypatch): from soup_cli.utils.agent_rollout import launch_rollout spec = _write_rollout_module( tmp_path, monkeypatch, "def rollout(prompts):\n return []\n" ) with pytest.raises(ValueError, match="no rows"): launch_rollout("openenv", prompts=["x"], rollout_func=spec) def test_non_iterable_rows_rejected(self, tmp_path, monkeypatch): from soup_cli.utils.agent_rollout import launch_rollout spec = _write_rollout_module( tmp_path, monkeypatch, "def rollout(prompts):\n return 42\n" ) with pytest.raises(ValueError, match="iterable"): launch_rollout("openenv", prompts=["x"], rollout_func=spec) def test_non_mapping_row_rejected(self, tmp_path, monkeypatch): from soup_cli.utils.agent_rollout import launch_rollout spec = _write_rollout_module( tmp_path, monkeypatch, "def rollout(prompts):\n return ['x']\n" ) with pytest.raises(ValueError, match="mapping"): launch_rollout("openenv", prompts=["x"], rollout_func=spec) @pytest.mark.parametrize("bad_steps", [True, 0, -3, "many"]) def test_max_steps_validation(self, tmp_path, monkeypatch, bad_steps): from soup_cli.utils.agent_rollout import launch_rollout spec = _write_rollout_module( tmp_path, monkeypatch, "def rollout(prompts):\n return [{'prompt': 'x'}]\n", ) with pytest.raises((TypeError, ValueError), match="max_steps"): launch_rollout( "openenv", prompts=["x"], rollout_func=spec, max_steps=bad_steps, ) def test_prompts_string_rejected(self, tmp_path, monkeypatch): from soup_cli.utils.agent_rollout import launch_rollout spec = _write_rollout_module( tmp_path, monkeypatch, "def rollout(prompts):\n return [{'prompt': 'x'}]\n", ) with pytest.raises(TypeError, match="prompts"): launch_rollout("openenv", prompts="hi", rollout_func=spec) class TestLaunchRolloutExternal: @pytest.mark.parametrize( ("name", "pkg"), [("art", "openpipe-art"), ("ruler", "ruler-eval"), ("nemo_gym", "nemo-gym")], ) def test_missing_package_friendly(self, monkeypatch, name, pkg): from soup_cli.utils import agent_rollout monkeypatch.setattr(agent_rollout, "_spec_exists", lambda _n: False) with pytest.raises(ImportError, match=pkg): agent_rollout.launch_rollout(name, prompts=["x"]) def test_present_package_honest_gate(self, monkeypatch): from soup_cli.utils import agent_rollout monkeypatch.setattr(agent_rollout, "_spec_exists", lambda _n: True) with pytest.raises(RuntimeError, match="openenv"): agent_rollout.launch_rollout("art", prompts=["x"]) def test_runner_override_seam(self, monkeypatch): from soup_cli.utils import agent_rollout def _fake_runner(**kwargs): return [{"prompt": "from-art", "answer": "ok"}] monkeypatch.setitem( agent_rollout._EXTERNAL_ROLLOUT_RUNNERS, "art", _fake_runner ) result = agent_rollout.launch_rollout("art", prompts=["x"]) assert result.backend == "art" assert result.rows[0]["prompt"] == "from-art" def test_unknown_backend_still_validated_first(self): from soup_cli.utils.agent_rollout import launch_rollout with pytest.raises(ValueError, match="not supported"): launch_rollout("trlx", prompts=["x"]) class TestRolloutLiveWiredFlags: def test_openenv_live_wired(self): from soup_cli.utils.agent_rollout import get_rollout_backend_spec assert get_rollout_backend_spec("openenv").live_wired is True @pytest.mark.parametrize("name", ["art", "ruler", "nemo_gym"]) def test_external_backends_stay_gated(self, name): from soup_cli.utils.agent_rollout import get_rollout_backend_spec assert get_rollout_backend_spec(name).live_wired is False class TestRolloutSchema: def test_openenv_with_func_happy(self): from soup_cli.config.loader import load_config_from_string cfg = load_config_from_string( "base: a/b\n" "task: grpo\n" "data: {train: x.jsonl}\n" "training: {rollout_backend: openenv, rollout_func: 'my_mod:my_fn'}\n" ) assert cfg.training.rollout_func == "my_mod:my_fn" def test_func_without_backend_rejected(self): from soup_cli.config.loader import load_config_from_string with pytest.raises(ValueError, match="openenv"): load_config_from_string( "base: a/b\n" "task: grpo\n" "data: {train: x.jsonl}\n" "training: {rollout_func: 'my_mod:my_fn'}\n" ) def test_openenv_without_func_rejected(self): from soup_cli.config.loader import load_config_from_string with pytest.raises(ValueError, match="rollout_func"): load_config_from_string( "base: a/b\n" "task: grpo\n" "data: {train: x.jsonl}\n" "training: {rollout_backend: openenv}\n" ) def test_func_with_art_rejected(self): from soup_cli.config.loader import load_config_from_string with pytest.raises(ValueError, match="openenv"): load_config_from_string( "base: a/b\n" "task: grpo\n" "data: {train: x.jsonl}\n" "training: {rollout_backend: art, rollout_func: 'my_mod:my_fn'}\n" ) def test_bad_func_shape_rejected(self): from soup_cli.config.loader import load_config_from_string with pytest.raises(ValueError, match="rollout_func"): load_config_from_string( "base: a/b\n" "task: grpo\n" "data: {train: x.jsonl}\n" "training: {rollout_backend: openenv, rollout_func: 'no-colon'}\n" ) def test_default_none(self): from soup_cli.config.schema import TrainingConfig assert TrainingConfig().rollout_func is None class TestRolloutGrpoWiring: def test_grpo_trainer_launches_rollout(self): source = (_SRC / "trainer" / "grpo.py").read_text(encoding="utf-8") assert "launch_rollout" in source assert "rollout_backend" in source # --------------------------------------------------------------------------- # #228 — apple-adapter live converter # --------------------------------------------------------------------------- class TestLoraKeyMapping: def test_hf_to_mlx_strips_prefix_and_lowers(self): from soup_cli.utils.apple_adapter import hf_key_to_mlx key = "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight" assert hf_key_to_mlx(key) == "model.layers.0.self_attn.q_proj.lora_a" def test_hf_to_mlx_lora_b(self): from soup_cli.utils.apple_adapter import hf_key_to_mlx key = "base_model.model.model.layers.2.mlp.gate_proj.lora_B.weight" assert hf_key_to_mlx(key) == "model.layers.2.mlp.gate_proj.lora_b" def test_non_lora_key_returns_none(self): from soup_cli.utils.apple_adapter import hf_key_to_mlx assert hf_key_to_mlx("base_model.model.model.embed_tokens.weight") is None def test_mlx_to_hf_round_trip(self): from soup_cli.utils.apple_adapter import hf_key_to_mlx, mlx_key_to_hf original = "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight" assert mlx_key_to_hf(hf_key_to_mlx(original)) == original def test_mlx_non_lora_returns_none(self): from soup_cli.utils.apple_adapter import mlx_key_to_hf assert mlx_key_to_hf("model.layers.0.self_attn.q_proj.weight") is None class TestConvertArrays: def _hf_arrays(self): rng = np.random.default_rng(0) return { "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": rng.standard_normal((4, 16)).astype(np.float32), "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight": rng.standard_normal((16, 4)).astype(np.float32), } def test_hf_to_mlx_transposes(self): from soup_cli.utils.apple_adapter import convert_hf_to_mlx_arrays converted, skipped = convert_hf_to_mlx_arrays(self._hf_arrays()) assert skipped == () lora_a = converted["model.layers.0.self_attn.q_proj.lora_a"] lora_b = converted["model.layers.0.self_attn.q_proj.lora_b"] assert lora_a.shape == (16, 4) # [in, r] assert lora_b.shape == (4, 16) # [r, out] def test_non_lora_keys_skipped(self): from soup_cli.utils.apple_adapter import convert_hf_to_mlx_arrays arrays = self._hf_arrays() arrays["base_model.model.model.embed_tokens.weight"] = np.zeros( (4, 4), dtype=np.float32 ) converted, skipped = convert_hf_to_mlx_arrays(arrays) assert len(converted) == 2 assert skipped == ("base_model.model.model.embed_tokens.weight",) def test_zero_lora_keys_rejected(self): from soup_cli.utils.apple_adapter import convert_hf_to_mlx_arrays with pytest.raises(ValueError, match="no LoRA"): convert_hf_to_mlx_arrays( {"embed_tokens.weight": np.zeros((2, 2), dtype=np.float32)} ) def test_mlx_to_hf_round_trip_values(self): from soup_cli.utils.apple_adapter import ( convert_hf_to_mlx_arrays, convert_mlx_to_hf_arrays, ) original = self._hf_arrays() mlx_arrays, _ = convert_hf_to_mlx_arrays(original) back, skipped = convert_mlx_to_hf_arrays(mlx_arrays) assert skipped == () for key, value in original.items(): np.testing.assert_allclose(back[key], value) def _write_peft_adapter(adapter_dir: Path) -> dict: """Write a synthetic PEFT LoRA adapter; return its arrays.""" from safetensors.numpy import save_file rng = np.random.default_rng(7) arrays = { "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight": rng.standard_normal((4, 16)).astype(np.float32), "base_model.model.model.layers.0.self_attn.q_proj.lora_B.weight": rng.standard_normal((16, 4)).astype(np.float32), "base_model.model.model.layers.0.self_attn.v_proj.lora_A.weight": rng.standard_normal((4, 16)).astype(np.float32), "base_model.model.model.layers.0.self_attn.v_proj.lora_B.weight": rng.standard_normal((16, 4)).astype(np.float32), } adapter_dir.mkdir(parents=True, exist_ok=True) save_file(arrays, str(adapter_dir / "adapter_model.safetensors")) (adapter_dir / "adapter_config.json").write_text( json.dumps({ "r": 4, "lora_alpha": 8, "peft_type": "LORA", "base_model_name_or_path": "tiny/base", }), encoding="utf-8", ) return arrays class TestConvertAppleAdapterLive: def test_hf_to_mlx_writes_safetensors(self, tmp_path, monkeypatch): """mlx-lm's load_adapters reads adapters.safetensors + num_layers (review fix — the legacy npz artifact is unloadable by current mlx-lm).""" from safetensors.numpy import load_file from soup_cli.utils.apple_adapter import ( build_apple_adapter_plan, convert_apple_adapter, ) monkeypatch.chdir(tmp_path) arrays = _write_peft_adapter(tmp_path / "adapter") plan = build_apple_adapter_plan( source_dir="adapter", output_dir="out", direction="hf-to-mlx", ) report = convert_apple_adapter(plan) assert report.converted_keys == 4 st_path = tmp_path / "out" / "adapters.safetensors" assert st_path.is_file() loaded = load_file(str(st_path)) key = "model.layers.0.self_attn.q_proj.lora_a" assert key in loaded np.testing.assert_allclose( loaded[key], arrays[ "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight" ].T, ) config = json.loads( (tmp_path / "out" / "adapter_config.json").read_text(encoding="utf-8") ) assert config["fine_tune_type"] == "lora" assert config["num_layers"] == 1 # max layer index + 1 assert config["lora_parameters"]["rank"] == 4 assert config["lora_parameters"]["scale"] == 2.0 # alpha 8 / r 4 def test_full_round_trip(self, tmp_path, monkeypatch): from safetensors.numpy import load_file from soup_cli.utils.apple_adapter import ( build_apple_adapter_plan, convert_apple_adapter, ) monkeypatch.chdir(tmp_path) arrays = _write_peft_adapter(tmp_path / "adapter") convert_apple_adapter( build_apple_adapter_plan( source_dir="adapter", output_dir="mlx", direction="hf-to-mlx", ) ) report = convert_apple_adapter( build_apple_adapter_plan( source_dir="mlx", output_dir="hf2", direction="mlx-to-hf", ) ) assert report.converted_keys == 4 back = load_file(str(tmp_path / "hf2" / "adapter_model.safetensors")) for key, value in arrays.items(): np.testing.assert_allclose(back[key], value) config = json.loads( (tmp_path / "hf2" / "adapter_config.json").read_text(encoding="utf-8") ) assert config["peft_type"] == "LORA" assert config["r"] == 4 # scale (2.0) * rank (4) reconstructs the source lora_alpha — and # the round trip exercises the adapters.safetensors input branch # (hf-to-mlx now emits safetensors). assert config["lora_alpha"] == 8 def test_sign_writes_signature(self, tmp_path, monkeypatch): from soup_cli.utils.apple_adapter import ( build_apple_adapter_plan, convert_apple_adapter, ) monkeypatch.chdir(tmp_path) _write_peft_adapter(tmp_path / "adapter") report = convert_apple_adapter( build_apple_adapter_plan( source_dir="adapter", output_dir="out", direction="hf-to-mlx", sign=True, ) ) assert report.signed is True assert (tmp_path / "out" / ".soup-signature.json").is_file() @pytest.mark.parametrize("direction", ["hf-to-apple", "mlx-to-apple"]) def test_apple_directions_upstream_gated(self, tmp_path, monkeypatch, direction): from soup_cli.utils.apple_adapter import ( build_apple_adapter_plan, convert_apple_adapter, ) monkeypatch.chdir(tmp_path) _write_peft_adapter(tmp_path / "adapter") plan = build_apple_adapter_plan( source_dir="adapter", output_dir="out", direction=direction, ) with pytest.raises(RuntimeError, match="FoundationModels"): convert_apple_adapter(plan) def test_missing_safetensors_friendly(self, tmp_path, monkeypatch): from soup_cli.utils.apple_adapter import ( build_apple_adapter_plan, convert_apple_adapter, ) monkeypatch.chdir(tmp_path) (tmp_path / "adapter").mkdir() plan = build_apple_adapter_plan( source_dir="adapter", output_dir="out", direction="hf-to-mlx", ) with pytest.raises(FileNotFoundError, match="adapter_model"): convert_apple_adapter(plan) def test_bin_adapter_rejected(self, tmp_path, monkeypatch): from soup_cli.utils.apple_adapter import ( build_apple_adapter_plan, convert_apple_adapter, ) monkeypatch.chdir(tmp_path) adapter = tmp_path / "adapter" adapter.mkdir() (adapter / "adapter_model.bin").write_bytes(b"\x80\x02") plan = build_apple_adapter_plan( source_dir="adapter", output_dir="out", direction="hf-to-mlx", ) with pytest.raises(ValueError, match="safetensors"): convert_apple_adapter(plan) def test_output_outside_cwd_rejected(self, tmp_path, monkeypatch): from soup_cli.utils.apple_adapter import ( build_apple_adapter_plan, convert_apple_adapter, ) workdir = tmp_path / "work" workdir.mkdir() monkeypatch.chdir(workdir) _write_peft_adapter(workdir / "adapter") plan = build_apple_adapter_plan( source_dir="adapter", output_dir=str(tmp_path / "outside"), direction="hf-to-mlx", ) with pytest.raises(ValueError, match="cwd"): convert_apple_adapter(plan) def test_mlx_source_missing_files_friendly(self, tmp_path, monkeypatch): from soup_cli.utils.apple_adapter import ( build_apple_adapter_plan, convert_apple_adapter, ) monkeypatch.chdir(tmp_path) (tmp_path / "mlx").mkdir() plan = build_apple_adapter_plan( source_dir="mlx", output_dir="out", direction="mlx-to-hf", ) with pytest.raises(FileNotFoundError, match="adapters"): convert_apple_adapter(plan) def test_non_plan_rejected(self): from soup_cli.utils.apple_adapter import convert_apple_adapter with pytest.raises(TypeError, match="AppleAdapterPlan"): convert_apple_adapter({}) # type: ignore[arg-type] def test_report_frozen(self): import dataclasses from soup_cli.utils.apple_adapter import ConversionReport report = ConversionReport( direction="hf-to-mlx", output_dir="out", converted_keys=2, skipped_keys=(), signed=False, ) with pytest.raises(dataclasses.FrozenInstanceError): report.direction = "mlx-to-hf" # type: ignore[misc] class TestAppleAdapterCli: def test_live_conversion_exit_0(self, tmp_path, monkeypatch): from soup_cli.cli import app monkeypatch.chdir(tmp_path) _write_peft_adapter(tmp_path / "adapter") result = runner.invoke( app, [ "apple-adapter", "adapter", "--direction", "hf-to-mlx", "--output", "out", ], ) assert result.exit_code == 0, (result.output, repr(result.exception)) assert (tmp_path / "out" / "adapters.safetensors").is_file() def test_apple_direction_exit_3(self, tmp_path, monkeypatch): from soup_cli.cli import app monkeypatch.chdir(tmp_path) _write_peft_adapter(tmp_path / "adapter") result = runner.invoke( app, [ "apple-adapter", "adapter", "--direction", "hf-to-apple", "--output", "out", ], ) assert result.exit_code == 3, (result.output, repr(result.exception)) def test_missing_weights_exit_2(self, tmp_path, monkeypatch): from soup_cli.cli import app monkeypatch.chdir(tmp_path) (tmp_path / "adapter").mkdir() result = runner.invoke( app, [ "apple-adapter", "adapter", "--direction", "hf-to-mlx", "--output", "out", ], ) assert result.exit_code == 2, (result.output, repr(result.exception)) # --------------------------------------------------------------------------- # #97 — delinearize-llama4 runtime # --------------------------------------------------------------------------- class TestIsExpertWeightKey: @pytest.mark.parametrize( "key", [ "language_model.model.layers.0.feed_forward.experts.gate_up_proj", "language_model.model.layers.3.feed_forward.experts.down_proj", "model.layers.0.feed_forward.experts.gate_up_proj.weight", ], ) def test_fused_expert_keys(self, key): from soup_cli.utils.delinearize_llama4 import is_expert_weight_key assert is_expert_weight_key(key) @pytest.mark.parametrize( "key", [ # Per-expert numbered keys are already unfused (Mixtral-style). "model.layers.0.feed_forward.experts.0.gate_proj.weight", "model.layers.0.feed_forward.router.weight", "model.layers.0.self_attn.q_proj.weight", "", None, 123, ], ) def test_non_fused_keys(self, key): from soup_cli.utils.delinearize_llama4 import is_expert_weight_key assert not is_expert_weight_key(key) class TestDelinearizeTensor: def test_reshapes_2d_to_3d(self): torch = _torch_or_skip() from soup_cli.utils.delinearize_llama4 import delinearize_tensor tensor = torch.arange(24, dtype=torch.float32).reshape(6, 4) out, status = delinearize_tensor(tensor, num_experts=3) assert status == "reshaped" assert tuple(out.shape) == (3, 2, 4) assert torch.equal(out[0], tensor[0:2]) def test_3d_passthrough(self): torch = _torch_or_skip() from soup_cli.utils.delinearize_llama4 import delinearize_tensor tensor = torch.zeros(3, 2, 4) out, status = delinearize_tensor(tensor, num_experts=3) assert status == "already_3d" assert out is tensor def test_not_divisible_rejected(self): torch = _torch_or_skip() from soup_cli.utils.delinearize_llama4 import delinearize_tensor tensor = torch.zeros(7, 4) with pytest.raises(ValueError, match="divisible"): delinearize_tensor(tensor, num_experts=3) def test_1d_rejected(self): torch = _torch_or_skip() from soup_cli.utils.delinearize_llama4 import delinearize_tensor with pytest.raises(ValueError, match="2-D"): delinearize_tensor(torch.zeros(8), num_experts=2) class TestReadNumExperts: def test_text_config_nested(self, tmp_path): from soup_cli.utils.delinearize_llama4 import read_num_experts (tmp_path / "config.json").write_text( json.dumps({"text_config": {"num_local_experts": 16}}), encoding="utf-8", ) assert read_num_experts(str(tmp_path)) == 16 def test_top_level(self, tmp_path): from soup_cli.utils.delinearize_llama4 import read_num_experts (tmp_path / "config.json").write_text( json.dumps({"num_local_experts": 8}), encoding="utf-8" ) assert read_num_experts(str(tmp_path)) == 8 def test_missing_returns_none(self, tmp_path): from soup_cli.utils.delinearize_llama4 import read_num_experts (tmp_path / "config.json").write_text("{}", encoding="utf-8") assert read_num_experts(str(tmp_path)) is None def test_no_config_returns_none(self, tmp_path): from soup_cli.utils.delinearize_llama4 import read_num_experts assert read_num_experts(str(tmp_path)) is None def test_bool_value_rejected(self, tmp_path): from soup_cli.utils.delinearize_llama4 import read_num_experts (tmp_path / "config.json").write_text( json.dumps({"num_local_experts": True}), encoding="utf-8" ) assert read_num_experts(str(tmp_path)) is None def _write_llama4_stub(source: Path, *, num_experts: int = 4) -> dict: """Write a stub Llama-4-shaped checkpoint; return the tensor dict.""" torch = _torch_or_skip() from safetensors.torch import save_file tensors = { "language_model.model.layers.0.feed_forward.experts.gate_up_proj": torch.arange(48, dtype=torch.float32).reshape(8, 6), "language_model.model.layers.0.feed_forward.experts.down_proj": torch.arange(48, dtype=torch.float32).reshape(12, 4), "language_model.model.layers.0.self_attn.q_proj.weight": torch.ones(4, 4), "language_model.model.layers.1.feed_forward.experts.gate_up_proj": torch.zeros(num_experts, 2, 6), } source.mkdir(parents=True, exist_ok=True) save_file(tensors, str(source / "model.safetensors")) (source / "config.json").write_text( json.dumps({"text_config": {"num_local_experts": num_experts}}), encoding="utf-8", ) return tensors class TestRunDelinearize: def test_happy_path(self, tmp_path, monkeypatch): torch = _torch_or_skip() from safetensors.torch import load_file from soup_cli.utils.delinearize_llama4 import ( plan_delinearize, run_delinearize, ) monkeypatch.chdir(tmp_path) tensors = _write_llama4_stub(tmp_path / "src") (tmp_path / "out").mkdir() plan = plan_delinearize("src", "out") result = run_delinearize(plan) assert result.reshaped_keys == 2 assert result.passthrough_keys == 1 assert result.already_3d_keys == 1 out = load_file(str(tmp_path / "out" / "model.safetensors")) gate_up = out[ "language_model.model.layers.0.feed_forward.experts.gate_up_proj" ] assert tuple(gate_up.shape) == (4, 2, 6) original = tensors[ "language_model.model.layers.0.feed_forward.experts.gate_up_proj" ] assert torch.equal(gate_up, original.reshape(4, 2, 6)) # Sidecar config copied for a loadable target checkpoint. assert (tmp_path / "out" / "config.json").is_file() def test_explicit_num_experts_overrides(self, tmp_path, monkeypatch): from safetensors.torch import load_file from soup_cli.utils.delinearize_llama4 import ( plan_delinearize, run_delinearize, ) monkeypatch.chdir(tmp_path) _write_llama4_stub(tmp_path / "src", num_experts=4) (tmp_path / "out").mkdir() plan = plan_delinearize("src", "out") result = run_delinearize(plan, num_experts=2) assert result.reshaped_keys == 2 out = load_file(str(tmp_path / "out" / "model.safetensors")) gate_up = out[ "language_model.model.layers.0.feed_forward.experts.gate_up_proj" ] assert tuple(gate_up.shape) == (2, 4, 6) def test_missing_num_experts_friendly(self, tmp_path, monkeypatch): torch = _torch_or_skip() from safetensors.torch import save_file from soup_cli.utils.delinearize_llama4 import ( plan_delinearize, run_delinearize, ) monkeypatch.chdir(tmp_path) src = tmp_path / "src" src.mkdir() save_file({"x": torch.zeros(2, 2)}, str(src / "model.safetensors")) (tmp_path / "out").mkdir() plan = plan_delinearize("src", "out") with pytest.raises(ValueError, match="--num-experts"): run_delinearize(plan) def test_non_divisible_names_key(self, tmp_path, monkeypatch): torch = _torch_or_skip() from safetensors.torch import save_file from soup_cli.utils.delinearize_llama4 import ( plan_delinearize, run_delinearize, ) monkeypatch.chdir(tmp_path) src = tmp_path / "src" src.mkdir() save_file( { "model.layers.0.feed_forward.experts.gate_up_proj": torch.zeros(7, 4), }, str(src / "model.safetensors"), ) (tmp_path / "out").mkdir() plan = plan_delinearize("src", "out") with pytest.raises(ValueError, match="gate_up_proj"): run_delinearize(plan, num_experts=4) @pytest.mark.parametrize("bad", [True, 0, -2, "four", 1_000_000]) def test_num_experts_bounds(self, tmp_path, monkeypatch, bad): from soup_cli.utils.delinearize_llama4 import ( plan_delinearize, run_delinearize, ) monkeypatch.chdir(tmp_path) _write_llama4_stub(tmp_path / "src") (tmp_path / "out").mkdir() plan = plan_delinearize("src", "out") with pytest.raises((TypeError, ValueError), match="num_experts"): run_delinearize(plan, num_experts=bad) def test_non_plan_rejected(self): from soup_cli.utils.delinearize_llama4 import run_delinearize with pytest.raises(TypeError, match="DelinearizePlan"): run_delinearize({}) # type: ignore[arg-type] def test_result_frozen(self): import dataclasses from soup_cli.utils.delinearize_llama4 import DelinearizeResult result = DelinearizeResult( source_dir="a", target_dir="b", files_written=("x",), reshaped_keys=1, passthrough_keys=0, already_3d_keys=0, ) with pytest.raises(dataclasses.FrozenInstanceError): result.reshaped_keys = 9 # type: ignore[misc] class TestDelinearizeCli: def test_live_run_exit_0(self, tmp_path, monkeypatch): from soup_cli.cli import app monkeypatch.chdir(tmp_path) _write_llama4_stub(tmp_path / "src") (tmp_path / "out").mkdir() result = runner.invoke( app, ["delinearize-llama4", "src", "--target", "out"], ) assert result.exit_code == 0, (result.output, repr(result.exception)) assert (tmp_path / "out" / "model.safetensors").is_file() def test_plan_only_writes_nothing(self, tmp_path, monkeypatch): from soup_cli.cli import app monkeypatch.chdir(tmp_path) _write_llama4_stub(tmp_path / "src") (tmp_path / "out").mkdir() result = runner.invoke( app, ["delinearize-llama4", "src", "--target", "out", "--plan-only"], ) assert result.exit_code == 0, (result.output, repr(result.exception)) assert not (tmp_path / "out" / "model.safetensors").exists() def test_missing_num_experts_exit_2(self, tmp_path, monkeypatch): torch = _torch_or_skip() from safetensors.torch import save_file from soup_cli.cli import app monkeypatch.chdir(tmp_path) src = tmp_path / "src" src.mkdir() save_file({"x": torch.zeros(2, 2)}, str(src / "model.safetensors")) (tmp_path / "out").mkdir() result = runner.invoke( app, ["delinearize-llama4", "src", "--target", "out"], ) assert result.exit_code == 2, (result.output, repr(result.exception)) assert "--num-experts" in result.output # --------------------------------------------------------------------------- # Cross-cutting invariants # --------------------------------------------------------------------------- class TestPatchInvariants: def test_version_bumped(self): import soup_cli parts = tuple(int(p) for p in soup_cli.__version__.split(".")[:3]) assert parts >= (0, 71, 21) @pytest.mark.parametrize( "rel", [ "utils/advanced_precision.py", "utils/grpo_long_context.py", "utils/agent_rollout.py", "utils/apple_adapter.py", "utils/delinearize_llama4.py", ], ) def test_no_heavy_top_level_imports(self, rel): # Full-source scan — TYPE_CHECKING-block imports are indented so # the column-0 patterns can never match them; splitting the source # there would blind the guard for everything below the block # (review fix). source = (_SRC / rel).read_text(encoding="utf-8") for heavy in ( "torch", "transformers", "vllm", "torchao", "numpy", "safetensors", "mlx", "peft", ): assert f"\nimport {heavy}\n" not in source assert f"\nimport {heavy} " not in source assert f"\nfrom {heavy}" not in source # --------------------------------------------------------------------------- # Review-fix follow-ups (v0.71.21 review wave) # --------------------------------------------------------------------------- class TestReviewFollowupsPrecision: """#141 review fixes — torchao probe, recipe NUL, partial-conversion.""" def test_out_proj_and_wqkv_in_allowlist(self): from soup_cli.utils.advanced_precision import is_attention_projection assert is_attention_projection("encoder.layers.0.self_attn.out_proj") assert is_attention_projection("transformer.blocks.0.attn.Wqkv") def test_is_blackwell_gpu_runtime_error_false(self, monkeypatch): torch = _torch_or_skip() from soup_cli.utils.advanced_precision import is_blackwell_gpu monkeypatch.setattr(torch.cuda, "is_available", lambda: True) def _boom(_i=0): raise RuntimeError("driver mismatch") monkeypatch.setattr(torch.cuda, "get_device_capability", _boom) assert is_blackwell_gpu() is False def test_already_converted_fast_path_no_float8_import(self, monkeypatch): """All projections pre-converted -> count returned WITHOUT importing torchao.float8 (the bare fake root has no float8 attr).""" torch = _torch_or_skip() nn = torch.nn from soup_cli.utils.advanced_precision import apply_fp8_attention class _F8(nn.Linear): pass _F8.__name__ = "Float8Linear" _enable_fp8_gates(monkeypatch) model = _tiny_attn_model() attn = model.layers[0].self_attn for name in ("q_proj", "k_proj", "v_proj", "o_proj"): setattr(attn, name, _F8(8, 8)) assert apply_fp8_attention(model) == 4 def test_partial_conversion_failure_is_honest(self, monkeypatch): """torchao crashing mid-conversion surfaces 'PARTIALLY converted'.""" from soup_cli.utils.advanced_precision import apply_fp8_attention record: dict = {} _install_fake_torchao_float8(monkeypatch, record) _enable_fp8_gates(monkeypatch) def _boom(model, config=None, module_filter_fn=None): raise ValueError("unsupported dim") sys.modules["torchao.float8"].convert_to_float8_training = _boom with pytest.raises(RuntimeError, match="PARTIALLY"): apply_fp8_attention(_tiny_attn_model()) def test_nvfp4_missing_quantize_friendly(self, monkeypatch): from soup_cli.utils.advanced_precision import apply_nvfp4 fake_q = types.ModuleType("torchao.quantization") class _NVFP4Config: pass fake_q.NVFP4Config = _NVFP4Config # no quantize_ fake_root = types.ModuleType("torchao") fake_root.quantization = fake_q monkeypatch.setitem(sys.modules, "torchao", fake_root) monkeypatch.setitem(sys.modules, "torchao.quantization", fake_q) monkeypatch.setattr( "soup_cli.utils.advanced_precision.is_blackwell_gpu", lambda: True ) with pytest.raises(RuntimeError, match="quantize_"): apply_nvfp4(_tiny_attn_model()) def test_v028_degrade_is_hermetic(self, monkeypatch): """Force the degrade path by patching the converters directly — environment-independent (review fix: the original test relied on the host lacking torchao/Hopper).""" from soup_cli.config.loader import load_config_from_string from soup_cli.utils import advanced_precision, v028_features def _gate(model, **kwargs): raise RuntimeError("gate fired") # v028_features lazy-imports the converters at call time, so the # patch targets the defining module. monkeypatch.setattr(advanced_precision, "apply_fp8_attention", _gate) monkeypatch.setattr(advanced_precision, "apply_nvfp4", _gate) cfg = load_config_from_string( "base: test-llama\n" "task: sft\n" "training:\n" " quantization_aware: fp8\n" " fp8_attention: true\n" " nvfp4: true\n" "data:\n" " train: data.jsonl\n" "output: ./out\n" ) applied = v028_features.apply_v028_speed_memory( model=object(), tcfg=cfg.training, base_model=cfg.base, console=None, device="cpu", backend="transformers", ) assert applied.get("fp8_attention") is False assert applied.get("nvfp4") is False class TestReviewFollowupsSleepMode: """#124 review fixes — level validation + version edge cases.""" def test_parse_version_tuple_none(self): from soup_cli.utils.grpo_long_context import _parse_version_tuple assert _parse_version_tuple(None) == () # type: ignore[arg-type] def test_vllm_without_version_attr_false(self, monkeypatch): from soup_cli.utils.grpo_long_context import vllm_supports_sleep_mode fake = types.ModuleType("vllm") # no __version__ monkeypatch.setitem(sys.modules, "vllm", fake) assert vllm_supports_sleep_mode() is False def test_exact_floor_0_7_0_supported(self, monkeypatch): from soup_cli.utils.grpo_long_context import vllm_supports_sleep_mode _install_fake_vllm(monkeypatch, "0.7.0") assert vllm_supports_sleep_mode() is True def test_sleep_cycle_level_2_forwarded(self): from soup_cli.utils.grpo_long_context import vllm_sleep_cycle calls: list = [] engine = SimpleNamespace( sleep=lambda level: calls.append(("sleep", level)), wake_up=lambda: calls.append(("wake", None)), ) with vllm_sleep_cycle(engine, level=2): pass assert calls == [("sleep", 2), ("wake", None)] @pytest.mark.parametrize("level", [True, 0, 3, "1"]) def test_sleep_cycle_bad_level_rejected(self, level): from soup_cli.utils.grpo_long_context import vllm_sleep_cycle engine = SimpleNamespace(sleep=lambda level: None, wake_up=lambda: None) with pytest.raises((TypeError, ValueError), match="level"): with vllm_sleep_cycle(engine, level=level): pass def test_sleep_without_wake_up_warns_and_runs(self, caplog): import logging from soup_cli.utils.grpo_long_context import vllm_sleep_cycle engine = SimpleNamespace(sleep=lambda level: None) # no wake_up ran = [] with caplog.at_level( logging.WARNING, logger="soup_cli.utils.grpo_long_context" ): with vllm_sleep_cycle(engine): ran.append(True) assert ran == [True] assert any("sleep" in rec.message for rec in caplog.records) class TestReviewFollowupsRollout: """#125 review fixes — smuggle-strip, loud answers, boundaries.""" def test_extra_keys_stripped(self, tmp_path, monkeypatch): from soup_cli.utils.agent_rollout import launch_rollout spec = _write_rollout_module( tmp_path, monkeypatch, "def rollout(prompts):\n" " return [{'prompt': 'x', 'evil': 'payload', 'answer': 'a'}]\n", ) result = launch_rollout("openenv", prompts=["p"], rollout_func=spec) assert set(result.rows[0]) == {"prompt", "answer"} def test_non_str_answer_rejected_loudly(self, tmp_path, monkeypatch): from soup_cli.utils.agent_rollout import launch_rollout spec = _write_rollout_module( tmp_path, monkeypatch, "def rollout(prompts):\n" " return [{'prompt': 'x', 'answer': 42}]\n", ) with pytest.raises(ValueError, match="answer"): launch_rollout("openenv", prompts=["p"], rollout_func=spec) def test_empty_prompt_rejected(self, tmp_path, monkeypatch): from soup_cli.utils.agent_rollout import launch_rollout spec = _write_rollout_module( tmp_path, monkeypatch, "def rollout(prompts):\n return [{'prompt': ''}]\n", ) with pytest.raises(ValueError, match="empty prompt"): launch_rollout("openenv", prompts=["p"], rollout_func=spec) def test_empty_rows_result_rejected(self): from soup_cli.utils.agent_rollout import RolloutResult with pytest.raises(ValueError, match="no rows"): RolloutResult(backend="openenv", rows=()) def test_result_stores_canonical_backend(self): from soup_cli.utils.agent_rollout import RolloutResult result = RolloutResult(backend="ART", rows=({"prompt": "x"},)) assert result.backend == "art" def test_message_list_prompts_alias_broken(self): """Mutating the rollout callable's retained message dict must not leak into the normalised training rows (immutability policy).""" from soup_cli.utils.agent_rollout import _normalise_rollout_rows message = {"role": "user", "content": "hi"} rows = _normalise_rollout_rows([{"prompt": [message]}], "openenv") message["content"] = "MUTATED" assert rows[0]["prompt"][0]["content"] == "hi" def test_validate_rollout_func_bool_rejected(self): from soup_cli.utils.agent_rollout import validate_rollout_func # ValueError (not TypeError) so the Pydantic field_validator wraps # it into a ValidationError (mode="before" validators re-raise # TypeError raw — v2 convention). with pytest.raises(ValueError, match="rollout_func"): validate_rollout_func(True) # type: ignore[arg-type] @pytest.mark.parametrize( ("max_steps", "ok"), [(100_000, True), (100_001, False)] ) def test_max_steps_exact_boundary( self, tmp_path, monkeypatch, max_steps, ok ): from soup_cli.utils.agent_rollout import launch_rollout spec = _write_rollout_module( tmp_path, monkeypatch, "def rollout(prompts):\n return [{'prompt': 'x'}]\n", ) if ok: result = launch_rollout( "openenv", prompts=["p"], rollout_func=spec, max_steps=max_steps, ) assert len(result.rows) == 1 else: with pytest.raises(ValueError, match="max_steps"): launch_rollout( "openenv", prompts=["p"], rollout_func=spec, max_steps=max_steps, ) def test_external_runner_receives_all_kwargs(self, monkeypatch): from soup_cli.utils import agent_rollout seen: dict = {} def _runner(**kwargs): seen.update(kwargs) return [{"prompt": "from-runner"}] monkeypatch.setitem( agent_rollout._EXTERNAL_ROLLOUT_RUNNERS, "art", _runner ) model, tokenizer, reward = object(), object(), object() agent_rollout.launch_rollout( "art", prompts=["p"], model=model, tokenizer=tokenizer, reward_fn=reward, max_steps=7, ) assert seen["prompts"] == ["p"] assert seen["model"] is model assert seen["tokenizer"] is tokenizer assert seen["reward_fn"] is reward assert seen["max_steps"] == 7 class TestReviewFollowupsAppleAdapter: """#228 review fixes — symlink TOCTOU, corrupt files, config carry.""" def test_infer_num_layers_gpt2_style(self): """GPT-2-style ``transformer.h.N`` paths also derive num_layers (caught by the real bf16 PEFT adapter smoke).""" from soup_cli.utils.apple_adapter import _infer_num_layers assert _infer_num_layers({ "transformer.h.4.attn.c_attn.lora_a": None, "transformer.h.0.attn.c_attn.lora_b": None, }) == 5 assert _infer_num_layers({"model.layers.2.q_proj.lora_a": None}) == 3 assert _infer_num_layers({"no_layer_key.lora_a": None}) is None def test_prefixless_hf_key_mapped(self): from soup_cli.utils.apple_adapter import hf_key_to_mlx, mlx_key_to_hf assert ( hf_key_to_mlx("model.layers.0.self_attn.q_proj.lora_A.weight") == "model.layers.0.self_attn.q_proj.lora_a" ) # The reverse direction re-adds the canonical prefix (documented # asymmetry — prefix-less inputs do not byte-round-trip). assert mlx_key_to_hf("model.layers.0.self_attn.q_proj.lora_a") == ( "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight" ) def test_lora_dropout_carried_through(self, tmp_path, monkeypatch): from soup_cli.utils.apple_adapter import ( build_apple_adapter_plan, convert_apple_adapter, ) monkeypatch.chdir(tmp_path) _write_peft_adapter(tmp_path / "adapter") config_path = tmp_path / "adapter" / "adapter_config.json" config = json.loads(config_path.read_text(encoding="utf-8")) config["lora_dropout"] = 0.05 config_path.write_text(json.dumps(config), encoding="utf-8") convert_apple_adapter( build_apple_adapter_plan( source_dir="adapter", output_dir="out", direction="hf-to-mlx", ) ) out_config = json.loads( (tmp_path / "out" / "adapter_config.json").read_text( encoding="utf-8" ) ) assert out_config["lora_parameters"]["dropout"] == 0.05 def test_legacy_npz_input_still_loads(self, tmp_path, monkeypatch): from safetensors.numpy import load_file from soup_cli.utils.apple_adapter import ( build_apple_adapter_plan, convert_apple_adapter, ) monkeypatch.chdir(tmp_path) mlx_dir = tmp_path / "mlx" mlx_dir.mkdir() rng = np.random.default_rng(3) np.savez( str(mlx_dir / "adapters.npz"), **{ "model.layers.0.self_attn.q_proj.lora_a": rng.standard_normal((16, 4)).astype(np.float32), "model.layers.0.self_attn.q_proj.lora_b": rng.standard_normal((4, 16)).astype(np.float32), }, ) report = convert_apple_adapter( build_apple_adapter_plan( source_dir="mlx", output_dir="hf", direction="mlx-to-hf", ) ) assert report.converted_keys == 2 back = load_file(str(tmp_path / "hf" / "adapter_model.safetensors")) assert ( "base_model.model.model.layers.0.self_attn.q_proj.lora_A.weight" in back ) def test_corrupt_safetensors_value_error(self, tmp_path, monkeypatch): pytest.importorskip("torch") from soup_cli.utils.apple_adapter import ( build_apple_adapter_plan, convert_apple_adapter, ) monkeypatch.chdir(tmp_path) adapter = tmp_path / "adapter" adapter.mkdir() (adapter / "adapter_model.safetensors").write_bytes(b"not-safetensors") plan = build_apple_adapter_plan( source_dir="adapter", output_dir="out", direction="hf-to-mlx", ) with pytest.raises(ValueError, match="not a valid safetensors"): convert_apple_adapter(plan) def test_corrupt_npz_value_error(self, tmp_path, monkeypatch): from soup_cli.utils.apple_adapter import ( build_apple_adapter_plan, convert_apple_adapter, ) monkeypatch.chdir(tmp_path) mlx_dir = tmp_path / "mlx" mlx_dir.mkdir() (mlx_dir / "adapters.npz").write_bytes(b"not-a-zip") plan = build_apple_adapter_plan( source_dir="mlx", output_dir="out", direction="mlx-to-hf", ) with pytest.raises(ValueError, match="not a valid npz"): convert_apple_adapter(plan) def test_npz_decompression_cap(self, tmp_path, monkeypatch): """The 4 GiB cap re-applies to DECOMPRESSED arrays (zip bomb).""" from soup_cli.utils import apple_adapter monkeypatch.chdir(tmp_path) mlx_dir = tmp_path / "mlx" mlx_dir.mkdir() zeros = np.zeros((128, 128), dtype=np.float32) # 64 KiB, compresses tiny with open(mlx_dir / "adapters.npz", "wb") as handle: np.savez_compressed( handle, **{"model.layers.0.self_attn.q_proj.lora_a": zeros} ) on_disk = (mlx_dir / "adapters.npz").stat().st_size monkeypatch.setattr( apple_adapter, "_MAX_ADAPTER_FILE_BYTES", on_disk + 1024 ) plan = apple_adapter.build_apple_adapter_plan( source_dir="mlx", output_dir="out", direction="mlx-to-hf", ) with pytest.raises(ValueError, match="decompresses past"): apple_adapter.convert_apple_adapter(plan) def test_adapter_size_cap_branch(self, tmp_path, monkeypatch): from soup_cli.utils import apple_adapter monkeypatch.chdir(tmp_path) _write_peft_adapter(tmp_path / "adapter") monkeypatch.setattr(apple_adapter, "_MAX_ADAPTER_FILE_BYTES", 16) plan = apple_adapter.build_apple_adapter_plan( source_dir="adapter", output_dir="out", direction="hf-to-mlx", ) with pytest.raises(ValueError, match="cap"): apple_adapter.convert_apple_adapter(plan) @pytest.mark.skipif(sys.platform == "win32", reason="POSIX symlink") def test_symlinked_weights_rejected(self, tmp_path, monkeypatch): import os from soup_cli.utils.apple_adapter import ( build_apple_adapter_plan, convert_apple_adapter, ) monkeypatch.chdir(tmp_path) real = tmp_path / "real" _write_peft_adapter(real) adapter = tmp_path / "adapter" adapter.mkdir() os.symlink( str(real / "adapter_model.safetensors"), str(adapter / "adapter_model.safetensors"), ) plan = build_apple_adapter_plan( source_dir="adapter", output_dir="out", direction="hf-to-mlx", ) with pytest.raises(ValueError, match="symlink"): convert_apple_adapter(plan) @pytest.mark.parametrize( ("field", "value", "exc"), [ ("converted_keys", -1, ValueError), ("converted_keys", True, TypeError), ("skipped_keys", ["a"], TypeError), ("signed", "yes", TypeError), ], ) def test_conversion_report_post_init(self, field, value, exc): from soup_cli.utils.apple_adapter import ConversionReport kwargs = { "direction": "hf-to-mlx", "output_dir": "out", "converted_keys": 2, "skipped_keys": (), "signed": False, } kwargs[field] = value with pytest.raises(exc): ConversionReport(**kwargs) class TestReviewFollowupsDelinearize: """#97 review fixes — plan containment, corrupt shards, boundaries.""" @pytest.mark.parametrize( ("value", "ok"), [(4096, True), (4097, False), (1, True), (0, False)] ) def test_num_experts_exact_boundary(self, value, ok): from soup_cli.utils.delinearize_llama4 import _validate_num_experts if ok: assert _validate_num_experts(value) == value else: with pytest.raises(ValueError, match="num_experts"): _validate_num_experts(value) def test_plan_direct_construction_outside_cwd_rejected( self, tmp_path, monkeypatch ): from soup_cli.utils.delinearize_llama4 import DelinearizePlan workdir = tmp_path / "work" workdir.mkdir() monkeypatch.chdir(workdir) with pytest.raises(ValueError, match="outside cwd"): DelinearizePlan( source_dir=str(workdir / "src"), target_dir=str(tmp_path / "outside"), weight_files=("model.safetensors",), ) def test_corrupt_shard_value_error(self, tmp_path, monkeypatch): pytest.importorskip("torch") from soup_cli.utils.delinearize_llama4 import ( plan_delinearize, run_delinearize, ) monkeypatch.chdir(tmp_path) src = tmp_path / "src" src.mkdir() (src / "model.safetensors").write_bytes(b"garbage") (tmp_path / "out").mkdir() plan = plan_delinearize("src", "out") with pytest.raises(ValueError, match="not a valid safetensors"): run_delinearize(plan, num_experts=4) def test_weight_file_size_cap_branch(self, tmp_path, monkeypatch): pytest.importorskip("torch") from soup_cli.utils import delinearize_llama4 monkeypatch.chdir(tmp_path) _write_llama4_stub(tmp_path / "src") (tmp_path / "out").mkdir() plan = delinearize_llama4.plan_delinearize("src", "out") monkeypatch.setattr(delinearize_llama4, "_MAX_WEIGHT_FILE_BYTES", 16) with pytest.raises(ValueError, match="cap"): delinearize_llama4.run_delinearize(plan) def test_multi_shard_run(self, tmp_path, monkeypatch): torch = _torch_or_skip() from safetensors.torch import save_file from soup_cli.utils.delinearize_llama4 import ( plan_delinearize, run_delinearize, ) monkeypatch.chdir(tmp_path) _write_llama4_stub(tmp_path / "src") save_file( { "language_model.model.layers.2.feed_forward.experts.down_proj": torch.zeros(8, 4), }, str(tmp_path / "src" / "model-00002.safetensors"), ) (tmp_path / "out").mkdir() result = run_delinearize(plan_delinearize("src", "out")) assert len(result.files_written) == 2 assert result.sidecars_copied >= 1 # config.json @pytest.mark.parametrize( ("config", "expected"), [ ({"num_experts": 16}, 16), # top-level fallback key ({"text_config": {"num_local_experts": 5000}}, None), # OOB ], ) def test_read_num_experts_fallbacks( self, tmp_path, monkeypatch, config, expected ): from soup_cli.utils.delinearize_llama4 import read_num_experts monkeypatch.chdir(tmp_path) src = tmp_path / "src" src.mkdir() (src / "config.json").write_text(json.dumps(config), encoding="utf-8") assert read_num_experts(str(src)) == expected @pytest.mark.skipif(sys.platform == "win32", reason="POSIX symlink") def test_symlinked_shard_rejected(self, tmp_path, monkeypatch): import os pytest.importorskip("torch") from soup_cli.utils.delinearize_llama4 import ( plan_delinearize, run_delinearize, ) monkeypatch.chdir(tmp_path) real = tmp_path / "real" _write_llama4_stub(real) src = tmp_path / "src" src.mkdir() os.symlink( str(real / "model.safetensors"), str(src / "model.safetensors") ) (tmp_path / "out").mkdir() plan = plan_delinearize("src", "out") with pytest.raises(ValueError, match="symlink"): run_delinearize(plan, num_experts=4) @pytest.mark.skipif(sys.platform == "win32", reason="POSIX symlink") def test_symlinked_config_json_skipped(self, tmp_path, monkeypatch): import os from soup_cli.utils.delinearize_llama4 import read_num_experts monkeypatch.chdir(tmp_path) real = tmp_path / "real" real.mkdir() (real / "config.json").write_text( json.dumps({"num_local_experts": 4}), encoding="utf-8" ) src = tmp_path / "src" src.mkdir() os.symlink(str(real / "config.json"), str(src / "config.json")) assert read_num_experts(str(src)) is None def test_cli_num_experts_flag_passthrough(self, tmp_path, monkeypatch): pytest.importorskip("torch") from soup_cli.cli import app monkeypatch.chdir(tmp_path) src = tmp_path / "src" _write_llama4_stub(src) (src / "config.json").unlink() # force the flag to matter (tmp_path / "out").mkdir() result = runner.invoke( app, [ "delinearize-llama4", "src", "--target", "out", "--num-experts", "4", ], ) assert result.exit_code == 0, (result.output, repr(result.exception)) assert (tmp_path / "out" / "model.safetensors").is_file()