diff --git a/studio/backend/core/__init__.py b/studio/backend/core/__init__.py index 46ee6c14b1..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 +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,4 +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..e957c668e3 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -15,13 +15,18 @@ 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 +import utils.hardware.hardware as _hw_module 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) @@ -32,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) @@ -78,23 +84,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 +105,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/__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..afb4bab65a --- /dev/null +++ b/studio/backend/tests/test_utils.py @@ -0,0 +1,350 @@ +""" +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, 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/mlx internals for mocking are skipped when unavailable. + +Run with: + cd studio/backend + python -m pytest tests/test_utils.py -v +""" +import platform +from unittest.mock import patch, MagicMock + +import pytest + +# --- 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, + 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 + + +# ========== Helpers ========== + +def _actual_device() -> str: + """Return the real device string for the current machine.""" + if HAS_TORCH and torch.cuda.is_available(): + return "cuda" + if is_apple_silicon() and HAS_MLX: + return "mlx" + 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): + self._saved_device = _hw_module.DEVICE + + def teardown_method(self): + _hw_module.DEVICE = self._saved_device + + def test_returns_valid_device_type(self): + result = get_device() + assert result in (DeviceType.CUDA, DeviceType.MLX, DeviceType.CPU) + + def test_matches_actual_hardware(self): + assert get_device().value == _actual_device() + + # --- Mocked paths --- + + @needs_torch + def test_returns_cuda_when_cuda_available(self): + 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_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), \ + patch("utils.hardware.hardware._has_mlx", return_value=True): + assert _reset_and_detect() == DeviceType.MLX + + def test_returns_cpu_when_nothing_available(self): + 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 + + +# ========== 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("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("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("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 + + +# ========== clear_gpu_cache() ========== + +class TestClearGpuCache: + """clear_gpu_cache() must never raise, regardless of platform.""" + + def test_does_not_raise(self): + clear_gpu_cache() + + @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: + clear_gpu_cache() + 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): + clear_gpu_cache() + + def test_noop_on_cpu(self): + with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU): + 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): + assert "available" in get_gpu_memory_info() + + def test_has_backend_key(self): + assert "backend" in get_gpu_memory_info() + + def test_backend_matches_device(self): + result = get_gpu_memory_info() + assert result["backend"] == get_device().value + + # --- 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 --- + + @needs_torch + def test_cuda_path_returns_correct_fields(self): + mock_props = MagicMock() + mock_props.total_memory = 16 * (1024 ** 3) + mock_props.name = "NVIDIA Test GPU" + + 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)), \ + 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 + + # --- MLX-specific mocked test --- + + @needs_mlx + def test_mlx_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 + + 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"] == "mlx" + assert "Apple Silicon" in result["device_name"] + assert abs(result["total_gb"] - 32.0) < 0.01 + + # --- CPU-only path --- + + 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 + + +# ========== 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, + } + 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 + assert "FakeGPU" in caplog.text + + def test_logs_cpu_fallback_when_no_gpu(self, caplog): + fake_info = {"available": False, "backend": "cpu"} + 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 + + +# ========== 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 --- + + @needs_torch + def test_cuda_oom(self): + err = Exception("CUDA out of memory") + 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 + assert "model" in msg + + # --- 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): + 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.hardware.get_device", return_value=DeviceType.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/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..25320a25e1 --- /dev/null +++ b/studio/backend/utils/hardware/hardware.py @@ -0,0 +1,208 @@ +""" +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" + MLX = "mlx" + 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" + + +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: + """ + 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. MLX (Apple Silicon via MLX framework) + 3. CPU (fallback) + """ + global DEVICE + + # --- 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 + + # --- 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 + + +# ========== 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 + + +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.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), MLX (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)} + + # ---- MLX path (Apple Silicon) ---- + if device == DeviceType.MLX: + try: + import mlx.core as mx + import psutil + + # 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, + "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, + "free_gb": (total - allocated) / (1024**3), + "utilization_pct": (allocated / total) * 100 if total else 0, + } + except Exception as 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"} + + +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)") diff --git a/studio/backend/utils/utils.py b/studio/backend/utils/utils.py index 74d5c05abf..0acdbb3313 100644 --- a/studio/backend/utils/utils.py +++ b/studio/backend/utils/utils.py @@ -1,18 +1,17 @@ """ Shared backend utilities """ -import gradio as gr import os import logging from contextlib import contextmanager from pathlib import Path -from typing import Optional, Dict, Any import shutil import tempfile logger = logging.getLogger(__name__) + @contextmanager def without_hf_auth(): """ @@ -96,113 +95,12 @@ 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 "mlx" in error_str or "out of memory" in error_str: + from utils.hardware import get_device + device = get_device() + 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 return str(error) pass - -def get_gpu_memory_info() -> Dict[str, Any]: - """Get GPU memory information.""" - import torch - - if not torch.cuda.is_available(): - return {"available": False} - - 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, - "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)} -pass - -def log_gpu_memory(context: str): - """Log GPU memory usage with context.""" - memory_info = get_gpu_memory_info() - if memory_info.get("available"): - logger.info( - f"GPU Memory [{context}]: " - 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") -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=[])