mirror of
https://github.com/unslothai/unsloth.git
synced 2026-08-13 02:44:21 +00:00
0b147fdd2a
* fix(studio): harden training setup and model loading Scan every model load root before approving remote code, and pin third-party codec sources to verified revisions. Align dataset option validation across the UI and backend, preserve manual drafts, and include edits in training start identity. Require job-scoped stop requests and retain bounded early-cancel tombstones without unsafe eviction. * fix(studio): harden training lifecycle and audio loading Preserve early start cancellations with bounded tombstones and explicit capacity handling for concurrent requests. Pin and verify third-party audio sources and codec artifacts with safe archive extraction and offline cache migration. Keep automatic evaluation data separate from every split included in a combined training instruction. Keep training summaries and GPU progress state aligned with the active configuration through public feature exports. Make route and lifecycle tests deterministic by isolating Hub reachability and executor behavior. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio): correct completion masking and offline training Apply Alpaca response markers safely when completion-only training is enabled. Skip optional runtime dependency installs while Studio is offline. Limit causal-conv1d hooks to recognized model families without dropping supported architectures. Disable known-broken TileLang dispatch when offline repair is unavailable. Use cached GGUF and model-size metadata without offline Hub retries. * fix(studio): resolve causal conv kernels from model configs Detect causal-conv1d requirements from resolved model architectures before loading model code. Keep name matching as a fallback while excluding unrelated renamed checkpoints. * fix(studio): preserve registered training cancellation Allow registered start cancellations to reclaim the oldest expiring tombstone when unknown cancellation capacity is full. Preserve the hard capacity limit and 429 response for unregistered request IDs. * Fix reset job scoping, DAC fast path, scan target and subset splits for PR #8103 - /api/train/reset: an unscoped reset could force-terminate a run mid-cancel. The guard now refuses to touch a live run it cannot prove it owns. The field stays optional so pre-rework clients, which POST /reset with no body, keep working. - ensure_dac_speech_weights: install the download into the pinned destination, so later loads hit the fast path instead of re-downloading and re-hashing 295 MB under the install lock. - _requires_security_review_for_model: apply the same load_scan_target alias normalization the sibling remote-code check gained, else the Spark-TTS alias 404s and fails open to "no review needed". - _purge_package_bytecode: best effort. It runs without the install lock over a cache shared by the inference and training workers; 7 of 8 concurrent imports died on it. Also moved inside the try so a failure cannot strand the cache dir on sys.path, and snapshot sys.modules before the origin audit. - commitSubset: clear the backing splits too, else the render-phase draft sync reads the previous subset's split back into the boxes it just reset. - Fix two tests that fail on the branch: the DAC assertion pinned the old return value, and the causal-conv1d assertion matched call formatting. * Make the SSM runtime tests Windows aware ensure_ssm_runtime deliberately skips causal-conv1d on win32 (no prebuilt wheel), so the two install-order assertions only hold off Windows. Caught on a real windows-latest runner. * Update two frontend source contracts the branch moved - captureTrainingStartInputs now delegates to createTrainingStartInputIdentity, so the normalize/flags assertions belong against training-start-inputs.ts. - resetTraining takes a RequiredTrainingJobScope and always sends the body, which is stronger than the hasScope branch the contract pinned. Both fail on the branch today; caught by tests/studio, which the studio backend job does not cover. * Fix unscoped reset compat, pyc purge fail-open and DAC fallback for PR #8103 Corrects four things in my earlier commits on this branch. Unscoped /api/train/reset returned "superseded" (HTTP 200) for a live run. The pre-rework cancel dialog chains stopTrainingRun then a bodyless reset, so an older client read that 200 as success and cleared its UI while training kept running. Return "active" (409) instead: same answer a live run already gives, and one those clients already handle. It still never force-terminates, so a bodyless reset landing between current_job_id being set and _cancel_requested being cleared cannot kill the run that just started. The bytecode purge was made best-effort, but it is the only thing stopping a stale or planted .pyc shadowing a verified .py: the manifest skips __pycache__ and the origin audit reads __file__, which still names the .py. Tolerate only FileNotFoundError, the real concurrent-purge race, and let PermissionError fail the load again. The DAC fast path copies 295 MB inside the hub cache with only Timeout caught, so a full disk turned a hash-verified download into a hard failure. Fall back to the verified hub path on OSError. commitSubset cleared both splits, but setDatasetSubset already does that; the extra setDatasetSplit(null) only cost a runDatasetCheck against an assumed "train" split. * Restore the eval split reset and the cancelled-run dismiss for PR #8103 Two corrections to 4e5389c21. commitSubset: I removed setDatasetEvalSplit(null) because setDatasetSubset already nulls datasetEvalSplit. It does, but it never resets evalSteps, and setDatasetEvalSplit is not a plain setter: it zeroes evalSteps and runs streamingCompatiblePatch. Without it, changing the subset left evaluation armed with no split, which routes/training.py rejects with 422 once streaming is on, and which silently auto-detects an eval split otherwise. Restored. setDatasetSplit(null) stays out, since its only unique effect was a runDatasetCheck against an assumed "train". Unscoped reset: returning "active" for every live run was too broad. The pre-rework cancel dialog only dismisses after stopTrainingRun succeeded, so _cancel_requested is already set and clearing the UI is right; 409 there just wedged the overlay behind a "Training still active" toast. Now 409 only when no stop was requested, which is the stale-tab case the change was for. Still no force_terminate on an unscoped reset. Also covers the __pycache__ branch of the purge, which is the route a planted .pyc actually takes; the existing test only reached the top-level .pyc loop. * Keep live start cancellations and fall back on a full disk for PR #8103 Two fixes plus the regression tests the earlier lifecycle fixes shipped without. Cancelling the active start at tombstone capacity reclaimed a slot by deleting the soonest-expiring entry. Expired ones are already pruned a few lines above, so that entry was always live, and dropping it let a delayed /start spawn the job it had cancelled. Reserve capacity instead: only the owner of the active start reaches that branch and there is at most one, so the table lands at cap + 1 rather than forgetting a cancellation. Unregistered ids still hit the hard cap. ensure_dac_speech_weights migrated a pre-existing legacy file with an unguarded copy, so a hub cache that cannot absorb a second 295 MB copy failed weights that had already passed size and sha256. Same fallback the download branch below it already uses. Tests: renamed and local SSM checkpoints resolving from config rather than name, the owner of an active start staying cancellable at capacity, live cancellations surviving that cancel, and the full-disk legacy DAC fallback. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Opt the pinned source checkouts into Git long paths for PR #8103 Cross-OS CI caught 14 failures on windows-latest, all of them the pinned checkout dying with "error: unable to write file ...: Filename too long". Git for Windows still enforces MAX_PATH unless told otherwise, and the cache nests a 40-char revision, a staging dir and .git/objects under the studio home. A venv-inferred home already measures about 253 of the 260 characters, so a slightly longer user or install path fails on a normal Windows machine, not just under the deeper pytest tmp dir. Passed per invocation with -c so no user or system Git config is touched, and it is a no-op off Windows. * Bound pending cancels, reach legacy DAC weights, and delete read-only checkouts for PR #8103 Three fixes, two of them on my own previous commit. Moving the owner cancel from evict-oldest to a one-slot overshoot also changed what the hardcoded reclaim_capacity=True on the pending non-owner branch did: it used to evict, so the table stayed at the cap, and it started overshooting instead. Start plus cancel could then be repeated to grow it without bound (1224 entries against a cap of 1024). That branch now takes the plain reservation, so the extra slot belongs to the owner of the active run alone. The DAC legacy fallback sat behind destination.parent.mkdir() and the install lock, both of which need a writable cache, so a read-only or full hub cache raised before weights we can already verify were ever looked at. Fall back at both points, and only to an artifact that passes the same size and sha256 check. Cross-OS CI then caught replacing a pinned checkout failing on windows-latest with WinError 5: Git marks .git/objects read-only and Windows will not delete a read-only file, so any repair or revision change died there. Clear the attribute and retry, only when the path is genuinely not writable, so an open handle still surfaces. * Delete the cached checkout the Windows-safe way in the migration test for PR #8103 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen <moonshotaisubstack@gmail.com>
2295 lines
90 KiB
Python
2295 lines
90 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,
|
|
get_vulkan_inference_gpu_info,
|
|
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)
|
|
|
|
def test_discrete_vulkan_inference_gpu_info(self):
|
|
with (
|
|
patch(
|
|
"core.inference.llama_cpp.LlamaCppBackend._is_vulkan_backend",
|
|
return_value = True,
|
|
),
|
|
patch(
|
|
"core.inference.llama_cpp.LlamaCppBackend.vulkan_device_inventory",
|
|
return_value = [
|
|
{
|
|
"index": 0,
|
|
"name": "Vulkan0",
|
|
"free_mib": 7402,
|
|
"total_mib": 8192,
|
|
"is_igpu": False,
|
|
}
|
|
],
|
|
),
|
|
):
|
|
result = get_vulkan_inference_gpu_info()
|
|
|
|
self.assertTrue(result["available"])
|
|
self.assertEqual(result["backend"], "vulkan")
|
|
# ggml Vulkan ordinals are the space `--device Vulkan<i>` pins, so they are selectable.
|
|
self.assertEqual(result["index_kind"], "vulkan")
|
|
self.assertEqual(result["parent_visible_gpu_ids"], [])
|
|
self.assertEqual(
|
|
result["devices"],
|
|
[
|
|
{
|
|
"index": 0,
|
|
"index_kind": "vulkan",
|
|
"visible_ordinal": 0,
|
|
"name": "Vulkan0",
|
|
"memory_total_gb": 8.0,
|
|
"vram_used_gb": 0.77,
|
|
"vram_free_gb": 7.23,
|
|
"vram_utilization_pct": 9.6,
|
|
"shared_memory": False,
|
|
}
|
|
],
|
|
)
|
|
|
|
def test_vulkan_igpu_info_uses_capped_free_budget(self):
|
|
with (
|
|
patch(
|
|
"core.inference.llama_cpp.LlamaCppBackend._is_vulkan_backend",
|
|
return_value = True,
|
|
),
|
|
patch(
|
|
"core.inference.llama_cpp.LlamaCppBackend.vulkan_device_inventory",
|
|
return_value = [
|
|
{
|
|
"index": 0,
|
|
"name": "Vulkan0",
|
|
"free_mib": 12288,
|
|
"total_mib": 32768,
|
|
"is_igpu": True,
|
|
}
|
|
],
|
|
),
|
|
patch(
|
|
"core.inference.llama_cpp._apply_igpu_host_reserve_mib",
|
|
return_value = 12288,
|
|
),
|
|
):
|
|
result = get_vulkan_inference_gpu_info()
|
|
|
|
device = result["devices"][0]
|
|
self.assertEqual(device["memory_total_gb"], 12.0)
|
|
self.assertEqual(device["vram_free_gb"], 12.0)
|
|
self.assertIsNone(device["vram_used_gb"])
|
|
self.assertIsNone(device["vram_utilization_pct"])
|
|
self.assertTrue(device["shared_memory"])
|
|
|
|
def test_forced_vulkan_overrides_torch_gpu_visibility_for_inference(self):
|
|
with (
|
|
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch(
|
|
"core.inference.llama_cpp.LlamaCppBackend._is_vulkan_backend",
|
|
return_value = True,
|
|
),
|
|
patch(
|
|
"core.inference.llama_cpp.LlamaCppBackend.vulkan_device_inventory",
|
|
return_value = [
|
|
{
|
|
"index": 1,
|
|
"name": "Vulkan1",
|
|
"free_mib": 6144,
|
|
"total_mib": 8192,
|
|
"is_igpu": False,
|
|
}
|
|
],
|
|
),
|
|
patch(
|
|
"utils.hardware.nvidia.get_backend_visible_gpu_info",
|
|
return_value = {
|
|
"available": True,
|
|
"backend": "cuda",
|
|
"devices": [{"index": 0, "name": "CUDA0", "memory_total_gb": 24.0}],
|
|
},
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._get_parent_visible_gpu_spec",
|
|
return_value = {"raw": None, "numeric_ids": None},
|
|
),
|
|
):
|
|
training_result = get_backend_visible_gpu_info()
|
|
inference_result = get_vulkan_inference_gpu_info()
|
|
|
|
self.assertEqual(training_result["backend"], "cuda")
|
|
self.assertEqual(inference_result["backend"], "vulkan")
|
|
self.assertEqual(inference_result["devices"][0]["index"], 1)
|
|
|
|
def test_vulkan_install_without_devices_reports_unavailable(self):
|
|
with (
|
|
patch(
|
|
"core.inference.llama_cpp.LlamaCppBackend._is_vulkan_backend",
|
|
return_value = True,
|
|
),
|
|
patch(
|
|
"core.inference.llama_cpp.LlamaCppBackend.vulkan_device_inventory",
|
|
return_value = [],
|
|
),
|
|
):
|
|
result = get_vulkan_inference_gpu_info()
|
|
|
|
self.assertFalse(result["available"])
|
|
self.assertEqual(result["backend"], "vulkan")
|
|
self.assertEqual(result["devices"], [])
|
|
|
|
|
|
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_offline_safetensors_probe_uses_config_without_hub_access(self):
|
|
config = object()
|
|
for offline_variable in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"):
|
|
with self.subTest(offline_variable = offline_variable):
|
|
with (
|
|
patch.dict(os.environ, {offline_variable: "true"}, clear = True),
|
|
patch("huggingface_hub.model_info") as hub_info,
|
|
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 = config,
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._estimate_fp16_model_size_bytes_from_config",
|
|
return_value = 1234,
|
|
),
|
|
patch(
|
|
"utils.hardware.hardware._get_local_weight_size_bytes",
|
|
return_value = None,
|
|
),
|
|
):
|
|
model_size_bytes, source = _hw_module.estimate_fp16_model_size_bytes(
|
|
"unsloth/test"
|
|
)
|
|
|
|
self.assertEqual((model_size_bytes, source), (1234, "config"))
|
|
hub_info.assert_not_called()
|
|
|
|
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 IDs use the normal resolver instead of a blanket rejection.
|
|
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("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_unreachable", 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("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_unreachable", 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, uses_vulkan_ordinals = asyncio.run(
|
|
inference_route._resolve_gguf_gpu_ids_for_request(config, [1, 0])
|
|
)
|
|
|
|
self.assertEqual(resolved, [0, 1])
|
|
self.assertTrue(uses_vulkan_ordinals)
|
|
|
|
def test_diffusion_gpu_ids_accept_rocm_physical_path(self):
|
|
import utils.hardware as hardware_pkg
|
|
import utils.hardware.hardware as hardware_mod
|
|
|
|
inference_route = _load_route_module(
|
|
"inference_route_module_for_diffusion_rocm_path_test",
|
|
"routes/inference.py",
|
|
)
|
|
config = SimpleNamespace(is_gguf = True)
|
|
fake_backend = SimpleNamespace(
|
|
is_vulkan_build = lambda: False,
|
|
_backend_lacks_gpu_lib = lambda *a, **k: False,
|
|
)
|
|
|
|
with (
|
|
patch.object(hardware_mod, "IS_ROCM", True),
|
|
patch.object(hardware_pkg, "get_device", return_value = DeviceType.CUDA),
|
|
patch.object(
|
|
hardware_mod,
|
|
"resolve_requested_gpu_ids",
|
|
return_value = [0, 1],
|
|
),
|
|
patch.object(
|
|
inference_route,
|
|
"get_llama_cpp_backend",
|
|
return_value = fake_backend,
|
|
),
|
|
patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
|
|
):
|
|
self.assertEqual(hardware_mod._backend_label(DeviceType.CUDA), "rocm")
|
|
resolved, uses_vulkan_ordinals = asyncio.run(
|
|
inference_route._resolve_gguf_gpu_ids_for_request(
|
|
config,
|
|
[1, 0],
|
|
diffusion_kind = True,
|
|
)
|
|
)
|
|
|
|
self.assertEqual(resolved, [0, 1])
|
|
self.assertFalse(uses_vulkan_ordinals)
|
|
|
|
def test_inference_route_validates_gpu_ids_for_gguf(self):
|
|
import utils.hardware.hardware as hardware_mod
|
|
import utils.hardware as hardware_pkg
|
|
|
|
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(
|
|
"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_unreachable", 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("gpu_ids", exc_info.exception.detail.lower())
|
|
|
|
def test_inference_route_rejects_gpu_ids_on_cpu_only_llama_build(self):
|
|
# A CPU-only llama.cpp build cannot honor a CUDA visibility pin.
|
|
import utils.hardware as hardware_pkg
|
|
|
|
inference_route = _load_route_module(
|
|
"inference_route_module_for_cpu_only_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,
|
|
)
|
|
fake_backend = SimpleNamespace(
|
|
is_loaded = False,
|
|
model_identifier = None,
|
|
is_vulkan_build = lambda: False,
|
|
_backend_lacks_gpu_lib = lambda *a, **k: True,
|
|
)
|
|
with (
|
|
patch.object(
|
|
inference_route,
|
|
"ModelConfig",
|
|
SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
|
|
),
|
|
patch.object(inference_route, "_classify_diffusion_gguf", return_value = 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_unreachable", nullcontext),
|
|
patch.object(inference_route, "get_llama_cpp_backend", return_value = fake_backend),
|
|
patch.object(hardware_pkg, "get_device", return_value = hardware_pkg.DeviceType.CUDA),
|
|
):
|
|
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("cpu-only build", exc_info.exception.detail.lower())
|
|
|
|
def test_diffusion_gguf_on_vulkan_build_rejects_ordinal_pin(self):
|
|
# The GGUF picker supplies Vulkan ordinals, not the CUDA physical IDs the diffusion runner uses.
|
|
import utils.hardware as hardware_pkg
|
|
|
|
inference_route = _load_route_module(
|
|
"inference_route_module_for_diffusion_cuda_path_test",
|
|
"routes/inference.py",
|
|
)
|
|
request = LoadRequest(model_path = "unsloth/diffusion.gguf", gpu_ids = [0, 1])
|
|
model_config = SimpleNamespace(
|
|
is_gguf = True,
|
|
is_lora = False,
|
|
gguf_hf_repo = None,
|
|
gguf_file = "/tmp/diffusion.gguf",
|
|
gguf_mmproj_file = None,
|
|
gguf_variant = None,
|
|
identifier = "unsloth/diffusion.gguf",
|
|
display_name = "unsloth/diffusion.gguf",
|
|
is_vision = False,
|
|
is_audio = False,
|
|
audio_type = None,
|
|
has_audio_input = False,
|
|
)
|
|
fake_backend = SimpleNamespace(
|
|
is_loaded = False,
|
|
model_identifier = None,
|
|
is_vulkan_build = lambda: True,
|
|
_backend_lacks_gpu_lib = lambda *a, **k: False,
|
|
)
|
|
with (
|
|
patch.object(
|
|
inference_route,
|
|
"ModelConfig",
|
|
SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
|
|
),
|
|
patch.object(inference_route, "_classify_diffusion_gguf", return_value = True),
|
|
patch.object(
|
|
_hw_module,
|
|
"resolve_requested_gpu_ids",
|
|
side_effect = AssertionError("Vulkan ordinal reached the CUDA resolver"),
|
|
),
|
|
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_unreachable", nullcontext),
|
|
patch.object(inference_route, "get_llama_cpp_backend", return_value = fake_backend),
|
|
patch.object(hardware_pkg, "get_device", return_value = hardware_pkg.DeviceType.CUDA),
|
|
):
|
|
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("no defined mapping", exc_info.exception.detail)
|
|
|
|
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): the point is the ORDER, validation before the handoff.
|
|
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_unreachable_for", 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 must 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_unreachable_for", 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 through acquire_for's
|
|
# `register` hook (which runs under the arbiter lock). Passing no register left a competing 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_unreachable_for", 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.object(
|
|
training_route,
|
|
"_remote_untrainable_model_format",
|
|
return_value = None,
|
|
),
|
|
patch.object(training_route.asyncio, "to_thread", new = _inline_to_thread),
|
|
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.object(
|
|
training_route,
|
|
"_remote_untrainable_model_format",
|
|
return_value = None,
|
|
),
|
|
patch.object(training_route.asyncio, "to_thread", new = _inline_to_thread),
|
|
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_unreachable", 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_unreachable", 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()
|