# 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 UUID/MIG" ): 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_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_cuda_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 devices", str(exc_info.exception)) def test_inference_route_rejects_gpu_ids_for_gguf(self): 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, ) 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), ): 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("GGUF", exc_info.exception.detail) 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 TestXpuRejection(_GpuCacheResetMixin, unittest.TestCase): def test_auto_select_returns_non_cuda_for_xpu(self): with patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU): selected, metadata = auto_select_gpu_ids("unsloth/test") self.assertIsNone(selected) self.assertEqual(metadata["selection_mode"], "non_cuda") def test_prepare_gpu_selection_rejects_explicit_ids_on_xpu(self): with patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU): with self.assertRaisesRegex(ValueError, "only supported on CUDA"): prepare_gpu_selection([0], model_name = "unsloth/test") 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()