ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

integrate global hardware detection at lifespan entrypoint

Browse source
This commit is contained in:
Roland Tannous 2026-02-11 15:34:26 +00:00
commit 59d5f24eb5
6 changed files with 276 additions and 238 deletions
Download patch file Download diff file Expand all files Collapse all files

6
studio/backend/core/__init__.py
View file

@ -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',
]

32
studio/backend/main.py
View file

@ -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),

70
studio/backend/tests/test_utils.py
View file

@ -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()

24
studio/backend/utils/hardware/__init__.py Normal file
View file

@ -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',
]

195
studio/backend/utils/hardware/hardware.py Normal file
View file

@ -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

187
studio/backend/utils/utils.py
View file

@ -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=[])