c8c371610b
TestLogGpuMemory tests used pytest's caplog fixture to capture log output, but log_gpu_memory() uses structlog which writes to stdout, not Python's standard logging module. caplog only captures standard logging messages, so assertions always failed on empty text. Switch to capsys which captures stdout directly.
379 lines
12 KiB
Python
379 lines
12 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""
|
|
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, capsys):
|
|
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.hardware.hardware.get_gpu_memory_info", return_value = fake_info
|
|
):
|
|
log_gpu_memory("unit-test")
|
|
|
|
captured = capsys.readouterr().out
|
|
assert "unit-test" in captured
|
|
assert "CUDA" in captured
|
|
assert "FakeGPU" in captured
|
|
|
|
def test_logs_cpu_fallback_when_no_gpu(self, capsys):
|
|
fake_info = {"available": False, "backend": "cpu"}
|
|
|
|
with patch(
|
|
"utils.hardware.hardware.get_gpu_memory_info", return_value = fake_info
|
|
):
|
|
log_gpu_memory("cpu-test")
|
|
|
|
captured = capsys.readouterr().out
|
|
assert "No GPU available" in captured
|
|
|
|
|
|
# ========== 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"
|