1813035d50
Follow-up review on the CPU-only / NUMA hardening: - Reserve a flat MTP fraction in the CPU context fit when MTP engages but the draft KV cannot be byte-sized (mtp_overhead_fn is None): budget_frac otherwise skips the flat reserve, so the fit could pick a context that OOMs once the draft allocates. - Bound the NUMA local-placement decision on the smallest node's free RAM, not the largest. The loader is not pinned, so first-touch may land on any node; interleave a footprint that fits only the larger node instead of gambling on placement. - Capture the first launch's scheduler abort before the no-spec MTP fallback resets the stdout tail, then fold it into the terminal memo decision, so a fallback that fails for another reason does not drop the memo and let the UI replay the load. Unit tests and the offline simulations updated and passing.
294 lines
12 KiB
Python
294 lines
12 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
|
|
|
|
"""Guards for CPU-only safe llama-server defaults in load_model.
|
|
|
|
On a CPU-only host (no discrete GPU, not Apple Metal) none of the GPU/Apple fit
|
|
branches run, so the launch used to keep GPU defaults: `--flash-attn on`, `--fit on`,
|
|
and the model's full native context. For large MLA/sparse-attention/MTP GGUFs (e.g.
|
|
GLM-5.2) `--fit`'s graph-reserve estimator aborts (llama.cpp #21932) and CPU flash
|
|
attention is unsafe. These tests pin that the launch now derives `_cpu_only` and uses
|
|
it to disable --fit and default flash-attn off, while leaving GPU/Apple behaviour and
|
|
user overrides intact. Source/AST level: load_model is too entangled to drive E2E.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
import inspect
|
|
import sys
|
|
import textwrap
|
|
import types as _types
|
|
from pathlib import Path
|
|
|
|
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
|
|
if _BACKEND_DIR not in sys.path:
|
|
sys.path.insert(0, _BACKEND_DIR)
|
|
|
|
try:
|
|
import structlog # noqa: F401
|
|
except ImportError:
|
|
_s = _types.ModuleType("structlog")
|
|
_s.get_logger = lambda *a, **k: __import__("logging").getLogger("stub")
|
|
_s.BoundLogger = type("BoundLogger", (), {})
|
|
sys.modules["structlog"] = _s
|
|
# Importing core.inference.* runs core/inference/__init__.py (orchestrator + loggers +
|
|
# httpx); stub those when absent so a dependency-light run can still collect this file.
|
|
try:
|
|
import loggers # noqa: F401
|
|
except ImportError:
|
|
_loggers_stub = _types.ModuleType("loggers")
|
|
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
|
|
sys.modules["loggers"] = _loggers_stub
|
|
try:
|
|
import httpx # noqa: F401
|
|
except ImportError:
|
|
_httpx_stub = _types.ModuleType("httpx")
|
|
for _exc in (
|
|
"ConnectError",
|
|
"TimeoutException",
|
|
"ReadTimeout",
|
|
"ReadError",
|
|
"RemoteProtocolError",
|
|
"CloseError",
|
|
"HTTPError",
|
|
"RequestError",
|
|
):
|
|
setattr(_httpx_stub, _exc, type(_exc, (Exception,), {}))
|
|
_httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None})
|
|
_httpx_stub.Response = type("Response", (), {})
|
|
_httpx_stub.Client = type(
|
|
"C",
|
|
(),
|
|
{
|
|
"__init__": lambda s, **kw: None,
|
|
"__enter__": lambda s: s,
|
|
"__exit__": lambda s, *a: None,
|
|
},
|
|
)
|
|
sys.modules["httpx"] = _httpx_stub
|
|
|
|
from core.inference.llama_cpp import LlamaCppBackend # noqa: E402
|
|
|
|
|
|
def _load_model_src() -> str:
|
|
return textwrap.dedent(inspect.getsource(LlamaCppBackend.load_model))
|
|
|
|
|
|
def test_cpu_only_flag_derived_from_no_gpu_and_not_apple():
|
|
src = _load_model_src()
|
|
assert "_cpu_only = (not gpus) and not _is_apple_silicon()" in src
|
|
|
|
|
|
def test_flash_attn_default_is_off_on_cpu_only():
|
|
"""The base cmd must use a computed flash default, off for CPU-only."""
|
|
src = _load_model_src()
|
|
assert '_flash_default = "off" if _cpu_only else "on"' in src
|
|
# And the base cmd must NOT hardcode flash-attn on anymore.
|
|
assert '"on", # Force flash attention for speed' not in src
|
|
# The --flash-attn argument in the base cmd list is the computed default.
|
|
assert "_flash_default, # CPU-only" in src
|
|
|
|
|
|
def test_fit_disabled_on_cpu_only():
|
|
src = _load_model_src()
|
|
assert "if _cpu_only and use_fit:" in src
|
|
fn = ast.parse(src).body[0]
|
|
# There is an `if _cpu_only and use_fit:` whose body sets use_fit = False.
|
|
found = False
|
|
for node in ast.walk(fn):
|
|
if (
|
|
isinstance(node, ast.If)
|
|
and isinstance(node.test, ast.BoolOp)
|
|
and any(isinstance(v, ast.Name) and v.id == "_cpu_only" for v in ast.walk(node.test))
|
|
):
|
|
for n in node.body:
|
|
if (
|
|
isinstance(n, ast.Assign)
|
|
and any(isinstance(t, ast.Name) and t.id == "use_fit" for t in n.targets)
|
|
and isinstance(n.value, ast.Constant)
|
|
and n.value.value is False
|
|
):
|
|
found = True
|
|
assert found, "expected `if _cpu_only and use_fit: ... use_fit = False`"
|
|
|
|
|
|
def test_gpu_path_still_emits_fit_on():
|
|
"""Backwards compat: the GPU branch still emits --fit on (only CPU changes)."""
|
|
src = _load_model_src()
|
|
assert 'cmd.extend(["--fit", "on"])' in src
|
|
|
|
|
|
def test_cpu_only_emits_explicit_fit_off():
|
|
"""--fit defaults to on in llama.cpp, so CPU-only must pass --fit off explicitly,
|
|
not just skip --fit on (PR review fix)."""
|
|
src = _load_model_src()
|
|
assert "elif _cpu_only:" in src
|
|
assert 'cmd.extend(["--fit", "off"])' in src
|
|
|
|
|
|
# ---- Phase 3: CPU context cap + RAM preflight ------------------------------
|
|
|
|
|
|
def test_cpu_context_ceiling_constant_is_sane():
|
|
from core.inference import llama_cpp as m
|
|
|
|
assert hasattr(m, "_CPU_CTX_AUTO_CEILING")
|
|
# A safe chat ceiling, far below a model's million-token native context.
|
|
assert 4096 <= m._CPU_CTX_AUTO_CEILING <= 131072
|
|
assert 0.5 < m._CPU_RAM_BUDGET_FRAC <= 1.0
|
|
|
|
|
|
def test_cpu_only_caps_auto_context_but_honors_explicit():
|
|
src = _load_model_src()
|
|
# The fit runs for any auto context (requested_ctx <= 0), against a ceiling of
|
|
# min(native, 32k); an explicit -c (requested_ctx > 0) is honored untouched.
|
|
assert "requested_ctx <= 0 and effective_ctx > 0" in src
|
|
assert "_ctx_ceiling = min(effective_ctx, _CPU_CTX_AUTO_CEILING)" in src
|
|
assert "effective_ctx = _cpu_cap" in src
|
|
|
|
|
|
def test_cpu_context_fit_runs_below_ceiling_too():
|
|
"""A large GGUF with a native context already <= 32k must still be RAM-fit, not
|
|
skipped, so it can be reduced toward 4096 instead of OS-killed (PR review fix)."""
|
|
src = _load_model_src()
|
|
# The gate is `> 0`, not `> _CPU_CTX_AUTO_CEILING`, and the fit ceiling is clamped.
|
|
assert "effective_ctx > _CPU_CTX_AUTO_CEILING" not in src
|
|
assert "requested_ctx = _ctx_ceiling" in src
|
|
|
|
|
|
def _nows(s: str) -> str:
|
|
"""Whitespace-stripped source, so assertions survive the formatter wrapping a line."""
|
|
return "".join(s.split())
|
|
|
|
|
|
def test_cpu_context_fit_accounts_for_mtp():
|
|
"""mtp_engaged alone is a no-op once budget_frac is set, so the CPU fit must pass the
|
|
byte-accurate MTP overhead fn to actually reserve MTP KV (PR review fix)."""
|
|
src = _nows(_load_model_src())
|
|
assert _nows("mtp_overhead_fn = (_mtp_bytes if _mtp_will_engage_cpu else None)") in src
|
|
|
|
|
|
def test_cpu_context_cap_reuses_fit_helper_against_ram():
|
|
"""The cap reuses _fit_context_to_vram with the system-RAM budget, not VRAM."""
|
|
src = _load_model_src()
|
|
assert "_available_system_memory_mib()" in src
|
|
assert "self._fit_context_to_vram(" in src
|
|
assert "_cpu_budget = _CPU_RAM_BUDGET_FRAC" in src
|
|
assert "budget_frac = _cpu_budget" in src
|
|
|
|
|
|
def test_cpu_ram_preflight_warns_when_weights_exceed_ram():
|
|
src = _load_model_src()
|
|
assert "CPU-only memory preflight" in src
|
|
|
|
|
|
def test_cpu_context_floors_to_min_when_weights_exceed_budget():
|
|
"""When the fixed footprint exceeds the RAM budget, _fit_context_to_vram returns
|
|
the ceiling unchanged; the cap must floor to the minimum instead (PR review fix)."""
|
|
src = _load_model_src()
|
|
# The check uses the fitted footprint (weights + compute buffer), not raw weights.
|
|
assert "_fixed = model_size_fit or model_size" in src
|
|
assert "_fixed >= _budget_b" in src
|
|
assert "_cpu_cap = 4096" in src
|
|
|
|
|
|
def test_cpu_context_fit_uses_fitted_footprint():
|
|
"""The CPU RAM fit must pass the fitted footprint (weights + compute buffer), so a
|
|
context that only fits when the buffer is ignored can't slip through (PR review fix)."""
|
|
src = _load_model_src()
|
|
assert "model_size_bytes = _fixed" in src
|
|
|
|
|
|
def test_numa_decision_uses_footprint_not_just_weights():
|
|
"""The NUMA interleave decision must use the full resident footprint (weights +
|
|
compute buffer + KV at the launched parallel slots + MTP reserve), so a model whose
|
|
weights fit one node but whose footprint does not still interleaves (PR review fixes)."""
|
|
src = _load_model_src()
|
|
assert "_numa_footprint" in src
|
|
assert "decide_interleave(_numa_footprint" in src
|
|
# Footprint = fitted weights (incl. compute buffer) + KV + MTP reserve.
|
|
assert "_resident = model_size_fit or model_size" in src
|
|
# MTP is recomputed at the post-cap context, not the stale pre-cap reserve.
|
|
assert _nows("_numa_mtp = _mtp_bytes(effective_ctx) if _mtp_will_engage_cpu else 0") in _nows(
|
|
src
|
|
)
|
|
# KV must be sized for the launched --parallel slots, not the n_parallel=1 default.
|
|
assert "effective_ctx, cache_type_kv, n_parallel = n_parallel" in src
|
|
|
|
|
|
def test_numa_surfaces_total_ram_failure():
|
|
"""When the footprint exceeds total RAM across all nodes, decide_interleave returns
|
|
an actionable 'interleave cannot help' reason; the caller must surface it, not only
|
|
the missing-numactl case (PR review fix)."""
|
|
src = _load_model_src()
|
|
assert '"interleave cannot help" in _numa.reason' in src
|
|
|
|
|
|
def test_explicit_user_numa_skips_auto_interleave_prefix():
|
|
"""An explicit user --numa must skip the numactl argv prefix (which user extra args
|
|
can't override), not just the --numa distribute flag (PR review fix)."""
|
|
src = _load_model_src()
|
|
assert 'if _numa.interleave and _extra_args_set_any_flag(extra_args, {"--numa"}):' in src
|
|
assert "leaving auto-interleave off" in src
|
|
|
|
|
|
def test_extra_args_forces_cpu_offload_helper():
|
|
"""The CPU-force detector: zero GPU layers or --device none (PR review fixes)."""
|
|
from core.inference.llama_cpp import _extra_args_forces_cpu_offload as f
|
|
|
|
assert f(["-ngl", "0"])
|
|
assert f(["--n-gpu-layers", "0"])
|
|
assert f(["--gpu-layers", "0"])
|
|
assert f(["-ngl=0"])
|
|
assert not f(["-ngl", "99"])
|
|
assert not f([])
|
|
assert not f(None)
|
|
assert not f(["--flash-attn", "on"])
|
|
# Each flag's last occurrence wins, matching llama-server's own parsing.
|
|
assert f(["-ngl", "99", "-ngl", "0"])
|
|
assert not f(["-ngl", "0", "-ngl", "99"])
|
|
# --device/-dev none also forces CPU, independently of -ngl.
|
|
assert f(["--device", "none"])
|
|
assert f(["-dev", "none"])
|
|
assert f(["--device=none"])
|
|
assert not f(["--device", "CUDA0"])
|
|
# The two controls are independent: -ngl 0 stays CPU even with a device named.
|
|
assert f(["-ngl", "0", "--device", "CUDA0"])
|
|
assert f(["--device", "none", "-ngl", "99"])
|
|
|
|
|
|
def test_cpu_cap_lowers_advertised_ceiling():
|
|
"""When the CPU cap reduces the launched context, max_available_ctx must drop too,
|
|
so /status and the UI safe-zone reflect the real window, not native (PR review fix)."""
|
|
src = _load_model_src()
|
|
assert "max_available_ctx = min(max_available_ctx, _cpu_cap)" in src
|
|
|
|
|
|
def test_cpu_fit_skips_mtp_reserve_when_mla_auto_drops():
|
|
"""Auto drops embedded MTP for MLA models, so the CPU cap / NUMA footprint must not
|
|
reserve a target-KV copy for a drafter that won't launch (PR review fix)."""
|
|
src = _nows(_load_model_src())
|
|
assert _nows("_mtp_will_engage_cpu = _mtp_will_engage and not (") in src
|
|
assert _nows("not _mla_mtp_auto_enabled()") in src
|
|
# The CPU cap and NUMA recompute use the gated flag, not the raw _mtp_will_engage.
|
|
assert _nows("mtp_overhead_fn = (_mtp_bytes if _mtp_will_engage_cpu else None)") in src
|
|
assert _nows("_numa_mtp = _mtp_bytes(effective_ctx) if _mtp_will_engage_cpu else 0") in src
|
|
|
|
|
|
def test_cpu_fit_reserves_flat_mtp_when_draft_unsized():
|
|
"""When MTP engages but the draft KV can't be byte-sized (mtp_overhead_fn is None),
|
|
budget_frac skips the flat reserve, so the CPU budget is trimmed to still hold MTP
|
|
RAM back instead of fitting a context that OOMs once the draft allocates (PR review)."""
|
|
src = _load_model_src()
|
|
assert "if _mtp_will_engage_cpu and mtp_overhead_fn is None:" in src
|
|
assert "_cpu_budget -= _MTP_VRAM_RESERVE_FRAC" in src
|
|
assert "budget_frac = _cpu_budget" in src
|
|
|
|
|
|
def test_zero_offload_folds_into_cpu_only():
|
|
"""A visible GPU plus a user -ngl 0 must be treated as CPU-only: the GPU list is
|
|
dropped before _cpu_only is computed so the CPU safe defaults apply (PR review fix)."""
|
|
src = _load_model_src()
|
|
assert "_extra_args_forces_cpu_offload(extra_args)" in src
|
|
assert "gpus, total_by_idx = [], {}" in src
|