test_deepspeed.py

"""Tests for Multi-GPU / DeepSpeed support."""

import json
import os
from unittest.mock import MagicMock, patch

import pytest


class TestDeepSpeedConfigs:
    """Test DeepSpeed configuration templates."""

    def test_zero2_config_structure(self):
        """ZeRO Stage 2 config should have correct structure."""
        from soup_cli.utils.deepspeed import get_deepspeed_config

        config = get_deepspeed_config("zero2")
        assert config["zero_optimization"]["stage"] == 2
        assert config["bf16"]["enabled"] is True
        assert config["gradient_accumulation_steps"] == "auto"

    def test_zero3_config_structure(self):
        """ZeRO Stage 3 config should have correct structure."""
        from soup_cli.utils.deepspeed import get_deepspeed_config

        config = get_deepspeed_config("zero3")
        assert config["zero_optimization"]["stage"] == 3
        assert config["zero_optimization"]["stage3_gather_16bit_weights_on_model_save"] is True

    def test_zero2_offload_config(self):
        """ZeRO Stage 2 with offload should enable CPU offloading."""
        from soup_cli.utils.deepspeed import get_deepspeed_config

        config = get_deepspeed_config("zero2_offload")
        assert config["zero_optimization"]["stage"] == 2
        offload = config["zero_optimization"]["offload_optimizer"]
        assert offload["device"] == "cpu"
        assert offload["pin_memory"] is True

    def test_invalid_config_name(self):
        """Should raise ValueError for unknown config name."""
        from soup_cli.utils.deepspeed import get_deepspeed_config

        with pytest.raises(ValueError, match="Unknown DeepSpeed config"):
            get_deepspeed_config("zero99")

    def test_get_config_returns_copy(self):
        """Should return a copy, not the original."""
        from soup_cli.utils.deepspeed import get_deepspeed_config

        config1 = get_deepspeed_config("zero2")
        config2 = get_deepspeed_config("zero2")
        config1["bf16"]["enabled"] = False
        assert config2["bf16"]["enabled"] is True

    def test_all_configs_have_auto_fields(self):
        """All configs should have 'auto' for batch sizes."""
        from soup_cli.utils.deepspeed import CONFIGS

        for name, config in CONFIGS.items():
            assert config["train_batch_size"] == "auto", f"{name} missing auto train_batch_size"
            assert config["train_micro_batch_size_per_gpu"] == "auto", (
                f"{name} missing auto micro batch"
            )


class TestWriteDeepSpeedConfig:
    """Test writing DeepSpeed config to temp file."""

    def test_write_creates_file(self):
        """Should create a valid JSON file."""
        from soup_cli.utils.deepspeed import write_deepspeed_config

        path = write_deepspeed_config("zero2")
        assert os.path.exists(path)

        with open(path) as f:
            config = json.load(f)
        assert config["zero_optimization"]["stage"] == 2

        # Cleanup
        os.unlink(path)

    def test_write_file_is_valid_json(self):
        """Written file should be parseable JSON."""
        from soup_cli.utils.deepspeed import write_deepspeed_config

        for stage in ["zero2", "zero3", "zero2_offload"]:
            path = write_deepspeed_config(stage)
            with open(path) as f:
                config = json.load(f)
            assert "zero_optimization" in config
            os.unlink(path)


class TestDetectMultiGPU:
    """Test multi-GPU detection."""

    def test_detect_no_gpu(self):
        """Should return 0 GPUs when CUDA not available."""
        from soup_cli.utils.deepspeed import detect_multi_gpu

        with patch("torch.cuda.is_available", return_value=False):
            result = detect_multi_gpu()
            assert result["gpu_count"] == 0
            assert result["gpus"] == []

    def test_detect_single_gpu(self):
        """Should detect a single GPU."""
        from soup_cli.utils.deepspeed import detect_multi_gpu

        mock_props = MagicMock()
        mock_props.name = "NVIDIA RTX 4090"
        mock_props.total_memory = 24 * (1024 ** 3)  # 24GB

        with patch("torch.cuda.is_available", return_value=True), \
             patch("torch.cuda.device_count", return_value=1), \
             patch("torch.cuda.get_device_properties", return_value=mock_props):
            result = detect_multi_gpu()
            assert result["gpu_count"] == 1
            assert len(result["gpus"]) == 1
            assert result["gpus"][0]["name"] == "NVIDIA RTX 4090"
            assert result["gpus"][0]["memory_gb"] == pytest.approx(24.0)

    def test_detect_multiple_gpus(self):
        """Should detect multiple GPUs."""
        from soup_cli.utils.deepspeed import detect_multi_gpu

        mock_props = MagicMock()
        mock_props.name = "NVIDIA A100"
        mock_props.total_memory = 80 * (1024 ** 3)

        with patch("torch.cuda.is_available", return_value=True), \
             patch("torch.cuda.device_count", return_value=4), \
             patch("torch.cuda.get_device_properties", return_value=mock_props):
            result = detect_multi_gpu()
            assert result["gpu_count"] == 4
            assert len(result["gpus"]) == 4

    def test_detect_without_torch(self):
        """Should handle missing torch gracefully."""
        from soup_cli.utils.deepspeed import detect_multi_gpu

        with patch.dict("sys.modules", {"torch": None}):
            # Import error should be caught
            result = detect_multi_gpu()
            assert result["gpu_count"] == 0


class TestResolveDeepSpeed:
    """Test DeepSpeed config resolution in train command."""

    def test_resolve_named_preset(self):
        """Should resolve named presets like 'zero2'."""
        from soup_cli.commands.train import _resolve_deepspeed

        path = _resolve_deepspeed("zero2")
        assert os.path.exists(path)

        with open(path) as f:
            config = json.load(f)
        assert config["zero_optimization"]["stage"] == 2
        os.unlink(path)

    def test_resolve_json_file(self, tmp_path):
        """Should resolve path to JSON file."""
        from soup_cli.commands.train import _resolve_deepspeed

        config_file = tmp_path / "ds_config.json"
        config_file.write_text(json.dumps({"zero_optimization": {"stage": 2}}))

        result = _resolve_deepspeed(str(config_file))
        assert result == str(config_file)

    def test_resolve_invalid_name(self):
        """Should raise exit for invalid name."""
        from click.exceptions import Exit

        from soup_cli.commands.train import _resolve_deepspeed

        with pytest.raises(Exit):
            _resolve_deepspeed("invalid_config")


class TestTrainerDeepSpeedParam:
    """Test that trainers accept deepspeed_config parameter."""

    def test_sft_trainer_accepts_deepspeed(self):
        """SFTTrainerWrapper should accept deepspeed_config."""
        from soup_cli.config.schema import SoupConfig
        from soup_cli.trainer.sft import SFTTrainerWrapper

        cfg = SoupConfig(
            base="test-model",
            data={"train": "test.jsonl"},
        )
        wrapper = SFTTrainerWrapper(cfg, device="cpu", deepspeed_config="/tmp/ds.json")
        assert wrapper.deepspeed_config == "/tmp/ds.json"

    def test_dpo_trainer_accepts_deepspeed(self):
        """DPOTrainerWrapper should accept deepspeed_config."""
        from soup_cli.config.schema import SoupConfig
        from soup_cli.trainer.dpo import DPOTrainerWrapper

        cfg = SoupConfig(
            base="test-model",
            task="dpo",
            data={"train": "test.jsonl"},
        )
        wrapper = DPOTrainerWrapper(cfg, device="cpu", deepspeed_config="/tmp/ds.json")
        assert wrapper.deepspeed_config == "/tmp/ds.json"

    def test_sft_trainer_default_no_deepspeed(self):
        """SFTTrainerWrapper should default to no DeepSpeed."""
        from soup_cli.config.schema import SoupConfig
        from soup_cli.trainer.sft import SFTTrainerWrapper

        cfg = SoupConfig(
            base="test-model",
            data={"train": "test.jsonl"},
        )
        wrapper = SFTTrainerWrapper(cfg, device="cpu")
        assert wrapper.deepspeed_config is None


class TestTrainDeepSpeedFlag:
    """Test --deepspeed flag in train command."""

    def test_train_help_shows_deepspeed(self):
        """Train help should mention --deepspeed option."""
        from typer.testing import CliRunner

        from soup_cli.cli import app

        runner = CliRunner()
        result = runner.invoke(app, ["train", "--help"])
        assert "deepspeed" in result.output.lower()