unslothai--unsloth/studio/backend/tests/test_gpu_selection_sandbox.py

#!/usr/bin/env python3
"""
Sandbox test for multi-GPU selection logic.

Tests GPU selection, memory estimation, and device_map logic in
isolation. Runs on Linux, macOS, and Windows without real GPUs -- all
hardware calls are mocked.

Usage:
    python -m pytest studio/backend/tests/test_gpu_selection_sandbox.py -v
    # or directly:
    python studio/backend/tests/test_gpu_selection_sandbox.py
"""

import os
import sys
import unittest
from pathlib import Path
from unittest.mock import patch, MagicMock

# Ensure backend is on sys.path.
_backend_root = Path(__file__).resolve().parent.parent
if str(_backend_root) not in sys.path:
    sys.path.insert(0, str(_backend_root))


def _make_fake_config(
    vocab_size = 32000,
    hidden_size = 4096,
    intermediate_size = 11008,
    num_hidden_layers = 32,
    num_attention_heads = 32,
    num_key_value_heads = 8,
    tie_word_embeddings = False,
):
    """Fake HF config-like object for estimation tests."""
    from types import SimpleNamespace
    return SimpleNamespace(
        vocab_size = vocab_size,
        hidden_size = hidden_size,
        intermediate_size = intermediate_size,
        num_hidden_layers = num_hidden_layers,
        num_attention_heads = num_attention_heads,
        num_key_value_heads = num_key_value_heads,
        tie_word_embeddings = tie_word_embeddings,
    )


class TestEstimateFP16ModelSizeFromConfig(unittest.TestCase):
    """Config-based model size estimation."""

    def test_llama_8b_size_reasonable(self):
        from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config

        config = _make_fake_config(
            vocab_size = 128256,
            hidden_size = 4096,
            intermediate_size = 14336,
            num_hidden_layers = 32,
            num_attention_heads = 32,
            num_key_value_heads = 8,
            tie_word_embeddings = False,
        )
        size = _estimate_fp16_model_size_bytes_from_config(config)
        self.assertIsNotNone(size)
        size_gb = size / (1024**3)
        # Llama 3.1 8B should be ~15GB in fp16
        self.assertGreater(size_gb, 12)
        self.assertLess(size_gb, 20)

    def test_small_model(self):
        from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config

        config = _make_fake_config(
            vocab_size = 32000,
            hidden_size = 2048,
            intermediate_size = 5504,
            num_hidden_layers = 22,
            num_attention_heads = 32,
            num_key_value_heads = 4,
        )
        size = _estimate_fp16_model_size_bytes_from_config(config)
        self.assertIsNotNone(size)
        size_gb = size / (1024**3)
        # ~1B model should be ~2GB in fp16
        self.assertGreater(size_gb, 1)
        self.assertLess(size_gb, 5)

    def test_returns_none_for_incomplete_config(self):
        from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config
        from types import SimpleNamespace

        config = SimpleNamespace(vocab_size = 32000)  # most fields missing
        size = _estimate_fp16_model_size_bytes_from_config(config)
        self.assertIsNone(size)

    def test_moe_model(self):
        from utils.hardware.hardware import _estimate_fp16_model_size_bytes_from_config
        from types import SimpleNamespace

        config = SimpleNamespace(
            vocab_size = 152064,
            hidden_size = 3584,
            intermediate_size = 18944,
            num_hidden_layers = 28,
            num_attention_heads = 28,
            num_key_value_heads = 4,
            tie_word_embeddings = False,
            num_local_experts = 64,
            moe_intermediate_size = 2560,
        )
        size = _estimate_fp16_model_size_bytes_from_config(config)
        self.assertIsNotNone(size)
        size_gb = size / (1024**3)
        # MoE model with 64 experts should be large
        self.assertGreater(size_gb, 50)


class TestEstimateRequiredModelMemory(unittest.TestCase):
    """Memory requirement estimation."""

    def test_inference_fp16_uses_1_3x(self):
        from utils.hardware.hardware import estimate_required_model_memory_gb
        with patch(
            "utils.hardware.hardware.estimate_fp16_model_size_bytes",
            return_value = (10 * (1024**3), "config"),  # 10GB model
        ):
            required, meta = estimate_required_model_memory_gb(
                "test/model",
                training_type = None,  # inference
                load_in_4bit = False,
            )
            self.assertIsNotNone(required)
            self.assertAlmostEqual(required, 13.0, places = 0)
            self.assertEqual(meta["mode"], "inference")

    def test_inference_4bit_uses_reduced_estimate(self):
        from utils.hardware.hardware import estimate_required_model_memory_gb
        with patch(
            "utils.hardware.hardware.estimate_fp16_model_size_bytes",
            return_value = (30 * (1024**3), "config"),  # 30GB fp16 model
        ):
            required, meta = estimate_required_model_memory_gb(
                "test/model",
                training_type = None,  # inference
                load_in_4bit = True,
            )
            self.assertIsNotNone(required)
            # 4bit base = 30/3.2 = 9.375GB, required = 9.375 + max(9.375*0.3, 2) = 12.19GB
            self.assertAlmostEqual(required, 12.2, places = 0)

    def test_4bit_training_reduces_base(self):
        from utils.hardware.hardware import estimate_required_model_memory_gb
        with patch(
            "utils.hardware.hardware.estimate_fp16_model_size_bytes",
            return_value = (30 * (1024**3), "config"),  # 30GB fp16 model
        ):
            required, meta = estimate_required_model_memory_gb(
                "test/model",
                training_type = "LoRA/QLoRA",
                load_in_4bit = True,
            )
            self.assertIsNotNone(required)
            # fallback: base=30/3.2=9.375, lora=30*0.04=1.2, act=30*0.15=4.5, cuda=1.4
            self.assertAlmostEqual(required, 16.5, places = 0)

    def test_full_finetune_uses_3_5x(self):
        from utils.hardware.hardware import estimate_required_model_memory_gb
        with patch(
            "utils.hardware.hardware.estimate_fp16_model_size_bytes",
            return_value = (10 * (1024**3), "config"),  # 10GB model
        ):
            required, meta = estimate_required_model_memory_gb(
                "test/model",
                training_type = "Full Finetuning",
            )
            self.assertIsNotNone(required)
            # fallback: 10 * 3.5 + 1.4 cuda overhead = 36.4
            self.assertAlmostEqual(required, 36.4, places = 0)

    def test_returns_none_when_unavailable(self):
        from utils.hardware.hardware import estimate_required_model_memory_gb
        with patch(
            "utils.hardware.hardware.estimate_fp16_model_size_bytes",
            return_value = (None, "unavailable"),
        ):
            required, meta = estimate_required_model_memory_gb("test/model")
            self.assertIsNone(required)


class TestAutoSelectGpuIds(unittest.TestCase):
    """Automatic GPU selection by model size and free memory."""

    def _make_utilization(self, devices):
        """Fake utilization response."""
        return {
            "available": True,
            "devices": [
                {
                    "index": idx,
                    "vram_total_gb": total,
                    "vram_used_gb": total - free,
                }
                for idx, total, free in devices
            ],
        }

    def test_single_gpu_sufficient(self):
        from utils.hardware.hardware import auto_select_gpu_ids
        import utils.hardware.hardware as hw
        with (
            patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA),
            patch.object(
                hw,
                "estimate_required_model_memory_gb",
                return_value = (
                    10.0,
                    {
                        "mode": "inference",
                        "required_gb": 10.0,
                        "model_size_source": "config",
                        "model_size_gb": 7.7,
                    },
                ),
            ),
            patch.object(
                hw,
                "_get_parent_visible_gpu_spec",
                return_value = {
                    "raw": "0,1,2,3",
                    "numeric_ids": [0, 1, 2, 3],
                    "supports_explicit_gpu_ids": True,
                },
            ),
            patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1, 2, 3]),
            patch.object(
                hw,
                "get_visible_gpu_utilization",
                return_value = self._make_utilization(
                    [
                        (0, 80.0, 75.0),
                        (1, 80.0, 78.0),
                        (2, 80.0, 70.0),
                        (3, 80.0, 72.0),
                    ]
                ),
            ),
        ):
            selected, meta = auto_select_gpu_ids("test/model")
            # Should pick GPU 1 (most free memory: 78GB) -- enough for 10GB
            self.assertEqual(len(selected), 1)
            self.assertEqual(selected[0], 1)

    def test_two_gpus_needed(self):
        from utils.hardware.hardware import auto_select_gpu_ids
        import utils.hardware.hardware as hw
        with (
            patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA),
            patch.object(
                hw,
                "estimate_required_model_memory_gb",
                return_value = (
                    50.0,
                    {
                        "mode": "inference",
                        "required_gb": 50.0,
                        "model_size_source": "config",
                        "model_size_gb": 38.0,
                    },
                ),
            ),
            patch.object(
                hw,
                "_get_parent_visible_gpu_spec",
                return_value = {
                    "raw": "0,1",
                    "numeric_ids": [0, 1],
                    "supports_explicit_gpu_ids": True,
                },
            ),
            patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1]),
            patch.object(
                hw,
                "get_visible_gpu_utilization",
                return_value = self._make_utilization(
                    [
                        (0, 40.0, 30.0),  # 30GB free
                        (1, 40.0, 35.0),  # 35GB free
                    ]
                ),
            ),
        ):
            selected, meta = auto_select_gpu_ids("test/model")
            # 35GB (first) + 30*0.85 (second) = 60.5GB > 50GB
            self.assertEqual(len(selected), 2)

    def test_non_cuda_returns_none(self):
        from utils.hardware.hardware import auto_select_gpu_ids
        import utils.hardware.hardware as hw
        with patch.object(hw, "get_device", return_value = hw.DeviceType.CPU):
            selected, meta = auto_select_gpu_ids("test/model")
            self.assertIsNone(selected)
            self.assertEqual(meta["selection_mode"], "non_cuda")


class TestGetDeviceMap(unittest.TestCase):
    """device_map string generation."""

    def test_single_gpu_returns_sequential(self):
        from utils.hardware.hardware import get_device_map
        import utils.hardware.hardware as hw
        with (
            patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA),
            patch.object(
                hw,
                "_get_parent_visible_gpu_spec",
                return_value = {
                    "raw": "0",
                    "numeric_ids": [0],
                    "supports_explicit_gpu_ids": True,
                },
            ),
            patch.object(hw, "get_visible_gpu_count", return_value = 1),
        ):
            dm = get_device_map(gpu_ids = [0])
            self.assertEqual(dm, "sequential")

    def test_multi_gpu_returns_balanced(self):
        from utils.hardware.hardware import get_device_map
        import utils.hardware.hardware as hw
        with patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA):
            dm = get_device_map(gpu_ids = [0, 1])
            self.assertEqual(dm, "balanced")

    def test_cpu_returns_sequential(self):
        from utils.hardware.hardware import get_device_map
        import utils.hardware.hardware as hw
        with patch.object(hw, "get_device", return_value = hw.DeviceType.CPU):
            dm = get_device_map(gpu_ids = None)
            self.assertEqual(dm, "sequential")


class TestResolveRequestedGpuIds(unittest.TestCase):
    """GPU ID validation."""

    def test_none_returns_parent_visible(self):
        from utils.hardware.hardware import resolve_requested_gpu_ids
        with (
            patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3"}, clear = False),
            patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
        ):
            result = resolve_requested_gpu_ids(None)
            self.assertEqual(result, [2, 3])

    def test_empty_list_returns_parent_visible(self):
        from utils.hardware.hardware import resolve_requested_gpu_ids
        with (
            patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3"}, clear = False),
            patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
        ):
            result = resolve_requested_gpu_ids([])
            self.assertEqual(result, [2, 3])

    def test_duplicates_rejected(self):
        from utils.hardware.hardware import resolve_requested_gpu_ids
        with (
            patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1,2"}, clear = False),
            patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
        ):
            with self.assertRaises(ValueError):
                resolve_requested_gpu_ids([1, 1])

    def test_out_of_range_rejected(self):
        from utils.hardware.hardware import resolve_requested_gpu_ids
        with (
            patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0,1"}, clear = False),
            patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 4),
        ):
            with self.assertRaises(ValueError):
                resolve_requested_gpu_ids([5])

    def test_uuid_env_var_rejects_explicit_ids(self):
        from utils.hardware.hardware import resolve_requested_gpu_ids
        with (
            patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-abc,GPU-def"}, clear = False),
            patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
        ):
            with self.assertRaises(ValueError):
                resolve_requested_gpu_ids([0])


class TestApplyGpuIds(unittest.TestCase):
    """CUDA_VISIBLE_DEVICES environment variable setting."""

    def test_apply_list(self):
        from utils.hardware.hardware import apply_gpu_ids
        with patch.dict(os.environ, {}, clear = False):
            apply_gpu_ids([3, 5])
            self.assertEqual(os.environ.get("CUDA_VISIBLE_DEVICES"), "3,5")

    def test_apply_none_does_nothing(self):
        from utils.hardware.hardware import apply_gpu_ids

        original = os.environ.get("CUDA_VISIBLE_DEVICES")
        apply_gpu_ids(None)
        self.assertEqual(os.environ.get("CUDA_VISIBLE_DEVICES"), original)


class TestMultiGpuOverheadAccounting(unittest.TestCase):
    """Multi-GPU overhead is applied correctly.

    The first GPU keeps its full free memory; the overhead factor applies
    only to additional GPUs.
    """

    def _make_utilization(self, devices):
        return {
            "available": True,
            "devices": [
                {
                    "index": idx,
                    "vram_total_gb": total,
                    "vram_used_gb": total - free,
                }
                for idx, total, free in devices
            ],
        }

    def test_first_gpu_not_penalized(self):
        """A model that just fits on 1 GPU should not require 2 GPUs."""
        from utils.hardware.hardware import auto_select_gpu_ids
        import utils.hardware.hardware as hw

        # Model requires 79GB, GPU has 80GB free
        with (
            patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA),
            patch.object(
                hw,
                "estimate_required_model_memory_gb",
                return_value = (
                    79.0,
                    {
                        "mode": "inference",
                        "required_gb": 79.0,
                        "model_size_source": "config",
                        "model_size_gb": 60.0,
                    },
                ),
            ),
            patch.object(
                hw,
                "_get_parent_visible_gpu_spec",
                return_value = {
                    "raw": "0,1",
                    "numeric_ids": [0, 1],
                    "supports_explicit_gpu_ids": True,
                },
            ),
            patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1]),
            patch.object(
                hw,
                "get_visible_gpu_utilization",
                return_value = self._make_utilization(
                    [
                        (0, 80.0, 80.0),
                        (1, 80.0, 80.0),
                    ]
                ),
            ),
        ):
            selected, meta = auto_select_gpu_ids("test/model")
            # Should fit on 1 GPU (80GB >= 79GB)
            self.assertEqual(len(selected), 1)

    def test_second_gpu_has_overhead(self):
        """When 2 GPUs are needed, the second one's contribution is reduced."""
        from utils.hardware.hardware import auto_select_gpu_ids
        import utils.hardware.hardware as hw

        # Model requires 110GB. First GPU has 80GB, second has 40GB.
        # With overhead: 80 + 40*0.85 = 114GB -- just enough
        with (
            patch.object(hw, "get_device", return_value = hw.DeviceType.CUDA),
            patch.object(
                hw,
                "estimate_required_model_memory_gb",
                return_value = (
                    110.0,
                    {
                        "mode": "inference",
                        "required_gb": 110.0,
                        "model_size_source": "config",
                        "model_size_gb": 85.0,
                    },
                ),
            ),
            patch.object(
                hw,
                "_get_parent_visible_gpu_spec",
                return_value = {
                    "raw": "0,1",
                    "numeric_ids": [0, 1],
                    "supports_explicit_gpu_ids": True,
                },
            ),
            patch.object(hw, "get_parent_visible_gpu_ids", return_value = [0, 1]),
            patch.object(
                hw,
                "get_visible_gpu_utilization",
                return_value = self._make_utilization(
                    [
                        (0, 80.0, 80.0),
                        (1, 80.0, 40.0),
                    ]
                ),
            ),
        ):
            selected, meta = auto_select_gpu_ids("test/model")
            # Should use both GPUs
            self.assertEqual(len(selected), 2)


if __name__ == "__main__":
    unittest.main()