From 107bd2be4cf454fea1e9329d609e4042caf1cc83 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 11 Feb 2026 14:00:39 +0000 Subject: [PATCH] feat: add Apple Silicon (MPS) compatibility to backend utils + tests --- studio/backend/core/__init__.py | 5 +- studio/backend/tests/__init__.py | 0 studio/backend/tests/conftest.py | 12 ++ studio/backend/tests/test_utils.py | 327 +++++++++++++++++++++++++++++ studio/backend/utils/utils.py | 131 +++++++++--- 5 files changed, 449 insertions(+), 26 deletions(-) create mode 100644 studio/backend/tests/__init__.py create mode 100644 studio/backend/tests/conftest.py create mode 100644 studio/backend/tests/test_utils.py diff --git a/studio/backend/core/__init__.py b/studio/backend/core/__init__.py index 46ee6c14b1..5fdb57ae7f 100644 --- a/studio/backend/core/__init__.py +++ b/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', ] diff --git a/studio/backend/tests/__init__.py b/studio/backend/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/studio/backend/tests/conftest.py b/studio/backend/tests/conftest.py new file mode 100644 index 0000000000..82cbeb3da5 --- /dev/null +++ b/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)) diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py new file mode 100644 index 0000000000..32c82cacb0 --- /dev/null +++ b/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" diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index 74d5c05abf..7bae2dd086 100644 --- a/studio/backend/utils/utils.py +++ b/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 """