diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 8ec11fa79f..bc31b58c16 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -811,6 +811,11 @@ _APPLE_UNIFIED_MEMORY_FRACTION = 0.85 # reserve (_estimate_mtp_overhead_bytes). Applied to both the fit budget and pin. _MTP_VRAM_RESERVE_FRAC = 0.05 +# CPU-only: cap an auto context to this ceiling (a full native context puts tens to +# hundreds of GB of KV + MTP reserve in RAM), then fit it to RAM. Explicit -c wins. +_CPU_CTX_AUTO_CEILING = 32768 +_CPU_RAM_BUDGET_FRAC = 0.9 # RAM headroom for compute buffers + OS + def _kv_bytes_per_elem(cache_type: Optional[str]) -> float: """Bytes per KV-cache element for a llama.cpp cache type (f16 default).""" @@ -1109,6 +1114,37 @@ def _extra_args_n_ubatch( return None +def _extra_args_forces_cpu_offload( + extra_args: Optional[Iterable[str]], env: Optional[Mapping[str, str]] = None +) -> bool: + """True when the launch runs fully on CPU despite a visible GPU, so the CPU-only + safe defaults apply: zero GPU layers (-ngl 0 / --n-gpu-layers 0 / --gpu-layers 0) + or no device (--device/-dev none). CLI extras win (each control's last value); else + the inherited LLAMA_ARG_N_GPU_LAYERS / LLAMA_ARG_DEVICE env the child would honor.""" + args = [str(a) for a in extra_args] if extra_args else [] + ngl_zero: Optional[bool] = None + device_none: Optional[bool] = None + for i, raw in enumerate(args): + flag, eq, inline = raw.partition("=") + value = inline if eq else (args[i + 1] if i + 1 < len(args) else "") + if flag in ("-ngl", "--n-gpu-layers", "--gpu-layers"): + try: + ngl_zero = int(value) == 0 + except (TypeError, ValueError): + continue + elif flag in ("--device", "-dev"): + device_none = value.strip().lower() == "none" + _env = os.environ if env is None else env + if ngl_zero is None and _env.get("LLAMA_ARG_N_GPU_LAYERS") is not None: + try: + ngl_zero = int(_env["LLAMA_ARG_N_GPU_LAYERS"]) == 0 + except (TypeError, ValueError): + pass + if device_none is None and _env.get("LLAMA_ARG_DEVICE") is not None: + device_none = _env["LLAMA_ARG_DEVICE"].strip().lower() == "none" + return bool(ngl_zero) or bool(device_none) + + def _build_ngram_mod_flags( caps: Optional[dict], n_match: int = 24, @@ -1234,6 +1270,8 @@ class LlamaCppBackend: def __init__(self): self._process: Optional[subprocess.Popen] = None + # Spawn-time argv prefix (e.g. numactl --interleave=all); recomputed per load. + self._numa_prefix: list[str] = [] self._port: Optional[int] = None self._model_identifier: Optional[str] = None self._gguf_path: Optional[str] = None @@ -2469,6 +2507,24 @@ class LlamaCppBackend: if key is not None: cls._tensor_split_abort_keys.add(key) + # (binary, mtime, model) that aborted in the ggml graph scheduler this process + # (GGML_ASSERT(*cur_backend_id != -1)): an unsupported op, so reloading just + # repeats the crash. Keyed like the tensor-split memo; mtime drops it on update. + _sched_reserve_abort_keys: set[tuple[str, int, str]] = set() + + @classmethod + def _sched_reserve_aborts(cls, binary: Optional[str], model: Optional[str]) -> bool: + """True if (binary, model) aborted in graph-scheduler reserve this session.""" + key = cls._tensor_split_cache_key(binary, model) + return key is not None and key in cls._sched_reserve_abort_keys + + @classmethod + def _record_sched_reserve_abort(cls, binary: Optional[str], model: Optional[str]) -> None: + """Remember a (binary, model) that aborts in graph-scheduler reserve.""" + key = cls._tensor_split_cache_key(binary, model) + if key is not None: + cls._sched_reserve_abort_keys.add(key) + @staticmethod def _windows_pip_nvidia_dll_dirs(prefix: str) -> list[str]: """Return DLL dirs from pip-installed CUDA wheels under @@ -4119,6 +4175,10 @@ class LlamaCppBackend: "settings and reload." ) + # ggml graph-scheduler abort: surface the real cause, not the generic fallback. + if LlamaCppBackend._is_sched_reserve_abort(output or ""): + return LlamaCppBackend._sched_reserve_abort_message() + # Detect Ollama source up front so the arch branch can keep the # Ollama hint instead of the generic "unsupported arch" message. gguf = gguf_path or "" @@ -4395,6 +4455,41 @@ class LlamaCppBackend: # the split-axis enum token, unique to this assert (not the source file). return "split_axis" in text + @staticmethod + def _is_sched_reserve_abort(output: str) -> bool: + """True for the ggml graph-scheduler abort GGML_ASSERT(*cur_backend_id != -1): + no backend can run an op in the graph (e.g. an MLA/sparse-attention/MTP op + unimplemented on this build). Needs a ggml-abort marker plus a scheduler marker + (matched on the backtrace frames, which outlive the [New LWP] dump in a short + tail). Excludes the #6415 split-axis abort. stderr is merged into output.""" + text = (output or "").lower() + if "ggml_assert" not in text and "ggml_abort" not in text: + return False + # #6415 split-axis abort shares the frame but is handled by the tensor latch. + if "split_axis" in text: + return False + return ( + "cur_backend_id" in text + or "ggml_backend_sched_split_graph" in text + or "sched_reserve" in text + or "graph_reserve" in text + ) + + @staticmethod + def _sched_reserve_abort_message() -> str: + """Actionable message for the graph-scheduler abort (classifier + fail-fast guard).""" + return ( + "llama.cpp aborted while reserving the compute graph " + "(GGML_ASSERT(*cur_backend_id != -1) in ggml_backend_sched_split_graph): " + "the active backend cannot run an operation in this model's graph. This " + "usually means a newer attention variant (MLA / sparse-attention / MTP) " + "is not implemented in this llama.cpp build for the backend you're using " + "(commonly CPU-only). Disabling flash attention did not help. Try: run " + "`unsloth studio update` for a newer llama.cpp, use a different " + "quantization, run on a supported GPU/backend, or disable speculative " + "decoding / MTP and lower the context length." + ) + @staticmethod def _is_signal_crash(returncode: Optional[int]) -> bool: """True only on a hard fault (SIGSEGV/SIGABRT/SIGILL/SIGFPE/SIGBUS or a @@ -4516,12 +4611,13 @@ class LlamaCppBackend: logger.debug(f"Could not open llama-server log file: {e}") self._llama_log_path = None - # Log the argv per attempt (the text-only mmproj retry re-enters here - # with --mmproj stripped), redacting the API key. - logger.info(f"Starting llama-server: {' '.join(self._redacted_cmd_for_log(cmd))}") + # Log the argv per attempt (mmproj retry re-enters with --mmproj stripped), + # redacting the API key. Prepend _numa_prefix so the log matches what runs. + _run_cmd = [*self._numa_prefix, *cmd] + logger.info(f"Starting llama-server: {' '.join(self._redacted_cmd_for_log(_run_cmd))}") self._process = subprocess.Popen( - cmd, + _run_cmd, stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True, @@ -4644,14 +4740,39 @@ 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 BEFORE killing the live server (so a known-bad reload is + # non-destructive) if this exact launch 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). The key includes the + # variant AND the launch settings (context, spec) so an identical replay is + # blocked but a user changing quant / lowering -c / disabling spec -- the + # exact recovery the error message recommends -- is allowed to retry. + _abort_memo_model = "\x00".join( + [ + model_identifier or "", + hf_variant or "", + gguf_path or "", + str(n_ctx), + str(speculative_type or ""), + " ".join(str(a) for a in (extra_args 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).", + model_identifier, + ) + 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. @@ -4865,9 +4986,16 @@ class LlamaCppBackend: "image input will be disabled for this session" ) model_size = None # set in the fit try; used by the APU RAM guard + model_size_fit = None # weights + compute buffer; set in the fit try # Layer-fallback min GPUs; raised below on a tensor downgrade. Bound # before the try so the --fit-on except path still has it (no UnboundLocal). _layer_min_gpus = 1 + # CPU-only (no discrete GPU, not Apple Metal): drives safe defaults + # below since no GPU/Apple fit branch runs. Bound before the try (the + # --fit except path needs it); refined once gpus is known. + from utils.hardware import is_apple_silicon as _is_apple_silicon + + _cpu_only = False try: gguf_size = self._get_gguf_size_bytes(model_path) # Include GPU-loaded mmproj in the fit budget (#5825). @@ -4880,6 +5008,13 @@ class LlamaCppBackend: _gpu_mem = self._get_gpu_memory() gpus = [(idx, free) for idx, free, _t in _gpu_mem] total_by_idx = {idx: total for idx, _f, total in _gpu_mem} + # A user -ngl 0 / --n-gpu-layers 0 pins every layer on CPU even with a + # visible GPU; drop the GPU list so the CPU-only safe defaults and the + # RAM-aware fit apply, not a GPU launch that then runs on CPU anyway. + if gpus and _extra_args_forces_cpu_offload(extra_args): + logger.info("User set zero GPU offload (-ngl 0): treating as CPU-only.") + gpus, total_by_idx = [], {} + _cpu_only = (not gpus) and not _is_apple_silicon() def _gpu_usable(g, frac = _CTX_FIT_VRAM_FRACTION): # Per-GPU usable budget for ranking: free - (1-frac)*total. @@ -4985,6 +5120,16 @@ class LlamaCppBackend: _mtp_will_engage = bool( _user_mtp_via_extras or _user_draft_via_extras or _auto_studio_mtp ) + # Auto drops embedded MTP for MLA models (GLM-5.2/DeepSeek/Kimi) unless + # forced; mirror that gate so the CPU cap / NUMA footprint don't reserve + # a target-KV copy for a drafter the launch will not start. + _mtp_will_engage_cpu = _mtp_will_engage and not ( + _auto_studio_mtp + and bool(self._nextn_predict_layers) + and self._kv_lora_rank is not None + and not bool(mtp_draft_path) + and not _mla_mtp_auto_enabled() + ) # The duplicated full target-KV copy (ctx_tgt) is an MTP-only # cost: the MTP head runs a second context over the target # model's own KV geometry. The separate-drafter spec modes @@ -5581,6 +5726,81 @@ class LlamaCppBackend: except Exception as e: logger.debug(f"mmproj audio-capability read failed: {e}") + # CPU-only safe defaults (user extra_args win, appended last): no --fit + # (its graph-reserve estimator hits the same abort, llama.cpp #21932, and + # mis-counts MTP KV #23472/#24117; offloads nothing on CPU) and flash-attn + # off (CPU FA is unsafe for large MLA/sparse/MTP graphs). + _flash_default = "off" if _cpu_only else "on" + if _cpu_only and use_fit: + use_fit = False + logger.info("CPU-only host: launching with --fit off and --flash-attn off.") + + if _cpu_only: + _avail_mib = self._available_system_memory_mib() + # Preflight: weights over total RAM get OS-killed mid-load (interleave + # spreads across nodes but can't beat the total). + if _avail_mib and model_size and model_size > _avail_mib * 1024 * 1024: + logger.warning( + "CPU-only memory preflight: model weights ~%.0f GB exceed " + "available RAM ~%.0f GB; loading may be OS-killed.", + model_size / (1024**3), + _avail_mib / 1024, + ) + # Fit an auto context to a RAM-aware ceiling (explicit -c is honored). + # Runs for every auto context, not only above the ceiling: a large + # small-context GGUF can still exceed RAM and must be reduced too. + if requested_ctx <= 0 and effective_ctx > 0: + _ctx_ceiling = min(effective_ctx, _CPU_CTX_AUTO_CEILING) + _cpu_cap = _ctx_ceiling + try: + if _avail_mib and model_size and self._can_estimate_kv(): + _budget_b = _avail_mib * _CPU_RAM_BUDGET_FRAC * 1024 * 1024 + # Fixed footprint = weights + compute buffer (the same lump + # the load allocates), so a context that only fits ignoring + # the buffer can't slip through and get OS-killed at startup. + _fixed = model_size_fit or model_size + if _fixed >= _budget_b: + # Footprint alone over budget: _fit_context_to_vram returns + # the ceiling unchanged, but KV could still OOM. Floor to + # 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, + model_size_bytes = _fixed, + 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 alone is a no-op once budget_frac is set; + # pass the byte-accurate overhead so MTP KV is reserved. + mtp_engaged = _mtp_will_engage_cpu, + mtp_overhead_fn = ( + _mtp_bytes if _mtp_will_engage_cpu else None + ), + budget_frac = _cpu_budget, + ) + _cpu_cap = max(4096, min(_ctx_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: + logger.info( + "CPU-only: capping context %d -> %d (set -c to override).", + effective_ctx, + _cpu_cap, + ) + effective_ctx = _cpu_cap + # Advertise the capped window as the ceiling too, so /status and + # the UI safe-zone don't steer back to the unlaunched native size. + max_available_ctx = min(max_available_ctx, _cpu_cap) + cmd = [ binary, "-m", @@ -5592,7 +5812,7 @@ class LlamaCppBackend: "--parallel", str(n_parallel), "--flash-attn", - "on", # Force flash attention for speed + _flash_default, # CPU-only: off; GPU: on for speed # Error out at n_ctx instead of silently rotating the KV cache; frontend catches it and points the user at "Context Length". "--no-context-shift", ] @@ -5613,6 +5833,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 @@ -5828,7 +6052,59 @@ class LlamaCppBackend: cmd.extend(str(a) for a in extra_args) logger.info(f"Appending user extra args to llama-server: {list(extra_args)}") - logger.info(f"Starting llama-server: {' '.join(self._redacted_cmd_for_log(cmd))}") + # NUMA auto-interleave: when the model overflows one node but fits + # across all (else first-touch thrashes/OOMs), wrap with + # `numactl --interleave=all` (applied at the Popen sites) AND pass + # `--numa distribute` -- both are needed, numactl alone doesn't spread + # pages evenly (llama.cpp #19102). User --numa wins. + self._numa_prefix = [] + try: + from core.inference.numa import decide_interleave + + # Decide on the full resident footprint (weights + compute buffer + KV + # at the capped context + MTP reserve), not weights alone, so a model + # whose weights fit one node but whose footprint does not still + # interleaves instead of first-touch thrashing on one node. + _resident = model_size_fit or model_size + _numa_footprint = _resident + if _resident and effective_ctx > 0 and self._can_estimate_kv(): + try: + # Recompute MTP at the post-cap context; the pre-cap + # _mtp_reserve_bytes would overstate a capped million-token load. + _numa_mtp = _mtp_bytes(effective_ctx) if _mtp_will_engage_cpu else 0 + _numa_footprint = ( + _resident + + self._estimate_kv_cache_bytes( + effective_ctx, cache_type_kv, n_parallel = n_parallel + ) + + _numa_mtp + ) + except Exception: + _numa_footprint = _resident + _numa = decide_interleave(_numa_footprint, cpu_only = _cpu_only) + if _numa.interleave and _extra_args_set_any_flag(extra_args, {"--numa"}): + # User set an explicit --numa policy: respect it. The numactl wrap + # is an argv prefix they can't override, so skip it (and --numa + # distribute) rather than force interleaving over their choice. + logger.info("NUMA: user --numa set; leaving auto-interleave off") + elif _numa.interleave: + self._numa_prefix = list(_numa.prefix) + cmd.extend(["--numa", "distribute"]) + logger.info("NUMA: %s", _numa.reason) + elif _cpu_only and ( + "numactl` is not installed" in _numa.reason + or "interleave cannot help" in _numa.reason + ): + # Actionable: numactl missing, or the footprint exceeds total RAM + # across all nodes (the weights-only preflight can't catch this). + logger.warning("NUMA: %s", _numa.reason) + except Exception as _numa_exc: # never block a load on the NUMA probe + logger.debug("NUMA interleave probe failed: %s", _numa_exc) + + logger.info( + "Starting llama-server: " + f"{' '.join(self._redacted_cmd_for_log([*self._numa_prefix, *cmd]))}" + ) # Library paths so llama-server finds its shared libs and CUDA DLLs. env = self._llama_server_env_for_binary(binary) @@ -5969,9 +6245,11 @@ class LlamaCppBackend: # Best-effort; never block the load on logging. logger.debug(f"Could not open llama-server log file: {e}") self._llama_log_path = None + # _last_spawn_cmd stays un-prefixed (retry helpers slice it); + # the NUMA prefix goes on the actual argv only, so retries don't double it. _last_spawn_cmd = list(run_cmd) self._process = subprocess.Popen( - run_cmd, + [*self._numa_prefix, *run_cmd], stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True, @@ -6047,6 +6325,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, @@ -6146,6 +6427,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. @@ -6205,6 +6492,25 @@ class LlamaCppBackend: out = "\n".join(self._stdout_lines[-50:]) # Read the crash code before _kill_process() clears _process. _crash_rc = self._process.poll() if self._process is not None else None + # Graph-scheduler abort (flash-off retry already failed)? Capture now (a + # text-only retry overwrites the stdout tail) but memo it only once the + # 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 ( + # 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). # Skip if a cancel/unload is pending (mirrors the MTP guard). @@ -6234,6 +6540,11 @@ class LlamaCppBackend: # an OS-killed text-only retry still gets the OOM message. _retry_rc = self._process.poll() if self._process is not None else None self._kill_process() + # Fallback exhausted: now memo the original scheduler abort. + if _was_sched_abort: + LlamaCppBackend._record_sched_reserve_abort( + binary, _abort_memo_model + ) raise RuntimeError( "Vision projector incompatible with this llama.cpp " "build, and the text-only retry also failed: " @@ -6245,6 +6556,10 @@ class LlamaCppBackend: ) ) else: + # No mmproj fallback available/eligible: memo the scheduler abort + # (if that is what crashed) so the next /load fails fast, then raise. + if _was_sched_abort: + LlamaCppBackend._record_sched_reserve_abort(binary, _abort_memo_model) raise RuntimeError( self._classify_llama_start_failure( out, diff --git a/studio/backend/core/inference/numa.py b/studio/backend/core/inference/numa.py new file mode 100644 index 0000000000..49a88e69e1 --- /dev/null +++ b/studio/backend/core/inference/numa.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""NUMA topology detection and an auto-interleave decision for CPU-only launches. + +Linux's first-touch policy faults a model's pages onto one node; if the GGUF is larger +than that node's free RAM the load thrashes/fails despite ample total RAM. +`numactl --interleave=all` spreads pages across nodes (llama.cpp discussion #19102). +Pure stdlib, Linux-only (no-op decision elsewhere), side-effect free. +""" + +from __future__ import annotations + +import shutil +from dataclasses import dataclass, field +from pathlib import Path + +_NODE_ROOT = Path("/sys/devices/system/node") + + +@dataclass(frozen = True) +class NumaTopology: + """Per-node free memory (MemFree, matching `numactl --hardware`'s 'node N free').""" + + node_free_mib: dict[int, int] = field(default_factory = dict) + + @property + def node_count(self) -> int: + return len(self.node_free_mib) + + @property + def total_free_mib(self) -> int: + return sum(self.node_free_mib.values()) + + @property + 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.""" + out: list[int] = [] + for part in spec.strip().split(","): + part = part.strip() + if not part: + continue + if "-" in part: + lo, _, hi = part.partition("-") + try: + out.extend(range(int(lo), int(hi) + 1)) + except ValueError: + continue + else: + try: + out.append(int(part)) + except ValueError: + continue + return out + + +def _mems_allowed() -> set[int] | None: + """NUMA nodes the process may allocate on (cpuset), from /proc/self/status + Mems_allowed_list. None when unavailable, so callers fall back to the online set. + numactl --interleave=all only spans these, so a cpuset-limited container must not + count host nodes it cannot use.""" + try: + text = Path("/proc/self/status").read_text() + except OSError: + return None + for line in text.splitlines(): + if line.startswith("Mems_allowed_list:"): + return set(_parse_online(line.split(":", 1)[1])) + return None + + +def _node_memfree_mib(node: int) -> int | None: + """MemFree for one node from /sys/.../nodeN/meminfo, in MiB (kB -> MiB).""" + try: + text = (_NODE_ROOT / f"node{node}" / "meminfo").read_text() + except OSError: + return None + for line in text.splitlines(): + # "Node 0 MemFree: 465594 kB" + if "MemFree:" in line: + parts = line.split() + try: + kb = int(parts[parts.index("MemFree:") + 1]) + except (ValueError, IndexError): + return None + return kb // 1024 + return None + + +def read_numa_topology() -> NumaTopology: + """Read NUMA node free memory from sysfs. Empty topology on non-Linux / no sysfs / + single-node (a single node is reported but callers treat node_count <= 1 as 'no + interleave needed').""" + try: + online = (_NODE_ROOT / "online").read_text() + except OSError: + return NumaTopology() + # Restrict to cpuset-allowed nodes: under a Docker/systemd cpuset the child can only + # allocate on these, so counting other host nodes would overstate the fittable RAM. + allowed = _mems_allowed() + free: dict[int, int] = {} + for node in _parse_online(online): + if allowed is not None and node not in allowed: + continue + mib = _node_memfree_mib(node) + if mib is not None: + free[node] = mib + return NumaTopology(node_free_mib = free) + + +def numactl_available() -> bool: + return shutil.which("numactl") is not None + + +@dataclass(frozen = True) +class InterleaveDecision: + interleave: bool + reason: str + prefix: tuple[str, ...] = () # argv prefix to apply, e.g. ("numactl", "--interleave=all") + + +def decide_interleave( + model_size_bytes: int | None, + *, + cpu_only: bool, + topology: NumaTopology | None = None, + has_numactl: bool | None = None, +) -> InterleaveDecision: + """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") + if not model_size_bytes or model_size_bytes <= 0: + return InterleaveDecision(False, "model size unknown; not forcing interleave") + + topo = topology if topology is not None else read_numa_topology() + if topo.node_count <= 1: + return InterleaveDecision(False, "single NUMA node; interleave not needed") + + model_mib = model_size_bytes // (1024 * 1024) + smallest = topo.smallest_node_free_mib + total = topo.total_free_mib + + # 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 every node's free RAM " + f"(smallest ~{smallest} MiB); keeping local placement", + ) + + # Impossible across all nodes regardless of numactl: surface the smaller-quant / + # free-memory path before the numactl hint, so a too-big model is not told to + # install numactl when interleaving could never make it fit. + if model_mib > total: + return InterleaveDecision( + False, + f"model ~{model_mib} MiB exceeds total free RAM across all nodes " + f"(~{total} MiB); interleave cannot help -- free memory or use a smaller quant", + ) + + avail = numactl_available() if has_numactl is None else has_numactl + if not avail: + # Needed but unavailable: surface it; caller decides whether to block. + return InterleaveDecision( + False, + 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 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 new file mode 100644 index 0000000000..9e639d2d04 --- /dev/null +++ b/studio/backend/tests/test_cpu_only_defaults.py @@ -0,0 +1,304 @@ +# 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 diff --git a/studio/backend/tests/test_numa_interleave.py b/studio/backend/tests/test_numa_interleave.py new file mode 100644 index 0000000000..7d856327a0 --- /dev/null +++ b/studio/backend/tests/test_numa_interleave.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Unit tests for the NUMA auto-interleave decision (core/inference/numa.py). + +Models the user's dual-NUMA Xeon (HF screenshots #23): node 0 ~465 GB free, node 1 +~223 GB free; a 583 GB GGUF exceeds the largest single node but fits across both, so +`numactl --interleave=all` is the right call. The decision is pure and topology is +injected, so these are deterministic on any host (no /sys, no numactl needed). +""" + +from __future__ import annotations + +import sys +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) + +# Importing core.inference.numa runs core/inference/__init__.py (orchestrator + structlog +# + loggers + httpx); stub those when absent so a dependency-light run can collect this. +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 +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.numa import ( # noqa: E402 + InterleaveDecision, + NumaTopology, + _parse_online, + decide_interleave, +) + +_GiB = 1024**3 +_MiB = 1024**2 + +# The user's box, in MiB free per node (from `numactl --hardware`). +_USER_TOPO = NumaTopology(node_free_mib = {0: 465594, 1: 223814}) +_SINGLE = NumaTopology(node_free_mib = {0: 900_000}) + + +def test_parse_online_ranges(): + assert _parse_online("0-1") == [0, 1] + assert _parse_online("0,2-3") == [0, 2, 3] + assert _parse_online("3") == [3] + assert _parse_online("") == [] + + +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 + + +def test_interleaves_when_model_exceeds_largest_node_but_fits_across(): + # 583 GB model: > 465 GB (node 0) but < ~689 GB total. + d = decide_interleave(583 * _GiB, cpu_only = True, topology = _USER_TOPO, has_numactl = True) + assert d.interleave is True + assert d.prefix == ("numactl", "--interleave=all") + assert "interleave=all" in d.reason + + +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 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(): + d = decide_interleave(583 * _GiB, cpu_only = False, topology = _USER_TOPO, has_numactl = True) + assert d.interleave is False + assert "not cpu-only" in d.reason + + +def test_no_interleave_single_node(): + d = decide_interleave(583 * _GiB, cpu_only = True, topology = _SINGLE, has_numactl = True) + assert d.interleave is False + assert "single NUMA node" in d.reason + + +def test_model_too_big_for_all_nodes_blocks_with_message(): + # 800 GB > ~689 GB total free across both nodes -> interleave can't help. + d = decide_interleave(800 * _GiB, cpu_only = True, topology = _USER_TOPO, has_numactl = True) + assert d.interleave is False + assert "exceeds total free RAM" in d.reason + + +def test_numactl_missing_surfaces_actionable_warning(): + d = decide_interleave(583 * _GiB, cpu_only = True, topology = _USER_TOPO, has_numactl = False) + assert d.interleave is False + assert "numactl` is not installed" in d.reason + + +def test_too_big_and_numactl_missing_prefers_total_ram_message(): + # Impossible across all nodes AND no numactl: the total-RAM guidance must win, so the + # user is not told to install numactl when interleaving could never help (PR review fix). + d = decide_interleave(800 * _GiB, cpu_only = True, topology = _USER_TOPO, has_numactl = False) + assert d.interleave is False + assert "exceeds total free RAM" in d.reason + assert "numactl` is not installed" not in d.reason + + +def test_unknown_model_size_does_not_force_interleave(): + for size in (None, 0, -1): + d = decide_interleave(size, cpu_only = True, topology = _USER_TOPO, has_numactl = True) + assert d.interleave is False + + +def test_topology_restricted_to_cpuset_allowed_nodes(tmp_path, monkeypatch): + """A cpuset that allows only node 0 must drop host node 1 from the topology, so a + container is not told a node's RAM is usable when the child can't allocate there + (PR review fix).""" + import core.inference.numa as m + + (tmp_path / "online").write_text("0-1") + for nid, free_kb in ((0, 100 * 1024), (1, 200 * 1024)): + d = tmp_path / f"node{nid}" + d.mkdir() + (d / "meminfo").write_text(f"Node {nid} MemFree: {free_kb} kB\n") + monkeypatch.setattr(m, "_NODE_ROOT", tmp_path) + + monkeypatch.setattr(m, "_mems_allowed", lambda: {0}) + assert m.read_numa_topology().node_free_mib == {0: 100} + # No cpuset info (None) -> trust the online set, both nodes present. + monkeypatch.setattr(m, "_mems_allowed", lambda: None) + assert m.read_numa_topology().node_free_mib == {0: 100, 1: 200} + + +def test_decision_is_frozen_dataclass(): + d = InterleaveDecision(True, "x", ("numactl", "--interleave=all")) + try: + d.interleave = False # type: ignore[misc] + except AttributeError: + return + raise AssertionError("InterleaveDecision should be immutable") diff --git a/studio/backend/tests/test_sched_reserve_abort.py b/studio/backend/tests/test_sched_reserve_abort.py new file mode 100644 index 0000000000..6ed98c2adf --- /dev/null +++ b/studio/backend/tests/test_sched_reserve_abort.py @@ -0,0 +1,325 @@ +# 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 the ggml graph-scheduler abort handling in load_model. + +A CPU-only user loading GLM-5.2 (MLA + sparse-attention "indexer" + embedded MTP) +hit `GGML_ASSERT(*cur_backend_id != -1)` in `ggml_backend_sched_split_graph` during +`sched_reserve`: the CPU backend cannot run an op in the graph. Studio's startup +crash raised a generic "invalid GGUF or out of memory" message and the UI replayed +`POST /load`, re-reading the ~583 GB weights into the identical crash every ~3.5 min. + +These tests pin: (1) the abort matcher recognises the real backtrace tail and only +that, (2) the classifier surfaces the actionable message, (3) the per-(binary,model) +memo round-trips and is mtime-invalidated, and (4) load_model fails fast on a memoed +abort and records on crash. No GPU, no network, fully deterministic. +""" + +from __future__ import annotations + +import ast +import inspect +import os +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) + +# External-dep stubs so importing the backend doesn't require structlog / loggers / +# httpx -- only installed when the real module is missing (mirrors test_tp_vision_regression). +try: + import structlog # noqa: F401 +except ImportError: + _structlog_stub = _types.ModuleType("structlog") + _structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub") + sys.modules["structlog"] = _structlog_stub +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 + + +# The real crash, faithfully reproduced from the user's llama-server log (HF +# screenshots discussion #23). The GGML_ASSERT line is followed by ~130 [New LWP] +# lines and then the gdb backtrace; the assert line scrolls out of a short tail. +_FULL_ABORT = "\n".join( + [ + "0.09.350.752 W llama_context: n_ctx_seq (4096) < n_ctx_train (1048576)", + "/tmp/llama.cpp/ggml/src/ggml-backend.cpp:1242: GGML_ASSERT(*cur_backend_id != -1) failed", + ] + + [f"[New LWP {2147067 - i}]" for i in range(130)] + + [ + "[Thread debugging using libthread_db enabled]", + "#1 0x... in ggml_print_backtrace () from /tmp/llama.cpp/build/bin/libggml-base.so.0", + "#2 0x... in ggml_abort () from /tmp/llama.cpp/build/bin/libggml-base.so.0", + "#3 0x... in ggml_backend_sched_split_graph () from /tmp/llama.cpp/build/bin/libggml-base.so.0", + "#4 0x... in llama_context::graph_reserve(...) () from /tmp/llama.cpp/build/bin/libllama.so.0", + "#5 0x... in llama_context::sched_reserve() () from /tmp/llama.cpp/build/bin/libllama.so.0", + ] +) + +# What a 50-line tail actually contains: the GGML_ASSERT line is gone, but the +# backtrace markers (ggml_abort, ggml_backend_sched_split_graph) remain. The matcher +# must still fire on this -- that's the realistic input at the recording site. +_ABORT_TAIL = "\n".join(_FULL_ABORT.splitlines()[-50:]) + +# The other ggml abort Studio already handles (#6415 split-axis): must NOT be +# misclassified as a scheduler-reserve abort. +_SPLIT_AXIS_ABORT = ( + "ggml/src/ggml-backend-meta.cpp:541: " + "GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0) failed\n" + "#3 ggml_backend_sched_split_graph ()" +) + +_OOM_OUTPUT = "llama_model_load: error loading model: unable to allocate buffer\nkilled" +_CLEAN_OUTPUT = "main: server is listening on http://127.0.0.1:8080 - starting the main loop" + + +# ---- matcher --------------------------------------------------------------- + + +def test_matcher_fires_on_full_abort_and_short_tail(): + assert LlamaCppBackend._is_sched_reserve_abort(_FULL_ABORT) + # The headline guarantee: the matcher survives the [New LWP] scroll. + assert "GGML_ASSERT(*cur_backend_id != -1)".lower() not in _ABORT_TAIL.lower() + assert LlamaCppBackend._is_sched_reserve_abort(_ABORT_TAIL) + + +def test_matcher_ignores_unrelated_crashes(): + assert not LlamaCppBackend._is_sched_reserve_abort(_SPLIT_AXIS_ABORT) + assert not LlamaCppBackend._is_sched_reserve_abort(_OOM_OUTPUT) + assert not LlamaCppBackend._is_sched_reserve_abort(_CLEAN_OUTPUT) + assert not LlamaCppBackend._is_sched_reserve_abort("") + + +def test_matcher_requires_both_a_ggml_marker_and_a_scheduler_marker(): + # scheduler word without any ggml abort/assert -> not our abort. + assert not LlamaCppBackend._is_sched_reserve_abort("graph_reserve completed in 3ms") + # ggml abort without a scheduler marker -> some other assert, not ours. + assert not LlamaCppBackend._is_sched_reserve_abort("ggml_abort: tensor type mismatch") + + +# ---- classifier ------------------------------------------------------------ + + +def test_classifier_surfaces_actionable_message(): + msg = LlamaCppBackend._classify_llama_start_failure( + _ABORT_TAIL, + gguf_path = "/x/GLM-5.2-UD-Q6_K-00001-of-00014.gguf", + model_identifier = "unsloth/GLM-5.2-GGUF", + returncode = -6, + ) + assert msg == LlamaCppBackend._sched_reserve_abort_message() + # The message names the real cause, not the generic invalid-GGUF/OOM fallback. + assert "ggml_backend_sched_split_graph" in msg + assert "enough memory" not in msg # i.e. not the generic fallback + + +def test_classifier_generic_fallback_unchanged_for_unknown_crash(): + msg = LlamaCppBackend._classify_llama_start_failure( + "some unrelated failure", + gguf_path = None, + model_identifier = None, + returncode = -6, + ) + assert "failed to start" in msg and "enough memory" in msg + + +# ---- memo round-trip / invalidation --------------------------------------- + + +def test_memo_round_trip_and_isolation(tmp_path): + binary = tmp_path / "llama-server" + binary.write_text("x") + b, model = str(binary), "unsloth/GLM-5.2-GGUF" + + LlamaCppBackend._sched_reserve_abort_keys.clear() + assert not LlamaCppBackend._sched_reserve_aborts(b, model) + LlamaCppBackend._record_sched_reserve_abort(b, model) + assert LlamaCppBackend._sched_reserve_aborts(b, model) + # A different model on the same binary is unaffected. + assert not LlamaCppBackend._sched_reserve_aborts(b, "unsloth/Qwen3.5-4B-MTP-GGUF") + LlamaCppBackend._sched_reserve_abort_keys.clear() + + +def test_memo_invalidated_by_binary_mtime_change(tmp_path): + """A `unsloth studio update` swaps the binary -> the memo must not persist.""" + binary = tmp_path / "llama-server" + binary.write_text("v1") + b, model = str(binary), "unsloth/GLM-5.2-GGUF" + + LlamaCppBackend._sched_reserve_abort_keys.clear() + LlamaCppBackend._record_sched_reserve_abort(b, model) + assert LlamaCppBackend._sched_reserve_aborts(b, model) + # Bump mtime to a distinct ns value (simulate a rebuilt binary). + st = binary.stat() + os.utime(binary, ns = (st.st_atime_ns + 10**9, st.st_mtime_ns + 10**9)) + assert not LlamaCppBackend._sched_reserve_aborts(b, model) + LlamaCppBackend._sched_reserve_abort_keys.clear() + + +def test_memo_safe_with_missing_binary_or_model(): + # None binary/model -> no key -> never aborts, never raises. + assert not LlamaCppBackend._sched_reserve_aborts(None, "m") + assert not LlamaCppBackend._sched_reserve_aborts("/x", None) + LlamaCppBackend._record_sched_reserve_abort(None, None) # no-op, no raise + + +# ---- load_model wiring (source-level, no GPU/network needed) --------------- + + +def _load_model_src() -> str: + return textwrap.dedent(inspect.getsource(LlamaCppBackend.load_model)) + + +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 == 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_and_launch_settings(): + """The memo key must include the variant AND the launch settings (context, spec), + so a failed quant does not block a different quant, and changing -c / spec (the + recommended recovery) is allowed to retry while an identical replay stays blocked.""" + src = _load_model_src() + assert "_abort_memo_model" in src + for tok in ("hf_variant", "gguf_path", "str(n_ctx)", "speculative_type", "extra_args"): + assert tok in src, tok + + def key( + model = "repo", + variant = "", + gguf = "", + n_ctx = 4096, + spec = "", + extra = "", + ): + return "\x00".join([model, variant, gguf, str(n_ctx), spec, extra]) + + base = key(variant = "UD-Q6_K") + assert base != key(variant = "UD-Q4_K_XL") # different quant -> retry allowed + assert base != key(variant = "UD-Q6_K", n_ctx = 2048) # lower context -> retry allowed + assert base != key(variant = "UD-Q6_K", spec = "off") # disable spec -> retry allowed + assert base == key(variant = "UD-Q6_K") # identical replay -> still blocked + + +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, _abort_memo_model)" in src + assert "_is_sched_reserve_abort(" in src + + +def test_abort_memo_deferred_until_mmproj_fallback_ruled_out(): + """The signature is captured up front but the memo is recorded only on a terminal + 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 "if _was_sched_abort:" in src + fn = ast.parse(src).body[0] + strip_line = _call_line(fn, "_strip_mmproj_args") + record_line = _call_line(fn, "_record_sched_reserve_abort") + # 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