ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: address review feedback on the CPU-only / NUMA hardening

Browse source 
- Emit --fit off explicitly on CPU. The flag defaults to on in recent llama.cpp, so
  setting use_fit = False only skipped --fit on and still ran the graph-reserve
  fitting step (same abort path) on the first attempt.
- Key the scheduler-abort memo per variant (model plus hf_variant / gguf_path) so a
  failed quant does not block trying a different quant of the same repo, which the
  error message itself suggests.
- Move the memo fail-fast guard ahead of _kill_process so a known-bad reload does not
  tear down a working server before returning the error.
- Floor the CPU context to the minimum when weights alone exceed the RAM budget, where
  the fit helper returns the requested ceiling unchanged and KV could OS-kill the load.
- Decide NUMA interleave on the resident footprint (weights plus KV at the capped
  context), not weights alone, so a model whose weights fit a node but whose footprint
  does not still interleaves.

Adds tests for each fix.
This commit is contained in:
Daniel Han 2026-06-29 05:56:36 +00:00
commit 8ebaf955da
4 changed files with 117 additions and 42 deletions
Download patch file Download diff file Expand all files Collapse all files

72
studio/backend/core/inference/llama_cpp.py
View file

@ -4709,18 +4709,19 @@ class LlamaCppBackend:
self._cancel_event.clear()
# ── Phase 1: kill old process (under lock, fast) ──────────
with self._lock:
self._kill_process()
# Resolve llama-server now but defer a not-found error: a block-diffusion
# GGUF uses the diffusion runner, and its arch is only known after the header.
binary = self._find_llama_server_binary()
# Fail fast if this (binary, model) already aborted in the graph scheduler
# this session: reloading just re-reads the weights into the same crash
# (a startup crash 500s and the UI replays /load).
if LlamaCppBackend._sched_reserve_aborts(binary, model_identifier):
# Fail fast BEFORE killing the live server (so a known-bad reload is
# non-destructive) if this (binary, model, variant) already aborted in the
# graph scheduler this session: reloading re-reads the weights into the same
# crash (a startup crash 500s and the UI replays /load). Keyed per variant so
# a failed quant does not block trying a different quant in the same repo.
_abort_memo_model = "\x00".join(
[model_identifier or "", hf_variant or "", gguf_path or ""]
)
if LlamaCppBackend._sched_reserve_aborts(binary, _abort_memo_model):
logger.warning(
"Skipping reload of '%s': it already aborted in the llama.cpp "
"graph scheduler this session (unsupported op on this backend).",
@ -4728,6 +4729,10 @@ class LlamaCppBackend:
)
raise RuntimeError(self._sched_reserve_abort_message())
# ── Phase 1: kill old process (under lock, fast) ──────────
with self._lock:
self._kill_process()
# ── Phase 2: download (NO lock held, so cancel can proceed) ──
# mtp_draft_path arrives set for local Gemma loads (detected
# sibling); for -hf loads it's None here and resolved just below.
@ -5689,18 +5694,25 @@ class LlamaCppBackend:
_cpu_cap = _CPU_CTX_AUTO_CEILING
try:
if _avail_mib and model_size and self._can_estimate_kv():
_fit = self._fit_context_to_vram(
requested_ctx = _CPU_CTX_AUTO_CEILING,
available_mib = _avail_mib,
model_size_bytes = model_size,
cache_type_kv = cache_type_kv,
min_ctx = 4096,
n_parallel = n_parallel,
kv_on_gpu = True, # KV lives in the RAM budget we fit
mtp_engaged = True, # flat reserve; no GPU draft here
budget_frac = _CPU_RAM_BUDGET_FRAC,
)
_cpu_cap = max(4096, min(_CPU_CTX_AUTO_CEILING, _fit))
_budget_b = _avail_mib * _CPU_RAM_BUDGET_FRAC * 1024 * 1024
if model_size >= _budget_b:
# Weights alone over budget: _fit_context_to_vram returns
# the ceiling unchanged, but KV at 32k could OOM. Floor to
# the minimum so the tightest fit gets the smallest context.
_cpu_cap = 4096
else:
_fit = self._fit_context_to_vram(
requested_ctx = _CPU_CTX_AUTO_CEILING,
available_mib = _avail_mib,
model_size_bytes = model_size,
cache_type_kv = cache_type_kv,
min_ctx = 4096,
n_parallel = n_parallel,
kv_on_gpu = True, # KV lives in the RAM budget we fit
mtp_engaged = True, # flat reserve; no GPU draft here
budget_frac = _CPU_RAM_BUDGET_FRAC,
)
_cpu_cap = max(4096, min(_CPU_CTX_AUTO_CEILING, _fit))
except Exception as _cap_exc: # best-effort; fall back to ceiling
logger.debug("CPU context-fit failed; using ceiling: %s", _cap_exc)
if _cpu_cap < effective_ctx:
@ -5743,6 +5755,10 @@ class LlamaCppBackend:
# Fits on selected GPU(s) -- offload all layers
cmd.extend(["-ngl", "-1"])
fully_gpu_offloaded = True
elif _cpu_only:
# --fit defaults to on in recent llama.cpp, so omitting it still runs
# the graph-reserve fitting step (same abort path); disable explicitly.
cmd.extend(["--fit", "off"])
server_caps = self.probe_server_capabilities(binary)
# Expose Prometheus /metrics for the engine-stats logger, only
@ -5966,7 +5982,19 @@ class LlamaCppBackend:
self._numa_prefix = []
try:
from core.inference.numa import decide_interleave
_numa = decide_interleave(model_size, cpu_only = _cpu_only)
# Decide on the resident footprint (weights + KV at the capped
# context), not weights alone, so a model whose weights fit one node
# but whose footprint does not still interleaves.
_numa_footprint = model_size
if model_size and effective_ctx > 0 and self._can_estimate_kv():
try:
_numa_footprint = model_size + self._estimate_kv_cache_bytes(
effective_ctx, cache_type_kv
)
except Exception:
_numa_footprint = model_size
_numa = decide_interleave(_numa_footprint, cpu_only = _cpu_only)
if _numa.interleave:
self._numa_prefix = list(_numa.prefix)
if not _extra_args_set_any_flag(extra_args, {"--numa"}):
@ -6367,7 +6395,7 @@ class LlamaCppBackend:
and (self._is_signal_crash(_crash_rc) or self._is_abort_exit(_crash_rc))
and self._is_sched_reserve_abort("\n".join(self._stdout_lines[-200:]))
):
LlamaCppBackend._record_sched_reserve_abort(binary, model_identifier)
LlamaCppBackend._record_sched_reserve_abort(binary, _abort_memo_model)
self._kill_process()
# The #6415 split-axis abort is latched earlier (first spawn).
# Skip if a cancel/unload is pending (mirrors the MTP guard).

6
studio/backend/core/inference/numa.py
View file

@ -110,9 +110,11 @@ def decide_interleave(
topology: NumaTopology | None = None,
has_numactl: bool | None = None,
) -> InterleaveDecision:
"""Interleave only when CPU-only, multi-node, and the model exceeds the largest
"""Interleave only when CPU-only, multi-node, and the footprint exceeds the largest
node's free RAM but fits across all nodes; otherwise leave placement local.
model_size_bytes (GGUF weights) is the conservative proxy for the footprint."""
model_size_bytes should be the resident footprint (weights + KV), not weights
alone, so a model whose weights fit a node but whose footprint does not still
interleaves."""
if not cpu_only:
return InterleaveDecision(False, "not cpu-only; leaving NUMA placement to the OS")
if not model_size_bytes or model_size_bytes <= 0:

25
studio/backend/tests/test_cpu_only_defaults.py
View file

@ -84,6 +84,14 @@ def test_gpu_path_still_emits_fit_on():
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 ------------------------------
@ -115,3 +123,20 @@ def test_cpu_context_cap_reuses_fit_helper_against_ram():
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 weights alone exceed 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()
assert "model_size >= _budget_b" in src
assert "_cpu_cap = 4096" in src
def test_numa_decision_uses_footprint_not_just_weights():
"""The NUMA interleave decision must use weights + KV, so a model whose weights fit
one node but whose footprint does not still interleaves (PR review fix)."""
src = _load_model_src()
assert "_numa_footprint" in src
assert "_estimate_kv_cache_bytes(" in src
assert "decide_interleave(_numa_footprint" in src

56
studio/backend/tests/test_sched_reserve_abort.py
View file

@ -207,36 +207,56 @@ def _load_model_src() -> str:
return textwrap.dedent(inspect.getsource(LlamaCppBackend.load_model))
def test_load_model_fails_fast_on_memoed_abort():
"""load_model must consult the memo and raise the actionable message before the
download/spawn, so a replayed /load doesn't re-read the weights."""
src = _load_model_src()
assert "_sched_reserve_aborts(binary, model_identifier)" in src
assert "_sched_reserve_abort_message()" in src
# The guard must precede the model download (the expensive reload it prevents).
fn = ast.parse(src).body[0]
guard_line = next(
node.lineno
for node in ast.walk(fn)
if isinstance(node, ast.Call)
and isinstance(node.func, ast.Attribute)
and node.func.attr == "_sched_reserve_aborts"
)
download_line = next(
def _call_line(fn, attr):
"""First line where load_model calls method `attr`, or None."""
return next(
(
node.lineno
for node in ast.walk(fn)
if isinstance(node, ast.Call)
and isinstance(node.func, ast.Attribute)
and node.func.attr == "_download_gguf"
and node.func.attr == attr
),
None,
)
def test_load_model_fails_fast_on_memoed_abort():
"""load_model must consult the memo and raise before the download/spawn, so a
replayed /load doesn't re-read the weights."""
src = _load_model_src()
assert "_sched_reserve_aborts(binary, _abort_memo_model)" in src
assert "_sched_reserve_abort_message()" in src
fn = ast.parse(src).body[0]
guard_line = _call_line(fn, "_sched_reserve_aborts")
download_line = _call_line(fn, "_download_gguf")
assert download_line is None or guard_line < download_line
def test_failfast_guard_runs_before_killing_the_live_server():
"""The memo guard must precede _kill_process so a known-bad reload does not tear
down a working server (PR review fix)."""
fn = ast.parse(_load_model_src()).body[0]
guard_line = _call_line(fn, "_sched_reserve_aborts")
kill_line = _call_line(fn, "_kill_process")
assert guard_line is not None and kill_line is not None
assert guard_line < kill_line
def test_abort_memo_key_includes_variant():
"""The memo key must include hf_variant / gguf_path so a failed quant does not
block a different quant of the same repo (PR review fix)."""
src = _load_model_src()
assert "_abort_memo_model" in src
assert "hf_variant" in src and "gguf_path" in src
# Sanity: distinct variants produce distinct keys for the same model.
def key(model, variant, gguf):
return "\x00".join([model or "", variant or "", gguf or ""])
assert key("repo", "UD-Q6_K", None) != key("repo", "UD-Q4_K_XL", None)
def test_load_model_records_abort_on_crash():
"""On a startup crash matching the signature, load_model must record the memo."""
src = _load_model_src()
assert "_record_sched_reserve_abort(binary, model_identifier)" in src
assert "_record_sched_reserve_abort(binary, _abort_memo_model)" in src
assert "_is_sched_reserve_abort(" in src