mirror of https://github.com/razor-ai/soup.git
289 lines
9.3 KiB
Python
289 lines
9.3 KiB
Python
"""Tests for Apple Silicon MLX backend — Part E of v0.25.0.
|
|
|
|
These tests mock MLX entirely so they run on CI (Linux / Windows / macOS).
|
|
"""
|
|
|
|
|
|
import pytest
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MLX detection
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestMLXDetection:
|
|
def test_detect_mlx_not_installed(self, monkeypatch):
|
|
"""detect_mlx returns False if mlx import fails."""
|
|
import builtins
|
|
|
|
real_import = builtins.__import__
|
|
|
|
def fake_import(name, *args, **kwargs):
|
|
if name.startswith("mlx"):
|
|
raise ImportError("no mlx")
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(builtins, "__import__", fake_import)
|
|
from soup_cli.utils import mlx as mlx_utils
|
|
|
|
# Force re-check via direct call
|
|
assert mlx_utils.detect_mlx() is False
|
|
|
|
def test_detect_mlx_installed_mock(self, monkeypatch):
|
|
"""detect_mlx returns True when mlx modules are importable (mocked)."""
|
|
import sys
|
|
import types
|
|
|
|
fake_mlx = types.ModuleType("mlx")
|
|
fake_mlx.__version__ = "0.20.0"
|
|
fake_core = types.ModuleType("mlx.core")
|
|
fake_core.metal = types.SimpleNamespace(is_available=lambda: True)
|
|
fake_mlx.core = fake_core
|
|
monkeypatch.setitem(sys.modules, "mlx", fake_mlx)
|
|
monkeypatch.setitem(sys.modules, "mlx.core", fake_core)
|
|
|
|
from soup_cli.utils import mlx as mlx_utils
|
|
|
|
assert mlx_utils.detect_mlx() is True
|
|
|
|
def test_get_mlx_info_not_installed(self, monkeypatch):
|
|
import builtins
|
|
|
|
real_import = builtins.__import__
|
|
|
|
def fake_import(name, *args, **kwargs):
|
|
if name.startswith("mlx"):
|
|
raise ImportError("no mlx")
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(builtins, "__import__", fake_import)
|
|
from soup_cli.utils import mlx as mlx_utils
|
|
|
|
info = mlx_utils.get_mlx_info()
|
|
assert info["available"] is False
|
|
|
|
def test_estimate_mlx_batch_size_small_model(self):
|
|
from soup_cli.utils.mlx import estimate_mlx_batch_size
|
|
|
|
# 7B model on 16GB unified memory
|
|
batch = estimate_mlx_batch_size(
|
|
model_params_b=7.0,
|
|
unified_memory_bytes=16 * 1024**3,
|
|
max_length=2048,
|
|
quantization="4bit",
|
|
)
|
|
assert batch >= 1
|
|
|
|
def test_estimate_mlx_batch_size_large_model_tiny_mem(self):
|
|
from soup_cli.utils.mlx import estimate_mlx_batch_size
|
|
|
|
# 70B on 16GB is not going to fit — should return 1 minimum
|
|
batch = estimate_mlx_batch_size(
|
|
model_params_b=70.0,
|
|
unified_memory_bytes=16 * 1024**3,
|
|
max_length=2048,
|
|
quantization="4bit",
|
|
)
|
|
assert batch >= 1
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Backend enum
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestMLXBackendConfig:
|
|
def test_backend_mlx_accepted(self):
|
|
from soup_cli.config.loader import load_config_from_string
|
|
|
|
yaml_str = """
|
|
base: mlx-community/Llama-3.1-8B-Instruct-4bit
|
|
task: sft
|
|
backend: mlx
|
|
data:
|
|
train: ./data/train.jsonl
|
|
format: chatml
|
|
training:
|
|
epochs: 1
|
|
lr: 1e-4
|
|
output: ./output
|
|
"""
|
|
cfg = load_config_from_string(yaml_str)
|
|
assert cfg.backend == "mlx"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MLX SFT trainer wrapper (mocked)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestMLXSFTTrainer:
|
|
def test_trainer_import(self):
|
|
"""Import the MLX SFT trainer."""
|
|
from soup_cli.trainer.mlx_sft import MLXSFTTrainerWrapper
|
|
|
|
assert MLXSFTTrainerWrapper is not None
|
|
|
|
def test_trainer_setup_mocked(self, tmp_path):
|
|
from soup_cli.config.schema import DataConfig, SoupConfig, TrainingConfig
|
|
from soup_cli.trainer.mlx_sft import MLXSFTTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="mlx-community/Llama-3.1-8B-Instruct-4bit",
|
|
task="sft",
|
|
backend="mlx",
|
|
data=DataConfig(train="./data/train.jsonl", format="chatml"),
|
|
training=TrainingConfig(epochs=1),
|
|
output=str(tmp_path),
|
|
)
|
|
wrapper = MLXSFTTrainerWrapper(cfg)
|
|
assert wrapper.config is cfg
|
|
assert wrapper.model is None
|
|
|
|
def test_trainer_raises_when_mlx_missing(self, tmp_path, monkeypatch):
|
|
import builtins
|
|
|
|
real_import = builtins.__import__
|
|
|
|
def fake_import(name, *args, **kwargs):
|
|
if name.startswith("mlx"):
|
|
raise ImportError("no mlx")
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
monkeypatch.setattr(builtins, "__import__", fake_import)
|
|
|
|
from soup_cli.config.schema import DataConfig, SoupConfig, TrainingConfig
|
|
from soup_cli.trainer.mlx_sft import MLXSFTTrainerWrapper
|
|
|
|
cfg = SoupConfig(
|
|
base="mlx-community/Llama-3.1-8B-Instruct-4bit",
|
|
task="sft",
|
|
backend="mlx",
|
|
data=DataConfig(train="./data/train.jsonl", format="chatml"),
|
|
training=TrainingConfig(),
|
|
output=str(tmp_path),
|
|
)
|
|
wrapper = MLXSFTTrainerWrapper(cfg)
|
|
with pytest.raises((ImportError, RuntimeError)):
|
|
wrapper.setup({"train": [], "val": []})
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MLX DPO + GRPO trainers — smoke import
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestMLXOtherTrainers:
|
|
def test_mlx_dpo_import(self):
|
|
from soup_cli.trainer.mlx_dpo import MLXDPOTrainerWrapper
|
|
|
|
assert MLXDPOTrainerWrapper is not None
|
|
|
|
def test_mlx_grpo_import(self):
|
|
from soup_cli.trainer.mlx_grpo import MLXGRPOTrainerWrapper
|
|
|
|
assert MLXGRPOTrainerWrapper is not None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# train command routing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestMLXRouting:
|
|
def test_mlx_routing_map(self):
|
|
"""Routing dict should map backend=mlx tasks to MLX trainers."""
|
|
from soup_cli.trainer import mlx_routing
|
|
|
|
assert mlx_routing.MLX_TRAINER_REGISTRY["sft"].__name__ == "MLXSFTTrainerWrapper"
|
|
assert mlx_routing.MLX_TRAINER_REGISTRY["dpo"].__name__ == "MLXDPOTrainerWrapper"
|
|
assert mlx_routing.MLX_TRAINER_REGISTRY["grpo"].__name__ == "MLXGRPOTrainerWrapper"
|
|
|
|
def test_mlx_unsupported_task_rejected(self):
|
|
from soup_cli.trainer import mlx_routing
|
|
|
|
assert "ppo" not in mlx_routing.MLX_TRAINER_REGISTRY
|
|
assert "pretrain" not in mlx_routing.MLX_TRAINER_REGISTRY
|
|
assert "embedding" not in mlx_routing.MLX_TRAINER_REGISTRY
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Recipes
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestMLXRecipes:
|
|
def test_llama3_1_8b_sft_mlx(self):
|
|
from soup_cli.config.loader import load_config_from_string
|
|
from soup_cli.recipes.catalog import get_recipe
|
|
|
|
recipe = get_recipe("llama3.1-8b-sft-mlx")
|
|
assert recipe is not None
|
|
cfg = load_config_from_string(recipe.yaml_str)
|
|
assert cfg.backend == "mlx"
|
|
assert cfg.task == "sft"
|
|
|
|
def test_qwen3_8b_sft_mlx(self):
|
|
from soup_cli.config.loader import load_config_from_string
|
|
from soup_cli.recipes.catalog import get_recipe
|
|
|
|
recipe = get_recipe("qwen3-8b-sft-mlx")
|
|
assert recipe is not None
|
|
cfg = load_config_from_string(recipe.yaml_str)
|
|
assert cfg.backend == "mlx"
|
|
|
|
def test_gemma3_9b_sft_mlx(self):
|
|
from soup_cli.recipes.catalog import get_recipe
|
|
|
|
recipe = get_recipe("gemma3-9b-sft-mlx")
|
|
assert recipe is not None
|
|
|
|
def test_mlx_dpo_config_rejected_at_load(self):
|
|
"""backend=mlx + task=dpo is rejected by the SoupConfig validator."""
|
|
from soup_cli.config.loader import load_config_from_string
|
|
|
|
yaml_str = """
|
|
base: mlx-community/Llama-3.1-8B-Instruct-4bit
|
|
task: dpo
|
|
backend: mlx
|
|
data:
|
|
train: ./x.jsonl
|
|
format: dpo
|
|
training:
|
|
epochs: 1
|
|
lr: 1e-6
|
|
output: ./output
|
|
"""
|
|
with pytest.raises(ValueError, match="MLX backend only ships SFT"):
|
|
load_config_from_string(yaml_str)
|
|
|
|
def test_mlx_grpo_config_rejected_at_load(self):
|
|
from soup_cli.config.loader import load_config_from_string
|
|
|
|
yaml_str = """
|
|
base: mlx-community/Llama-3.1-8B-Instruct-4bit
|
|
task: grpo
|
|
backend: mlx
|
|
data:
|
|
train: ./x.jsonl
|
|
format: chatml
|
|
training:
|
|
epochs: 1
|
|
lr: 1e-6
|
|
output: ./output
|
|
"""
|
|
with pytest.raises(ValueError, match="MLX backend only ships SFT"):
|
|
load_config_from_string(yaml_str)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# doctor command reports MLX
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestMLXDoctor:
|
|
def test_doctor_has_mlx_info(self):
|
|
"""`soup doctor` helpers surface MLX info (no crash on non-Apple)."""
|
|
from soup_cli.commands.doctor import _get_mlx_info
|
|
|
|
info = _get_mlx_info()
|
|
assert isinstance(info, dict)
|
|
assert "available" in info
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|