diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index 1c2a80b8bd..03e221d336 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -2,12 +2,12 @@ Tests for utils/hardware and utils/utils — device detection, GPU memory, error formatting. These tests are designed to pass on ANY platform: - • NVIDIA GPU (CUDA backend) - • Apple Silicon (MPS backend) + • NVIDIA GPU (CUDA backend, requires torch) + • Apple Silicon (MLX backend, requires mlx) • CPU-only (no GPU at all) No ML framework is imported at the top level. -Tests that need torch internals for mocking are skipped when torch is unavailable. +Tests that need torch/mlx internals for mocking are skipped when unavailable. Run with: cd studio/backend @@ -18,14 +18,21 @@ from unittest.mock import patch, MagicMock import pytest -# --- Conditional torch import --- +# --- Conditional framework imports --- try: import torch HAS_TORCH = True except ImportError: HAS_TORCH = False +try: + import mlx.core as mx + HAS_MLX = True +except ImportError: + HAS_MLX = False + needs_torch = pytest.mark.skipif(not HAS_TORCH, reason="PyTorch not installed") +needs_mlx = pytest.mark.skipif(not HAS_MLX, reason="MLX not installed") from utils.hardware import ( get_device, @@ -44,11 +51,10 @@ from utils.utils import format_error_message def _actual_device() -> str: """Return the real device string for the current machine.""" - if HAS_TORCH: - if torch.cuda.is_available(): - return "cuda" - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - return "mps" + if HAS_TORCH and torch.cuda.is_available(): + return "cuda" + if is_apple_silicon() and HAS_MLX: + return "mlx" return "cpu" @@ -64,39 +70,36 @@ class TestGetDevice: """Tests for get_device() — should agree with the real hardware.""" def setup_method(self): - """Save DEVICE before each test.""" self._saved_device = _hw_module.DEVICE def teardown_method(self): - """Restore DEVICE after each test so mocked tests don't poison later ones.""" _hw_module.DEVICE = self._saved_device def test_returns_valid_device_type(self): result = get_device() - assert result in (DeviceType.CUDA, DeviceType.MPS, DeviceType.CPU) + assert result in (DeviceType.CUDA, DeviceType.MLX, DeviceType.CPU) def test_matches_actual_hardware(self): assert get_device().value == _actual_device() - # --- Mocked paths (require torch for patching) --- + # --- Mocked paths --- @needs_torch def test_returns_cuda_when_cuda_available(self): - with patch("torch.cuda.is_available", return_value=True): + with patch("utils.hardware.hardware._has_torch", return_value=True), \ + patch("torch.cuda.is_available", return_value=True): assert _reset_and_detect() == DeviceType.CUDA - @needs_torch - def test_returns_mps_when_only_mps_available(self): - mock_mps = MagicMock() - mock_mps.is_available.return_value = True - with patch("torch.cuda.is_available", return_value=False), \ - patch.object(torch.backends, "mps", mock_mps, create=True): - assert _reset_and_detect() == DeviceType.MPS + def test_returns_mlx_when_on_apple_silicon_with_mlx(self): + with patch("utils.hardware.hardware._has_torch", return_value=False), \ + patch("utils.hardware.hardware.is_apple_silicon", return_value=True), \ + patch("utils.hardware.hardware._has_mlx", return_value=True): + assert _reset_and_detect() == DeviceType.MLX - @needs_torch def test_returns_cpu_when_nothing_available(self): - with patch("torch.cuda.is_available", return_value=False), \ - patch("builtins.hasattr", side_effect=lambda obj, name: False if name == "mps" else hasattr(obj, name)): + with patch("utils.hardware.hardware._has_torch", return_value=False), \ + patch("utils.hardware.hardware.is_apple_silicon", return_value=False), \ + patch("utils.hardware.hardware._has_mlx", return_value=False): assert _reset_and_detect() == DeviceType.CPU @@ -144,15 +147,10 @@ class TestClearGpuCache: mock_empty.assert_called_once() mock_ipc.assert_called_once() - @needs_torch - def test_calls_mps_cache_when_mps(self): - mock_mps_module = MagicMock() - mock_mps_module.empty_cache = MagicMock() - - with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MPS), \ - patch.object(torch, "mps", mock_mps_module, create=True): + def test_mlx_does_not_raise(self): + """MLX cache clear is a no-op — should just succeed.""" + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MLX): clear_gpu_cache() - mock_mps_module.empty_cache.assert_called_once() def test_noop_on_cpu(self): with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU): @@ -215,29 +213,24 @@ class TestGetGpuMemoryInfo: assert abs(result["free_gb"] - 12.0) < 0.01 assert abs(result["utilization_pct"] - 25.0) < 0.1 - # --- MPS-specific mocked test --- + # --- MLX-specific mocked test --- - @needs_torch - def test_mps_path_returns_correct_fields(self): + @needs_mlx + def test_mlx_path_returns_correct_fields(self): mock_psutil_mem = MagicMock() - mock_psutil_mem.total = 32 * (1024 ** 3) + mock_psutil_mem.total = 32 * (1024 ** 3) # 32 GB unified mock_psutil = MagicMock() mock_psutil.virtual_memory.return_value = mock_psutil_mem - mock_mps = MagicMock() - mock_mps.current_allocated_memory.return_value = 2 * (1024 ** 3) - - with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MPS), \ - patch.dict("sys.modules", {"psutil": mock_psutil}), \ - patch.object(torch, "mps", mock_mps, create=True): + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MLX), \ + patch.dict("sys.modules", {"psutil": mock_psutil}): result = get_gpu_memory_info() assert result["available"] is True - assert result["backend"] == "mps" + assert result["backend"] == "mlx" assert "Apple Silicon" in result["device_name"] assert abs(result["total_gb"] - 32.0) < 0.01 - assert abs(result["allocated_gb"] - 2.0) < 0.01 # --- CPU-only path --- @@ -329,11 +322,11 @@ class TestFormatErrorMessage: assert "big/model" not in msg assert "model" in msg - # --- OOM on MPS --- + # --- OOM on MLX --- - def test_mps_oom(self): - err = Exception("MPS backend out of memory") - with patch("utils.hardware.get_device", return_value=DeviceType.MPS): + def test_mlx_oom(self): + err = Exception("MLX backend out of memory") + with patch("utils.hardware.get_device", return_value=DeviceType.MLX): msg = format_error_message(err, "unsloth/huge-model") assert "Apple Silicon" in msg diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index ea223096b1..25320a25e1 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -25,7 +25,7 @@ logger = logging.getLogger(__name__) class DeviceType(str, Enum): """Supported compute backends. Inherits from str so it serializes cleanly in JSON.""" CUDA = "cuda" - MPS = "mps" + MLX = "mlx" CPU = "cpu" @@ -39,7 +39,24 @@ DEVICE: Optional[DeviceType] = None def is_apple_silicon() -> bool: """Check if running on Apple Silicon hardware (pure platform check, no ML imports).""" return platform.system() == "Darwin" and platform.machine() == "arm64" -pass + + +def _has_torch() -> bool: + """Check if PyTorch is importable.""" + try: + import torch + return True + except ImportError: + return False + + +def _has_mlx() -> bool: + """Check if MLX is importable.""" + try: + import mlx.core + return True + except ImportError: + return False def detect_hardware() -> DeviceType: @@ -51,34 +68,31 @@ def detect_hardware() -> DeviceType: Detection order: 1. CUDA (NVIDIA GPU, requires torch) - 2. MPS (Apple Silicon via PyTorch MPS backend) + 2. MLX (Apple Silicon via MLX framework) 3. CPU (fallback) """ global DEVICE - # --- Try PyTorch first (covers CUDA and MPS) --- - try: + # --- CUDA: try PyTorch --- + if _has_torch(): import torch - if torch.cuda.is_available(): DEVICE = DeviceType.CUDA device_name = torch.cuda.get_device_properties(0).name logger.info(f"Hardware detected: CUDA — {device_name}") return DEVICE - if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): - DEVICE = DeviceType.MPS - chip = platform.processor() or platform.machine() - logger.info(f"Hardware detected: MPS — Apple Silicon ({chip})") - return DEVICE - - except ImportError: - logger.warning("PyTorch not installed — falling back to CPU") + # --- MLX: Apple Silicon --- + if is_apple_silicon() and _has_mlx(): + DEVICE = DeviceType.MLX + chip = platform.processor() or platform.machine() + logger.info(f"Hardware detected: MLX — Apple Silicon ({chip})") + return DEVICE + # --- Fallback --- DEVICE = DeviceType.CPU logger.info("Hardware detected: CPU (no GPU backend available)") return DEVICE -pass # ========== Convenience helpers ========== @@ -92,7 +106,6 @@ def get_device() -> DeviceType: if DEVICE is None: detect_hardware() return DEVICE -pass def clear_gpu_cache(): @@ -109,17 +122,16 @@ def clear_gpu_cache(): import torch torch.cuda.empty_cache() torch.cuda.ipc_collect() - elif device == DeviceType.MPS: - import torch - if hasattr(torch.mps, "empty_cache"): - torch.mps.empty_cache() -pass + elif device == DeviceType.MLX: + # MLX manages memory automatically; no explicit cache clear needed. + # mlx.core has no empty_cache equivalent — gc.collect() above is enough. + pass def get_gpu_memory_info() -> Dict[str, Any]: """ Get GPU memory information. - Supports CUDA (NVIDIA), MPS (Apple Silicon), and CPU-only environments. + Supports CUDA (NVIDIA), MLX (Apple Silicon), and CPU-only environments. """ device = get_device() @@ -149,13 +161,16 @@ def get_gpu_memory_info() -> Dict[str, Any]: logger.error(f"Error getting CUDA GPU info: {e}") return {"available": False, "backend": device.value, "error": str(e)} - # ---- MPS path (Apple Silicon) ---- - if device == DeviceType.MPS: + # ---- MLX path (Apple Silicon) ---- + if device == DeviceType.MLX: try: - import torch + import mlx.core as mx import psutil - allocated = torch.mps.current_allocated_memory() if hasattr(torch.mps, "current_allocated_memory") else 0 + + # MLX uses unified memory — report system memory as the pool total = psutil.virtual_memory().total + # MLX doesn't expose per-process GPU allocation; report 0 as allocated + allocated = 0 return { "available": True, @@ -164,17 +179,16 @@ def get_gpu_memory_info() -> Dict[str, Any]: "device_name": f"Apple Silicon ({platform.processor() or platform.machine()})", "total_gb": total / (1024**3), "allocated_gb": allocated / (1024**3), - "reserved_gb": 0, # MPS doesn't have a separate reserved pool + "reserved_gb": 0, "free_gb": (total - allocated) / (1024**3), "utilization_pct": (allocated / total) * 100 if total else 0, } except Exception as e: - logger.error(f"Error getting MPS GPU info: {e}") + logger.error(f"Error getting MLX GPU info: {e}") return {"available": False, "backend": device.value, "error": str(e)} # ---- CPU-only ---- return {"available": False, "backend": "cpu"} -pass def log_gpu_memory(context: str): @@ -192,4 +206,3 @@ def log_gpu_memory(context: str): ) else: logger.info(f"GPU Memory [{context}]: No GPU available (CPU-only)") -pass diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index ee88406965..0acdbb3313 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -95,10 +95,10 @@ def format_error_message(error: Exception, model_name: str) -> str: if "invalid user token" in error_str: return "Invalid HF token. Please check your token and try again." - if "memory" in error_str or "cuda" in error_str or "mps" in error_str or "out of memory" in error_str: + if "memory" in error_str or "cuda" in error_str or "mlx" in error_str or "out of memory" in error_str: from utils.hardware import get_device device = get_device() - device_label = {"cuda": "GPU", "mps": "Apple Silicon GPU", "cpu": "system"}.get(device.value, "GPU") + device_label = {"cuda": "GPU", "mlx": "Apple Silicon GPU", "cpu": "system"}.get(device.value, "GPU") return f"Not enough {device_label} memory to load '{model_short}'. Try a smaller model or free memory." # Generic fallback