From 107bd2be4cf454fea1e9329d609e4042caf1cc83 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 11 Feb 2026 14:00:39 +0000 Subject: [PATCH 1/8] 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 """ From 59d5f24eb5b0eb4b1506864d95ade1840cf1fe0a Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 11 Feb 2026 15:34:26 +0000 Subject: [PATCH 2/8] integrate global hardware detection at lifespan entrypoint --- studio/backend/core/__init__.py | 6 +- studio/backend/main.py | 32 ++-- studio/backend/tests/test_utils.py | 70 ++++---- studio/backend/utils/hardware/__init__.py | 24 +++ studio/backend/utils/hardware/hardware.py | 195 ++++++++++++++++++++++ studio/backend/utils/utils.py | 187 +-------------------- 6 files changed, 276 insertions(+), 238 deletions(-) create mode 100644 studio/backend/utils/hardware/__init__.py create mode 100644 studio/backend/utils/hardware/hardware.py diff --git a/studio/backend/core/__init__.py b/studio/backend/core/__init__.py index 5fdb57ae7f..7b562c6cb0 100644 --- a/studio/backend/core/__init__.py +++ b/studio/backend/core/__init__.py @@ -13,7 +13,8 @@ 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, get_device, is_apple_silicon, clear_gpu_cache +from utils.utils import without_hf_auth, format_error_message +from utils.hardware import get_device, is_apple_silicon, clear_gpu_cache, get_gpu_memory_info, log_gpu_memory, DeviceType from utils.datasets import format_and_template_dataset __all__ = [ @@ -37,7 +38,6 @@ __all__ = [ 'get_base_model_from_lora', # Utils - 'search_hf_models', 'format_and_template_dataset', 'normalize_path', 'is_local_path', @@ -45,7 +45,9 @@ __all__ = [ 'without_hf_auth', 'format_error_message', 'get_gpu_memory_info', + 'log_gpu_memory', 'get_device', 'is_apple_silicon', 'clear_gpu_cache', + 'DeviceType', ] diff --git a/studio/backend/main.py b/studio/backend/main.py index b08ef3430e..de468c183a 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -15,13 +15,17 @@ from datetime import datetime # Import routers from routes import training_router, models_router, inference_router, datasets_router, auth_router from auth import storage +from utils.hardware import detect_hardware UNSLOTH_CACHE_DIR = Path(__file__).parent / "unsloth_compiled_cache" @asynccontextmanager async def lifespan(app: FastAPI): - """Startup: print setup token if needed. Shutdown: clean up compiled cache.""" + """Startup: detect hardware, print setup token if needed. Shutdown: clean up compiled cache.""" + # Detect hardware first — sets DEVICE global used everywhere + detect_hardware() + if not storage.is_initialized(): setup_token = secrets.token_urlsafe(32) storage.save_setup_token(setup_token) @@ -78,23 +82,20 @@ async def health_check(): @app.get("/api/system") async def get_system_info(): """Get system information""" - import torch import platform import psutil + from utils.hardware import get_device, get_gpu_memory_info, DeviceType - # GPU Info - gpu_info = {"available": False, "devices": []} - if torch.cuda.is_available(): - gpu_info["available"] = True - for i in range(torch.cuda.device_count()): - props = torch.cuda.get_device_properties(i) - gpu_info["devices"].append( - { - "index": i, - "name": props.name, - "memory_total_gb": round(props.total_memory / 1e9, 2), - } - ) + # GPU Info — uses the hardware module (works on CUDA, MPS, CPU) + mem_info = get_gpu_memory_info() + gpu_info = {"available": mem_info.get("available", False), "devices": []} + + if mem_info.get("available"): + gpu_info["devices"].append({ + "index": mem_info.get("device", 0), + "name": mem_info.get("device_name", "Unknown"), + "memory_total_gb": round(mem_info.get("total_gb", 0), 2), + }) # CPU & Memory memory = psutil.virtual_memory() @@ -102,6 +103,7 @@ async def get_system_info(): return { "platform": platform.platform(), "python_version": platform.python_version(), + "device_backend": get_device().value, "cpu_count": psutil.cpu_count(), "memory": { "total_gb": round(memory.total / 1e9, 2), diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index 32c82cacb0..89b16dcd0e 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -1,5 +1,5 @@ """ -Tests for utils/utils.py — device detection, GPU memory, error formatting. +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) @@ -16,14 +16,15 @@ from unittest.mock import patch, MagicMock import pytest import torch -from utils.utils import ( +from utils.hardware import ( get_device, is_apple_silicon, clear_gpu_cache, get_gpu_memory_info, log_gpu_memory, - format_error_message, + DeviceType, ) +from utils.utils import format_error_message # ========== Helpers ========== @@ -42,31 +43,30 @@ def _actual_device() -> str: class TestGetDevice: """Tests for get_device() — should agree with the real hardware.""" - def test_returns_valid_string(self): + def test_returns_valid_device_type(self): result = get_device() - assert result in ("cuda", "mps", "cpu") + assert result in (DeviceType.CUDA, DeviceType.MPS, DeviceType.CPU) def test_matches_actual_hardware(self): - assert get_device() == _actual_device() + assert get_device().value == _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" + assert get_device() == DeviceType.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" + assert get_device() == DeviceType.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" + assert get_device() == DeviceType.CPU # ========== is_apple_silicon() ========== @@ -77,19 +77,22 @@ class TestIsAppleSilicon: 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"): + with patch("utils.hardware.hardware.platform") as mock_plat: + mock_plat.system.return_value = "Darwin" + mock_plat.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"): + with patch("utils.hardware.hardware.platform") as mock_plat: + mock_plat.system.return_value = "Linux" + mock_plat.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"): + with patch("utils.hardware.hardware.platform") as mock_plat: + mock_plat.system.return_value = "Darwin" + mock_plat.machine.return_value = "x86_64" assert is_apple_silicon() is False @@ -103,7 +106,7 @@ class TestClearGpuCache: clear_gpu_cache() def test_calls_cuda_cache_when_cuda_available(self): - with patch("torch.cuda.is_available", return_value=True), \ + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CUDA), \ patch("torch.cuda.empty_cache") as mock_empty, \ patch("torch.cuda.ipc_collect") as mock_ipc: clear_gpu_cache() @@ -111,21 +114,16 @@ class TestClearGpuCache: 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), \ + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MPS), \ 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)): + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU): # Must not raise clear_gpu_cache() @@ -150,7 +148,7 @@ class TestGetGpuMemoryInfo: """The reported backend must agree with get_device().""" result = get_gpu_memory_info() device = get_device() - assert result["backend"] == device + assert result["backend"] == device.value # --- When a GPU IS available --- @@ -174,7 +172,7 @@ class TestGetGpuMemoryInfo: mock_props.total_memory = 16 * (1024 ** 3) # 16 GB mock_props.name = "NVIDIA Test GPU" - with patch("utils.utils.get_device", return_value="cuda"), \ + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.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)), \ @@ -201,7 +199,7 @@ class TestGetGpuMemoryInfo: mock_mps = MagicMock() mock_mps.current_allocated_memory.return_value = 2 * (1024 ** 3) - with patch("utils.utils.get_device", return_value="mps"), \ + 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): result = get_gpu_memory_info() @@ -215,7 +213,7 @@ class TestGetGpuMemoryInfo: # --- CPU-only path --- def test_cpu_path_returns_unavailable(self): - with patch("utils.utils.get_device", return_value="cpu"): + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU): result = get_gpu_memory_info() assert result["available"] is False @@ -224,7 +222,7 @@ class TestGetGpuMemoryInfo: # --- Error resilience --- def test_cuda_error_returns_unavailable(self): - with patch("utils.utils.get_device", return_value="cuda"), \ + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CUDA), \ patch("torch.cuda.current_device", side_effect=RuntimeError("CUDA init failed")): result = get_gpu_memory_info() @@ -249,9 +247,9 @@ class TestLogGpuMemory: "utilization_pct": 12.5, "free_gb": 14.0, } - with patch("utils.utils.get_gpu_memory_info", return_value=fake_info): + with patch("utils.hardware.hardware.get_gpu_memory_info", return_value=fake_info): import logging - with caplog.at_level(logging.INFO, logger="utils.utils"): + with caplog.at_level(logging.INFO, logger="utils.hardware.hardware"): log_gpu_memory("unit-test") assert "unit-test" in caplog.text @@ -260,9 +258,9 @@ class TestLogGpuMemory: 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): + with patch("utils.hardware.hardware.get_gpu_memory_info", return_value=fake_info): import logging - with caplog.at_level(logging.INFO, logger="utils.utils"): + with caplog.at_level(logging.INFO, logger="utils.hardware.hardware"): log_gpu_memory("cpu-test") assert "No GPU available" in caplog.text @@ -297,7 +295,7 @@ class TestFormatErrorMessage: def test_cuda_oom(self): err = Exception("CUDA out of memory") - with patch("utils.utils.get_device", return_value="cuda"): + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CUDA): msg = format_error_message(err, "big/model") assert "GPU" in msg assert "big/model" not in msg # should use short name @@ -307,7 +305,7 @@ class TestFormatErrorMessage: def test_mps_oom(self): err = Exception("MPS backend out of memory") - with patch("utils.utils.get_device", return_value="mps"): + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MPS): msg = format_error_message(err, "unsloth/huge-model") assert "Apple Silicon" in msg @@ -315,7 +313,7 @@ class TestFormatErrorMessage: def test_cpu_oom(self): err = Exception("not enough memory to allocate") - with patch("utils.utils.get_device", return_value="cpu"): + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU): msg = format_error_message(err, "any/model") assert "system" in msg.lower() diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py new file mode 100644 index 0000000000..b992fc191d --- /dev/null +++ b/studio/backend/utils/hardware/__init__.py @@ -0,0 +1,24 @@ +""" +Hardware detection and GPU utilities +""" +from .hardware import ( + DeviceType, + DEVICE, + detect_hardware, + get_device, + is_apple_silicon, + clear_gpu_cache, + get_gpu_memory_info, + log_gpu_memory, +) + +__all__ = [ + 'DeviceType', + 'DEVICE', + 'detect_hardware', + 'get_device', + 'is_apple_silicon', + 'clear_gpu_cache', + 'get_gpu_memory_info', + 'log_gpu_memory', +] diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py new file mode 100644 index 0000000000..ea223096b1 --- /dev/null +++ b/studio/backend/utils/hardware/hardware.py @@ -0,0 +1,195 @@ +""" +Hardware detection — run once at startup, read everywhere. + +Usage: + # At FastAPI lifespan startup: + from utils.hardware import detect_hardware + detect_hardware() + + # Anywhere else: + from utils.hardware import DEVICE, DeviceType, is_apple_silicon + if DEVICE == DeviceType.CUDA: + import torch + ... +""" +import platform +import logging +from enum import Enum +from typing import Optional, Dict, Any + +logger = logging.getLogger(__name__) + + +# ========== Device Enum ========== + +class DeviceType(str, Enum): + """Supported compute backends. Inherits from str so it serializes cleanly in JSON.""" + CUDA = "cuda" + MPS = "mps" + CPU = "cpu" + + +# ========== Global State (set once by detect_hardware) ========== + +DEVICE: Optional[DeviceType] = None + + +# ========== Detection ========== + +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 detect_hardware() -> DeviceType: + """ + Detect the best available compute device and set the module-level DEVICE global. + + Should be called exactly once during FastAPI lifespan startup. + Safe to call multiple times (idempotent). + + Detection order: + 1. CUDA (NVIDIA GPU, requires torch) + 2. MPS (Apple Silicon via PyTorch MPS backend) + 3. CPU (fallback) + """ + global DEVICE + + # --- Try PyTorch first (covers CUDA and MPS) --- + try: + 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") + + DEVICE = DeviceType.CPU + logger.info("Hardware detected: CPU (no GPU backend available)") + return DEVICE +pass + + +# ========== Convenience helpers ========== + +def get_device() -> DeviceType: + """ + Return the detected device. Auto-detects if detect_hardware() hasn't been called yet. + Prefer calling detect_hardware() explicitly at startup instead. + """ + global DEVICE + if DEVICE is None: + detect_hardware() + return DEVICE +pass + + +def clear_gpu_cache(): + """ + Clear GPU memory cache for the current device. + Safe to call on any platform — no-ops gracefully. + """ + import gc + gc.collect() + + device = get_device() + + if device == DeviceType.CUDA: + 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 + + +def get_gpu_memory_info() -> Dict[str, Any]: + """ + Get GPU memory information. + Supports CUDA (NVIDIA), MPS (Apple Silicon), and CPU-only environments. + """ + device = get_device() + + # ---- CUDA path ---- + if device == DeviceType.CUDA: + try: + import torch + idx = torch.cuda.current_device() + props = torch.cuda.get_device_properties(idx) + + total = props.total_memory + allocated = torch.cuda.memory_allocated(idx) + reserved = torch.cuda.memory_reserved(idx) + + return { + "available": True, + "backend": device.value, + "device": idx, + "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": device.value, "error": str(e)} + + # ---- MPS path (Apple Silicon) ---- + if device == DeviceType.MPS: + try: + import torch + import psutil + allocated = torch.mps.current_allocated_memory() if hasattr(torch.mps, "current_allocated_memory") else 0 + total = psutil.virtual_memory().total + + return { + "available": True, + "backend": device.value, + "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": device.value, "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}] {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 GPU available (CPU-only)") +pass diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index 7bae2dd086..ee88406965 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -1,61 +1,16 @@ """ Shared backend utilities """ -import gradio as gr import os -import platform import logging from contextlib import contextmanager from pathlib import Path -from typing import Optional, Dict, Any import shutil 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(): @@ -141,149 +96,11 @@ def format_error_message(error: Exception, model_name: str) -> 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: + from utils.hardware import get_device device = get_device() - device_label = {"cuda": "GPU", "mps": "Apple Silicon GPU", "cpu": "system"}.get(device, "GPU") + device_label = {"cuda": "GPU", "mps": "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 return str(error) pass - -def get_gpu_memory_info() -> Dict[str, Any]: - """ - Get GPU memory information. - Supports CUDA (NVIDIA), MPS (Apple Silicon), and CPU-only environments. - """ - import torch - - device_type = get_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) - - 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}] {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 GPU available (CPU-only)") -pass - -""" -Model utility functions - search, discovery, etc. -""" - - -def search_hf_models(search_query: str, hf_token: Optional[str] = None): - """ - Search HuggingFace model hub. - """ - import requests - - if not search_query or not search_query.strip(): - return gr.update(choices=[]) - - # Simple debouncing: only search if query is at least 2 characters - if len(search_query.strip()) < 2: - return gr.update(choices=[]) - - try: - headers = {} - if hf_token and hf_token.strip(): - headers["Authorization"] = f"Bearer {hf_token.strip()}" - - url = "https://huggingface.co/api/models" - params = { - "search": search_query, - "pipeline_tag": "text-generation", - "library": "transformers", - "limit": 15, - "sort": "downloads", - "direction": -1 - } - - response = requests.get(url, headers=headers, params=params, timeout=10) - - if response.status_code == 200: - models = response.json() - unsloth_results = [] - other_results = [] - - for model in models: - model_id = model.get("modelId", "") - if model_id and "gguf" not in model_id.lower(): - result = (f"{model_id}", model_id) - - if model_id.startswith("unsloth/"): - unsloth_results.append(result) - else: - other_results.append(result) - - # Combine with unsloth models first - search_results = unsloth_results + other_results - return gr.update(choices=search_results) - else: - logger.warning(f"HF API returned status {response.status_code}") - return gr.update(choices=[]) - - except Exception as e: - logger.warning(f"Model search failed: {e}") - return gr.update(choices=[]) From 85fc481afe34a046ee94c8e24e6bbfcca59884ce Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 11 Feb 2026 15:37:34 +0000 Subject: [PATCH 3/8] fixed tests to be hardware specific --- studio/backend/tests/test_utils.py | 72 +++++++++++++++++------------- 1 file changed, 42 insertions(+), 30 deletions(-) diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index 89b16dcd0e..0ac160ecb9 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -6,6 +6,9 @@ These tests are designed to pass on ANY platform: • Apple Silicon (MPS backend) • 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. + Run with: cd studio/backend python -m pytest tests/test_utils.py -v @@ -14,7 +17,15 @@ import platform from unittest.mock import patch, MagicMock import pytest -import torch + +# --- Conditional torch import --- +try: + import torch + HAS_TORCH = True +except ImportError: + HAS_TORCH = False + +needs_torch = pytest.mark.skipif(not HAS_TORCH, reason="PyTorch not installed") from utils.hardware import ( get_device, @@ -31,10 +42,11 @@ from utils.utils import format_error_message 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" + if HAS_TORCH: + if torch.cuda.is_available(): + return "cuda" + if hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): + return "mps" return "cpu" @@ -50,12 +62,14 @@ class TestGetDevice: def test_matches_actual_hardware(self): assert get_device().value == _actual_device() - # --- Mocked paths to cover all branches regardless of hardware --- + # --- Mocked paths (require torch for patching) --- + @needs_torch def test_returns_cuda_when_cuda_available(self): with patch("torch.cuda.is_available", return_value=True): assert get_device() == DeviceType.CUDA + @needs_torch def test_returns_mps_when_only_mps_available(self): mock_mps = MagicMock() mock_mps.is_available.return_value = True @@ -63,6 +77,7 @@ class TestGetDevice: patch.object(torch.backends, "mps", mock_mps, create=True): assert get_device() == DeviceType.MPS + @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)): @@ -102,10 +117,10 @@ 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): + @needs_torch + def test_calls_cuda_cache_when_cuda(self): with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CUDA), \ patch("torch.cuda.empty_cache") as mock_empty, \ patch("torch.cuda.ipc_collect") as mock_ipc: @@ -113,7 +128,8 @@ class TestClearGpuCache: mock_empty.assert_called_once() mock_ipc.assert_called_once() - def test_calls_mps_cache_when_mps_available(self): + @needs_torch + def test_calls_mps_cache_when_mps(self): mock_mps_module = MagicMock() mock_mps_module.empty_cache = MagicMock() @@ -124,7 +140,6 @@ class TestClearGpuCache: def test_noop_on_cpu(self): with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU): - # Must not raise clear_gpu_cache() @@ -137,18 +152,14 @@ class TestGetGpuMemoryInfo: assert isinstance(result, dict) def test_has_available_key(self): - result = get_gpu_memory_info() - assert "available" in result + assert "available" in get_gpu_memory_info() def test_has_backend_key(self): - result = get_gpu_memory_info() - assert "backend" in result + assert "backend" in get_gpu_memory_info() 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.value + assert result["backend"] == get_device().value # --- When a GPU IS available --- @@ -167,9 +178,10 @@ class TestGetGpuMemoryInfo: # --- CUDA-specific mocked test --- + @needs_torch def test_cuda_path_returns_correct_fields(self): mock_props = MagicMock() - mock_props.total_memory = 16 * (1024 ** 3) # 16 GB + mock_props.total_memory = 16 * (1024 ** 3) mock_props.name = "NVIDIA Test GPU" with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CUDA), \ @@ -189,9 +201,10 @@ class TestGetGpuMemoryInfo: # --- MPS-specific mocked test --- + @needs_torch def test_mps_path_returns_correct_fields(self): mock_psutil_mem = MagicMock() - mock_psutil_mem.total = 32 * (1024 ** 3) # 32 GB unified + mock_psutil_mem.total = 32 * (1024 ** 3) mock_psutil = MagicMock() mock_psutil.virtual_memory.return_value = mock_psutil_mem @@ -215,17 +228,16 @@ class TestGetGpuMemoryInfo: def test_cpu_path_returns_unavailable(self): with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU): result = get_gpu_memory_info() - assert result["available"] is False assert result["backend"] == "cpu" # --- Error resilience --- + @needs_torch def test_cuda_error_returns_unavailable(self): with patch("utils.hardware.hardware.get_device", return_value=DeviceType.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 @@ -247,10 +259,10 @@ class TestLogGpuMemory: "utilization_pct": 12.5, "free_gb": 14.0, } - with patch("utils.hardware.hardware.get_gpu_memory_info", return_value=fake_info): - import logging - with caplog.at_level(logging.INFO, logger="utils.hardware.hardware"): - log_gpu_memory("unit-test") + import logging + with patch("utils.hardware.hardware.get_gpu_memory_info", return_value=fake_info), \ + caplog.at_level(logging.INFO, logger="utils.hardware.hardware"): + log_gpu_memory("unit-test") assert "unit-test" in caplog.text assert "CUDA" in caplog.text @@ -258,10 +270,10 @@ class TestLogGpuMemory: def test_logs_cpu_fallback_when_no_gpu(self, caplog): fake_info = {"available": False, "backend": "cpu"} - with patch("utils.hardware.hardware.get_gpu_memory_info", return_value=fake_info): - import logging - with caplog.at_level(logging.INFO, logger="utils.hardware.hardware"): - log_gpu_memory("cpu-test") + import logging + with patch("utils.hardware.hardware.get_gpu_memory_info", return_value=fake_info), \ + caplog.at_level(logging.INFO, logger="utils.hardware.hardware"): + log_gpu_memory("cpu-test") assert "No GPU available" in caplog.text @@ -298,7 +310,7 @@ class TestFormatErrorMessage: with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CUDA): msg = format_error_message(err, "big/model") assert "GPU" in msg - assert "big/model" not in msg # should use short name + assert "big/model" not in msg assert "model" in msg # --- OOM on MPS --- From e7c3e7b48d4ac9132a6049b9d826842bcb6ba95b Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 11 Feb 2026 15:40:00 +0000 Subject: [PATCH 4/8] fixed tests to be hardware specific --- studio/backend/tests/test_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index 0ac160ecb9..ffc3599474 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -307,7 +307,7 @@ class TestFormatErrorMessage: def test_cuda_oom(self): err = Exception("CUDA out of memory") - with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CUDA): + with patch("utils.hardware.get_device", return_value=DeviceType.CUDA): msg = format_error_message(err, "big/model") assert "GPU" in msg assert "big/model" not in msg @@ -317,7 +317,7 @@ class TestFormatErrorMessage: def test_mps_oom(self): err = Exception("MPS backend out of memory") - with patch("utils.hardware.hardware.get_device", return_value=DeviceType.MPS): + with patch("utils.hardware.get_device", return_value=DeviceType.MPS): msg = format_error_message(err, "unsloth/huge-model") assert "Apple Silicon" in msg @@ -325,7 +325,7 @@ class TestFormatErrorMessage: def test_cpu_oom(self): err = Exception("not enough memory to allocate") - with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU): + with patch("utils.hardware.get_device", return_value=DeviceType.CPU): msg = format_error_message(err, "any/model") assert "system" in msg.lower() From 7db31723b9ce7310b2107c95dbf0e25be7cde1c9 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 11 Feb 2026 15:58:13 +0000 Subject: [PATCH 5/8] reset DEVICE type on fastapi lifespan exit --- studio/backend/main.py | 4 +++- studio/backend/tests/test_utils.py | 22 +++++++++++++++++++--- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/studio/backend/main.py b/studio/backend/main.py index de468c183a..e957c668e3 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -16,6 +16,7 @@ from datetime import datetime from routes import training_router, models_router, inference_router, datasets_router, auth_router from auth import storage from utils.hardware import detect_hardware +import utils.hardware.hardware as _hw_module UNSLOTH_CACHE_DIR = Path(__file__).parent / "unsloth_compiled_cache" @@ -36,7 +37,8 @@ async def lifespan(app: FastAPI): print("This token can only be used once.") print("=" * 60 + "\n") yield - # Cleanup: remove Unsloth compiled cache on shutdown + # Cleanup + _hw_module.DEVICE = None shutil.rmtree(UNSLOTH_CACHE_DIR, ignore_errors=True) diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index ffc3599474..1c2a80b8bd 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -29,12 +29,14 @@ needs_torch = pytest.mark.skipif(not HAS_TORCH, reason="PyTorch not installed") from utils.hardware import ( get_device, + detect_hardware, is_apple_silicon, clear_gpu_cache, get_gpu_memory_info, log_gpu_memory, DeviceType, ) +import utils.hardware.hardware as _hw_module from utils.utils import format_error_message @@ -50,11 +52,25 @@ def _actual_device() -> str: return "cpu" +def _reset_and_detect(): + """Reset the cached DEVICE global and re-run detection.""" + _hw_module.DEVICE = None + return detect_hardware() + + # ========== get_device() ========== 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) @@ -67,7 +83,7 @@ class TestGetDevice: @needs_torch def test_returns_cuda_when_cuda_available(self): with patch("torch.cuda.is_available", return_value=True): - assert get_device() == DeviceType.CUDA + assert _reset_and_detect() == DeviceType.CUDA @needs_torch def test_returns_mps_when_only_mps_available(self): @@ -75,13 +91,13 @@ class TestGetDevice: 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() == DeviceType.MPS + assert _reset_and_detect() == DeviceType.MPS @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)): - assert get_device() == DeviceType.CPU + assert _reset_and_detect() == DeviceType.CPU # ========== is_apple_silicon() ========== From 63c583c54f1cf6b250a78a499d15cbcfbcbb5a77 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 11 Feb 2026 16:04:35 +0000 Subject: [PATCH 6/8] replace torch MPS with MLX --- studio/backend/tests/test_utils.py | 89 +++++++++++------------ studio/backend/utils/hardware/hardware.py | 73 +++++++++++-------- studio/backend/utils/utils.py | 4 +- 3 files changed, 86 insertions(+), 80 deletions(-) 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 From 95038d61290edb03f2c218784a2faa97b5c5d18a Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 11 Feb 2026 16:10:22 +0000 Subject: [PATCH 7/8] add @needs_mlx decorator on tests --- studio/backend/tests/test_utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index 03e221d336..60a9969edc 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -90,6 +90,7 @@ class TestGetDevice: patch("torch.cuda.is_available", return_value=True): assert _reset_and_detect() == DeviceType.CUDA + @needs_mlx 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), \ @@ -147,6 +148,7 @@ class TestClearGpuCache: mock_empty.assert_called_once() mock_ipc.assert_called_once() + @needs_mlx 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): @@ -324,6 +326,7 @@ class TestFormatErrorMessage: # --- OOM on MLX --- + @needs_mlx def test_mlx_oom(self): err = Exception("MLX backend out of memory") with patch("utils.hardware.get_device", return_value=DeviceType.MLX): From 1a6bfe51b620da2c491756441b42d594fcc841ed Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 11 Feb 2026 16:12:37 +0000 Subject: [PATCH 8/8] added @needs_torch to test_cuda_oom --- studio/backend/tests/test_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index 60a9969edc..afb4bab65a 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -316,6 +316,7 @@ class TestFormatErrorMessage: # --- OOM on CUDA --- + @needs_torch def test_cuda_oom(self): err = Exception("CUDA out of memory") with patch("utils.hardware.get_device", return_value=DeviceType.CUDA):