unsloth/tests/studio/test_hardware_dispatch_matrix.py

# SPDX-License-Identifier: AGPL-3.0-only
"""
Comprehensive hardware dispatch matrix for Studio.

Drives every supported hardware profile from a single test host by
spoofing platform / torch.cuda / torch.xpu / sys.modules['mlx'] so we
can exercise the CUDA, ROCm, XPU, MLX, and CPU dispatch paths
deterministically without real hardware.

Profiles checked:

    nvidia_cuda          Linux x86_64 + torch.cuda.is_available()=True,
                         torch.version.hip=None
    amd_rocm             Linux x86_64 + torch.cuda.is_available()=True,
                         torch.version.hip="6.1"   (PyTorch ROCm aliases
                         torch.cuda.* over HIP)
    intel_xpu            Linux x86_64 + torch.cuda off, torch.xpu.is_available()=True
    apple_silicon_mlx    Darwin arm64 + cuda off + xpu off + mlx importable
    apple_silicon_no_mlx Darwin arm64 + everything off (no mlx pkg)
    linux_arm64_with_mlx Linux arm64 + mlx importable -- gate must NOT activate
                         (canary against accidental Linux-arm64 hijack)
    cpu_only             Linux x86_64 + nothing -- pure CPU fallback

For each profile we assert three contracts:

  1. ``unsloth._IS_MLX`` (re-evaluated under the spoof).
  2. ``utils.hardware.detect_hardware()`` ``DeviceType`` and ``IS_ROCM``.
  3. ``utils.hardware.is_apple_silicon()``.

Add a row to ``PROFILES`` to extend coverage; tests parametrize over it
automatically. No real hardware required.
"""

from __future__ import annotations

import importlib
import importlib.machinery
import importlib.util
import sys
import types
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

import pytest


REPO_ROOT = Path(__file__).resolve().parents[2]
STUDIO_BACKEND = REPO_ROOT / "studio" / "backend"


# ---------------------------------------------------------------------------
# Profile definition
# ---------------------------------------------------------------------------


@dataclass
class HardwareProfile:
    name: str
    system: str  # platform.system() value
    machine: str  # platform.machine() value
    cuda_available: bool  # torch.cuda.is_available() value
    hip_version: Optional[str]  # torch.version.hip; None for NVIDIA, "6.1" etc. for ROCm
    xpu_available: bool  # torch.xpu.is_available() value
    has_mlx: bool  # whether to inject a fake mlx into sys.modules
    mps_available: bool  # torch.backends.mps.is_available() value

    expect_is_mlx: bool  # unsloth._IS_MLX
    expect_device_type: str  # Studio DeviceType (uppercased name: "CUDA"/"XPU"/"MLX"/"CPU")
    expect_is_rocm: bool  # Studio IS_ROCM
    expect_apple_silicon: bool  # Studio is_apple_silicon()
    extra_notes: str = ""


PROFILES = [
    HardwareProfile(
        name = "nvidia_cuda",
        system = "Linux",
        machine = "x86_64",
        cuda_available = True,
        hip_version = None,
        xpu_available = False,
        has_mlx = False,
        mps_available = False,
        expect_is_mlx = False,
        expect_device_type = "CUDA",
        expect_is_rocm = False,
        expect_apple_silicon = False,
    ),
    HardwareProfile(
        name = "amd_rocm",
        system = "Linux",
        machine = "x86_64",
        cuda_available = True,
        hip_version = "6.1",
        xpu_available = False,
        has_mlx = False,
        mps_available = False,
        expect_is_mlx = False,
        expect_device_type = "CUDA",
        expect_is_rocm = True,
        expect_apple_silicon = False,
        extra_notes = "PyTorch ROCm reuses torch.cuda.* over HIP; "
        "Studio still uses DeviceType.CUDA but flips IS_ROCM=True.",
    ),
    HardwareProfile(
        name = "intel_xpu",
        system = "Linux",
        machine = "x86_64",
        cuda_available = False,
        hip_version = None,
        xpu_available = True,
        has_mlx = False,
        mps_available = False,
        expect_is_mlx = False,
        expect_device_type = "XPU",
        expect_is_rocm = False,
        expect_apple_silicon = False,
    ),
    HardwareProfile(
        name = "apple_silicon_mlx",
        system = "Darwin",
        machine = "arm64",
        cuda_available = False,
        hip_version = None,
        xpu_available = False,
        has_mlx = True,
        mps_available = True,
        expect_is_mlx = True,
        expect_device_type = "MLX",
        expect_is_rocm = False,
        expect_apple_silicon = True,
    ),
    HardwareProfile(
        name = "apple_silicon_no_mlx",
        system = "Darwin",
        machine = "arm64",
        cuda_available = False,
        hip_version = None,
        xpu_available = False,
        has_mlx = False,
        mps_available = True,
        expect_is_mlx = False,
        expect_device_type = "CPU",
        expect_is_rocm = False,
        expect_apple_silicon = True,
        extra_notes = "Mac without mlx falls through to CPU (chat-only).",
    ),
    HardwareProfile(
        name = "linux_arm64_with_mlx",
        system = "Linux",
        machine = "arm64",
        cuda_available = False,
        hip_version = None,
        xpu_available = False,
        has_mlx = True,
        mps_available = False,
        expect_is_mlx = False,
        expect_device_type = "CPU",
        expect_is_rocm = False,
        expect_apple_silicon = False,
        extra_notes = "Canary: Linux ARM64 with mlx package installed must NOT "
        "trigger MLX dispatch; the system check is what guards it.",
    ),
    HardwareProfile(
        name = "cpu_only",
        system = "Linux",
        machine = "x86_64",
        cuda_available = False,
        hip_version = None,
        xpu_available = False,
        has_mlx = False,
        mps_available = False,
        expect_is_mlx = False,
        expect_device_type = "CPU",
        expect_is_rocm = False,
        expect_apple_silicon = False,
    ),
]

PROFILE_IDS = [p.name for p in PROFILES]


# ---------------------------------------------------------------------------
# Spoofing helpers
# ---------------------------------------------------------------------------


@pytest.fixture
def spoof_hardware(monkeypatch):
    """Return a function that applies a HardwareProfile to the live process.

    Idempotent: each call re-applies the profile. Cleanup happens
    automatically when the test exits via monkeypatch.
    """

    def _apply(profile: HardwareProfile) -> None:
        import platform
        import torch

        # platform spoof (used by both the unsloth gate and Studio's helpers)
        monkeypatch.setattr(platform, "system", lambda: profile.system)
        monkeypatch.setattr(platform, "machine", lambda: profile.machine)

        # torch.cuda.is_available
        monkeypatch.setattr(torch.cuda, "is_available", lambda: profile.cuda_available)
        # detect_hardware reads torch.cuda.get_device_properties(0).name when
        # cuda_available is True. On a CPU CI runner that triggers _cuda_init
        # and crashes with "No CUDA GPUs are available". Stub it so the
        # dispatch path under test runs end-to-end.
        if profile.cuda_available:
            stub_props = types.SimpleNamespace(
                name = "Stub GPU" if not profile.hip_version else "Stub AMD GPU",
            )
            monkeypatch.setattr(
                torch.cuda,
                "get_device_properties",
                lambda i = 0: stub_props,
                raising = False,
            )

        # torch.version.hip — None on NVIDIA, "6.1" etc. on ROCm
        torch_version = torch.version
        monkeypatch.setattr(torch_version, "hip", profile.hip_version, raising = False)

        # torch.xpu.is_available + get_device_name -- detect_hardware reads both.
        # Real torch.xpu.get_device_name requires the XPU-compiled torch build,
        # so always stub it under the spoof to keep tests hardware-agnostic.
        if hasattr(torch, "xpu"):
            monkeypatch.setattr(torch.xpu, "is_available", lambda: profile.xpu_available)
            monkeypatch.setattr(
                torch.xpu,
                "get_device_name",
                lambda i = 0: "Intel XPU (stub)",
                raising = False,
            )
        elif profile.xpu_available:
            xpu_stub = types.SimpleNamespace(
                is_available = lambda: True,
                get_device_name = lambda i = 0: "Intel XPU (stub)",
            )
            monkeypatch.setattr(torch, "xpu", xpu_stub, raising = False)

        # torch.backends.mps.is_available
        if hasattr(torch.backends, "mps"):
            monkeypatch.setattr(torch.backends.mps, "is_available", lambda: profile.mps_available)

        # mlx + mlx.core in sys.modules
        if profile.has_mlx:
            fake_mlx = types.ModuleType("mlx")
            fake_mlx.__spec__ = importlib.machinery.ModuleSpec("mlx", loader = None)
            fake_mlx.__path__ = []
            fake_mlx_core = types.ModuleType("mlx.core")
            fake_mlx.core = fake_mlx_core
            monkeypatch.setitem(sys.modules, "mlx", fake_mlx)
            monkeypatch.setitem(sys.modules, "mlx.core", fake_mlx_core)
        else:
            # Drop any cached mlx modules and patch find_spec so the
            # unsloth gate (which uses importlib.util.find_spec) sees
            # mlx as absent.
            monkeypatch.delitem(sys.modules, "mlx", raising = False)
            monkeypatch.delitem(sys.modules, "mlx.core", raising = False)
            real_find_spec = importlib.util.find_spec

            def _no_mlx(name, *args, **kwargs):
                if name == "mlx" or name.startswith("mlx."):
                    return None
                return real_find_spec(name, *args, **kwargs)

            monkeypatch.setattr(importlib.util, "find_spec", _no_mlx)

            # Studio's _has_mlx() literally does `import mlx.core`, not
            # find_spec, so on a real Apple Silicon host with mlx
            # genuinely installed the import would still succeed. Block
            # it via a meta_path finder that raises ImportError for any
            # `mlx` / `mlx.*` import while this profile is active.
            class _BlockMLXFinder:
                def find_spec(
                    self_inner,
                    name,
                    path = None,
                    target = None,
                ):
                    if name == "mlx" or name.startswith("mlx."):
                        raise ImportError(
                            f"mlx import blocked by spoof_hardware " f"(profile={profile.name})"
                        )
                    return None

            blocker = _BlockMLXFinder()
            # Replace meta_path with a NEW list so monkeypatch can fully
            # restore the original on teardown (mutating the list in
            # place would survive the test).
            monkeypatch.setattr(
                sys,
                "meta_path",
                [blocker, *sys.meta_path],
            )

    return _apply


def _evaluate_unsloth_is_mlx_gate() -> bool:
    """Re-evaluate the exact expression from unsloth/__init__.py:20-24."""
    import importlib.util
    import platform

    return (
        platform.system() == "Darwin"
        and platform.machine() == "arm64"
        and importlib.util.find_spec("mlx") is not None
    )


def _import_studio_hardware_module():
    """Lazy-load Studio's hardware module under the bare-imports layout."""
    if str(STUDIO_BACKEND) not in sys.path:
        sys.path.insert(0, str(STUDIO_BACKEND))
    # Force a fresh import so detect_hardware re-runs under the current spoofs.
    sys.modules.pop("utils.hardware.hardware", None)
    sys.modules.pop("utils.hardware", None)
    from utils.hardware import hardware as hw  # type: ignore

    return hw


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


@pytest.mark.parametrize("profile", PROFILES, ids = PROFILE_IDS)
def test_unsloth_is_mlx_gate_matches_profile(profile, spoof_hardware):
    """The _IS_MLX expression in unsloth/__init__.py flips correctly per profile."""
    spoof_hardware(profile)
    actual = _evaluate_unsloth_is_mlx_gate()
    assert actual is profile.expect_is_mlx, (
        f"profile {profile.name}: expected _IS_MLX={profile.expect_is_mlx}, "
        f"got {actual}. {profile.extra_notes}"
    )


@pytest.mark.parametrize("profile", PROFILES, ids = PROFILE_IDS)
def test_studio_detect_hardware_matches_profile(profile, spoof_hardware):
    """Studio's detect_hardware() routes to the right DeviceType per profile."""
    spoof_hardware(profile)
    hw = _import_studio_hardware_module()
    detected = hw.detect_hardware()
    expected = getattr(hw.DeviceType, profile.expect_device_type)
    assert detected == expected, (
        f"profile {profile.name}: expected {profile.expect_device_type}, "
        f"got {detected!r}. {profile.extra_notes}"
    )
    assert hw.IS_ROCM is profile.expect_is_rocm, (
        f"profile {profile.name}: expected IS_ROCM={profile.expect_is_rocm}, " f"got {hw.IS_ROCM}"
    )


@pytest.mark.parametrize("profile", PROFILES, ids = PROFILE_IDS)
def test_studio_is_apple_silicon_matches_profile(profile, spoof_hardware):
    """Studio's is_apple_silicon() helper agrees with platform spoof."""
    spoof_hardware(profile)
    hw = _import_studio_hardware_module()
    assert hw.is_apple_silicon() is profile.expect_apple_silicon, (
        f"profile {profile.name}: expected is_apple_silicon={profile.expect_apple_silicon}, "
        f"got {hw.is_apple_silicon()}"
    )


# ---------------------------------------------------------------------------
# Negative-space tests: catch regressions where the dispatch order changes.
# ---------------------------------------------------------------------------


def test_cuda_takes_priority_over_mlx_when_both_available(spoof_hardware):
    """If both CUDA and MLX are available, Studio MUST pick CUDA. This is the
    canary that protects every existing GPU user from being silently routed
    to MLX after future refactors.
    """
    profile = HardwareProfile(
        name = "cuda_plus_mlx",
        system = "Darwin",
        machine = "arm64",
        cuda_available = True,
        hip_version = None,
        xpu_available = False,
        has_mlx = True,
        mps_available = True,
        expect_is_mlx = True,
        expect_device_type = "CUDA",
        expect_is_rocm = False,
        expect_apple_silicon = True,
    )
    spoof_hardware(profile)
    hw = _import_studio_hardware_module()
    assert hw.detect_hardware() == hw.DeviceType.CUDA


def test_xpu_takes_priority_over_mlx_when_both_available(spoof_hardware):
    """XPU is selected over MLX in the dispatch order."""
    profile = HardwareProfile(
        name = "xpu_plus_mlx",
        system = "Darwin",
        machine = "arm64",
        cuda_available = False,
        hip_version = None,
        xpu_available = True,
        has_mlx = True,
        mps_available = True,
        expect_is_mlx = True,
        expect_device_type = "XPU",
        expect_is_rocm = False,
        expect_apple_silicon = True,
    )
    spoof_hardware(profile)
    hw = _import_studio_hardware_module()
    assert hw.detect_hardware() == hw.DeviceType.XPU