Files
unslothai--unsloth/studio/backend/tests/test_chat_load_during_training.py
T
wehub-resource-sync 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
chore: import upstream snapshot with attribution
2026-07-13 12:59:56 +08:00

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()