diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 433502d43d..b37d151dac 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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). diff --git a/studio/backend/core/inference/numa.py b/studio/backend/core/inference/numa.py index 9c9728fc59..ba41331726 100644 --- a/studio/backend/core/inference/numa.py +++ b/studio/backend/core/inference/numa.py @@ -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: diff --git a/studio/backend/tests/test_cpu_only_defaults.py b/studio/backend/tests/test_cpu_only_defaults.py index 3bd7021875..98dfd9930c 100644 --- a/studio/backend/tests/test_cpu_only_defaults.py +++ b/studio/backend/tests/test_cpu_only_defaults.py @@ -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 diff --git a/studio/backend/tests/test_sched_reserve_abort.py b/studio/backend/tests/test_sched_reserve_abort.py index 9f038c5c9e..a5bcbe17df 100644 --- a/studio/backend/tests/test_sched_reserve_abort.py +++ b/studio/backend/tests/test_sched_reserve_abort.py @@ -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