feat: add Apple Silicon (MPS) compatibility to backend utils + tests

studio/backend/core/__init__.py

@@ -13,7 +13,7 @@ from utils.models import is_vision_model, ModelConfig, scan_trained_loras, load_
 
 # Utilities (from utils)
 from utils.paths import normalize_path, is_local_path, is_model_cached
-from utils.utils import without_hf_auth, format_error_message, get_gpu_memory_info, search_hf_models
+from utils.utils import without_hf_auth, format_error_message, get_gpu_memory_info, search_hf_models, get_device, is_apple_silicon, clear_gpu_cache
 from utils.datasets import format_and_template_dataset
 
 __all__ = [
@@ -45,4 +45,7 @@ __all__ = [
     'without_hf_auth',
     'format_error_message',
     'get_gpu_memory_info',
+    'get_device',
+    'is_apple_silicon',
+    'clear_gpu_cache',
 ]

studio/backend/tests/conftest.py

@@ -0,0 +1,12 @@
+"""
+Shared pytest configuration for the backend test suite.
+Ensures that the backend root is on sys.path so that
+`import utils.utils` (and similar flat imports) resolve correctly.
+"""
+import sys
+from pathlib import Path
+
+# Add backend root to sys.path (mirrors how the app itself is launched)
+_backend_root = Path(__file__).resolve().parent.parent
+if str(_backend_root) not in sys.path:
+    sys.path.insert(0, str(_backend_root))

studio/backend/tests/test_utils.py

@@ -0,0 +1,327 @@
+"""
+Tests for utils/utils.py — device detection, GPU memory, error formatting.
+
+These tests are designed to pass on ANY platform:
+  • NVIDIA GPU  (CUDA backend)
+  • Apple Silicon (MPS backend)
+  • CPU-only     (no GPU at all)
+
+Run with:
+    cd studio/backend
+    python -m pytest tests/test_utils.py -v
+"""
+import platform
+from unittest.mock import patch, MagicMock
+
+import pytest
+import torch
+
+from utils.utils import (
+    get_device,
+    is_apple_silicon,
+    clear_gpu_cache,
+    get_gpu_memory_info,
+    log_gpu_memory,
+    format_error_message,
+)
+
+
+# ========== Helpers ==========
+
+def _actual_device() -> str:
+    """Return the real device string for the current machine."""
+    if torch.cuda.is_available():
+        return "cuda"
+    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
+# ========== get_device() ==========
+
+class TestGetDevice:
+    """Tests for get_device() — should agree with the real hardware."""
+
+    def test_returns_valid_string(self):
+        result = get_device()
+        assert result in ("cuda", "mps", "cpu")
+
+    def test_matches_actual_hardware(self):
+        assert get_device() == _actual_device()
+
+    # --- Mocked paths to cover all branches regardless of hardware ---
+
+    def test_returns_cuda_when_cuda_available(self):
+        with patch("torch.cuda.is_available", return_value=True):
+            assert get_device() == "cuda"
+
+    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 get_device() == "mps"
+
+    def test_returns_cpu_when_nothing_available(self):
+        # Patch out mps entirely so hasattr(..., "mps") returns False
+        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)):
+            assert get_device() == "cpu"
+
+
+# ========== is_apple_silicon() ==========
+
+class TestIsAppleSilicon:
+
+    def test_returns_bool(self):
+        assert isinstance(is_apple_silicon(), bool)
+
+    def test_true_on_darwin_arm64(self):
+        with patch("platform.system", return_value="Darwin"), \
+             patch("platform.machine", return_value="arm64"):
+            assert is_apple_silicon() is True
+
+    def test_false_on_linux_x86(self):
+        with patch("platform.system", return_value="Linux"), \
+             patch("platform.machine", return_value="x86_64"):
+            assert is_apple_silicon() is False
+
+    def test_false_on_darwin_x86(self):
+        """Intel Mac should return False."""
+        with patch("platform.system", return_value="Darwin"), \
+             patch("platform.machine", return_value="x86_64"):
+            assert is_apple_silicon() is False
+
+
+# ========== clear_gpu_cache() ==========
+
+class TestClearGpuCache:
+    """clear_gpu_cache() must never raise, regardless of platform."""
+
+    def test_does_not_raise(self):
+        # Should be safe on any hardware
+        clear_gpu_cache()
+
+    def test_calls_cuda_cache_when_cuda_available(self):
+        with patch("torch.cuda.is_available", return_value=True), \
+             patch("torch.cuda.empty_cache") as mock_empty, \
+             patch("torch.cuda.ipc_collect") as mock_ipc:
+            clear_gpu_cache()
+            mock_empty.assert_called_once()
+            mock_ipc.assert_called_once()
+
+    def test_calls_mps_cache_when_mps_available(self):
+        mock_mps_backend = MagicMock()
+        mock_mps_backend.is_available.return_value = True
+
+        mock_mps_module = MagicMock()
+        mock_mps_module.empty_cache = MagicMock()
+
+        with patch("torch.cuda.is_available", return_value=False), \
+             patch.object(torch.backends, "mps", mock_mps_backend, create=True), \
+             patch.object(torch, "mps", mock_mps_module, create=True):
+            clear_gpu_cache()
+            mock_mps_module.empty_cache.assert_called_once()
+
+    def test_noop_on_cpu(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)):
+            # Must not raise
+            clear_gpu_cache()
+
+
+# ========== get_gpu_memory_info() ==========
+
+class TestGetGpuMemoryInfo:
+
+    def test_returns_dict(self):
+        result = get_gpu_memory_info()
+        assert isinstance(result, dict)
+
+    def test_has_available_key(self):
+        result = get_gpu_memory_info()
+        assert "available" in result
+
+    def test_has_backend_key(self):
+        result = get_gpu_memory_info()
+        assert "backend" in result
+
+    def test_backend_matches_device(self):
+        """The reported backend must agree with get_device()."""
+        result = get_gpu_memory_info()
+        device = get_device()
+        assert result["backend"] == device
+
+    # --- When a GPU IS available ---
+
+    @pytest.mark.skipif(
+        _actual_device() == "cpu",
+        reason="No GPU available on this machine"
+    )
+    def test_gpu_available_fields(self):
+        result = get_gpu_memory_info()
+        assert result["available"] is True
+        assert result["total_gb"] > 0
+        assert result["allocated_gb"] >= 0
+        assert result["free_gb"] >= 0
+        assert 0 <= result["utilization_pct"] <= 100
+        assert "device_name" in result
+
+    # --- CUDA-specific mocked test ---
+
+    def test_cuda_path_returns_correct_fields(self):
+        mock_props = MagicMock()
+        mock_props.total_memory = 16 * (1024 ** 3)  # 16 GB
+        mock_props.name = "NVIDIA Test GPU"
+
+        with patch("utils.utils.get_device", return_value="cuda"), \
+             patch("torch.cuda.current_device", return_value=0), \
+             patch("torch.cuda.get_device_properties", return_value=mock_props), \
+             patch("torch.cuda.memory_allocated", return_value=4 * (1024 ** 3)), \
+             patch("torch.cuda.memory_reserved", return_value=6 * (1024 ** 3)):
+            result = get_gpu_memory_info()
+
+        assert result["available"] is True
+        assert result["backend"] == "cuda"
+        assert result["device_name"] == "NVIDIA Test GPU"
+        assert abs(result["total_gb"] - 16.0) < 0.01
+        assert abs(result["allocated_gb"] - 4.0) < 0.01
+        assert abs(result["free_gb"] - 12.0) < 0.01
+        assert abs(result["utilization_pct"] - 25.0) < 0.1
+
+    # --- MPS-specific mocked test ---
+
+    def test_mps_path_returns_correct_fields(self):
+        mock_psutil_mem = MagicMock()
+        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.utils.get_device", return_value="mps"), \
+             patch.dict("sys.modules", {"psutil": mock_psutil}), \
+             patch.object(torch, "mps", mock_mps, create=True):
+            result = get_gpu_memory_info()
+
+        assert result["available"] is True
+        assert result["backend"] == "mps"
+        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 ---
+
+    def test_cpu_path_returns_unavailable(self):
+        with patch("utils.utils.get_device", return_value="cpu"):
+            result = get_gpu_memory_info()
+
+        assert result["available"] is False
+        assert result["backend"] == "cpu"
+
+    # --- Error resilience ---
+
+    def test_cuda_error_returns_unavailable(self):
+        with patch("utils.utils.get_device", return_value="cuda"), \
+             patch("torch.cuda.current_device", side_effect=RuntimeError("CUDA init failed")):
+            result = get_gpu_memory_info()
+
+        assert result["available"] is False
+        assert "error" in result
+
+
+# ========== log_gpu_memory() ==========
+
+class TestLogGpuMemory:
+
+    def test_does_not_raise(self):
+        log_gpu_memory("test")
+
+    def test_logs_gpu_info_when_available(self, caplog):
+        fake_info = {
+            "available": True,
+            "backend": "cuda",
+            "device_name": "FakeGPU",
+            "allocated_gb": 2.0,
+            "total_gb": 16.0,
+            "utilization_pct": 12.5,
+            "free_gb": 14.0,
+        }
+        with patch("utils.utils.get_gpu_memory_info", return_value=fake_info):
+            import logging
+            with caplog.at_level(logging.INFO, logger="utils.utils"):
+                log_gpu_memory("unit-test")
+
+        assert "unit-test" in caplog.text
+        assert "CUDA" in caplog.text
+        assert "FakeGPU" in caplog.text
+
+    def test_logs_cpu_fallback_when_no_gpu(self, caplog):
+        fake_info = {"available": False, "backend": "cpu"}
+        with patch("utils.utils.get_gpu_memory_info", return_value=fake_info):
+            import logging
+            with caplog.at_level(logging.INFO, logger="utils.utils"):
+                log_gpu_memory("cpu-test")
+
+        assert "No GPU available" in caplog.text
+
+
+# ========== format_error_message() ==========
+
+class TestFormatErrorMessage:
+
+    def test_not_found(self):
+        err = Exception("Repository not found for unsloth/test")
+        msg = format_error_message(err, "unsloth/test")
+        assert "not found" in msg.lower()
+        assert "test" in msg
+
+    def test_unauthorized(self):
+        err = Exception("401 Unauthorized")
+        msg = format_error_message(err, "some/model")
+        assert "authentication" in msg.lower() or "unauthorized" in msg.lower()
+
+    def test_gated_model(self):
+        err = Exception("Access to model requires authentication")
+        msg = format_error_message(err, "meta/llama")
+        assert "authentication" in msg.lower()
+
+    def test_invalid_token(self):
+        err = Exception("Invalid user token")
+        msg = format_error_message(err, "any/model")
+        assert "invalid" in msg.lower()
+
+    # --- OOM on CUDA ---
+
+    def test_cuda_oom(self):
+        err = Exception("CUDA out of memory")
+        with patch("utils.utils.get_device", return_value="cuda"):
+            msg = format_error_message(err, "big/model")
+        assert "GPU" in msg
+        assert "big/model" not in msg  # should use short name
+        assert "model" in msg
+
+    # --- OOM on MPS ---
+
+    def test_mps_oom(self):
+        err = Exception("MPS backend out of memory")
+        with patch("utils.utils.get_device", return_value="mps"):
+            msg = format_error_message(err, "unsloth/huge-model")
+        assert "Apple Silicon" in msg
+
+    # --- OOM on CPU ---
+
+    def test_cpu_oom(self):
+        err = Exception("not enough memory to allocate")
+        with patch("utils.utils.get_device", return_value="cpu"):
+            msg = format_error_message(err, "any/model")
+        assert "system" in msg.lower()
+
+    # --- Generic fallback ---
+
+    def test_generic_error(self):
+        err = Exception("Something completely unexpected")
+        msg = format_error_message(err, "any/model")
+        assert msg == "Something completely unexpected"

studio/backend/utils/utils.py

@@ -3,6 +3,7 @@ Shared backend utilities
 """
 import gradio as gr
 import os
+import platform
 import logging
 from contextlib import contextmanager
 from pathlib import Path
@@ -13,6 +14,49 @@ import tempfile
 
 logger = logging.getLogger(__name__)
 
+# ========== Device Detection & Management ==========
+
+def get_device() -> str:
+    """
+    Detect the best available compute device.
+
+    Returns:
+        "cuda" on NVIDIA GPUs, "mps" on Apple Silicon, "cpu" otherwise.
+    """
+    import torch
+
+    if torch.cuda.is_available():
+        return "cuda"
+    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+pass
+
+
+def is_apple_silicon() -> bool:
+    """Check if running on Apple Silicon hardware."""
+    return platform.system() == "Darwin" and platform.machine() == "arm64"
+pass
+
+
+def clear_gpu_cache():
+    """
+    Clear GPU memory cache for the current device.
+    Safe to call on any platform — no-ops gracefully when the backend is unavailable.
+    """
+    import torch
+    import gc
+
+    gc.collect()
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        if hasattr(torch.mps, "empty_cache"):
+            torch.mps.empty_cache()
+pass
+
 @contextmanager
 def without_hf_auth():
     """
@@ -96,54 +140,91 @@ 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 "out of memory" in error_str:
-        return f"Not enough GPU memory to load '{model_short}'. Try a smaller model or free GPU memory."
+    if "memory" in error_str or "cuda" in error_str or "mps" in error_str or "out of memory" in error_str:
+        device = get_device()
+        device_label = {"cuda": "GPU", "mps": "Apple Silicon GPU", "cpu": "system"}.get(device, "GPU")
+        return f"Not enough {device_label} memory to load '{model_short}'. Try a smaller model or free memory."
 
     # Generic fallback
     return str(error)
 pass
 
 def get_gpu_memory_info() -> Dict[str, Any]:
-    """Get GPU memory information."""
+    """
+    Get GPU memory information.
+    Supports CUDA (NVIDIA), MPS (Apple Silicon), and CPU-only environments.
+    """
     import torch
 
-    if not torch.cuda.is_available():
-        return {"available": False}
+    device_type = get_device()
 
-    try:
-        device = torch.cuda.current_device()
-        props = torch.cuda.get_device_properties(device)
+    # ---- CUDA path ----
+    if device_type == "cuda":
+        try:
+            device = torch.cuda.current_device()
+            props = torch.cuda.get_device_properties(device)
 
-        total = props.total_memory
-        allocated = torch.cuda.memory_allocated(device)
-        reserved = torch.cuda.memory_reserved(device)
+            total = props.total_memory
+            allocated = torch.cuda.memory_allocated(device)
+            reserved = torch.cuda.memory_reserved(device)
 
-        return {
-            "available": True,
-            "device": device,
-            "total_gb": total / (1024**3),
-            "allocated_gb": allocated / (1024**3),
-            "reserved_gb": reserved / (1024**3),
-            "free_gb": (total - allocated) / (1024**3),
-            "utilization_pct": (allocated / total) * 100
-        }
-    except Exception as e:
-        logger.error(f"Error getting GPU info: {e}")
-        return {"available": False, "error": str(e)}
+            return {
+                "available": True,
+                "backend": "cuda",
+                "device": device,
+                "device_name": props.name,
+                "total_gb": total / (1024**3),
+                "allocated_gb": allocated / (1024**3),
+                "reserved_gb": reserved / (1024**3),
+                "free_gb": (total - allocated) / (1024**3),
+                "utilization_pct": (allocated / total) * 100,
+            }
+        except Exception as e:
+            logger.error(f"Error getting CUDA GPU info: {e}")
+            return {"available": False, "backend": "cuda", "error": str(e)}
+
+    # ---- MPS path (Apple Silicon) ----
+    if device_type == "mps":
+        try:
+            allocated = torch.mps.current_allocated_memory() if hasattr(torch.mps, "current_allocated_memory") else 0
+            # MPS doesn't expose total VRAM directly — use unified memory from psutil as a proxy
+            import psutil
+            total = psutil.virtual_memory().total
+
+            return {
+                "available": True,
+                "backend": "mps",
+                "device": 0,
+                "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
+                "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}")
+            return {"available": False, "backend": "mps", "error": str(e)}
+
+    # ---- CPU-only ----
+    return {"available": False, "backend": "cpu"}
 pass
 
 def log_gpu_memory(context: str):
     """Log GPU memory usage with context."""
     memory_info = get_gpu_memory_info()
     if memory_info.get("available"):
+        backend = memory_info.get("backend", "unknown").upper()
+        device_name = memory_info.get("device_name", "")
+        label = f"{backend}" + (f" ({device_name})" if device_name else "")
         logger.info(
-            f"GPU Memory [{context}]: "
+            f"GPU Memory [{context}] {label}: "
             f"{memory_info['allocated_gb']:.2f}GB/{memory_info['total_gb']:.2f}GB "
             f"({memory_info['utilization_pct']:.1f}% used, "
             f"{memory_info['free_gb']:.2f}GB free)"
         )
     else:
-        logger.info(f"GPU Memory [{context}]: No CUDA GPU available")
+        logger.info(f"GPU Memory [{context}]: No GPU available (CPU-only)")
 pass
 
 """