aa9847d337
Follow-up review on the CPU-only / NUMA hardening: - Restrict the NUMA topology to the cpuset-allowed nodes (/proc/self/status Mems_allowed_list). numactl --interleave=all only spans those, so under a Docker/systemd cpuset that exposes the host's nodes but allows one, the decision no longer sums RAM the child cannot use and wrap a launch that OOMs. - Honor env-forced CPU offload: _extra_args_forces_cpu_offload now also reads the inherited LLAMA_ARG_N_GPU_LAYERS / LLAMA_ARG_DEVICE the child would apply, so a GPU host running with LLAMA_ARG_N_GPU_LAYERS=0 or LLAMA_ARG_DEVICE=none gets the CPU-only safe defaults. CLI extras still win over the env. Unit tests and the offline simulations updated and passing.
304 lines
13 KiB
Python
304 lines
13 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, via CLI or the inherited
|
|
LLAMA_ARG_* env (PR review fixes)."""
|
|
from core.inference.llama_cpp import _extra_args_forces_cpu_offload as f
|
|
|
|
E: dict = {} # explicit empty env so cases ignore the ambient environment
|
|
assert f(["-ngl", "0"], env = E)
|
|
assert f(["--n-gpu-layers", "0"], env = E)
|
|
assert f(["--gpu-layers", "0"], env = E)
|
|
assert f(["-ngl=0"], env = E)
|
|
assert not f(["-ngl", "99"], env = E)
|
|
assert not f([], env = E)
|
|
assert not f(None, env = E)
|
|
assert not f(["--flash-attn", "on"], env = E)
|
|
# Each flag's last occurrence wins, matching llama-server's own parsing.
|
|
assert f(["-ngl", "99", "-ngl", "0"], env = E)
|
|
assert not f(["-ngl", "0", "-ngl", "99"], env = E)
|
|
# --device/-dev none also forces CPU, independently of -ngl.
|
|
assert f(["--device", "none"], env = E)
|
|
assert f(["-dev", "none"], env = E)
|
|
assert f(["--device=none"], env = E)
|
|
assert not f(["--device", "CUDA0"], env = E)
|
|
# The two controls are independent: -ngl 0 stays CPU even with a device named.
|
|
assert f(["-ngl", "0", "--device", "CUDA0"], env = E)
|
|
assert f(["--device", "none", "-ngl", "99"], env = E)
|
|
# Inherited env forces CPU when the CLI does not set the control.
|
|
assert f([], env = {"LLAMA_ARG_N_GPU_LAYERS": "0"})
|
|
assert f([], env = {"LLAMA_ARG_DEVICE": "none"})
|
|
assert not f([], env = {"LLAMA_ARG_N_GPU_LAYERS": "99"})
|
|
assert not f([], env = {"LLAMA_ARG_DEVICE": "CUDA0"})
|
|
# CLI wins over env.
|
|
assert not f(["-ngl", "99"], env = {"LLAMA_ARG_N_GPU_LAYERS": "0"})
|
|
assert f(["-ngl", "0"], env = {"LLAMA_ARG_N_GPU_LAYERS": "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
|