bc00a8e797
Six review findings, three of them evict-then-fail orderings: - The chat load reclaimed the GPU without telling the arbiter it existed. A chat load holds no llama-server process until its GGUF has downloaded, which is minutes, so a competing Images/Video acquire in that window found nothing to cancel, took the GPU, and the chat load then spawned onto the same device. It now registers an in-flight marker through acquire_for's register hook (under the arbiter lock, as the image and video loads do), the evictor cancels a marked load, and the route undoes itself if ownership moved while it loaded. - The Hub-download conflict check ran after that handoff, so a GGUF the download manager already owns destroyed the resident Images/Video pipeline and then 409'd, having loaded nothing. It moves above the handoff, together with the marker it handshakes with. - The image load released the engine router's transition lock before registering the load, so a second load choosing the other engine could unload the still-idle engine this one captured; the load then landed on a deactivated engine, where generate, status, unload and the arbiter's evictor can no longer reach it. Registration now happens under that lock and refuses if the engine changed. - Training a DiT family on a host with no GPU was accepted: nf4 is not a CPU fallback, its 4-bit load goes through bitsandbytes, which requires CUDA, XPU or MPS. The start unloaded the working Images pipeline, pulled the text encoders, and only then died in the child. Rejected before the teardown now, and /info stops advertising a precision that always 400s. SDXL keeps its documented fp32-on-CPU path. - Both diffusion pages kept the routed-pick marker forever, so re-picking the same checkpoint (after chat evicted it) neither loaded nor cleared the query string. The marker is released once the query is gone. The Images key also carried a stray NUL byte, which made the file read as binary to grep and other tooling. - diffusers dropped Python 3.9 in 0.38, so the unconditional >=0.39.0 pin left pip no candidate at all on 3.9 and made every install that composes the huggingface extras unresolvable there. The floor is conditional now. Also fixes tests that were already red on the branch: two hand-built request fakes had gone stale against fields this branch added, and the handoff-ordering test only failed on a host with fewer than two GPUs.
1986 lines
77 KiB
Python
1986 lines
77 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
|
|
|
|
import asyncio
|
|
import importlib.util
|
|
import os
|
|
import re
|
|
import sys
|
|
import unittest
|
|
from contextlib import nullcontext
|
|
from pathlib import Path
|
|
from types import ModuleType, SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from fastapi import HTTPException
|
|
|
|
from core.training.training import TrainingBackend
|
|
from models.inference import LoadRequest
|
|
from models.training import TrainingStartRequest
|
|
from utils.hardware import (
|
|
apply_gpu_ids,
|
|
DeviceType,
|
|
auto_select_gpu_ids,
|
|
estimate_required_model_memory_gb,
|
|
get_backend_visible_gpu_info,
|
|
get_device_map,
|
|
get_gpu_utilization,
|
|
get_offloaded_device_map_entries,
|
|
get_parent_visible_gpu_ids,
|
|
get_visible_gpu_utilization,
|
|
prepare_gpu_selection,
|
|
resolve_requested_gpu_ids,
|
|
)
|
|
import utils.hardware.hardware as _hw_module
|
|
|
|
_BACKEND_ROOT = Path(__file__).resolve().parent.parent
|
|
|
|
|
|
async def _inline_to_thread(func, /, *args, **kwargs):
|
|
return func(*args, **kwargs)
|
|
|
|
|
|
def _fake_unsloth_attention_modules(resolver):
|
|
unsloth_module = ModuleType("unsloth")
|
|
models_module = ModuleType("unsloth.models")
|
|
utils_module = ModuleType("unsloth.models._utils")
|
|
utils_module.resolve_attention_implementation = resolver
|
|
models_module._utils = utils_module
|
|
unsloth_module.models = models_module
|
|
return {
|
|
"unsloth": unsloth_module,
|
|
"unsloth.models": models_module,
|
|
"unsloth.models._utils": utils_module,
|
|
}
|
|
|
|
|
|
def _load_route_module(name: str, relative_path: str):
|
|
spec = importlib.util.spec_from_file_location(name, _BACKEND_ROOT / relative_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class _GpuCacheResetMixin:
|
|
"""Reset module-level GPU caches between tests to prevent state leaks."""
|
|
|
|
def tearDown(self):
|
|
_hw_module._physical_gpu_count = None
|
|
_hw_module._visible_gpu_count = None
|
|
|
|
|
|
class TestResolveRequestedGpuIds(_GpuCacheResetMixin, unittest.TestCase):
|
|
def test_parent_visibility_defaults_to_physical_enumeration(self):
|
|
with (
|
|
patch.dict(os.environ, {}, clear = True),
|
|
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 4),
|
|
):
|
|
self.assertEqual(get_parent_visible_gpu_ids(), [0, 1, 2, 3])
|
|
self.assertEqual(resolve_requested_gpu_ids(None), [0, 1, 2, 3])
|
|
|
|
def test_parent_visibility_uses_cuda_visible_devices(self):
|
|
with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "1,3"}, clear = True):
|
|
self.assertEqual(get_parent_visible_gpu_ids(), [1, 3])
|
|
self.assertEqual(resolve_requested_gpu_ids(None), [1, 3])
|
|
|
|
def test_parent_visibility_uses_empty_numeric_ids_for_uuid_masks(self):
|
|
with (
|
|
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True),
|
|
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
|
|
):
|
|
self.assertEqual(get_parent_visible_gpu_ids(), [])
|
|
|
|
def test_invalid_requests_raise_clear_value_errors(self):
|
|
cases = [
|
|
([1, 1], "duplicate GPU IDs"),
|
|
([-1], "Rejected IDs: [-1]"),
|
|
([99], "Rejected IDs: [99]"),
|
|
([0], "outside the parent-visible set [1, 3]"),
|
|
]
|
|
with (
|
|
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "1,3"}, clear = True),
|
|
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
|
|
):
|
|
for gpu_ids, message in cases:
|
|
with self.subTest(gpu_ids = gpu_ids):
|
|
with self.assertRaisesRegex(ValueError, re.escape(message)):
|
|
resolve_requested_gpu_ids(gpu_ids)
|
|
|
|
def test_explicit_ids_must_be_physical_not_relative(self):
|
|
with (
|
|
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "1,3"}, clear = True),
|
|
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
|
|
):
|
|
self.assertEqual(resolve_requested_gpu_ids([1, 3]), [1, 3])
|
|
|
|
def test_explicit_ids_are_rejected_for_uuid_parent_visibility(self):
|
|
with (
|
|
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True),
|
|
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
|
|
):
|
|
with self.assertRaisesRegex(
|
|
ValueError,
|
|
"unsupported when CUDA_VISIBLE_DEVICES uses non-numeric or subdevice",
|
|
):
|
|
resolve_requested_gpu_ids([1])
|
|
|
|
def test_empty_list_is_treated_as_auto(self):
|
|
with (
|
|
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "1,3"}, clear = True),
|
|
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 8),
|
|
):
|
|
self.assertEqual(resolve_requested_gpu_ids([]), [1, 3])
|
|
|
|
def test_vulkan_ordinals_bypass_cuda_parent_visible_validation(self):
|
|
# Vulkan build on a CPU-only torch host: no CUDA parent-visible set and a
|
|
# zero physical count, yet a valid Vulkan ordinal must not be rejected as
|
|
# a CUDA physical id (issue #7239).
|
|
with (
|
|
patch.dict(os.environ, {}, clear = True),
|
|
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 0),
|
|
):
|
|
# As a CUDA physical id, [0] is outside the empty parent-visible set.
|
|
with self.assertRaises(ValueError):
|
|
resolve_requested_gpu_ids([0])
|
|
# As Vulkan ordinals, [0] and [0, 1] pass through unchanged.
|
|
self.assertEqual(resolve_requested_gpu_ids([0], is_vulkan = True), [0])
|
|
self.assertEqual(resolve_requested_gpu_ids([0, 1], is_vulkan = True), [0, 1])
|
|
# Malformed ordinals are still rejected.
|
|
with self.assertRaisesRegex(ValueError, "duplicate GPU IDs"):
|
|
resolve_requested_gpu_ids([0, 0], is_vulkan = True)
|
|
with self.assertRaisesRegex(ValueError, "non-negative"):
|
|
resolve_requested_gpu_ids([-1], is_vulkan = True)
|
|
|
|
def test_apply_gpu_ids_only_updates_cuda_visible_devices(self):
|
|
with patch.dict(
|
|
os.environ,
|
|
{"CUDA_VISIBLE_DEVICES": "1,3", "TEST_PARENT_ENV": "keep-me"},
|
|
clear = True,
|
|
):
|
|
apply_gpu_ids([5, 6])
|
|
|
|
self.assertEqual(os.environ["CUDA_VISIBLE_DEVICES"], "5,6")
|
|
self.assertEqual(os.environ["TEST_PARENT_ENV"], "keep-me")
|
|
|
|
|
|
class TestVisibleGpuUtilization(_GpuCacheResetMixin, unittest.TestCase):
|
|
def test_gpu_utilization_preserves_primary_shape_with_devices(self):
|
|
devices = [
|
|
{
|
|
"index": 5,
|
|
"visible_ordinal": 0,
|
|
"gpu_utilization_pct": 11.0,
|
|
"temperature_c": 40.0,
|
|
"vram_used_gb": 4.0,
|
|
"vram_total_gb": 24.0,
|
|
"vram_utilization_pct": 16.7,
|
|
"power_draw_w": 80.0,
|
|
"power_limit_w": 300.0,
|
|
"power_utilization_pct": 26.7,
|
|
},
|
|
{
|
|
"index": 3,
|
|
"visible_ordinal": 1,
|
|
"gpu_utilization_pct": 22.0,
|
|
"temperature_c": 50.0,
|
|
"vram_used_gb": 8.0,
|
|
"vram_total_gb": 24.0,
|
|
"vram_utilization_pct": 33.3,
|
|
"power_draw_w": 120.0,
|
|
"power_limit_w": 300.0,
|
|
"power_utilization_pct": 40.0,
|
|
},
|
|
]
|
|
|
|
with (
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch.object(_hw_module, "IS_ROCM", False),
|
|
patch(
|
|
"utils.hardware.hardware._get_parent_visible_gpu_spec",
|
|
return_value = {"raw": "5,3", "numeric_ids": [5, 3]},
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._smi_query",
|
|
return_value = {
|
|
"available": True,
|
|
"devices": devices,
|
|
"backend_cuda_visible_devices": "5,3",
|
|
"parent_visible_gpu_ids": [5, 3],
|
|
"index_kind": "physical",
|
|
},
|
|
),
|
|
):
|
|
result = get_gpu_utilization()
|
|
|
|
self.assertIsInstance(result, dict)
|
|
self.assertTrue(result["available"])
|
|
self.assertEqual(result["backend"], "cuda")
|
|
self.assertEqual(result["index"], 5)
|
|
self.assertEqual(result["visible_ordinal"], 0)
|
|
self.assertEqual(result["vram_total_gb"], 24.0)
|
|
self.assertEqual(result["parent_visible_gpu_ids"], [5, 3])
|
|
self.assertEqual([device["index"] for device in result["devices"]], [5, 3])
|
|
|
|
def test_gpu_utilization_cpu_returns_legacy_unavailable_object(self):
|
|
with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CPU):
|
|
result = get_gpu_utilization()
|
|
|
|
self.assertEqual(result, {"available": False, "backend": "cpu", "devices": []})
|
|
|
|
def test_gpu_utilization_mlx_stays_available_without_agx_stats(self):
|
|
fake_psutil = ModuleType("psutil")
|
|
fake_psutil.virtual_memory = lambda: SimpleNamespace(total = 64 * 1024**3)
|
|
|
|
with (
|
|
patch.dict(sys.modules, {"psutil": fake_psutil}),
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.MLX),
|
|
patch("utils.hardware.hardware._read_apple_gpu_stats", return_value = {}),
|
|
patch(
|
|
"core.training.get_training_backend",
|
|
return_value = SimpleNamespace(_progress = None),
|
|
),
|
|
patch("utils.hardware.apple.read_gpu_temperature_c", return_value = None),
|
|
patch("utils.hardware.apple.read_gpu_power_w", return_value = None),
|
|
):
|
|
result = get_gpu_utilization()
|
|
|
|
self.assertTrue(result["available"])
|
|
self.assertEqual(result["backend"], "mlx")
|
|
self.assertIsNone(result["gpu_utilization_pct"])
|
|
self.assertEqual(result["vram_used_gb"], 0)
|
|
self.assertEqual(result["vram_total_gb"], 64.0)
|
|
self.assertEqual(len(result["devices"]), 1)
|
|
|
|
def test_gpu_utilization_xpu_uses_visible_devices(self):
|
|
with (
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU),
|
|
patch(
|
|
"utils.hardware.hardware.get_visible_gpu_utilization",
|
|
return_value = {
|
|
"available": True,
|
|
"backend": "xpu",
|
|
"parent_visible_gpu_ids": [2, 0],
|
|
"index_kind": "physical",
|
|
"devices": [
|
|
{
|
|
"index": 2,
|
|
"visible_ordinal": 1,
|
|
"gpu_utilization_pct": None,
|
|
"temperature_c": None,
|
|
"vram_used_gb": 3.0,
|
|
"vram_total_gb": 16.0,
|
|
"vram_utilization_pct": 18.8,
|
|
"power_draw_w": None,
|
|
"power_limit_w": None,
|
|
"power_utilization_pct": None,
|
|
},
|
|
{
|
|
"index": 0,
|
|
"visible_ordinal": 0,
|
|
"gpu_utilization_pct": None,
|
|
"temperature_c": None,
|
|
"vram_used_gb": 1.0,
|
|
"vram_total_gb": 16.0,
|
|
"vram_utilization_pct": 6.3,
|
|
"power_draw_w": None,
|
|
"power_limit_w": None,
|
|
"power_utilization_pct": None,
|
|
},
|
|
],
|
|
},
|
|
),
|
|
):
|
|
result = get_gpu_utilization()
|
|
|
|
self.assertEqual(result["backend"], "xpu")
|
|
self.assertEqual(result["index"], 0)
|
|
self.assertEqual(result["visible_ordinal"], 0)
|
|
self.assertEqual([device["index"] for device in result["devices"]], [0, 2])
|
|
|
|
def test_visible_gpu_utilization_filters_to_parent_visible_ids(self):
|
|
smi_output = "\n".join(
|
|
[
|
|
"0, 10, 30, 1000, 10000, 50, 100",
|
|
"1, 20, 40, 2000, 10000, 60, 120",
|
|
"3, 30, 50, 3000, 10000, 70, 140",
|
|
]
|
|
)
|
|
|
|
with (
|
|
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "1,3"}, clear = True),
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch("utils.hardware.nvidia.subprocess.run") as mock_run,
|
|
):
|
|
mock_run.return_value = SimpleNamespace(
|
|
returncode = 0,
|
|
stdout = smi_output,
|
|
)
|
|
result = get_visible_gpu_utilization()
|
|
|
|
self.assertTrue(result["available"])
|
|
self.assertEqual(result["parent_visible_gpu_ids"], [1, 3])
|
|
self.assertEqual(result["index_kind"], "physical")
|
|
self.assertEqual([device["index"] for device in result["devices"]], [1, 3])
|
|
self.assertEqual(result["devices"][0]["visible_ordinal"], 0)
|
|
self.assertEqual(result["devices"][1]["visible_ordinal"], 1)
|
|
self.assertEqual(result["devices"][0]["gpu_utilization_pct"], 20.0)
|
|
self.assertEqual(result["devices"][1]["power_utilization_pct"], 50.0)
|
|
|
|
def test_backend_visible_gpu_info_preserves_physical_indices(self):
|
|
smi_output = "\n".join(
|
|
[
|
|
"0, GPU Zero, 10000",
|
|
"1, GPU One, 20000",
|
|
"3, GPU Three, 30000",
|
|
]
|
|
)
|
|
|
|
with (
|
|
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "1,3"}, clear = True),
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch("utils.hardware.nvidia.subprocess.run") as mock_run,
|
|
):
|
|
mock_run.return_value = SimpleNamespace(
|
|
returncode = 0,
|
|
stdout = smi_output,
|
|
)
|
|
result = get_backend_visible_gpu_info()
|
|
|
|
self.assertTrue(result["available"])
|
|
self.assertEqual(result["parent_visible_gpu_ids"], [1, 3])
|
|
self.assertEqual(result["index_kind"], "physical")
|
|
self.assertEqual([device["index"] for device in result["devices"]], [1, 3])
|
|
self.assertEqual(result["devices"][0]["visible_ordinal"], 0)
|
|
self.assertEqual(result["devices"][1]["visible_ordinal"], 1)
|
|
self.assertEqual(result["devices"][0]["name"], "GPU One")
|
|
self.assertAlmostEqual(result["devices"][1]["memory_total_gb"], 29.3, places = 1)
|
|
|
|
def test_uuid_parent_visibility_falls_back_to_torch(self):
|
|
"""UUID/MIG masks fall through nvidia to the torch fallback and
|
|
still report visible devices using relative ordinals."""
|
|
fake_torch_devices = [
|
|
{
|
|
"index": 0,
|
|
"visible_ordinal": 0,
|
|
"name": "GPU-A",
|
|
"total_gb": 24.0,
|
|
"used_gb": 2.0,
|
|
},
|
|
{
|
|
"index": 1,
|
|
"visible_ordinal": 1,
|
|
"name": "GPU-B",
|
|
"total_gb": 24.0,
|
|
"used_gb": 3.0,
|
|
},
|
|
]
|
|
with (
|
|
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True),
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch("utils.hardware.hardware._torch_get_physical_gpu_count", return_value = 2),
|
|
patch(
|
|
"utils.hardware.hardware._torch_get_per_device_info",
|
|
return_value = fake_torch_devices,
|
|
),
|
|
):
|
|
result = get_backend_visible_gpu_info()
|
|
|
|
self.assertTrue(result["available"])
|
|
self.assertEqual(result["parent_visible_gpu_ids"], [])
|
|
self.assertEqual(len(result["devices"]), 2)
|
|
self.assertEqual(result["index_kind"], "relative")
|
|
|
|
def test_mlx_visible_gpu_info_is_best_effort_relative(self):
|
|
with (
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.MLX),
|
|
patch(
|
|
"utils.hardware.hardware.get_gpu_memory_info",
|
|
return_value = {
|
|
"available": True,
|
|
"device_name": "Apple Silicon",
|
|
"total_gb": 64.0,
|
|
"allocated_gb": 8.0,
|
|
"utilization_pct": 12.5,
|
|
},
|
|
),
|
|
):
|
|
result = get_backend_visible_gpu_info()
|
|
|
|
self.assertTrue(result["available"])
|
|
self.assertEqual(result["index_kind"], "relative")
|
|
self.assertEqual(result["devices"][0]["index"], 0)
|
|
self.assertEqual(result["devices"][0]["visible_ordinal"], 0)
|
|
|
|
|
|
class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase):
|
|
def test_get_device_map_uses_explicit_gpu_selection(self):
|
|
with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA):
|
|
self.assertEqual(get_device_map(None), "sequential")
|
|
self.assertEqual(get_device_map([0]), "sequential")
|
|
self.assertEqual(get_device_map([0, 1]), "balanced")
|
|
|
|
def test_get_device_map_uses_all_inherited_visible_gpus_for_uuid_masks(self):
|
|
with (
|
|
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True),
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
):
|
|
self.assertEqual(get_device_map(None), "balanced")
|
|
|
|
def test_get_offloaded_device_map_entries_returns_only_cpu_and_disk(self):
|
|
model = SimpleNamespace(
|
|
hf_device_map = {
|
|
"model.embed_tokens": 0,
|
|
"model.layers.0": 1,
|
|
"model.layers.1": "cpu",
|
|
"lm_head": "disk",
|
|
}
|
|
)
|
|
|
|
self.assertEqual(
|
|
get_offloaded_device_map_entries(model),
|
|
{
|
|
"model.layers.1": "cpu",
|
|
"lm_head": "disk",
|
|
},
|
|
)
|
|
|
|
def test_get_offloaded_device_map_entries_handles_models_without_device_map(self):
|
|
self.assertEqual(get_offloaded_device_map_entries(SimpleNamespace()), {})
|
|
|
|
@patch(
|
|
"utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
|
|
new = lambda model_name, **_: model_name,
|
|
)
|
|
@patch(
|
|
"utils.hardware.hardware._load_config_for_gpu_estimate",
|
|
new = lambda *_args, **_kwargs: None,
|
|
)
|
|
def test_estimate_required_memory_formulas(self):
|
|
eight_gb = 8 * (1024**3)
|
|
|
|
with patch(
|
|
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
|
return_value = (eight_gb, "config"),
|
|
):
|
|
# FP16 inference: 8GB * 1.3 = 10.4GB
|
|
required_gb, metadata = estimate_required_model_memory_gb(
|
|
"unsloth/test",
|
|
load_in_4bit = False,
|
|
)
|
|
self.assertAlmostEqual(required_gb, 10.4, places = 3)
|
|
self.assertEqual(metadata["model_size_source"], "config")
|
|
|
|
# 4bit inference: base_4bit = 8/3.2 = 2.5GB
|
|
# required = 2.5 + max(2.5*0.3, 2.0) = 2.5 + 2.0 = 4.5GB
|
|
required_gb, _ = estimate_required_model_memory_gb(
|
|
"unsloth/test",
|
|
load_in_4bit = True,
|
|
)
|
|
self.assertAlmostEqual(required_gb, 4.5, places = 2)
|
|
|
|
# Full FT fallback: model_size * 3.5 + overhead
|
|
required_gb, metadata = estimate_required_model_memory_gb(
|
|
"unsloth/test", training_type = "Full Finetuning"
|
|
)
|
|
self.assertEqual(metadata.get("estimation_mode"), "fallback")
|
|
self.assertGreater(required_gb, 25.0)
|
|
self.assertLess(required_gb, 40.0)
|
|
|
|
# LoRA fp16 fallback: model_size + lora_overhead + activations + overhead
|
|
required_gb, metadata = estimate_required_model_memory_gb(
|
|
"unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
load_in_4bit = False,
|
|
)
|
|
self.assertEqual(metadata.get("estimation_mode"), "fallback")
|
|
self.assertGreater(required_gb, 8.0)
|
|
self.assertLess(required_gb, 15.0)
|
|
|
|
# QLoRA 4-bit fallback: compressed weights + lora overhead + activations + overhead
|
|
required_gb, metadata = estimate_required_model_memory_gb(
|
|
"unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
load_in_4bit = True,
|
|
)
|
|
self.assertEqual(metadata.get("estimation_mode"), "fallback")
|
|
self.assertGreater(required_gb, 3.0)
|
|
self.assertLess(required_gb, 8.0)
|
|
|
|
# Larger model: 16GB fp16
|
|
sixteen_gb = 16 * (1024**3)
|
|
with patch(
|
|
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
|
return_value = (sixteen_gb, "config"),
|
|
):
|
|
required_gb, _ = estimate_required_model_memory_gb(
|
|
"unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
load_in_4bit = True,
|
|
)
|
|
# QLoRA for 16GB model should be < 12 GB
|
|
self.assertGreater(required_gb, 5.0)
|
|
self.assertLess(required_gb, 12.0)
|
|
|
|
def test_estimate_fp16_model_size_bytes_uses_vllm_fallback_last(self):
|
|
config = object()
|
|
with (
|
|
patch(
|
|
"utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
|
|
return_value = "unsloth/test",
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._get_hf_safetensors_total_params",
|
|
return_value = None,
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._load_config_for_gpu_estimate",
|
|
return_value = config,
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._estimate_fp16_model_size_bytes_from_config",
|
|
return_value = None,
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._get_local_weight_size_bytes",
|
|
return_value = None,
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._estimate_fp16_model_size_bytes_from_vllm_utils",
|
|
return_value = 1234,
|
|
),
|
|
):
|
|
model_size_bytes, source = _hw_module.estimate_fp16_model_size_bytes("unsloth/test")
|
|
|
|
self.assertEqual(model_size_bytes, 1234)
|
|
self.assertEqual(source, "vllm_utils")
|
|
|
|
def test_auto_select_gpu_ids_chooses_smallest_fitting_subset(self):
|
|
fake_devices = {
|
|
"devices": [
|
|
{"index": 0, "vram_total_gb": 16.0, "vram_used_gb": 4.0},
|
|
{"index": 1, "vram_total_gb": 16.0, "vram_used_gb": 6.0},
|
|
{"index": 2, "vram_total_gb": 16.0, "vram_used_gb": 7.0},
|
|
]
|
|
}
|
|
|
|
with (
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch(
|
|
"utils.hardware.hardware.estimate_required_model_memory_gb",
|
|
return_value = (
|
|
14.0,
|
|
{"required_gb": 14.0, "model_size_source": "config"},
|
|
),
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware.get_visible_gpu_utilization",
|
|
return_value = fake_devices,
|
|
),
|
|
):
|
|
selected, metadata = auto_select_gpu_ids("unsloth/test")
|
|
|
|
self.assertEqual(selected, [0, 1])
|
|
self.assertEqual(metadata["selection_mode"], "auto")
|
|
# First GPU full (12GB) + second GPU with overhead (10*0.85=8.5) = 20.5GB
|
|
self.assertAlmostEqual(metadata["usable_gb"], 20.5, places = 3)
|
|
|
|
def test_auto_select_gpu_ids_falls_back_to_all_visible(self):
|
|
fake_devices = {
|
|
"devices": [
|
|
{"index": 0, "vram_total_gb": 12.0, "vram_used_gb": 2.0},
|
|
{"index": 1, "vram_total_gb": 12.0, "vram_used_gb": 2.0},
|
|
]
|
|
}
|
|
|
|
with (
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch(
|
|
"utils.hardware.hardware.estimate_required_model_memory_gb",
|
|
return_value = (
|
|
30.0,
|
|
{"required_gb": 30.0, "model_size_source": "config"},
|
|
),
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware.get_visible_gpu_utilization",
|
|
return_value = fake_devices,
|
|
),
|
|
):
|
|
selected, metadata = auto_select_gpu_ids("unsloth/test")
|
|
|
|
self.assertEqual(selected, [0, 1])
|
|
self.assertEqual(metadata["selection_mode"], "fallback_all")
|
|
# First GPU full (10GB) + second GPU with overhead (10*0.85=8.5) = 18.5GB
|
|
self.assertAlmostEqual(metadata["usable_gb"], 18.5, places = 3)
|
|
|
|
def test_prepare_gpu_selection_preserves_explicit_ids_without_auto_selection(self):
|
|
with (
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch(
|
|
"utils.hardware.hardware.resolve_requested_gpu_ids",
|
|
return_value = [2, 3],
|
|
),
|
|
patch("utils.hardware.hardware.auto_select_gpu_ids") as mock_auto_select,
|
|
):
|
|
selected, metadata = prepare_gpu_selection(
|
|
[2, 3],
|
|
model_name = "unsloth/test",
|
|
)
|
|
|
|
self.assertEqual(selected, [2, 3])
|
|
self.assertEqual(metadata["selection_mode"], "explicit")
|
|
mock_auto_select.assert_not_called()
|
|
|
|
def test_prepare_gpu_selection_treats_empty_list_as_auto(self):
|
|
with patch(
|
|
"utils.hardware.hardware.auto_select_gpu_ids",
|
|
return_value = ([0, 1], {"selection_mode": "auto"}),
|
|
) as mock_auto_select:
|
|
selected, metadata = prepare_gpu_selection(
|
|
[],
|
|
model_name = "unsloth/test",
|
|
)
|
|
|
|
self.assertEqual(selected, [0, 1])
|
|
self.assertEqual(metadata["selection_mode"], "auto")
|
|
mock_auto_select.assert_called_once()
|
|
|
|
def test_prepare_gpu_selection_preserves_uuid_parent_visibility_in_auto_mode(self):
|
|
with (
|
|
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True),
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch(
|
|
"utils.hardware.hardware.estimate_required_model_memory_gb",
|
|
return_value = (
|
|
14.0,
|
|
{"required_gb": 14.0, "model_size_source": "config"},
|
|
),
|
|
),
|
|
):
|
|
selected, metadata = prepare_gpu_selection(
|
|
None,
|
|
model_name = "unsloth/test",
|
|
)
|
|
|
|
self.assertIsNone(selected)
|
|
self.assertEqual(metadata["selection_mode"], "inherit_parent_visible")
|
|
self.assertIsNone(metadata["selected_gpu_ids"])
|
|
|
|
|
|
class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase):
|
|
def test_training_backend_resolves_explicit_gpu_ids_before_spawn(self):
|
|
backend = TrainingBackend()
|
|
|
|
class DummyProcess:
|
|
pid = 12345
|
|
|
|
def start(self):
|
|
return None
|
|
|
|
class DummyThread:
|
|
def start(self):
|
|
return None
|
|
|
|
dummy_queue = object()
|
|
|
|
with (
|
|
patch(
|
|
"core.training.training.prepare_gpu_selection",
|
|
return_value = ([1, 2], {"selection_mode": "explicit"}),
|
|
),
|
|
patch(
|
|
"core.training.training._CTX.Queue",
|
|
side_effect = [dummy_queue, dummy_queue],
|
|
),
|
|
patch(
|
|
"core.training.training._CTX.Process", return_value = DummyProcess()
|
|
) as mock_process,
|
|
patch("core.training.training.threading.Thread", return_value = DummyThread()),
|
|
):
|
|
backend.start_training(
|
|
job_id = "test-job-1",
|
|
model_name = "unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
gpu_ids = [1, 2],
|
|
)
|
|
|
|
config = mock_process.call_args.kwargs["kwargs"]["config"]
|
|
self.assertEqual(config["gpu_ids"], [1, 2])
|
|
self.assertEqual(config["resolved_gpu_ids"], [1, 2])
|
|
self.assertEqual(config["gpu_selection"]["selection_mode"], "explicit")
|
|
|
|
def test_training_backend_auto_selects_gpu_ids_when_omitted(self):
|
|
backend = TrainingBackend()
|
|
|
|
class DummyProcess:
|
|
pid = 12345
|
|
|
|
def start(self):
|
|
return None
|
|
|
|
class DummyThread:
|
|
def start(self):
|
|
return None
|
|
|
|
dummy_queue = object()
|
|
|
|
with (
|
|
patch(
|
|
"core.training.training.prepare_gpu_selection",
|
|
return_value = ([0, 1], {"selection_mode": "auto"}),
|
|
),
|
|
patch(
|
|
"core.training.training._CTX.Queue",
|
|
side_effect = [dummy_queue, dummy_queue],
|
|
),
|
|
patch(
|
|
"core.training.training._CTX.Process", return_value = DummyProcess()
|
|
) as mock_process,
|
|
patch("core.training.training.threading.Thread", return_value = DummyThread()),
|
|
):
|
|
backend.start_training(
|
|
job_id = "test-job-2",
|
|
model_name = "unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
gpu_ids = None,
|
|
)
|
|
|
|
config = mock_process.call_args.kwargs["kwargs"]["config"]
|
|
self.assertIsNone(config["gpu_ids"])
|
|
self.assertEqual(config["resolved_gpu_ids"], [0, 1])
|
|
self.assertEqual(config["gpu_selection"]["selection_mode"], "auto")
|
|
|
|
def test_training_backend_preserves_uuid_parent_visibility_in_auto_mode(self):
|
|
backend = TrainingBackend()
|
|
|
|
class DummyProcess:
|
|
pid = 12345
|
|
|
|
def start(self):
|
|
return None
|
|
|
|
class DummyThread:
|
|
def start(self):
|
|
return None
|
|
|
|
dummy_queue = object()
|
|
|
|
with (
|
|
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True),
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch(
|
|
"core.training.training._CTX.Queue",
|
|
side_effect = [dummy_queue, dummy_queue],
|
|
),
|
|
patch(
|
|
"core.training.training._CTX.Process", return_value = DummyProcess()
|
|
) as mock_process,
|
|
patch("core.training.training.threading.Thread", return_value = DummyThread()),
|
|
patch(
|
|
"utils.hardware.hardware.estimate_required_model_memory_gb",
|
|
return_value = (
|
|
14.0,
|
|
{"required_gb": 14.0, "model_size_source": "config"},
|
|
),
|
|
),
|
|
):
|
|
backend.start_training(
|
|
job_id = "test-job-uuid-auto",
|
|
model_name = "unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
gpu_ids = None,
|
|
)
|
|
|
|
config = mock_process.call_args.kwargs["kwargs"]["config"]
|
|
self.assertIsNone(config["resolved_gpu_ids"])
|
|
self.assertEqual(config["gpu_selection"]["selection_mode"], "inherit_parent_visible")
|
|
|
|
def test_inference_orchestrator_resolves_explicit_gpu_ids_before_spawn(self):
|
|
class DummyThread:
|
|
def __init__(self, *args, **kwargs):
|
|
pass
|
|
|
|
def start(self):
|
|
return None
|
|
|
|
with patch("core.inference.orchestrator.threading.Thread", DummyThread):
|
|
from core.inference.orchestrator import InferenceOrchestrator
|
|
orchestrator = InferenceOrchestrator()
|
|
|
|
config = SimpleNamespace(identifier = "unsloth/test", gguf_variant = None)
|
|
|
|
with (
|
|
patch(
|
|
"core.inference.orchestrator.prepare_gpu_selection",
|
|
return_value = ([1], {"selection_mode": "explicit"}),
|
|
),
|
|
patch.object(orchestrator, "_ensure_subprocess_alive", return_value = False),
|
|
patch.object(orchestrator, "_spawn_subprocess") as mock_spawn,
|
|
patch.object(
|
|
orchestrator,
|
|
"_wait_response",
|
|
return_value = {"success": True, "model_info": {}},
|
|
),
|
|
patch("utils.transformers_version.needs_transformers_5", return_value = False),
|
|
):
|
|
self.assertTrue(orchestrator.load_model(config = config, gpu_ids = [1]))
|
|
|
|
sub_config = mock_spawn.call_args.args[0]
|
|
self.assertEqual(sub_config["gpu_ids"], [1])
|
|
self.assertEqual(sub_config["resolved_gpu_ids"], [1])
|
|
self.assertEqual(sub_config["gpu_selection"]["selection_mode"], "explicit")
|
|
|
|
def test_inference_orchestrator_auto_selects_gpu_ids_when_omitted(self):
|
|
class DummyThread:
|
|
def __init__(self, *args, **kwargs):
|
|
pass
|
|
|
|
def start(self):
|
|
return None
|
|
|
|
with patch("core.inference.orchestrator.threading.Thread", DummyThread):
|
|
from core.inference.orchestrator import InferenceOrchestrator
|
|
orchestrator = InferenceOrchestrator()
|
|
|
|
config = SimpleNamespace(identifier = "unsloth/test", gguf_variant = None)
|
|
|
|
with (
|
|
patch(
|
|
"core.inference.orchestrator.prepare_gpu_selection",
|
|
return_value = ([0], {"selection_mode": "auto"}),
|
|
),
|
|
patch.object(orchestrator, "_ensure_subprocess_alive", return_value = False),
|
|
patch.object(orchestrator, "_spawn_subprocess") as mock_spawn,
|
|
patch.object(
|
|
orchestrator,
|
|
"_wait_response",
|
|
return_value = {"success": True, "model_info": {}},
|
|
),
|
|
patch("utils.transformers_version.needs_transformers_5", return_value = False),
|
|
):
|
|
self.assertTrue(orchestrator.load_model(config = config, gpu_ids = None))
|
|
|
|
sub_config = mock_spawn.call_args.args[0]
|
|
self.assertIsNone(sub_config["gpu_ids"])
|
|
self.assertEqual(sub_config["resolved_gpu_ids"], [0])
|
|
self.assertEqual(sub_config["gpu_selection"]["selection_mode"], "auto")
|
|
|
|
|
|
class TestRouteErrors(unittest.TestCase):
|
|
def test_prepare_gpu_selection_rejects_gpu_ids_on_non_accelerator_backend(self):
|
|
with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CPU):
|
|
with self.assertRaises(ValueError) as exc_info:
|
|
prepare_gpu_selection([0], model_name = "unsloth/test")
|
|
|
|
self.assertIn("only supported on CUDA and Intel XPU", str(exc_info.exception))
|
|
|
|
def test_inference_route_resolves_gguf_gpu_ids(self):
|
|
# GGUF gpu_ids are now supported: /load routes them through the same
|
|
# resolution as non-GGUF loads (rejecting only genuinely invalid ids with
|
|
# the resolver's actionable message) rather than a blanket "not supported"
|
|
# reject, so /validate can stay consistent with /load (#7239).
|
|
import utils.hardware.hardware as hardware_mod
|
|
|
|
inference_route = _load_route_module(
|
|
"inference_route_module_for_gguf_gpu_ids_test",
|
|
"routes/inference.py",
|
|
)
|
|
request = LoadRequest(model_path = "unsloth/test.gguf", gpu_ids = [0, 1])
|
|
model_config = SimpleNamespace(
|
|
is_gguf = True,
|
|
is_lora = False,
|
|
gguf_hf_repo = None,
|
|
gguf_file = "/tmp/test.gguf",
|
|
gguf_mmproj_file = None,
|
|
gguf_variant = None,
|
|
identifier = "unsloth/test.gguf",
|
|
display_name = "unsloth/test.gguf",
|
|
is_vision = False,
|
|
is_audio = False,
|
|
audio_type = None,
|
|
has_audio_input = False,
|
|
)
|
|
|
|
def _fake_resolve(ids, is_vulkan = False):
|
|
raise ValueError("SENTINEL requested GPUs are outside the parent-visible set")
|
|
|
|
with (
|
|
patch.object(
|
|
inference_route,
|
|
"ModelConfig",
|
|
SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
|
|
),
|
|
# Patch both the package re-export and the defining module so the stub
|
|
# fires no matter which import path the route uses.
|
|
patch("utils.hardware.resolve_requested_gpu_ids", _fake_resolve),
|
|
patch.object(hardware_mod, "resolve_requested_gpu_ids", _fake_resolve),
|
|
patch.object(
|
|
inference_route,
|
|
"_guard_chat_load_against_training",
|
|
return_value = None,
|
|
),
|
|
patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
|
|
patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
|
|
):
|
|
with self.assertRaises(HTTPException) as exc_info:
|
|
asyncio.run(
|
|
inference_route._load_model_impl(
|
|
request,
|
|
SimpleNamespace(
|
|
app = SimpleNamespace(
|
|
state = SimpleNamespace(llama_parallel_slots = 1),
|
|
),
|
|
),
|
|
current_subject = "test-user",
|
|
)
|
|
)
|
|
|
|
# The selection was routed through resolution (not the old blanket reject).
|
|
self.assertEqual(exc_info.exception.status_code, 400)
|
|
self.assertIn("SENTINEL", exc_info.exception.detail)
|
|
self.assertNotIn("not supported for GGUF", exc_info.exception.detail)
|
|
|
|
def test_load_rejects_unavailable_vulkan_ordinal_before_training_guard(self):
|
|
inference_route = _load_route_module(
|
|
"inference_route_module_for_vulkan_preflight_test",
|
|
"routes/inference.py",
|
|
)
|
|
request = LoadRequest(model_path = "unsloth/test.gguf", gpu_ids = [99])
|
|
model_config = SimpleNamespace(
|
|
is_gguf = True,
|
|
is_lora = False,
|
|
gguf_hf_repo = None,
|
|
gguf_file = "/tmp/test.gguf",
|
|
gguf_mmproj_file = None,
|
|
gguf_variant = None,
|
|
identifier = "unsloth/test.gguf",
|
|
display_name = "unsloth/test.gguf",
|
|
is_vision = False,
|
|
is_audio = False,
|
|
audio_type = None,
|
|
has_audio_input = False,
|
|
)
|
|
|
|
with (
|
|
patch.object(
|
|
inference_route,
|
|
"ModelConfig",
|
|
SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
|
|
),
|
|
patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch.object(inference_route, "_classify_diffusion_gguf", return_value = None),
|
|
patch.object(
|
|
inference_route.LlamaCppBackend,
|
|
"_is_vulkan_backend",
|
|
return_value = True,
|
|
),
|
|
patch.object(
|
|
inference_route.LlamaCppBackend,
|
|
"_find_llama_server_binary",
|
|
return_value = "/tmp/llama-server",
|
|
),
|
|
patch.object(
|
|
inference_route.LlamaCppBackend,
|
|
"_get_gpu_memory",
|
|
return_value = [(0, 8 * 1024**3, 16 * 1024**3)],
|
|
),
|
|
patch.object(
|
|
inference_route,
|
|
"_guard_chat_load_against_training",
|
|
return_value = None,
|
|
) as training_guard,
|
|
patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
|
|
patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
|
|
):
|
|
with self.assertRaises(HTTPException) as exc_info:
|
|
asyncio.run(
|
|
inference_route._load_model_impl(
|
|
request,
|
|
SimpleNamespace(
|
|
app = SimpleNamespace(
|
|
state = SimpleNamespace(llama_parallel_slots = 1),
|
|
),
|
|
),
|
|
current_subject = "test-user",
|
|
)
|
|
)
|
|
|
|
self.assertEqual(exc_info.exception.status_code, 400)
|
|
self.assertIn("Vulkan GPU ordinal(s) [99]", exc_info.exception.detail)
|
|
training_guard.assert_not_called()
|
|
|
|
def test_vulkan_ordinals_are_allowed_on_xpu_hosts(self):
|
|
import utils.hardware.hardware as hardware_mod
|
|
|
|
inference_route = _load_route_module(
|
|
"inference_route_module_for_xpu_vulkan_test",
|
|
"routes/inference.py",
|
|
)
|
|
config = SimpleNamespace(is_gguf = True)
|
|
|
|
with (
|
|
patch("utils.hardware.get_device", return_value = DeviceType.XPU),
|
|
patch.object(
|
|
inference_route.LlamaCppBackend,
|
|
"_is_vulkan_backend",
|
|
return_value = True,
|
|
),
|
|
patch.object(inference_route, "_classify_diffusion_gguf", return_value = False),
|
|
patch.object(hardware_mod, "resolve_requested_gpu_ids", return_value = [0, 1]),
|
|
patch.object(
|
|
inference_route.LlamaCppBackend,
|
|
"_find_llama_server_binary",
|
|
return_value = None,
|
|
),
|
|
):
|
|
resolved = asyncio.run(
|
|
inference_route._resolve_gguf_gpu_ids_for_request(config, [1, 0])
|
|
)
|
|
|
|
self.assertEqual(resolved, [0, 1])
|
|
|
|
def test_inference_route_validates_gpu_ids_for_gguf(self):
|
|
# gpu_ids is now SUPPORTED for GGUF (the GPU picker), but still
|
|
# validated: a rejected pick surfaces as a clean 400, not the old
|
|
# "not supported for GGUF" rejection. Patch the validator so the test
|
|
# is deterministic regardless of the host's (or a prior test's) GPU env.
|
|
import utils.hardware.hardware as hardware_mod
|
|
|
|
inference_route = _load_route_module(
|
|
"inference_route_module_for_gguf_gpu_ids_test2",
|
|
"routes/inference.py",
|
|
)
|
|
request = LoadRequest(model_path = "unsloth/test.gguf", gpu_ids = [0, 1])
|
|
model_config = SimpleNamespace(
|
|
is_gguf = True,
|
|
is_lora = False,
|
|
gguf_hf_repo = None,
|
|
gguf_file = "/tmp/test.gguf",
|
|
gguf_mmproj_file = None,
|
|
gguf_variant = None,
|
|
identifier = "unsloth/test.gguf",
|
|
display_name = "unsloth/test.gguf",
|
|
is_vision = False,
|
|
is_audio = False,
|
|
audio_type = None,
|
|
has_audio_input = False,
|
|
)
|
|
|
|
with (
|
|
patch.object(
|
|
inference_route,
|
|
"ModelConfig",
|
|
SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
|
|
),
|
|
# Patch both the package re-export and the defining module so the stub
|
|
# fires no matter which import path the route uses.
|
|
patch(
|
|
"utils.hardware.resolve_requested_gpu_ids",
|
|
side_effect = ValueError("Invalid gpu_ids [0, 1]: rejected by test"),
|
|
),
|
|
patch.object(
|
|
hardware_mod,
|
|
"resolve_requested_gpu_ids",
|
|
side_effect = ValueError("Invalid gpu_ids [0, 1]: rejected by test"),
|
|
),
|
|
patch.object(
|
|
inference_route,
|
|
"_guard_chat_load_against_training",
|
|
return_value = None,
|
|
),
|
|
patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
|
|
patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
|
|
):
|
|
with self.assertRaises(HTTPException) as exc_info:
|
|
asyncio.run(
|
|
inference_route._load_model_impl(
|
|
request,
|
|
SimpleNamespace(
|
|
app = SimpleNamespace(
|
|
state = SimpleNamespace(llama_parallel_slots = 1),
|
|
),
|
|
),
|
|
current_subject = "test-user",
|
|
)
|
|
)
|
|
|
|
# The validator's ValueError becomes a clean 400 (not the removed
|
|
# "not supported for GGUF" rejection).
|
|
self.assertEqual(exc_info.exception.status_code, 400)
|
|
self.assertIn("gpu_ids", exc_info.exception.detail.lower())
|
|
self.assertNotIn("not supported", exc_info.exception.detail.lower())
|
|
|
|
def test_inference_route_defers_gpu_handoff_until_after_validation(self):
|
|
# A doomed chat load (GGUF + gpu_ids -> 400) must NOT reclaim the CHAT arbiter owner
|
|
# first: the handoff is deferred past validation, so a resident Images/Video pipeline is
|
|
# never evicted for a load that then errors.
|
|
import core.inference.gpu_arbiter as arb
|
|
|
|
inference_route = _load_route_module(
|
|
"inference_route_module_for_handoff_test",
|
|
"routes/inference.py",
|
|
)
|
|
request = LoadRequest(model_path = "unsloth/test.gguf", gpu_ids = [0, 1])
|
|
model_config = SimpleNamespace(
|
|
is_gguf = True,
|
|
is_lora = False,
|
|
gguf_hf_repo = None,
|
|
gguf_file = "/tmp/test.gguf",
|
|
gguf_mmproj_file = None,
|
|
gguf_mtp_file = None,
|
|
gguf_variant = None,
|
|
identifier = "unsloth/test.gguf",
|
|
display_name = "unsloth/test.gguf",
|
|
is_vision = False,
|
|
is_audio = False,
|
|
audio_type = None,
|
|
has_audio_input = False,
|
|
)
|
|
acquired = []
|
|
# Make [0, 1] invalid on any host (a duplicate id is rejected everywhere, GPUs or not):
|
|
# the point of the test is the ORDER -- validation before the handoff -- not this
|
|
# machine's device count.
|
|
request.gpu_ids = [0, 0]
|
|
with (
|
|
patch.object(
|
|
inference_route,
|
|
"ModelConfig",
|
|
SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
|
|
),
|
|
patch.object(inference_route, "_guard_chat_load_against_training", return_value = None),
|
|
patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
|
|
patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
|
|
# The chat handoff passes a `register` hook (the in-flight marker), so accept it.
|
|
patch.object(
|
|
arb, "acquire_for", lambda owner, register = None: acquired.append(owner)
|
|
),
|
|
):
|
|
with self.assertRaises(HTTPException) as exc_info:
|
|
asyncio.run(
|
|
inference_route._load_model_impl(
|
|
request,
|
|
SimpleNamespace(
|
|
app = SimpleNamespace(
|
|
state = SimpleNamespace(llama_parallel_slots = 1),
|
|
),
|
|
),
|
|
current_subject = "test-user",
|
|
)
|
|
)
|
|
self.assertEqual(exc_info.exception.status_code, 400)
|
|
self.assertEqual(acquired, []) # no CHAT handoff before the doomed load errored
|
|
|
|
def test_inference_route_checks_hub_download_conflict_before_the_handoff(self):
|
|
# A GGUF the download manager is fetching 409s and loads nothing, so that check has to
|
|
# run BEFORE the CHAT handoff: afterwards, it destroyed the resident Images/Video
|
|
# pipeline for a load that could never start.
|
|
import core.inference.gpu_arbiter as arb
|
|
import core.inference.llama_cpp as llama_cpp
|
|
|
|
inference_route = _load_route_module(
|
|
"inference_route_module_for_hub_conflict_test",
|
|
"routes/inference.py",
|
|
)
|
|
request = LoadRequest(model_path = "unsloth/Qwen3-4B-GGUF", gguf_variant = "Q4_K_M")
|
|
model_config = SimpleNamespace(
|
|
is_gguf = True,
|
|
is_lora = False,
|
|
gguf_hf_repo = "unsloth/Qwen3-4B-GGUF",
|
|
gguf_file = None,
|
|
gguf_mmproj_file = None,
|
|
gguf_variant = "Q4_K_M",
|
|
identifier = "unsloth/Qwen3-4B-GGUF",
|
|
display_name = "Qwen3-4B",
|
|
is_vision = False,
|
|
is_audio = False,
|
|
audio_type = None,
|
|
has_audio_input = False,
|
|
)
|
|
acquired = []
|
|
with (
|
|
patch.object(
|
|
inference_route,
|
|
"ModelConfig",
|
|
SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
|
|
),
|
|
patch.object(inference_route, "_guard_chat_load_against_training", return_value = None),
|
|
patch.object(inference_route, "_resolve_inherited_extra_args", lambda *a, **k: None),
|
|
patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
|
|
patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
|
|
patch.object(llama_cpp, "_hub_download_blocks_gguf_load", lambda *a, **k: True),
|
|
patch.object(arb, "acquire_for", lambda *a, **k: acquired.append(a[0])),
|
|
):
|
|
with self.assertRaises(HTTPException) as exc_info:
|
|
asyncio.run(
|
|
inference_route._load_model_impl(
|
|
request,
|
|
SimpleNamespace(
|
|
app = SimpleNamespace(
|
|
state = SimpleNamespace(llama_parallel_slots = 1),
|
|
),
|
|
),
|
|
current_subject = "test-user",
|
|
)
|
|
)
|
|
self.assertEqual(exc_info.exception.status_code, 409)
|
|
self.assertIn("download", exc_info.exception.detail.lower())
|
|
self.assertEqual(acquired, []) # nothing evicted for a load that cannot start
|
|
|
|
def test_inference_route_marks_the_chat_load_under_the_arbiter_lock(self):
|
|
# A chat load holds no llama-server process until its GGUF downloaded, so the arbiter is
|
|
# told about it through acquire_for's `register` hook (which runs under the arbiter lock).
|
|
# Passing no register left a competing Images/Video acquire with nothing to cancel.
|
|
import core.inference.gpu_arbiter as arb
|
|
import core.inference.llama_cpp as llama_cpp
|
|
|
|
inference_route = _load_route_module(
|
|
"inference_route_module_for_chat_marker_test",
|
|
"routes/inference.py",
|
|
)
|
|
request = LoadRequest(model_path = "unsloth/Qwen3-4B-GGUF", gguf_variant = "Q4_K_M")
|
|
model_config = SimpleNamespace(
|
|
is_gguf = True,
|
|
is_lora = False,
|
|
gguf_hf_repo = "unsloth/Qwen3-4B-GGUF",
|
|
gguf_file = None,
|
|
gguf_mmproj_file = None,
|
|
gguf_variant = "Q4_K_M",
|
|
identifier = "unsloth/Qwen3-4B-GGUF",
|
|
display_name = "Qwen3-4B",
|
|
is_vision = False,
|
|
is_audio = False,
|
|
audio_type = None,
|
|
has_audio_input = False,
|
|
)
|
|
marked = []
|
|
|
|
def _acquire(owner, register = None):
|
|
# Under the arbiter lock the evictor must already be able to see this load.
|
|
if register is not None:
|
|
register()
|
|
marked.append(llama_cpp.chat_load_active())
|
|
raise RuntimeError("stop the load here")
|
|
|
|
with (
|
|
patch.object(
|
|
inference_route,
|
|
"ModelConfig",
|
|
SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
|
|
),
|
|
patch.object(inference_route, "_guard_chat_load_against_training", return_value = None),
|
|
patch.object(inference_route, "_resolve_inherited_extra_args", lambda *a, **k: None),
|
|
patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
|
|
patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
|
|
patch.object(llama_cpp, "_hub_download_blocks_gguf_load", lambda *a, **k: False),
|
|
patch.object(arb, "acquire_for", _acquire),
|
|
):
|
|
with self.assertRaises(HTTPException):
|
|
asyncio.run(
|
|
inference_route._load_model_impl(
|
|
request,
|
|
SimpleNamespace(
|
|
app = SimpleNamespace(
|
|
state = SimpleNamespace(llama_parallel_slots = 1),
|
|
),
|
|
),
|
|
current_subject = "test-user",
|
|
)
|
|
)
|
|
self.assertEqual(marked, [True])
|
|
# The marker is scoped to the request: it must not outlive the failed load.
|
|
self.assertFalse(llama_cpp.chat_load_active())
|
|
|
|
def test_training_route_returns_400_for_invalid_gpu_ids(self):
|
|
training_route = _load_route_module(
|
|
"training_route_module_for_test",
|
|
"routes/training.py",
|
|
)
|
|
request = TrainingStartRequest(
|
|
model_name = "unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
format_type = "alpaca",
|
|
gpu_ids = [99],
|
|
)
|
|
|
|
class DummyBackend:
|
|
current_job_id = None
|
|
|
|
def is_training_active(self):
|
|
return False
|
|
|
|
def start_training(self, **kwargs):
|
|
raise ValueError("Invalid gpu_ids [99]")
|
|
|
|
with (
|
|
patch.object(training_route, "get_training_backend", return_value = DummyBackend()),
|
|
patch(
|
|
"routes.training_vram.summarize_resident_chat",
|
|
return_value = {"any": False, "hf": None, "gguf": None},
|
|
),
|
|
patch(
|
|
"core.export.get_export_backend",
|
|
return_value = SimpleNamespace(current_checkpoint = None),
|
|
),
|
|
):
|
|
with self.assertRaises(HTTPException) as exc_info:
|
|
asyncio.run(training_route.start_training(request, current_subject = "test-user"))
|
|
|
|
self.assertEqual(exc_info.exception.status_code, 400)
|
|
self.assertIn("gpu_ids [99]", exc_info.exception.detail)
|
|
|
|
def test_training_route_returns_400_for_uuid_parent_visibility_gpu_ids(self):
|
|
training_route = _load_route_module(
|
|
"training_route_module_for_uuid_parent_visibility_test",
|
|
"routes/training.py",
|
|
)
|
|
request = TrainingStartRequest(
|
|
model_name = "unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
format_type = "alpaca",
|
|
gpu_ids = [1],
|
|
)
|
|
|
|
class DummyBackend:
|
|
current_job_id = None
|
|
|
|
def is_training_active(self):
|
|
return False
|
|
|
|
def start_training(self, **kwargs):
|
|
raise ValueError(
|
|
"Invalid gpu_ids [1]: explicit physical GPU IDs are unsupported when CUDA_VISIBLE_DEVICES uses UUID/MIG entries"
|
|
)
|
|
|
|
with (
|
|
patch.object(training_route, "get_training_backend", return_value = DummyBackend()),
|
|
patch(
|
|
"routes.training_vram.summarize_resident_chat",
|
|
return_value = {"any": False, "hf": None, "gguf": None},
|
|
),
|
|
patch(
|
|
"core.export.get_export_backend",
|
|
return_value = SimpleNamespace(current_checkpoint = None),
|
|
),
|
|
):
|
|
with self.assertRaises(HTTPException) as exc_info:
|
|
asyncio.run(training_route.start_training(request, current_subject = "test-user"))
|
|
|
|
self.assertEqual(exc_info.exception.status_code, 400)
|
|
self.assertIn("UUID/MIG", exc_info.exception.detail)
|
|
|
|
def test_inference_route_returns_400_for_invalid_gpu_ids(self):
|
|
inference_route = _load_route_module(
|
|
"inference_route_module_for_test",
|
|
"routes/inference.py",
|
|
)
|
|
request = LoadRequest(model_path = "unsloth/test", gpu_ids = [99])
|
|
model_config = SimpleNamespace(
|
|
is_gguf = False,
|
|
is_lora = False,
|
|
path = None,
|
|
identifier = "unsloth/test",
|
|
display_name = "unsloth/test",
|
|
is_vision = False,
|
|
is_audio = False,
|
|
audio_type = None,
|
|
has_audio_input = False,
|
|
)
|
|
|
|
class DummyInferenceBackend:
|
|
active_model_name = None
|
|
models = {}
|
|
|
|
def load_model(self, **kwargs):
|
|
raise ValueError("Invalid gpu_ids [99]")
|
|
|
|
with (
|
|
patch.object(
|
|
inference_route,
|
|
"ModelConfig",
|
|
SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
|
|
),
|
|
patch.object(
|
|
inference_route,
|
|
"get_inference_backend",
|
|
return_value = DummyInferenceBackend(),
|
|
),
|
|
patch.object(
|
|
inference_route,
|
|
"get_llama_cpp_backend",
|
|
return_value = SimpleNamespace(is_loaded = False),
|
|
),
|
|
patch.object(
|
|
inference_route,
|
|
"_guard_chat_load_against_training",
|
|
return_value = None,
|
|
),
|
|
patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
|
|
patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
|
|
patch(
|
|
"core.export.get_export_backend",
|
|
return_value = SimpleNamespace(current_checkpoint = None),
|
|
),
|
|
):
|
|
with self.assertRaises(HTTPException) as exc_info:
|
|
asyncio.run(
|
|
inference_route._load_model_impl(
|
|
request,
|
|
SimpleNamespace(
|
|
app = SimpleNamespace(
|
|
state = SimpleNamespace(llama_parallel_slots = 1),
|
|
),
|
|
),
|
|
current_subject = "test-user",
|
|
)
|
|
)
|
|
|
|
self.assertEqual(exc_info.exception.status_code, 400)
|
|
self.assertIn("gpu_ids [99]", exc_info.exception.detail)
|
|
|
|
def test_inference_route_returns_400_for_uuid_parent_visibility_gpu_ids(self):
|
|
inference_route = _load_route_module(
|
|
"inference_route_module_for_uuid_parent_visibility_test",
|
|
"routes/inference.py",
|
|
)
|
|
request = LoadRequest(model_path = "unsloth/test", gpu_ids = [1])
|
|
model_config = SimpleNamespace(
|
|
is_gguf = False,
|
|
is_lora = False,
|
|
path = None,
|
|
identifier = "unsloth/test",
|
|
display_name = "unsloth/test",
|
|
is_vision = False,
|
|
is_audio = False,
|
|
audio_type = None,
|
|
has_audio_input = False,
|
|
)
|
|
|
|
class DummyInferenceBackend:
|
|
active_model_name = None
|
|
models = {}
|
|
|
|
def load_model(self, **kwargs):
|
|
raise ValueError(
|
|
"Invalid gpu_ids [1]: explicit physical GPU IDs are unsupported when CUDA_VISIBLE_DEVICES uses UUID/MIG entries"
|
|
)
|
|
|
|
with (
|
|
patch.object(
|
|
inference_route,
|
|
"ModelConfig",
|
|
SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
|
|
),
|
|
patch.object(
|
|
inference_route,
|
|
"get_inference_backend",
|
|
return_value = DummyInferenceBackend(),
|
|
),
|
|
patch.object(
|
|
inference_route,
|
|
"get_llama_cpp_backend",
|
|
return_value = SimpleNamespace(is_loaded = False),
|
|
),
|
|
patch.object(
|
|
inference_route,
|
|
"_guard_chat_load_against_training",
|
|
return_value = None,
|
|
),
|
|
patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
|
|
patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
|
|
patch(
|
|
"core.export.get_export_backend",
|
|
return_value = SimpleNamespace(current_checkpoint = None),
|
|
),
|
|
):
|
|
with self.assertRaises(HTTPException) as exc_info:
|
|
asyncio.run(
|
|
inference_route._load_model_impl(
|
|
request,
|
|
SimpleNamespace(
|
|
app = SimpleNamespace(
|
|
state = SimpleNamespace(llama_parallel_slots = 1),
|
|
),
|
|
),
|
|
current_subject = "test-user",
|
|
)
|
|
)
|
|
|
|
self.assertEqual(exc_info.exception.status_code, 400)
|
|
self.assertIn("UUID/MIG", exc_info.exception.detail)
|
|
|
|
|
|
class TestRaiseIfOffloaded(unittest.TestCase):
|
|
def test_no_offload_is_noop(self):
|
|
from utils.hardware import raise_if_offloaded
|
|
model = SimpleNamespace(hf_device_map = {"model.embed_tokens": 0, "lm_head": 1})
|
|
raise_if_offloaded(model, "balanced", "Test")
|
|
|
|
def test_cpu_offload_raises(self):
|
|
from utils.hardware import raise_if_offloaded
|
|
model = SimpleNamespace(hf_device_map = {"model.layers.0": 0, "model.layers.1": "cpu"})
|
|
with self.assertRaisesRegex(ValueError, "offloaded"):
|
|
raise_if_offloaded(model, "balanced", "Test")
|
|
|
|
def test_no_device_map_attr_is_noop(self):
|
|
from utils.hardware import raise_if_offloaded
|
|
raise_if_offloaded(SimpleNamespace(), "sequential", "Test")
|
|
|
|
|
|
class TestMinGpuVram(unittest.TestCase):
|
|
def test_min_gpu_vram_decreases_with_more_gpus(self):
|
|
from utils.hardware.vram_estimation import (
|
|
ModelArchConfig,
|
|
TrainingVramConfig,
|
|
estimate_training_vram,
|
|
)
|
|
|
|
arch = ModelArchConfig(
|
|
hidden_size = 4096,
|
|
num_hidden_layers = 32,
|
|
num_attention_heads = 32,
|
|
num_key_value_heads = 8,
|
|
intermediate_size = 14336,
|
|
vocab_size = 128256,
|
|
tie_word_embeddings = False,
|
|
)
|
|
config = TrainingVramConfig(
|
|
training_method = "qlora",
|
|
load_in_4bit = True,
|
|
)
|
|
breakdown = estimate_training_vram(arch, config)
|
|
v1 = breakdown.min_gpu_vram(1)
|
|
v2 = breakdown.min_gpu_vram(2)
|
|
v4 = breakdown.min_gpu_vram(4)
|
|
self.assertGreater(v1, v2)
|
|
self.assertGreater(v2, v4)
|
|
self.assertGreater(v4, 0)
|
|
|
|
def test_total_equals_min_gpu_vram_1(self):
|
|
from utils.hardware.vram_estimation import (
|
|
ModelArchConfig,
|
|
TrainingVramConfig,
|
|
estimate_training_vram,
|
|
)
|
|
|
|
arch = ModelArchConfig(
|
|
hidden_size = 4096,
|
|
num_hidden_layers = 32,
|
|
num_attention_heads = 32,
|
|
num_key_value_heads = 8,
|
|
intermediate_size = 14336,
|
|
vocab_size = 128256,
|
|
tie_word_embeddings = False,
|
|
)
|
|
config = TrainingVramConfig(
|
|
training_method = "qlora",
|
|
load_in_4bit = True,
|
|
)
|
|
breakdown = estimate_training_vram(arch, config)
|
|
self.assertEqual(breakdown.total, breakdown.min_gpu_vram(1))
|
|
|
|
|
|
class TestPerGpuFitGuardAllCounts(unittest.TestCase):
|
|
def test_training_estimate_resolves_attention_without_raising(self):
|
|
with (
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch(
|
|
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
|
return_value = (8 * (1024**3), "config"),
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
|
|
return_value = "unsloth/test",
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._load_config_for_gpu_estimate",
|
|
return_value = SimpleNamespace(
|
|
hidden_size = 4096,
|
|
num_hidden_layers = 32,
|
|
num_attention_heads = 32,
|
|
num_key_value_heads = 8,
|
|
intermediate_size = 14336,
|
|
vocab_size = 128256,
|
|
tie_word_embeddings = False,
|
|
),
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._determine_attention_impl_for_gpu_estimate",
|
|
return_value = "eager",
|
|
),
|
|
patch("utils.hardware.hardware.get_visible_gpu_count", return_value = 1),
|
|
):
|
|
_, metadata = estimate_required_model_memory_gb(
|
|
"unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
load_in_4bit = True,
|
|
)
|
|
|
|
self.assertEqual(metadata.get("estimation_mode"), "detailed")
|
|
self.assertEqual(metadata.get("attention_implementation"), "eager")
|
|
|
|
def test_training_estimate_falls_back_when_attention_resolution_fails(self):
|
|
with (
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch(
|
|
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
|
return_value = (8 * (1024**3), "config"),
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
|
|
return_value = "unsloth/test",
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._load_config_for_gpu_estimate",
|
|
return_value = SimpleNamespace(
|
|
hidden_size = 4096,
|
|
num_hidden_layers = 32,
|
|
num_attention_heads = 32,
|
|
num_key_value_heads = 8,
|
|
intermediate_size = 14336,
|
|
vocab_size = 128256,
|
|
tie_word_embeddings = False,
|
|
),
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._determine_attention_impl_for_gpu_estimate",
|
|
side_effect = RuntimeError("attention unavailable"),
|
|
),
|
|
patch("utils.hardware.hardware.get_visible_gpu_count", return_value = 1),
|
|
):
|
|
_, metadata = estimate_required_model_memory_gb(
|
|
"unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
load_in_4bit = True,
|
|
)
|
|
|
|
self.assertEqual(metadata.get("estimation_mode"), "detailed")
|
|
self.assertEqual(
|
|
metadata.get("attention_implementation"),
|
|
"eager",
|
|
)
|
|
|
|
def test_attention_resolver_does_not_mutate_loaded_config(self):
|
|
from utils.hardware import hardware as hardware_module
|
|
|
|
config = SimpleNamespace(
|
|
hidden_size = 1024,
|
|
num_hidden_layers = 2,
|
|
num_attention_heads = 8,
|
|
num_key_value_heads = 8,
|
|
intermediate_size = 2048,
|
|
vocab_size = 1024,
|
|
tie_word_embeddings = True,
|
|
)
|
|
|
|
def _stub_resolver(model_class, cfg):
|
|
cfg._attn_implementation = "eager"
|
|
return "eager"
|
|
|
|
with patch.dict(sys.modules, _fake_unsloth_attention_modules(_stub_resolver)):
|
|
hardware_module._determine_attention_impl_for_gpu_estimate(config)
|
|
|
|
self.assertFalse(hasattr(config, "_attn_implementation"))
|
|
|
|
def test_attention_resolver_handles_missing_model_mapping(self):
|
|
from utils.hardware import hardware as hardware_module
|
|
|
|
config = SimpleNamespace(
|
|
hidden_size = 1024,
|
|
num_hidden_layers = 2,
|
|
num_attention_heads = 8,
|
|
num_key_value_heads = 8,
|
|
intermediate_size = 2048,
|
|
vocab_size = 1024,
|
|
tie_word_embeddings = True,
|
|
)
|
|
captured = {}
|
|
|
|
def _stub_resolver(model_class, cfg):
|
|
captured["model_class"] = model_class
|
|
return "eager"
|
|
|
|
from transformers import AutoModel, AutoModelForCausalLM
|
|
|
|
with (
|
|
patch.object(AutoModelForCausalLM, "_model_mapping", new = None),
|
|
patch.object(AutoModel, "_model_mapping", new = None),
|
|
patch.dict(sys.modules, _fake_unsloth_attention_modules(_stub_resolver)),
|
|
):
|
|
result = hardware_module._determine_attention_impl_for_gpu_estimate(config)
|
|
|
|
self.assertEqual(result, "eager")
|
|
self.assertIsNone(captured["model_class"])
|
|
|
|
def test_attention_resolver_does_not_mutate_nested_text_config(self):
|
|
from utils.hardware import hardware as hardware_module
|
|
|
|
text_config = SimpleNamespace(
|
|
hidden_size = 1024,
|
|
num_hidden_layers = 2,
|
|
num_attention_heads = 8,
|
|
num_key_value_heads = 8,
|
|
intermediate_size = 2048,
|
|
vocab_size = 1024,
|
|
tie_word_embeddings = True,
|
|
)
|
|
config = SimpleNamespace(
|
|
hidden_size = 1024,
|
|
num_hidden_layers = 2,
|
|
num_attention_heads = 8,
|
|
num_key_value_heads = 8,
|
|
intermediate_size = 2048,
|
|
vocab_size = 1024,
|
|
tie_word_embeddings = True,
|
|
text_config = text_config,
|
|
)
|
|
|
|
def _stub_resolver(model_class, cfg):
|
|
cfg._attn_implementation = "eager"
|
|
inner = getattr(cfg, "text_config", None)
|
|
if inner is not None:
|
|
inner._attn_implementation = "eager"
|
|
return "eager"
|
|
|
|
with patch.dict(sys.modules, _fake_unsloth_attention_modules(_stub_resolver)):
|
|
hardware_module._determine_attention_impl_for_gpu_estimate(config)
|
|
|
|
self.assertFalse(hasattr(config, "_attn_implementation"))
|
|
self.assertFalse(hasattr(text_config, "_attn_implementation"))
|
|
|
|
def test_min_per_gpu_generated_for_all_visible_counts(self):
|
|
with (
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch(
|
|
"utils.hardware.hardware.estimate_fp16_model_size_bytes",
|
|
return_value = (8 * (1024**3), "config"),
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
|
|
return_value = "unsloth/test",
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._load_config_for_gpu_estimate",
|
|
return_value = SimpleNamespace(
|
|
hidden_size = 4096,
|
|
num_hidden_layers = 32,
|
|
num_attention_heads = 32,
|
|
num_key_value_heads = 8,
|
|
intermediate_size = 14336,
|
|
vocab_size = 128256,
|
|
tie_word_embeddings = False,
|
|
),
|
|
),
|
|
patch("utils.hardware.hardware.get_visible_gpu_count", return_value = 6),
|
|
):
|
|
_, metadata = estimate_required_model_memory_gb(
|
|
"unsloth/test",
|
|
training_type = "LoRA/QLoRA",
|
|
load_in_4bit = True,
|
|
)
|
|
|
|
self.assertEqual(metadata.get("estimation_mode"), "detailed")
|
|
breakdown = metadata["vram_breakdown"]
|
|
for n in range(1, 7):
|
|
self.assertIn(f"min_per_gpu_{n}", breakdown)
|
|
|
|
|
|
class TestAutoSelectWithNoneRequired(_GpuCacheResetMixin, unittest.TestCase):
|
|
def test_auto_select_falls_back_when_estimate_unavailable(self):
|
|
with (
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch(
|
|
"utils.hardware.hardware.estimate_required_model_memory_gb",
|
|
return_value = (None, {"model_size_source": "unavailable"}),
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._get_parent_visible_gpu_spec",
|
|
return_value = {
|
|
"raw": "0,1",
|
|
"numeric_ids": [0, 1],
|
|
"supports_explicit_gpu_ids": True,
|
|
},
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware.get_parent_visible_gpu_ids",
|
|
return_value = [0, 1],
|
|
),
|
|
):
|
|
selected, metadata = auto_select_gpu_ids("unsloth/test")
|
|
|
|
self.assertEqual(selected, [0, 1])
|
|
self.assertEqual(metadata["selection_mode"], "fallback_all")
|
|
|
|
|
|
class TestXpuSelection(_GpuCacheResetMixin, unittest.TestCase):
|
|
def test_auto_select_supports_xpu(self):
|
|
with (
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU),
|
|
patch(
|
|
"utils.hardware.hardware.estimate_required_model_memory_gb",
|
|
return_value = (1.0, {}),
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware.get_visible_gpu_utilization",
|
|
return_value = {
|
|
"devices": [
|
|
{"index": 0, "vram_total_gb": 8, "vram_used_gb": 1},
|
|
]
|
|
},
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._get_parent_visible_gpu_spec",
|
|
return_value = {
|
|
"raw": None,
|
|
"numeric_ids": [0],
|
|
"supports_explicit_gpu_ids": True,
|
|
},
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware.get_parent_visible_gpu_ids",
|
|
return_value = [0],
|
|
),
|
|
):
|
|
selected, metadata = auto_select_gpu_ids("unsloth/test")
|
|
|
|
self.assertEqual(selected, [0])
|
|
self.assertEqual(metadata["selection_mode"], "auto")
|
|
|
|
def test_prepare_gpu_selection_accepts_explicit_ids_on_xpu(self):
|
|
with (
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU),
|
|
patch(
|
|
"utils.hardware.hardware._get_parent_visible_gpu_spec",
|
|
return_value = {
|
|
"raw": "0",
|
|
"numeric_ids": [0],
|
|
"supports_explicit_gpu_ids": True,
|
|
},
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware.get_parent_visible_gpu_ids",
|
|
return_value = [0],
|
|
),
|
|
patch("utils.hardware.hardware.get_physical_gpu_count", return_value = 1),
|
|
):
|
|
selected, metadata = prepare_gpu_selection([0], model_name = "unsloth/test")
|
|
|
|
self.assertEqual(selected, [0])
|
|
self.assertEqual(metadata["selection_mode"], "explicit")
|
|
|
|
|
|
class TestEstimateFp16ModelSizeBytesPrefersLocalWeights(unittest.TestCase):
|
|
def _run(
|
|
self,
|
|
model_path,
|
|
*,
|
|
config_bytes,
|
|
local_bytes,
|
|
safetensors_params = None,
|
|
config = object(),
|
|
):
|
|
from utils.hardware import hardware as hardware_module
|
|
with (
|
|
patch.object(
|
|
hardware_module,
|
|
"_resolve_model_identifier_for_gpu_estimate",
|
|
return_value = model_path,
|
|
),
|
|
patch.object(
|
|
hardware_module,
|
|
"_get_hf_safetensors_total_params",
|
|
return_value = safetensors_params,
|
|
),
|
|
patch.object(
|
|
hardware_module,
|
|
"_load_config_for_gpu_estimate",
|
|
return_value = config,
|
|
),
|
|
patch.object(
|
|
hardware_module,
|
|
"_estimate_fp16_model_size_bytes_from_config",
|
|
return_value = config_bytes,
|
|
),
|
|
patch.object(
|
|
hardware_module,
|
|
"_get_local_weight_size_bytes",
|
|
return_value = local_bytes,
|
|
),
|
|
):
|
|
return hardware_module.estimate_fp16_model_size_bytes(model_path)
|
|
|
|
def test_local_weight_bytes_preferred_when_larger_than_config(self):
|
|
bytes_, src = self._run(
|
|
"/local/vlm",
|
|
config_bytes = 2 * (1 << 30),
|
|
local_bytes = 20 * (1 << 30),
|
|
)
|
|
self.assertEqual(bytes_, 20 * (1 << 30))
|
|
self.assertEqual(src, "weight_bytes")
|
|
|
|
def test_config_bytes_preferred_when_larger_than_local(self):
|
|
bytes_, src = self._run(
|
|
"/local/text-only",
|
|
config_bytes = 20 * (1 << 30),
|
|
local_bytes = 2 * (1 << 30),
|
|
)
|
|
self.assertEqual(bytes_, 20 * (1 << 30))
|
|
self.assertEqual(src, "config")
|
|
|
|
def test_config_bytes_returned_when_no_local_weights(self):
|
|
bytes_, src = self._run(
|
|
"/local/no-weights",
|
|
config_bytes = 5 * (1 << 30),
|
|
local_bytes = None,
|
|
)
|
|
self.assertEqual(bytes_, 5 * (1 << 30))
|
|
self.assertEqual(src, "config")
|
|
|
|
def test_local_bytes_returned_when_config_resolution_fails(self):
|
|
bytes_, src = self._run(
|
|
"/local/no-config",
|
|
config_bytes = None,
|
|
local_bytes = 7 * (1 << 30),
|
|
config = None,
|
|
)
|
|
self.assertEqual(bytes_, 7 * (1 << 30))
|
|
self.assertEqual(src, "weight_bytes")
|
|
|
|
def test_equal_local_and_config_keeps_config_label(self):
|
|
# Tie-breaker is "local must be strictly larger", so an exact
|
|
# match keeps the config-derived path.
|
|
same = 8 * (1 << 30)
|
|
bytes_, src = self._run(
|
|
"/local/equal",
|
|
config_bytes = same,
|
|
local_bytes = same,
|
|
)
|
|
self.assertEqual(bytes_, same)
|
|
self.assertEqual(src, "config")
|
|
|
|
def test_remote_safetensors_path_unaffected_by_local_weights(self):
|
|
from utils.hardware import hardware as hardware_module
|
|
with (
|
|
patch.object(
|
|
hardware_module,
|
|
"_resolve_model_identifier_for_gpu_estimate",
|
|
return_value = "owner/repo",
|
|
),
|
|
patch.object(
|
|
hardware_module,
|
|
"_get_hf_safetensors_total_params",
|
|
return_value = 1_000_000_000,
|
|
),
|
|
patch.object(
|
|
hardware_module,
|
|
"_load_config_for_gpu_estimate",
|
|
) as mock_load,
|
|
patch.object(
|
|
hardware_module,
|
|
"_get_local_weight_size_bytes",
|
|
) as mock_local,
|
|
):
|
|
bytes_, src = hardware_module.estimate_fp16_model_size_bytes("owner/repo")
|
|
self.assertEqual(bytes_, 2 * 1_000_000_000)
|
|
self.assertEqual(src, "safetensors")
|
|
mock_load.assert_not_called()
|
|
mock_local.assert_not_called()
|