diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index cec37609ff..d72f261e7d 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -5754,6 +5754,12 @@ class LlamaCppBackend: # the minimum so the tightest fit gets the smallest context. _cpu_cap = 4096 else: + # MTP engages but the draft KV can't be byte-sized: + # _mtp_bytes returns 0 and budget_frac skips the flat + # reserve, so trim the budget to still hold back MTP RAM. + _cpu_budget = _CPU_RAM_BUDGET_FRAC + if _mtp_will_engage_cpu and mtp_overhead_fn is None: + _cpu_budget -= _MTP_VRAM_RESERVE_FRAC _fit = self._fit_context_to_vram( requested_ctx = _ctx_ceiling, available_mib = _avail_mib, @@ -5768,7 +5774,7 @@ class LlamaCppBackend: mtp_overhead_fn = ( _mtp_bytes if _mtp_will_engage_cpu else None ), - budget_frac = _CPU_RAM_BUDGET_FRAC, + budget_frac = _cpu_budget, ) _cpu_cap = max(4096, min(_ctx_ceiling, _fit)) except Exception as _cap_exc: # best-effort; fall back to ceiling @@ -6308,6 +6314,9 @@ class LlamaCppBackend: ) healthy = _spawn_and_wait(cmd) + # A scheduler abort on the first launch must still be memoed even if a later + # no-spec/text-only fallback overwrites the stdout tail; track it across them. + _pre_fallback_sched_abort = False # #6415 split-mode tensor warmup abort. Latch it on THIS first spawn: # the flash-attn-off retry below can't run tensor (needs flash_attn), # so its output drops the marker and recording later would miss it, @@ -6407,6 +6416,12 @@ class LlamaCppBackend: # cancel check stops an /unload-killed attempt respawning. A # decode-probe failure above also routes here. if not healthy and _spec_requested_mtp and not self._cancel_event.is_set(): + # The no-spec retry below resets the stdout tail; capture a scheduler + # abort from this first launch now so it can still be memoed if the + # retries then fail for another reason (else the UI replays the load). + _pre_fallback_sched_abort = _pre_fallback_sched_abort or ( + self._is_sched_reserve_abort("\n".join(self._stdout_lines[-200:])) + ) # Blame the binary only when the output shows MTP itself # failing (unknown arch / draft or context build); an # unrelated crash (e.g. OOM) gets a neutral message. @@ -6471,10 +6486,14 @@ class LlamaCppBackend: # mmproj fallback is ruled out, so a VLM that recovers text-only is not # blocked by the fail-fast guard on its next load. Wider slice as the # GGML_ASSERT line can scroll past the [New LWP] dump. - _was_sched_abort = ( - not self._cancel_event.is_set() - 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:])) + _was_sched_abort = not self._cancel_event.is_set() and ( + # A scheduler abort captured before the no-spec fallback reset stdout, + _pre_fallback_sched_abort + # or one still visible in the current tail. + or ( + (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:])) + ) ) self._kill_process() # The #6415 split-axis abort is latched earlier (first spawn). diff --git a/studio/backend/core/inference/numa.py b/studio/backend/core/inference/numa.py index d53fafc5a4..dfdae1e02e 100644 --- a/studio/backend/core/inference/numa.py +++ b/studio/backend/core/inference/numa.py @@ -36,6 +36,10 @@ class NumaTopology: def largest_node_free_mib(self) -> int: return max(self.node_free_mib.values(), default = 0) + @property + def smallest_node_free_mib(self) -> int: + return min(self.node_free_mib.values(), default = 0) + def _parse_online(spec: str) -> list[int]: """Parse a sysfs cpulist-style range, e.g. '0-1' or '0,2-3', into node ids.""" @@ -110,10 +114,11 @@ def decide_interleave( topology: NumaTopology | None = None, has_numactl: bool | None = None, ) -> InterleaveDecision: - """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 should be the resident footprint (weights + KV), not weights - alone, so a model whose weights fit a node but whose footprint does not still + """Interleave only when CPU-only, multi-node, and the footprint exceeds the smallest + node's free RAM but fits across all nodes; otherwise leave placement local. The + smallest node is the bound because the loader is not pinned, so first-touch may land + on any node. 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") @@ -125,14 +130,17 @@ def decide_interleave( return InterleaveDecision(False, "single NUMA node; interleave not needed") model_mib = model_size_bytes // (1024 * 1024) - largest = topo.largest_node_free_mib + smallest = topo.smallest_node_free_mib total = topo.total_free_mib - if model_mib <= largest: + # Keep local only when the footprint fits EVERY node (the smallest). We don't bind + # the loader, so first-touch may land on any node; local placement is safe only when + # any node can hold it. A footprint that fits only the larger node still interleaves. + if model_mib <= smallest: return InterleaveDecision( False, - f"model ~{model_mib} MiB fits the largest node's free RAM " - f"(~{largest} MiB); keeping local placement", + f"model ~{model_mib} MiB fits every node's free RAM " + f"(smallest ~{smallest} MiB); keeping local placement", ) # Impossible across all nodes regardless of numactl: surface the smaller-quant / @@ -150,16 +158,16 @@ def decide_interleave( # Needed but unavailable: surface it; caller decides whether to block. return InterleaveDecision( False, - f"model ~{model_mib} MiB exceeds the largest NUMA node's free RAM " - f"(~{largest} MiB) and needs interleaving across {topo.node_count} nodes, " + f"model ~{model_mib} MiB exceeds the smallest NUMA node's free RAM " + f"(~{smallest} MiB) and needs interleaving across {topo.node_count} nodes, " f"but `numactl` is not installed. Install numactl (e.g. `apt install " f"numactl`) or the model may fail to fit a single node.", ) return InterleaveDecision( True, - f"model ~{model_mib} MiB exceeds the largest NUMA node's free RAM " - f"(~{largest} MiB) but fits across {topo.node_count} nodes (~{total} MiB total " + f"model ~{model_mib} MiB exceeds the smallest NUMA node's free RAM " + f"(~{smallest} MiB) but fits across {topo.node_count} nodes (~{total} MiB total " f"free); wrapping with numactl --interleave=all", prefix = ("numactl", "--interleave=all"), ) diff --git a/studio/backend/tests/test_cpu_only_defaults.py b/studio/backend/tests/test_cpu_only_defaults.py index c3035e69ea..8c548353d4 100644 --- a/studio/backend/tests/test_cpu_only_defaults.py +++ b/studio/backend/tests/test_cpu_only_defaults.py @@ -174,7 +174,8 @@ def test_cpu_context_cap_reuses_fit_helper_against_ram(): src = _load_model_src() assert "_available_system_memory_mib()" in src assert "self._fit_context_to_vram(" in src - assert "budget_frac = _CPU_RAM_BUDGET_FRAC" 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(): @@ -275,6 +276,16 @@ def test_cpu_fit_skips_mtp_reserve_when_mla_auto_drops(): 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).""" diff --git a/studio/backend/tests/test_numa_interleave.py b/studio/backend/tests/test_numa_interleave.py index 179504badf..8cd4a3770c 100644 --- a/studio/backend/tests/test_numa_interleave.py +++ b/studio/backend/tests/test_numa_interleave.py @@ -87,6 +87,7 @@ def test_parse_online_ranges(): def test_topology_aggregates(): assert _USER_TOPO.node_count == 2 assert _USER_TOPO.largest_node_free_mib == 465594 + assert _USER_TOPO.smallest_node_free_mib == 223814 assert _USER_TOPO.total_free_mib == 465594 + 223814 @@ -98,12 +99,20 @@ def test_interleaves_when_model_exceeds_largest_node_but_fits_across(): assert "interleave=all" in d.reason -def test_no_interleave_when_model_fits_largest_node(): - # A 200 GB model fits node 0's 465 GB free -> keep local placement. +def test_no_interleave_when_model_fits_every_node(): + # A 200 GB model fits even the smaller node (223 GB) -> safe on any node, keep local. d = decide_interleave(200 * _GiB, cpu_only = True, topology = _USER_TOPO, has_numactl = True) assert d.interleave is False assert d.prefix == () - assert "fits" in d.reason + assert "fits every node" in d.reason + + +def test_interleaves_when_fits_larger_node_but_not_smaller(): + # 300 GB fits node 0 (465) but not node 1 (223); the loader is not bound, so first- + # touch could land on node 1. Interleave instead of gambling on placement (PR review). + d = decide_interleave(300 * _GiB, cpu_only = True, topology = _USER_TOPO, has_numactl = True) + assert d.interleave is True + assert d.prefix == ("numactl", "--interleave=all") def test_no_interleave_on_gpu_host(): diff --git a/studio/backend/tests/test_sched_reserve_abort.py b/studio/backend/tests/test_sched_reserve_abort.py index 9e031ba6b7..6ed98c2adf 100644 --- a/studio/backend/tests/test_sched_reserve_abort.py +++ b/studio/backend/tests/test_sched_reserve_abort.py @@ -281,7 +281,7 @@ def test_abort_memo_deferred_until_mmproj_fallback_ruled_out(): raise, after the text-only mmproj fallback is ruled out, so a VLM that recovers text-only is not blocked by the fail-fast guard next time (PR review fix).""" src = _load_model_src() - assert "_was_sched_abort = (" in src + assert "_was_sched_abort = " in src assert "if _was_sched_abort:" in src fn = ast.parse(src).body[0] strip_line = _call_line(fn, "_strip_mmproj_args") @@ -289,3 +289,37 @@ def test_abort_memo_deferred_until_mmproj_fallback_ruled_out(): # Recording happens after the projector strip, i.e. only once the fallback is tried. assert strip_line is not None and record_line is not None assert record_line > strip_line + + +def test_sched_abort_captured_before_mtp_fallback(): + """The no-spec MTP fallback resets the stdout tail, so the first launch's scheduler + abort must be captured before it runs and folded into the terminal decision; else a + differently-failing fallback drops the memo and the UI replays the load (PR review fix).""" + src = _load_model_src() + assert "_pre_fallback_sched_abort = False" in src + assert "_pre_fallback_sched_abort = _pre_fallback_sched_abort or (" in src + # The capture is set before the no-spec fallback spawns. + fn = ast.parse(src).body[0] + capture_line = next( + ( + n.lineno + for n in ast.walk(fn) + if isinstance(n, ast.Assign) + and any( + isinstance(t, ast.Name) and t.id == "_pre_fallback_sched_abort" for t in n.targets + ) + and isinstance(n.value, ast.BoolOp) + ), + None, + ) + fallback_line = next( + ( + n.lineno + for n in ast.walk(fn) + if isinstance(n, ast.Call) + and any(isinstance(a, ast.Name) and a.id == "fallback_cmd" for a in n.args) + ), + None, + ) + assert capture_line is not None and fallback_line is not None + assert capture_line < fallback_line