Final cleanup
This commit is contained in:
Roland Tannous
parent
7e336049d8
commit
985d2e43ee
123 changed files with 7474 additions and 5805 deletions
|
|
@ -16,6 +16,7 @@ Run with:
|
|||
cd studio/backend
|
||||
python -m pytest tests/test_utils.py -v
|
||||
"""
|
||||
|
||||
import platform
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
|
|
@ -24,18 +25,20 @@ 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")
|
||||
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,
|
||||
|
|
@ -52,6 +55,7 @@ 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():
|
||||
|
|
@ -69,6 +73,7 @@ def _reset_and_detect():
|
|||
|
||||
# ========== get_device() ==========
|
||||
|
||||
|
||||
class TestGetDevice:
|
||||
"""Tests for get_device() — should agree with the real hardware."""
|
||||
|
||||
|
|
@ -89,28 +94,34 @@ class TestGetDevice:
|
|||
|
||||
@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):
|
||||
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):
|
||||
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):
|
||||
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:
|
||||
|
||||
class TestIsAppleSilicon:
|
||||
def test_returns_bool(self):
|
||||
assert isinstance(is_apple_silicon(), bool)
|
||||
|
||||
|
|
@ -136,6 +147,7 @@ class TestIsAppleSilicon:
|
|||
|
||||
# ========== clear_gpu_cache() ==========
|
||||
|
||||
|
||||
class TestClearGpuCache:
|
||||
"""clear_gpu_cache() must never raise, regardless of platform."""
|
||||
|
||||
|
|
@ -144,9 +156,11 @@ class TestClearGpuCache:
|
|||
|
||||
@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:
|
||||
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()
|
||||
|
|
@ -154,18 +168,18 @@ class TestClearGpuCache:
|
|||
@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):
|
||||
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):
|
||||
with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CPU):
|
||||
clear_gpu_cache()
|
||||
|
||||
|
||||
# ========== get_gpu_memory_info() ==========
|
||||
|
||||
class TestGetGpuMemoryInfo:
|
||||
|
||||
class TestGetGpuMemoryInfo:
|
||||
def test_returns_dict(self):
|
||||
result = get_gpu_memory_info()
|
||||
assert isinstance(result, dict)
|
||||
|
|
@ -183,8 +197,7 @@ class TestGetGpuMemoryInfo:
|
|||
# --- When a GPU IS available ---
|
||||
|
||||
@pytest.mark.skipif(
|
||||
_actual_device() == "cpu",
|
||||
reason="No GPU available on this machine"
|
||||
_actual_device() == "cpu", reason = "No GPU available on this machine"
|
||||
)
|
||||
def test_gpu_available_fields(self):
|
||||
result = get_gpu_memory_info()
|
||||
|
|
@ -200,14 +213,16 @@ class TestGetGpuMemoryInfo:
|
|||
@needs_torch
|
||||
def test_cuda_path_returns_correct_fields(self):
|
||||
mock_props = MagicMock()
|
||||
mock_props.total_memory = 16 * (1024 ** 3)
|
||||
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)):
|
||||
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
|
||||
|
|
@ -223,13 +238,15 @@ class TestGetGpuMemoryInfo:
|
|||
@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_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}):
|
||||
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
|
||||
|
|
@ -240,7 +257,7 @@ class TestGetGpuMemoryInfo:
|
|||
# --- CPU-only path ---
|
||||
|
||||
def test_cpu_path_returns_unavailable(self):
|
||||
with patch("utils.hardware.hardware.get_device", return_value=DeviceType.CPU):
|
||||
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"
|
||||
|
|
@ -249,8 +266,13 @@ class TestGetGpuMemoryInfo:
|
|||
|
||||
@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")):
|
||||
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
|
||||
|
|
@ -258,8 +280,8 @@ class TestGetGpuMemoryInfo:
|
|||
|
||||
# ========== log_gpu_memory() ==========
|
||||
|
||||
class TestLogGpuMemory:
|
||||
|
||||
class TestLogGpuMemory:
|
||||
def test_does_not_raise(self):
|
||||
log_gpu_memory("test")
|
||||
|
||||
|
|
@ -275,8 +297,13 @@ class TestLogGpuMemory:
|
|||
}
|
||||
import structlog
|
||||
from loggers import get_logger
|
||||
with patch("utils.hardware.hardware.get_gpu_memory_info", return_value=fake_info), \
|
||||
caplog.at_level(logging.INFO, logger="utils.hardware.hardware"):
|
||||
|
||||
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
|
||||
|
|
@ -287,8 +314,13 @@ class TestLogGpuMemory:
|
|||
fake_info = {"available": False, "backend": "cpu"}
|
||||
import structlog
|
||||
from loggers import get_logger
|
||||
with patch("utils.hardware.hardware.get_gpu_memory_info", return_value=fake_info), \
|
||||
caplog.at_level(logging.INFO, logger="utils.hardware.hardware"):
|
||||
|
||||
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
|
||||
|
|
@ -296,8 +328,8 @@ class TestLogGpuMemory:
|
|||
|
||||
# ========== format_error_message() ==========
|
||||
|
||||
class TestFormatErrorMessage:
|
||||
|
||||
class TestFormatErrorMessage:
|
||||
def test_not_found(self):
|
||||
err = Exception("Repository not found for unsloth/test")
|
||||
msg = format_error_message(err, "unsloth/test")
|
||||
|
|
@ -324,7 +356,7 @@ class TestFormatErrorMessage:
|
|||
@needs_torch
|
||||
def test_cuda_oom(self):
|
||||
err = Exception("CUDA out of memory")
|
||||
with patch("utils.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
|
||||
|
|
@ -335,7 +367,7 @@ class TestFormatErrorMessage:
|
|||
@needs_mlx
|
||||
def test_mlx_oom(self):
|
||||
err = Exception("MLX backend out of memory")
|
||||
with patch("utils.hardware.get_device", return_value=DeviceType.MLX):
|
||||
with patch("utils.hardware.get_device", return_value = DeviceType.MLX):
|
||||
msg = format_error_message(err, "unsloth/huge-model")
|
||||
assert "Apple Silicon" in msg
|
||||
|
||||
|
|
@ -343,7 +375,7 @@ class TestFormatErrorMessage:
|
|||
|
||||
def test_cpu_oom(self):
|
||||
err = Exception("not enough memory to allocate")
|
||||
with patch("utils.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()
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue