e93507a09c
Lockfile supply-chain audit / lockfile supply-chain audit (push) Has been cancelled
Windows Studio GGUF CI / GPU prebuilt resolves without Visual Studio (push) Has been cancelled
Windows Studio GGUF CI / setup.ps1 unit tests (VS 2026 / CMake guard) (push) Has been cancelled
Windows Studio GGUF CI / real-VS detection (VS 2022) (push) Has been cancelled
Windows Studio GGUF CI / real-VS detection (VS 2026) (push) Has been cancelled
Windows Studio GGUF CI / VC++ runtime detect + install round-trip (windows-2025-vs2026) (push) Has been cancelled
Windows Studio GGUF CI / VC++ runtime detect + install round-trip (windows-latest) (push) Has been cancelled
Windows Studio Update CI / Studio Updating Tests (push) Has been cancelled
Wheel CI / Wheel build + content sanity + import smoke (push) Has been cancelled
Lint CI / Source lint (Python + shell + YAML + JSON + safety nets) (push) Has been cancelled
MLX CI on Mac M1 / dispatch (push) Has been cancelled
Security audit / advisory audit (pip + npm + cargo) (push) Has been cancelled
Security audit / pip scan-packages :: extras (push) Has been cancelled
Security audit / pip scan-packages :: studio (push) Has been cancelled
Security audit / pip scan-packages :: hf-stack (push) Has been cancelled
Security audit / npm scan-packages (Studio frontend tarballs) (push) Has been cancelled
Security audit / workflow-trigger lint (pull_request_target / cache-poisoning) (push) Has been cancelled
Security audit / pytest tests/security (push) Has been cancelled
Security audit / npm provenance + new install-script diff (push) Has been cancelled
Studio API CI / Studio API & Auth Tests (push) Has been cancelled
Backend CI / (Python 3.10) (push) Has been cancelled
Backend CI / (Python 3.11) (push) Has been cancelled
Backend CI / (Python 3.12) (push) Has been cancelled
Backend CI / (Python 3.13) (push) Has been cancelled
Backend CI / Repo tests (CPU) (push) Has been cancelled
Frontend CI / Frontend build + bundle sanity (push) Has been cancelled
Studio GGUF CI / OpenAI, Anthropic API tests (push) Has been cancelled
Studio GGUF CI / Tool calling Tests (push) Has been cancelled
Studio GGUF CI / JSON, images (push) Has been cancelled
Mac Studio GGUF CI / OpenAI, Anthropic API tests (push) Has been cancelled
Mac Studio GGUF CI / Tool calling Tests (push) Has been cancelled
Mac Studio GGUF CI / JSON, images (push) Has been cancelled
Mac Studio Install Matrix CI / Install + load (macos-14) (push) Has been cancelled
Mac Studio Install Matrix CI / Install + load (macos-15) (push) Has been cancelled
Mac Studio Install Matrix CI / Install + load (macos-26) (push) Has been cancelled
Mac Studio Install Matrix CI / Install + load (macos-15-intel) (push) Has been cancelled
Mac Studio API CI / Studio API & Auth Tests (push) Has been cancelled
Mac Studio Install Matrix CI / Install + load (macos-26-intel) (push) Has been cancelled
Mac Studio UI CI / Chat UI Tests (push) Has been cancelled
Studio Tauri CI / Tauri Linux debug build (no codesign) (push) Has been cancelled
Mac Studio Update CI / Studio Updating Tests (push) Has been cancelled
Studio UI CI / Chat UI Tests (push) Has been cancelled
Windows Studio API CI / Studio API & Auth Tests (push) Has been cancelled
Windows Studio UI CI / Chat UI Tests (push) Has been cancelled
Studio Update CI / Studio Updating Tests (push) Has been cancelled
Core / Core (HF=default + TRL=default) (push) Has been cancelled
Core / Core (HF=4.57.6 + TRL<1) (push) Has been cancelled
Core / Core (HF=latest + TRL=latest) (push) Has been cancelled
Core / llama.cpp build + smoke (push) Has been cancelled
Windows Studio GGUF CI / OpenAI, Anthropic API tests (push) Has been cancelled
Windows Studio GGUF CI / Tool calling Tests (push) Has been cancelled
Windows Studio GGUF CI / JSON, images (push) Has been cancelled
Windows Studio GGUF CI / Studio install + inference without Visual Studio (push) Has been cancelled
Studio export capability / capability (macos-latest) (push) Has been cancelled
Studio export capability / capability (ubuntu-latest) (push) Has been cancelled
Studio export capability / capability (windows-latest) (push) Has been cancelled
Cross-platform parity / parity (macos-latest) (push) Has been cancelled
Cross-platform parity / parity (windows-latest) (push) Has been cancelled
Scorecard supply-chain security / Scorecard analysis (push) Has been cancelled
Studio load-orchestrator CI / test (push) Has been cancelled
686 lines
27 KiB
Python
686 lines
27 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
|
|
|
|
"""Loading a NEW chat model while training runs: can_load_chat_during_training
|
|
(VRAM fit check), _guard_chat_load_against_training and _effective_load_in_4bit
|
|
(409 + sizing wiring). The guard sizes the same effective load the backend will
|
|
perform (HF auto reuses the loader's selector, HF explicit applies a per-GPU
|
|
floor, GGUF sizes from on-disk weights, LoRA 4-bit->16-bit flips resolved first)
|
|
and leaves non-training/external loads untouched."""
|
|
|
|
import asyncio
|
|
import importlib.util
|
|
import sys
|
|
import types
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from fastapi import HTTPException
|
|
|
|
from utils.hardware import DeviceType
|
|
import utils.hardware.hardware as _hw_module
|
|
|
|
_BACKEND_ROOT = Path(__file__).resolve().parent.parent
|
|
|
|
# Load training_vram.py standalone (avoids the heavy routes/__init__.py); its
|
|
# lazy hardware imports still resolve against the patched utils.hardware names.
|
|
_spec = importlib.util.spec_from_file_location(
|
|
"training_vram_load_test", _BACKEND_ROOT / "routes" / "training_vram.py"
|
|
)
|
|
tv = importlib.util.module_from_spec(_spec)
|
|
_spec.loader.exec_module(tv)
|
|
|
|
|
|
class _GpuCacheResetMixin:
|
|
def tearDown(self):
|
|
_hw_module._physical_gpu_count = None
|
|
_hw_module._visible_gpu_count = None
|
|
|
|
|
|
def _devices(*free_specs):
|
|
"""Build a device list from (index, total, used) tuples."""
|
|
return [
|
|
{"index": i, "vram_total_gb": total, "vram_used_gb": used}
|
|
for (i, total, used) in free_specs
|
|
]
|
|
|
|
|
|
# ── can_load_chat_during_training: HF auto (reuses auto_select_gpu_ids) ───────
|
|
|
|
|
|
class TestCanLoadAutoHF(_GpuCacheResetMixin, unittest.TestCase):
|
|
def _run(self, *, selection_mode, required, usable):
|
|
meta = {"selection_mode": selection_mode, "required_gb": required, "usable_gb": usable}
|
|
with (
|
|
patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch("utils.hardware.auto_select_gpu_ids", return_value = ([0], meta)) as auto_mock,
|
|
):
|
|
ok, info = tv.can_load_chat_during_training(
|
|
model_name = "unsloth/Qwen3-1.7B",
|
|
hf_token = None,
|
|
load_in_4bit = True,
|
|
max_seq_length = 0,
|
|
requested_gpu_ids = None,
|
|
is_gguf = False,
|
|
)
|
|
return ok, info, auto_mock
|
|
|
|
def test_fits_with_margin(self):
|
|
# free 60 >= 8*1.15+4 = 13.2
|
|
ok, info, auto_mock = self._run(selection_mode = "auto", required = 8.0, usable = 60.0)
|
|
self.assertTrue(ok)
|
|
self.assertEqual(info["mode"], "auto")
|
|
self.assertAlmostEqual(info["needed_gb"], 13.2, places = 3)
|
|
auto_mock.assert_called_once() # mirrors the loader's own selection
|
|
|
|
def test_too_tight_refuses(self):
|
|
# free 10 < 8*1.15+4 = 13.2 -> refuse even though raw 10 > 8
|
|
ok, _, _ = self._run(selection_mode = "auto", required = 8.0, usable = 10.0)
|
|
self.assertFalse(ok)
|
|
|
|
def test_fallback_all_refuses(self):
|
|
# Selector couldn't confirm placement -> default-deny to protect training.
|
|
ok, info = self._run(selection_mode = "fallback_all", required = 8.0, usable = 999.0)[:2]
|
|
self.assertFalse(ok)
|
|
|
|
|
|
# ── can_load_chat_during_training: HF explicit (per-GPU floor) ────────────────
|
|
|
|
|
|
class TestCanLoadExplicitHF(_GpuCacheResetMixin, unittest.TestCase):
|
|
def _run(
|
|
self,
|
|
*,
|
|
required,
|
|
devices,
|
|
gpu_ids,
|
|
resolved = None,
|
|
resolve_side_effect = None,
|
|
):
|
|
resolve_kwargs = (
|
|
{"side_effect": resolve_side_effect}
|
|
if resolve_side_effect
|
|
else {"return_value": resolved if resolved is not None else gpu_ids}
|
|
)
|
|
with (
|
|
patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch("utils.hardware.estimate_required_model_memory_gb", return_value = (required, {})),
|
|
patch("utils.hardware.get_visible_gpu_utilization", return_value = {"devices": devices}),
|
|
patch("utils.hardware.resolve_requested_gpu_ids", **resolve_kwargs),
|
|
patch("utils.hardware.auto_select_gpu_ids") as auto_mock,
|
|
):
|
|
ok, info = tv.can_load_chat_during_training(
|
|
model_name = "m",
|
|
hf_token = None,
|
|
load_in_4bit = True,
|
|
max_seq_length = 0,
|
|
requested_gpu_ids = gpu_ids,
|
|
is_gguf = False,
|
|
)
|
|
return ok, info, auto_mock
|
|
|
|
def test_single_gpu_fits(self):
|
|
ok, info, auto_mock = self._run(required = 8.0, devices = _devices((0, 80, 20)), gpu_ids = [0])
|
|
self.assertTrue(ok)
|
|
self.assertEqual(info["mode"], "explicit")
|
|
auto_mock.assert_not_called() # explicit never calls the auto selector
|
|
|
|
def test_per_gpu_floor_blocks_uneven_split(self):
|
|
# free [45, 10]; aggregate 45 + 10*0.85 = 53.5 >= needed 27, but the 10 GB
|
|
# GPU is below the even-share floor 27/2 = 13.5 -> refuse (would OOM it).
|
|
ok, info, _ = self._run(
|
|
required = 20.0, devices = _devices((0, 80, 35), (1, 80, 70)), gpu_ids = [0, 1]
|
|
)
|
|
self.assertFalse(ok)
|
|
self.assertAlmostEqual(info["min_free_gb"], 10.0, places = 3)
|
|
|
|
def test_per_gpu_floor_passes_when_even(self):
|
|
# free [30, 30]; both clear the 13.5 even-share floor -> allow.
|
|
ok, _, _ = self._run(
|
|
required = 20.0, devices = _devices((0, 80, 50), (1, 80, 50)), gpu_ids = [0, 1]
|
|
)
|
|
self.assertTrue(ok)
|
|
|
|
def test_missing_gpu_counts_as_zero(self):
|
|
ok, _, _ = self._run(required = 5.0, devices = _devices((0, 80, 5)), gpu_ids = [3], resolved = [3])
|
|
self.assertFalse(ok)
|
|
|
|
def test_invalid_ids_does_not_block(self):
|
|
ok, info, _ = self._run(
|
|
required = 5.0,
|
|
devices = [],
|
|
gpu_ids = [99],
|
|
resolve_side_effect = ValueError("Invalid gpu_ids [99]"),
|
|
)
|
|
self.assertTrue(ok)
|
|
self.assertEqual(info["reason"], "invalid_gpu_ids")
|
|
|
|
|
|
# ── can_load_chat_during_training: GGUF (sized from on-disk weights) ──────────
|
|
|
|
|
|
class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase):
|
|
def _run(
|
|
self,
|
|
*,
|
|
devices,
|
|
required_override = None,
|
|
estimate = None,
|
|
):
|
|
with (
|
|
patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch("utils.hardware.estimate_required_model_memory_gb", return_value = (estimate, {})),
|
|
patch("utils.hardware.get_visible_gpu_utilization", return_value = {"devices": devices}),
|
|
patch("utils.hardware.auto_select_gpu_ids") as auto_mock,
|
|
):
|
|
ok, info = tv.can_load_chat_during_training(
|
|
model_name = "unsloth/gemma-GGUF",
|
|
hf_token = None,
|
|
load_in_4bit = True,
|
|
max_seq_length = 0,
|
|
requested_gpu_ids = None,
|
|
is_gguf = True,
|
|
required_override_gb = required_override,
|
|
)
|
|
return ok, info, auto_mock
|
|
|
|
def test_override_fits(self):
|
|
ok, info, auto_mock = self._run(devices = _devices((0, 80, 20)), required_override = 10.0)
|
|
self.assertTrue(ok)
|
|
self.assertEqual(info["mode"], "gguf")
|
|
auto_mock.assert_not_called() # GGUF never uses the HF auto selector
|
|
|
|
def test_no_per_gpu_floor_for_gguf(self):
|
|
# free [45, 10], override 20 -> needed 27, aggregate 53.5 >= 27. GGUF self-
|
|
# places, so the per-GPU floor that would block HF doesn't apply -> allow.
|
|
ok, _, _ = self._run(devices = _devices((0, 80, 35), (1, 80, 70)), required_override = 20.0)
|
|
self.assertTrue(ok)
|
|
|
|
def test_estimate_unavailable_refuses(self):
|
|
# No override and the estimator can't size it -> default-deny.
|
|
ok, info, _ = self._run(devices = _devices((0, 80, 0)), required_override = None, estimate = None)
|
|
self.assertFalse(ok)
|
|
self.assertEqual(info["reason"], "estimate_unavailable")
|
|
|
|
|
|
# ── can_load_chat_during_training: device-independent paths ──────────────────
|
|
|
|
|
|
class TestCanLoadMisc(_GpuCacheResetMixin, unittest.TestCase):
|
|
def test_non_cuda_allows(self):
|
|
with patch("utils.hardware.get_device", return_value = DeviceType.MLX):
|
|
ok, info = tv.can_load_chat_during_training(
|
|
model_name = "m",
|
|
hf_token = None,
|
|
load_in_4bit = True,
|
|
max_seq_length = 0,
|
|
requested_gpu_ids = None,
|
|
)
|
|
self.assertTrue(ok)
|
|
self.assertEqual(info["mode"], "non_cuda")
|
|
|
|
def test_no_visible_gpus_refuses(self):
|
|
# GGUF with an empty device list -> no candidate GPU -> default-deny.
|
|
with (
|
|
patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
|
|
patch("utils.hardware.get_visible_gpu_utilization", return_value = {"devices": []}),
|
|
patch("utils.hardware.auto_select_gpu_ids"),
|
|
):
|
|
ok, info = tv.can_load_chat_during_training(
|
|
model_name = "m",
|
|
hf_token = None,
|
|
load_in_4bit = True,
|
|
max_seq_length = 0,
|
|
requested_gpu_ids = None,
|
|
is_gguf = True,
|
|
required_override_gb = 8.0,
|
|
)
|
|
self.assertFalse(ok)
|
|
self.assertEqual(info["reason"], "no_visible_gpus")
|
|
|
|
def test_probe_exception_refuses(self):
|
|
with patch("utils.hardware.get_device", side_effect = RuntimeError("boom")):
|
|
ok, info = tv.can_load_chat_during_training(
|
|
model_name = "m",
|
|
hf_token = None,
|
|
load_in_4bit = True,
|
|
max_seq_length = 0,
|
|
requested_gpu_ids = None,
|
|
)
|
|
self.assertFalse(ok)
|
|
self.assertEqual(info["reason"], "probe_error")
|
|
|
|
|
|
# ── _guard_chat_load_against_training + _effective_load_in_4bit (route) ───────
|
|
|
|
|
|
def _load_inference_route():
|
|
spec = importlib.util.spec_from_file_location(
|
|
"inference_route_chatload_test", _BACKEND_ROOT / "routes" / "inference.py"
|
|
)
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def _stub_guard_deps(
|
|
*,
|
|
training_active,
|
|
decision,
|
|
captured = None,
|
|
):
|
|
"""Inject the guard's two lazy imports (get_training_backend, can_load_chat_
|
|
during_training); `captured` records the can_load kwargs for assertions."""
|
|
core_training = types.ModuleType("core.training")
|
|
if isinstance(training_active, Exception):
|
|
|
|
def _raise():
|
|
raise training_active
|
|
|
|
core_training.get_training_backend = _raise
|
|
else:
|
|
core_training.get_training_backend = lambda: SimpleNamespace(
|
|
is_training_active = lambda: training_active
|
|
)
|
|
|
|
def _can_load(**kwargs):
|
|
if captured is not None:
|
|
captured.append(kwargs)
|
|
return decision
|
|
|
|
tv_stub = types.ModuleType("routes.training_vram")
|
|
tv_stub.can_load_chat_during_training = _can_load
|
|
return patch.dict(
|
|
sys.modules, {"core.training": core_training, "routes.training_vram": tv_stub}
|
|
)
|
|
|
|
|
|
class TestChatLoadGuardRoute(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.route = _load_inference_route()
|
|
|
|
def _guard(
|
|
self,
|
|
*,
|
|
config = None,
|
|
captured = None,
|
|
training_active,
|
|
decision,
|
|
):
|
|
config = config or SimpleNamespace(is_gguf = False, is_lora = False, path = None)
|
|
with _stub_guard_deps(
|
|
training_active = training_active, decision = decision, captured = captured
|
|
):
|
|
self.route._guard_chat_load_against_training(
|
|
config,
|
|
model_identifier = "unsloth/Qwen3-1.7B",
|
|
hf_token = None,
|
|
load_in_4bit = True,
|
|
max_seq_length = 0,
|
|
requested_gpu_ids = None,
|
|
)
|
|
|
|
def test_noop_when_training_inactive(self):
|
|
self._guard(training_active = False, decision = (False, {})) # must not raise
|
|
|
|
def test_noop_when_training_state_unknown(self):
|
|
self._guard(training_active = RuntimeError("no backend"), decision = (False, {}))
|
|
|
|
def test_allows_when_fits(self):
|
|
self._guard(training_active = True, decision = (True, {"mode": "auto"}))
|
|
|
|
def test_refuses_with_headroom_number(self):
|
|
info = {"required_gb": 30.0, "usable_gb": 6.0, "needed_gb": 39.0, "mode": "auto"}
|
|
with self.assertRaises(HTTPException) as exc:
|
|
self._guard(training_active = True, decision = (False, info))
|
|
self.assertEqual(exc.exception.status_code, 409)
|
|
self.assertIn("39 GB", exc.exception.detail) # reports needed_gb, not required_gb 30
|
|
self.assertNotIn("30 GB", exc.exception.detail)
|
|
self.assertIn("including safety headroom", exc.exception.detail)
|
|
self.assertNotIn("chat is disabled", exc.exception.detail.lower())
|
|
|
|
def test_refuses_generic_when_unsizable(self):
|
|
with self.assertRaises(HTTPException) as exc:
|
|
self._guard(training_active = True, decision = (False, {"reason": "estimate_unavailable"}))
|
|
self.assertEqual(exc.exception.status_code, 409)
|
|
self.assertIn("could not be verified", exc.exception.detail)
|
|
|
|
def test_gguf_config_passes_is_gguf_and_override(self):
|
|
captured = []
|
|
config = SimpleNamespace(is_gguf = True)
|
|
with patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5):
|
|
self._guard(
|
|
config = config,
|
|
captured = captured,
|
|
training_active = True,
|
|
decision = (True, {}),
|
|
)
|
|
self.assertEqual(captured[0]["is_gguf"], True)
|
|
self.assertEqual(captured[0]["required_override_gb"], 12.5)
|
|
|
|
|
|
class TestEffectiveLoadIn4bit(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.route = _load_inference_route()
|
|
|
|
def _write_adapter(self, tmpdir, payload):
|
|
import json
|
|
(Path(tmpdir) / "adapter_config.json").write_text(json.dumps(payload))
|
|
|
|
def test_non_lora_returns_request(self):
|
|
cfg = SimpleNamespace(is_lora = False, path = None, base_model = None)
|
|
self.assertTrue(self.route._effective_load_in_4bit(cfg, True))
|
|
|
|
def test_lora_method_flips_to_16bit(self):
|
|
import tempfile
|
|
with tempfile.TemporaryDirectory() as d:
|
|
self._write_adapter(d, {"unsloth_training_method": "lora"})
|
|
cfg = SimpleNamespace(is_lora = True, path = d, base_model = "x")
|
|
# requested 4-bit, but a 'lora' adapter loads 16-bit
|
|
self.assertFalse(self.route._effective_load_in_4bit(cfg, True))
|
|
|
|
def test_qlora_method_keeps_4bit(self):
|
|
import tempfile
|
|
with tempfile.TemporaryDirectory() as d:
|
|
self._write_adapter(d, {"unsloth_training_method": "qlora"})
|
|
cfg = SimpleNamespace(is_lora = True, path = d, base_model = "x")
|
|
self.assertTrue(self.route._effective_load_in_4bit(cfg, True))
|
|
|
|
def test_no_method_non_bnb_base_flips_to_16bit(self):
|
|
import tempfile
|
|
with tempfile.TemporaryDirectory() as d:
|
|
self._write_adapter(d, {})
|
|
cfg = SimpleNamespace(is_lora = True, path = d, base_model = "meta/Llama-3-8B")
|
|
self.assertFalse(self.route._effective_load_in_4bit(cfg, True))
|
|
|
|
def test_malformed_adapter_config_returns_request(self):
|
|
import tempfile
|
|
with tempfile.TemporaryDirectory() as d:
|
|
(Path(d) / "adapter_config.json").write_text("[1, 2, 3]") # not a dict
|
|
cfg = SimpleNamespace(is_lora = True, path = d, base_model = "x")
|
|
self.assertTrue(self.route._effective_load_in_4bit(cfg, True)) # no crash
|
|
|
|
|
|
# ── validate_model integration (early refusal, real settings) ────────────────
|
|
|
|
|
|
class TestValidateRefusesDuringTraining(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.route = _load_inference_route()
|
|
|
|
def _validate(
|
|
self,
|
|
*,
|
|
training_active,
|
|
decision,
|
|
captured = None,
|
|
load_in_4bit = True,
|
|
):
|
|
from models.inference import ValidateModelRequest
|
|
|
|
request = ValidateModelRequest(
|
|
model_path = "unsloth/Qwen3-1.7B", load_in_4bit = load_in_4bit, max_seq_length = 4096
|
|
)
|
|
cfg = SimpleNamespace(
|
|
identifier = "unsloth/Qwen3-1.7B",
|
|
display_name = "Qwen3-1.7B",
|
|
is_gguf = False,
|
|
is_lora = False,
|
|
is_vision = False,
|
|
path = None,
|
|
base_model = None,
|
|
)
|
|
with (
|
|
patch.object(
|
|
self.route,
|
|
"_resolve_model_identifier_for_request",
|
|
return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False),
|
|
),
|
|
patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
|
|
patch.object(self.route, "load_inference_config", return_value = {}),
|
|
_stub_guard_deps(training_active = training_active, decision = decision, captured = captured),
|
|
):
|
|
return asyncio.run(self.route.validate_model(request, current_subject = "test-user"))
|
|
|
|
def test_ok_when_training_inactive(self):
|
|
resp = self._validate(training_active = False, decision = (False, {}))
|
|
self.assertTrue(resp.valid)
|
|
|
|
def test_refuses_when_wont_fit(self):
|
|
info = {"required_gb": 40.0, "usable_gb": 5.0, "needed_gb": 50.0}
|
|
with self.assertRaises(HTTPException) as exc:
|
|
self._validate(training_active = True, decision = (False, info))
|
|
self.assertEqual(exc.exception.status_code, 409)
|
|
self.assertIn("training is running", exc.exception.detail)
|
|
|
|
def test_passes_real_load_settings_to_guard(self):
|
|
# validate must size with the request's settings, not hardcoded defaults.
|
|
captured = []
|
|
self._validate(
|
|
training_active = True, decision = (True, {}), captured = captured, load_in_4bit = False
|
|
)
|
|
self.assertEqual(captured[0]["load_in_4bit"], False)
|
|
self.assertEqual(captured[0]["max_seq_length"], 4096)
|
|
|
|
def test_rejects_gguf_with_gpu_ids_before_guard(self):
|
|
# /validate must mirror /load's GGUF + gpu_ids 400, before the VRAM guard.
|
|
from models.inference import ValidateModelRequest
|
|
|
|
request = ValidateModelRequest(model_path = "x.gguf", gpu_ids = [0])
|
|
cfg = SimpleNamespace(
|
|
identifier = "x.gguf",
|
|
display_name = "x",
|
|
is_gguf = True,
|
|
is_lora = False,
|
|
is_vision = False,
|
|
path = None,
|
|
base_model = None,
|
|
)
|
|
captured = []
|
|
with (
|
|
patch.object(
|
|
self.route,
|
|
"_resolve_model_identifier_for_request",
|
|
return_value = ("x.gguf", "x.gguf", False),
|
|
),
|
|
patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
|
|
patch.object(self.route, "load_inference_config", return_value = {}),
|
|
_stub_guard_deps(training_active = True, decision = (True, {}), captured = captured),
|
|
):
|
|
with self.assertRaises(HTTPException) as exc:
|
|
asyncio.run(self.route.validate_model(request, current_subject = "u"))
|
|
self.assertEqual(exc.exception.status_code, 400)
|
|
self.assertIn("gpu_ids is not supported for GGUF", exc.exception.detail)
|
|
self.assertEqual(captured, []) # guard never reached
|
|
|
|
|
|
# ── _estimate_gguf_required_gb (sizes the same weights the loader loads) ──────
|
|
|
|
|
|
class TestEstimateGgufRequiredGb(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.route = _load_inference_route()
|
|
|
|
def test_local_sums_split_shards(self):
|
|
import tempfile
|
|
with tempfile.TemporaryDirectory() as d:
|
|
p = Path(d)
|
|
(p / "model-00001-of-00002.gguf").write_bytes(b"x" * 1000)
|
|
(p / "model-00002-of-00002.gguf").write_bytes(b"y" * 2000)
|
|
cfg = SimpleNamespace(
|
|
gguf_file = str(p / "model-00001-of-00002.gguf"),
|
|
gguf_mmproj_file = None,
|
|
gguf_mtp_file = None,
|
|
gguf_hf_repo = None,
|
|
gguf_variant = None,
|
|
)
|
|
gb = self.route._estimate_gguf_required_gb(cfg)
|
|
self.assertAlmostEqual(gb, 3000 / (1024**3), places = 9) # both shards
|
|
|
|
def test_remote_threads_token_and_adds_companions(self):
|
|
import utils.models.model_config as mc
|
|
|
|
cfg = SimpleNamespace(
|
|
gguf_file = None,
|
|
gguf_mmproj_file = None,
|
|
gguf_mtp_file = None,
|
|
gguf_hf_repo = "org/repo",
|
|
gguf_variant = "Q4_K_M",
|
|
)
|
|
variant = SimpleNamespace(quant = "Q4_K_M", size_bytes = 10 * 1024**3)
|
|
captured = {}
|
|
|
|
def fake_list(repo, hf_token = None):
|
|
captured["token"] = hf_token
|
|
return ([variant], True) # has_vision -> include mmproj
|
|
|
|
with (
|
|
patch.object(mc, "list_gguf_variants", fake_list),
|
|
patch.object(
|
|
self.route, "_remote_gguf_companion_bytes", return_value = 2 * 1024**3
|
|
) as comp,
|
|
):
|
|
gb = self.route._estimate_gguf_required_gb(cfg, hf_token = "tok")
|
|
self.assertEqual(captured["token"], "tok") # token threaded for gated repos
|
|
self.assertAlmostEqual(gb, 12.0, places = 6) # 10 GB variant + 2 GB companions
|
|
self.assertTrue(comp.call_args.kwargs["include_mmproj"])
|
|
|
|
def test_remote_unknown_variant_returns_none(self):
|
|
import utils.models.model_config as mc
|
|
cfg = SimpleNamespace(
|
|
gguf_file = None,
|
|
gguf_mmproj_file = None,
|
|
gguf_mtp_file = None,
|
|
gguf_hf_repo = "org/repo",
|
|
gguf_variant = "Q8_0",
|
|
)
|
|
with patch.object(
|
|
mc,
|
|
"list_gguf_variants",
|
|
return_value = ([SimpleNamespace(quant = "Q4_K_M", size_bytes = 1)], False),
|
|
):
|
|
self.assertIsNone(self.route._estimate_gguf_required_gb(cfg))
|
|
|
|
def test_local_adds_kv_cache(self):
|
|
import tempfile
|
|
with tempfile.TemporaryDirectory() as d:
|
|
p = Path(d) / "model.gguf"
|
|
p.write_bytes(b"x" * 1000)
|
|
cfg = SimpleNamespace(
|
|
gguf_file = str(p),
|
|
gguf_mmproj_file = None,
|
|
gguf_mtp_file = None,
|
|
gguf_hf_repo = None,
|
|
gguf_variant = None,
|
|
)
|
|
with patch.object(self.route, "_estimate_gguf_kv_gb", return_value = 2.0):
|
|
gb = self.route._estimate_gguf_required_gb(cfg, max_seq_length = 8192)
|
|
self.assertAlmostEqual(gb, 1000 / (1024**3) + 2.0, places = 6) # weights + KV
|
|
|
|
def test_kv_helper_graceful_on_non_gguf(self):
|
|
import tempfile
|
|
with tempfile.TemporaryDirectory() as d:
|
|
p = Path(d) / "not-a.gguf"
|
|
p.write_bytes(b"not a gguf")
|
|
self.assertEqual(self.route._estimate_gguf_kv_gb(str(p), 4096), 0.0)
|
|
|
|
def test_kv_sizes_at_larger_of_max_seq_len_and_ctx_override(self):
|
|
# KV sized at the larger of max_seq_length and --ctx-size, else native.
|
|
seen = {}
|
|
|
|
class _FakeBackend:
|
|
_context_length = 2048
|
|
|
|
def _read_gguf_metadata(self, path):
|
|
pass
|
|
|
|
def _can_estimate_kv(self):
|
|
return True
|
|
|
|
def _estimate_kv_cache_bytes(
|
|
self,
|
|
ctx,
|
|
n_parallel = 1,
|
|
):
|
|
seen["ctx"] = ctx
|
|
seen["n_parallel"] = n_parallel
|
|
return ctx * n_parallel * (1024**2) # 1 MiB per ctx unit per slot
|
|
|
|
with patch.object(self.route, "LlamaCppBackend", _FakeBackend):
|
|
r = self.route
|
|
# --ctx-size override above max_seq_length -> override wins
|
|
self.assertAlmostEqual(
|
|
r._estimate_gguf_kv_gb("m", 4096, ["--ctx-size", "131072"]), 128.0
|
|
)
|
|
self.assertEqual(seen["ctx"], 131072)
|
|
self.assertEqual(seen["n_parallel"], 1) # default single slot
|
|
# override below max_seq_length -> larger (max_seq_length) wins
|
|
self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 4096, ["--ctx-size", "1024"]), 4.0)
|
|
self.assertEqual(seen["ctx"], 4096)
|
|
# no override, no max_seq_length -> native context fallback
|
|
self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 0, None), 2.0)
|
|
self.assertEqual(seen["ctx"], 2048)
|
|
# malformed extras are ignored (fall back to max_seq_length)
|
|
self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 4096, ["--ctx-size", "oops"]), 4.0)
|
|
# --parallel slots scale the cache the same way the launcher does
|
|
self.assertAlmostEqual(r._estimate_gguf_kv_gb("m", 4096, None, 4), 16.0)
|
|
self.assertEqual(seen["n_parallel"], 4)
|
|
|
|
|
|
# ── load_model integration: authoritative 409, and no unload before refusal ──
|
|
|
|
|
|
class TestLoadModelGuardIntegration(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
cls.route = _load_inference_route()
|
|
|
|
def test_refusal_409_and_no_unload(self):
|
|
import contextlib
|
|
from unittest.mock import MagicMock
|
|
from models.inference import LoadRequest
|
|
|
|
inf = SimpleNamespace(active_model_name = None)
|
|
inf.unload_model = MagicMock()
|
|
inf._shutdown_subprocess = MagicMock()
|
|
llama = SimpleNamespace(is_loaded = False, model_identifier = None, hf_variant = None)
|
|
llama.unload_model = MagicMock()
|
|
cfg = SimpleNamespace(is_gguf = False, is_lora = False, path = None, base_model = None)
|
|
request = LoadRequest(model_path = "unsloth/Qwen3-1.7B")
|
|
info = {"required_gb": 40.0, "usable_gb": 5.0, "needed_gb": 50.0, "mode": "auto"}
|
|
|
|
with (
|
|
patch.object(self.route, "validate_extra_args", return_value = None),
|
|
patch.object(
|
|
self.route,
|
|
"_resolve_model_identifier_for_request",
|
|
return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False),
|
|
),
|
|
patch.object(self.route, "resolve_effective_chat_template_override", return_value = None),
|
|
patch.object(self.route, "get_inference_backend", return_value = inf),
|
|
patch.object(self.route, "get_llama_cpp_backend", return_value = llama),
|
|
patch.object(self.route, "_hf_offline_if_dns_dead", lambda: contextlib.nullcontext()),
|
|
patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
|
|
_stub_guard_deps(training_active = True, decision = (False, info)),
|
|
):
|
|
with self.assertRaises(HTTPException) as exc:
|
|
asyncio.run(
|
|
self.route.load_model(request, fastapi_request = MagicMock(), current_subject = "u")
|
|
)
|
|
|
|
self.assertEqual(exc.exception.status_code, 409)
|
|
# Guard runs before the unload step, so a refused load tears down nothing.
|
|
inf.unload_model.assert_not_called()
|
|
inf._shutdown_subprocess.assert_not_called()
|
|
llama.unload_model.assert_not_called()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|