6d8c18cd1a* Replace standalone Studio wording with Unsloth Replace the single word Studio with Unsloth wherever it is used as shorthand for Unsloth Studio in docs, CLI output, UI strings, i18n locales, workflow display names, comments and docstrings. Kept unchanged: the full name Unsloth Studio, third party product names (LM Studio, Visual Studio, Mac Studio), feature names (Recipe Studio, Fine-tuning Studio and its translations), and all identifiers such as env vars, commands, paths and filenames. * Address review feedback on the Studio wording rename Use "an" before Unsloth where the rename left the article as "a". Restore the split brand where Unsloth and Studio render as two halves of the full product name: the onboarding sidebar subtitle and the IPv6 localhost warning. Scope two messages to the full name Unsloth Studio where plain Unsloth was misleading: the AMD README bullet and the CLI studio setup error.
360 lines
13 KiB
Python
360 lines
13 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
"""Unsloth hardware dispatch matrix: spoofs platform/torch/mlx per PROFILES to exercise CUDA/ROCm/XPU/MLX/CPU paths without real hardware."""
|
|
|
|
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"
|
|
|
|
|
|
@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 # Unsloth DeviceType (uppercased name: "CUDA"/"XPU"/"MLX"/"CPU")
|
|
expect_is_rocm: bool # Unsloth IS_ROCM
|
|
expect_apple_silicon: bool # Unsloth 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; "
|
|
"Unsloth 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]
|
|
|
|
|
|
@pytest.fixture
|
|
def spoof_hardware(monkeypatch):
|
|
"""Return a function that applies a HardwareProfile to the live process; monkeypatch cleans up on exit."""
|
|
|
|
def _apply(profile: HardwareProfile) -> None:
|
|
import platform
|
|
import torch
|
|
|
|
# platform spoof (used by both the unsloth gate and Unsloth's helpers)
|
|
monkeypatch.setattr(platform, "system", lambda: profile.system)
|
|
monkeypatch.setattr(platform, "machine", lambda: profile.machine)
|
|
|
|
monkeypatch.setattr(torch.cuda, "is_available", lambda: profile.cuda_available)
|
|
# Stub get_device_properties: detect_hardware reads .name, which crashes on a CPU CI runner.
|
|
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)
|
|
|
|
# Stub torch.xpu.* always; real get_device_name needs the XPU torch build.
|
|
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)
|
|
# detect_hardware now gates MLX on the full stack via
|
|
# utils.mlx_repair.mlx_stack_available() (it imports mlx_lm/mlx_vlm and
|
|
# checks dist versions), which faking only mlx.core cannot satisfy. An
|
|
# mlx profile means a complete, healthy stack, so model that here;
|
|
# mlx_stack_available's own internals are covered by test_mlx_repair.py.
|
|
if str(STUDIO_BACKEND) not in sys.path:
|
|
sys.path.insert(0, str(STUDIO_BACKEND))
|
|
import utils.mlx_repair as _mlx_repair # type: ignore
|
|
|
|
monkeypatch.setattr(_mlx_repair, "mlx_stack_available", lambda: True)
|
|
else:
|
|
# Drop cached mlx and patch find_spec so the unsloth gate 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)
|
|
|
|
# Unsloth's _has_mlx() does `import mlx.core`, not find_spec; block it
|
|
# with a meta_path finder that raises ImportError for mlx.*.
|
|
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()
|
|
# New list so monkeypatch fully restores on teardown.
|
|
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 Unsloth's hardware module under the bare-imports layout."""
|
|
if str(STUDIO_BACKEND) not in sys.path:
|
|
sys.path.insert(0, str(STUDIO_BACKEND))
|
|
# 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
|
|
|
|
|
|
@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):
|
|
"""Unsloth'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):
|
|
"""Unsloth'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):
|
|
"""CUDA wins over MLX when both available: canary against GPU users being routed to MLX after 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
|