diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 31969afb14..8ec11fa79f 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1271,6 +1271,9 @@ class LlamaCppBackend: self._cache_type_kv: Optional[str] = None # Whether --split-mode tensor was applied on the active load. self._tensor_parallel: bool = False + # Layer load kept multi-GPU only to honor a downgraded tensor request, so a + # later explicit tensor-off reloads instead of deduping to it (#6659). + self._layer_preserves_tensor_intent: bool = False self._reasoning_default: bool = True self._speculative_type: Optional[str] = None # Canonical UI-facing mode the user requested @@ -1643,6 +1646,11 @@ class LlamaCppBackend: """Whether --split-mode tensor is active on the loaded server.""" return self._tensor_parallel + @property + def layer_preserves_tensor_intent(self) -> bool: + """True when a downgraded tensor request kept this layer load multi-GPU.""" + return self._layer_preserves_tensor_intent + @property def speculative_type(self) -> Optional[str]: return self._speculative_type @@ -2430,6 +2438,37 @@ class LlamaCppBackend: # aborts a --split-mode tensor load, so it's dropped for the tensor attempt. _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"}) + # (binary, mtime, model) that aborted on --split-mode tensor this process (#6415 + # geometry limit, e.g. MQA n_head_kv=1). Model-keyed so one model's abort doesn't + # skip tensor for others; tensor is tried by default, recorded only on a real abort. + _tensor_split_abort_keys: set[tuple[str, int, str]] = set() + + @classmethod + def _tensor_split_cache_key( + cls, binary: Optional[str], model: Optional[str] + ) -> Optional[tuple[str, int, str]]: + """(path, mtime_ns, model) key; ns mtime re-probes a same-second binary swap.""" + if not binary or not model: + return None + try: + mtime = Path(binary).stat().st_mtime_ns + except OSError: + mtime = 0 + return (binary, mtime, model) + + @classmethod + def _tensor_split_aborts(cls, binary: Optional[str], model: Optional[str]) -> bool: + """True if (binary, model) aborted on --split-mode tensor this session.""" + key = cls._tensor_split_cache_key(binary, model) + return key is not None and key in cls._tensor_split_abort_keys + + @classmethod + def _record_tensor_split_abort(cls, binary: Optional[str], model: Optional[str]) -> None: + """Remember a (binary, model) that aborts on --split-mode tensor.""" + key = cls._tensor_split_cache_key(binary, model) + if key is not None: + cls._tensor_split_abort_keys.add(key) + @staticmethod def _windows_pip_nvidia_dll_dirs(prefix: str) -> list[str]: """Return DLL dirs from pip-installed CUDA wheels under @@ -2569,9 +2608,13 @@ class LlamaCppBackend: usable_fraction: Optional[float] = None, total_by_idx: Optional[dict[int, int]] = None, per_device_overhead_bytes: int = 0, + min_gpus: int = 1, ) -> tuple[Optional[list[int]], bool]: """Pick GPU(s) for a model from estimated VRAM and free memory. + ``min_gpus`` (default 1, capped at ``len(gpus)``) keeps a downgraded + tensor/multi-GPU request spread instead of collapsing to one card. + ``model_size_bytes`` should include weights and estimated KV cache. ``usable_fraction`` (default ``_GPU_PIN_VRAM_FRACTION``) provides headroom for compute buffers, CUDA context, and other runtime @@ -2590,9 +2633,11 @@ class LlamaCppBackend: if not gpus: return None, True + min_gpus = max(1, min(min_gpus, len(gpus))) model_size_mib = model_size_bytes / (1024 * 1024) if usable_fraction is None: usable_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION + overhead_mib = per_device_overhead_bytes / (1024 * 1024) # Per-GPU usable budget: free - (1-frac)*total when total is known, else # the legacy free*frac (also covers a total-0 two-column probe). @@ -2606,19 +2651,26 @@ class LlamaCppBackend: # card can have less usable room than a less-used small one. ranked = sorted(gpus, key = lambda g: _usable(g[0], g[1]), reverse = True) - # Try 1 GPU at the usable-VRAM threshold. - if _usable(ranked[0][0], ranked[0][1]) >= model_size_mib: + # Cap a downgraded multi-GPU request to the usable count so it doesn't pull + # in a near-full card to hit min_gpus. No-op for the default min_gpus == 1. + usable_count = sum(1 for idx, free_mib in ranked if _usable(idx, free_mib) > overhead_mib) + min_gpus = max(1, min(min_gpus, usable_count or 1)) + + # Try 1 GPU at the usable-VRAM threshold (only when one device is allowed). + if min_gpus <= 1 and _usable(ranked[0][0], ranked[0][1]) >= model_size_mib: return [ranked[0][0]], False - # Try N GPUs (accumulate usable memory from most-free). Each GPU past the - # first adds a fixed per-device overhead the pool must hold. - overhead_mib = per_device_overhead_bytes / (1024 * 1024) + # Try N GPUs (most-free first); each past the first adds per-device overhead. + # Require at least min_gpus devices before accepting a fit. cumulative = 0.0 selected = [] for idx, free_mib in ranked: selected.append(idx) cumulative += _usable(idx, free_mib) - if cumulative >= model_size_mib + (len(selected) - 1) * overhead_mib: + if ( + len(selected) >= min_gpus + and cumulative >= model_size_mib + (len(selected) - 1) * overhead_mib + ): return sorted(selected), False # Too large even for all GPUs; let --fit handle it @@ -3868,7 +3920,7 @@ class LlamaCppBackend: logger.debug(f"Could not list repo files for {label}: {e}") break logger.debug( - f"Could not list repo files for {label} " f"(attempt {attempt + 1}/3): {e}" + f"Could not list repo files for {label} (attempt {attempt + 1}/3): {e}" ) if attempt < 2: self._cancel_event.wait(2**attempt) @@ -4332,6 +4384,17 @@ class LlamaCppBackend: ) ) + @staticmethod + def _is_tensor_split_assert(output: str) -> bool: + """True only for the #6415 split-axis warmup assert (GGML_BACKEND_SPLIT_AXIS_*), + not any ggml assert/abort, so an unrelated invariant isn't cached. stderr is + merged into output.""" + text = (output or "").lower() + if "ggml_assert" not in text and "ggml_abort" not in text: + return False + # the split-axis enum token, unique to this assert (not the source file). + return "split_axis" in text + @staticmethod def _is_signal_crash(returncode: Optional[int]) -> bool: """True only on a hard fault (SIGSEGV/SIGABRT/SIGILL/SIGFPE/SIGBUS or a @@ -4344,6 +4407,20 @@ class LlamaCppBackend: return True return -returncode in (4, 6, 7, 8, 11) # SIGILL SIGABRT SIGBUS SIGFPE SIGSEGV + @staticmethod + def _is_abort_exit(returncode: Optional[int]) -> bool: + """Windows CRT abort() exit code (3) from GGML_ASSERT on MSVC -- not a POSIX + signal or 0xC0000000+ NTSTATUS.""" + return returncode == 3 + + @classmethod + def _should_record_tensor_split_abort(cls, returncode: Optional[int], output: str) -> bool: + """The #6415 split-axis abort: the marker plus a hard crash (POSIX signal or + Windows abort exit). Marker required so a generic crash isn't cached.""" + return cls._is_tensor_split_assert(output) and ( + cls._is_signal_crash(returncode) or cls._is_abort_exit(returncode) + ) + @staticmethod def _with_flash_attn_off(cmd: list[str]) -> Optional[list[str]]: """Return cmd with flash attention forced off, or None when its effective @@ -4488,6 +4565,8 @@ class LlamaCppBackend: n_gpu_layers: Optional[int] = None, # caller compat, unused n_parallel: int = 1, extra_args: Optional[List[str]] = None, + # Route-level tensor->layer fallback retry: keep the layer split multi-GPU. + preserve_multi_gpu_on_layer: bool = False, ) -> bool: """Start llama-server with a GGUF model. @@ -4518,6 +4597,8 @@ class LlamaCppBackend: "n_gpu_layers": n_gpu_layers, "n_parallel": n_parallel, "extra_args": list(extra_args) if extra_args is not None else None, + # Replayed by _respawn_if_dead so a downgraded model stays multi-GPU. + "preserve_multi_gpu_on_layer": preserve_multi_gpu_on_layer, } # Serialise the whole load so concurrent /load calls never leave two # llama-server processes alive (#5401 / #5161). Doesn't block /unload. @@ -4541,6 +4622,7 @@ class LlamaCppBackend: chat_template_override = chat_template_override, extra_args = extra_args, is_vision = is_vision, + preserve_multi_gpu_on_layer = preserve_multi_gpu_on_layer, ): logger.info( f"load_model: backend already in target state for " @@ -4626,6 +4708,9 @@ class LlamaCppBackend: # Block-diffusion GGUFs (DiffusionGemma) cannot run on llama-server; # serve them with the diffusion runner (same OpenAI-compat interface). if self._is_diffusion: + # Not a tensor/layer GGUF: clear any preserved-fallback flag from a + # prior load (this path skips the command builder that clears it). + self._layer_preserves_tensor_intent = False with self._lock: if self._cancel_event.is_set(): logger.info("Load cancelled before diffusion server start") @@ -4780,6 +4865,9 @@ class LlamaCppBackend: "image input will be disabled for this session" ) model_size = None # set in the fit try; used by the APU RAM guard + # 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 try: gguf_size = self._get_gguf_size_bytes(model_path) # Include GPU-loaded mmproj in the fit budget (#5825). @@ -5064,10 +5152,8 @@ class LlamaCppBackend: _apple_budget_mib = self._apple_metal_memory_budget_bytes() // (1024 * 1024) def _restore_after_tensor_downgrade(): - # Tensor mode dropped a quantized KV and stripped the cache - # extras (it rejects quantized); layer split supports them, so - # restore the original type + extras (minus --split-mode) and - # clear the env flag so the layer launch re-emits them. + # Restore the quantized KV + extras tensor dropped (layer + # split supports them), minus --split-mode. nonlocal cache_type_kv, _cache_type_from_env, extra_args if _tensor_dropped_cache_type_kv is not None: cache_type_kv = _tensor_dropped_cache_type_kv @@ -5078,13 +5164,22 @@ class LlamaCppBackend: else extra_args ) - if tensor_parallel and effective_is_vision: + # The route fallback retry is tensor-off; keep it multi-GPU. + if preserve_multi_gpu_on_layer: + _layer_min_gpus = max(_layer_min_gpus, len(gpus)) + + if tensor_parallel and self._tensor_split_aborts(binary, model_identifier): + # Aborted on tensor for this model this session (#6415); skip + # tensor upfront, layer split serves it. logger.info( - "Tensor parallelism skipped for vision model: " - "--split-mode tensor is incompatible with --mmproj " - "in the current llama.cpp build; using layer split." + "Tensor parallelism skipped: this llama.cpp build aborted " + "on --split-mode tensor for this model earlier this " + "session; using layer split across %d GPU(s).", + len(gpus), ) tensor_parallel = False + # Keep the multi-GPU request (gated on it, not the cache). + _layer_min_gpus = max(_layer_min_gpus, len(gpus)) _restore_after_tensor_downgrade() # Tensor mode replicates a compute buffer on every GPU, so drop @@ -5124,6 +5219,11 @@ class LlamaCppBackend: len(gpus), ) tensor_parallel = False + # GPUs below tensor's compute-buffer reserve can still do layer + # split, so keep multi-GPU (mirrors the budget/geometry drops); + # _select_gpus caps unusable cards. + if len(gpus) >= 2: + _layer_min_gpus = max(_layer_min_gpus, len(gpus)) # Layer split supports a quantized KV the tensor attempt # dropped; restore the original cache type + extras (minus # --split-mode) so the layer launch re-emits them. @@ -5160,8 +5260,12 @@ class LlamaCppBackend: "per-device compute buffers; falling back to layer split." ) tensor_parallel = False - # Restore the dropped quantized KV + original cache extras - # (minus --split-mode); layer split supports them. + # Weights needed >1 card, so keep multi-GPU across the + # usable tensor GPUs. + if len(tp_gpus) >= 2: + _layer_min_gpus = max(_layer_min_gpus, len(tp_gpus)) + # Restore the dropped quantized KV + cache extras (minus + # --split-mode); layer split supports them. _restore_after_tensor_downgrade() if tensor_parallel and tp_gpus: @@ -5263,6 +5367,7 @@ class LlamaCppBackend: usable_fraction = _pin_fraction, total_by_idx = total_by_idx, per_device_overhead_bytes = _pipeline_overhead_bytes, + min_gpus = _layer_min_gpus, ) # No silent shrink: effective_ctx stays == requested_ctx. else: @@ -5273,7 +5378,22 @@ class LlamaCppBackend: ranked = sorted( gpus, key = lambda g: _gpu_usable(g, pin_fraction), reverse = True ) - for n_gpus in range(1, len(ranked) + 1): + # Skips _select_gpus, so apply its cap: count only cards + # whose usable VRAM clears the per-device layer overhead. + _pipeline_overhead_mib = _pipeline_overhead_bytes / (1024 * 1024) + _auto_min_gpus = max( + 1, + min( + _layer_min_gpus, + sum( + 1 + for g in ranked + if _gpu_usable(g, pin_fraction) > _pipeline_overhead_mib + ) + or 1, + ), + ) + for n_gpus in range(_auto_min_gpus, len(ranked) + 1): subset = ranked[:n_gpus] pool_budget = _pool_budget_mib(subset, pin_fraction) _ms = _subset_model_size(n_gpus) @@ -5303,7 +5423,7 @@ class LlamaCppBackend: # at 131k may pin fine with a 4096 KV (#5106). effective_ctx = min(4096, effective_ctx) if effective_ctx > 0: - for n_gpus in range(1, len(ranked) + 1): + for n_gpus in range(_auto_min_gpus, len(ranked) + 1): subset = ranked[:n_gpus] kv = self._estimate_kv_cache_bytes( effective_ctx, @@ -5339,6 +5459,7 @@ class LlamaCppBackend: usable_fraction = _pin_fraction, total_by_idx = total_by_idx, per_device_overhead_bytes = _pipeline_overhead_bytes, + min_gpus = _layer_min_gpus, ) if use_fit and not explicit_ctx: # Weights don't fit on any subset; default UI to 4096 @@ -5578,12 +5699,15 @@ class LlamaCppBackend: ] ) self._tensor_parallel = True + self._layer_preserves_tensor_intent = False logger.info( "Tensor parallelism: --split-mode tensor, --tensor-split %s", tp_tensor_split, ) else: self._tensor_parallel = False + # > 1 only when a tensor request was downgraded but kept multi-GPU. + self._layer_preserves_tensor_intent = _layer_min_gpus > 1 # Speculative decoding. See _build_speculative_flags for the # mode resolution, benchmarks, and llama.cpp references. @@ -5867,7 +5991,17 @@ class LlamaCppBackend: _startup_crashed = ( self._process.poll() is not None and self._process.returncode != 0 ) - if _spawn_attempt == 0 and _fit_retry_allowed and _startup_crashed: + # A split-axis abort (#6415) is fit-independent: skip the + # --fit off retry and let the caller latch it. + _split_axis_crash = self._is_tensor_split_assert( + "\n".join(self._stdout_lines[-50:]) + ) + if ( + _spawn_attempt == 0 + and _fit_retry_allowed + and _startup_crashed + and not _split_axis_crash + ): logger.warning( "llama-server crashed during startup (exit code %s) " "with the default memory-fit step enabled; Studio " @@ -5913,6 +6047,21 @@ class LlamaCppBackend: ) healthy = _spawn_and_wait(cmd) + # #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, + # looping every load. Record and raise to the route's layer fallback, + # skipping the futile flash-attn/MTP retries. + if not healthy and self._tensor_parallel and not self._cancel_event.is_set(): + _ts_out = "\n".join(self._stdout_lines[-50:]) + _ts_rc = self._process.poll() if self._process is not None else None + if self._should_record_tensor_split_abort(_ts_rc, _ts_out): + LlamaCppBackend._record_tensor_split_abort(binary, model_identifier) + self._kill_process() + raise RuntimeError( + "llama-server aborted on --split-mode tensor " + "(split-axis geometry); retrying with layer split." + ) # Flash-attention kernels hard-crash at startup on some ROCm/GPU # builds (frequently inside the vision tower). Disabling FA keeps # both vision and MTP, so retry that way before dropping either. @@ -6057,6 +6206,7 @@ class LlamaCppBackend: # Read the crash code before _kill_process() clears _process. _crash_rc = self._process.poll() if self._process is not None else None self._kill_process() + # The #6415 split-axis abort is latched earlier (first spawn). # Skip if a cancel/unload is pending (mirrors the MTP guard). if ( launched_with_mmproj @@ -6488,6 +6638,7 @@ class LlamaCppBackend: spec_draft_n_max: Optional[int] = None, tensor_parallel: bool = False, mtp_draft_path: Optional[str] = None, + preserve_multi_gpu_on_layer: bool = False, ) -> bool: """True iff the live server already satisfies these load kwargs. @@ -6530,6 +6681,17 @@ class LlamaCppBackend: # server. An identical request would downgrade the same way. if not _tensor_parallel_matches_loaded(extra_args, tensor_parallel, self._tensor_parallel): return False + # Preserved tensor->layer fallback + an EXPLICIT tensor drop: reload so + # placement re-selects instead of keeping the all-GPU mask (mirrors the route, + # #6659). preserve_multi_gpu_on_layer carries the route's carry-forward decision + # (True for an implicit same-settings reload), so those still dedupe -- the HF + # auto-pick / local-dir flows skip the route guard and only reach here. + if ( + self._layer_preserves_tensor_intent + and not _effective_tensor_parallel(extra_args, tensor_parallel) + and not preserve_multi_gpu_on_layer + ): + return False # Compare on the canonical requested mode. With --spec-type in # extra_args the backend stores None; mirror that here. @@ -6641,6 +6803,7 @@ class LlamaCppBackend: self._supports_tools = False self._cache_type_kv = None self._tensor_parallel = False + self._layer_preserves_tensor_intent = False self._speculative_type = None self._requested_spec_mode = None self._spec_draft_n_max = None diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index caf2a262a3..d3f56edec5 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -683,7 +683,9 @@ try: detect_reasoning_flags, ) from core.inference.llama_server_args import ( + _effective_tensor_parallel, _tensor_parallel_matches_loaded, + parse_split_mode_override, resolve_tensor_parallel, strip_shadowing_flags, validate_extra_args, @@ -718,7 +720,9 @@ except ImportError: detect_reasoning_flags, ) from core.inference.llama_server_args import ( + _effective_tensor_parallel, _tensor_parallel_matches_loaded, + parse_split_mode_override, resolve_tensor_parallel, strip_shadowing_flags, validate_extra_args, @@ -2078,6 +2082,32 @@ def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[ ) +def _carry_preserved_tensor_intent( + *, preserved: bool, same_model: bool, explicit_drop: bool +) -> bool: + """Carry a preserved multi-GPU layer fallback forward only for a reload of the + SAME loaded model that doesn't explicitly drop tensor intent, so a fitting model + isn't collapsed to one GPU on a ctx-only change -- but an unrelated model switch + (without /unload) or an explicit tensor-off doesn't inherit it (#6659).""" + return preserved and same_model and not explicit_drop + + +def _is_explicit_tensor_drop(request: LoadRequest) -> bool: + """True only when the request explicitly selects a non-tensor --split-mode (e.g. + layer/row/none), a deliberate departure from a preserved tensor->layer fallback. + + A bare tensor_parallel field is NOT a drop: the Studio UI always sends it and echoes + the /load response's resolved value back, so after a fallback every reload carries + tensor_parallel=false even though the user never changed it -- treating that as a drop + would collapse the preserved multi-GPU placement on the next ctx/settings reload. An + empty clear is not a drop either (a fallback always stores --split-mode layer, never a + tensor split mode, so a clear never wipes tensor intent), nor is an unrelated extra + (--top-k) or inherit (None). tensor_parallel=true / --split-mode tensor re-engage + tensor. Shared by the already-loaded dedup and the load carry-forward (#6659).""" + override = parse_split_mode_override(request.llama_extra_args) + return override is not None and override.strip().lower() != "tensor" + + def _request_matches_loaded_settings( request: LoadRequest, llama_backend: LlamaCppBackend, @@ -2116,6 +2146,13 @@ def _request_matches_loaded_settings( effective_extra, request.tensor_parallel, llama_backend.tensor_parallel ): return False + # Preserved tensor->layer fallback (both report tensor=off, so the check above + # matches): if the user now explicitly drops tensor intent, reload so placement + # re-selects instead of keeping the all-GPU mask (#6659). The effective check + # includes the env, so an env-only tensor (LLAMA_ARG_SPLIT_MODE=tensor) that + # can't actually be dropped falls through to the env-downgrade match, not a loop. + if llama_backend.layer_preserves_tensor_intent and _is_explicit_tensor_drop(request): + return False # Spec decoding works on vision models too (MTP is mmproj-compatible, # llama.cpp #22673; the old ``not is_vision`` gate is gone), so compare # the real requested mode -- coercing vision to ``off`` here used to @@ -2810,6 +2847,48 @@ async def load_model( hf_variant = config.gguf_variant, ) + # Tensor intent for this load: the request itself, or a preserved + # multi-GPU layer fallback carried across a reload of the SAME model that + # doesn't drop it (e.g. a ctx-only change), so a fitting model doesn't + # silently collapse to one GPU. Only an explicit non-tensor --split-mode + # override counts as the drop -- the tensor field echo / unrelated extras keep + # the preserved placement; the same-model guard stops a switch-without-unload + # inheriting the prior model's intent. + _explicit_tensor_drop = _is_explicit_tensor_drop(request) + # Compare the resolved config.identifier (what load_model stores), not the + # raw request id: from_identifier normalizes shorthands (adds unsloth/, fixes + # case), so a reload with the shorthand would otherwise miss the match and + # drop the carry-forward. #6659 + _same_model_loaded = ( + llama_backend.is_loaded + and (llama_backend.model_identifier or "").lower() + == (config.identifier or "").lower() + ) + # model_identifier is variant-agnostic for HF repos and dir-level for a + # local multi-variant directory, so also require the loaded quant to match + # (path else variant, mirroring _already_in_target_state) -- otherwise a + # different variant inherits the prior one's preserved intent. #6659 + if _same_model_loaded: + if config.gguf_file and llama_backend.gguf_path: + try: + _same_model_loaded = ( + Path(llama_backend.gguf_path).resolve() + == Path(config.gguf_file).resolve() + ) + except OSError: + _same_model_loaded = False + else: + _same_model_loaded = (llama_backend.hf_variant or "").lower() == ( + config.gguf_variant or "" + ).lower() + _tensor_intent_overall = _effective_tensor_parallel( + extra_llama_args, request.tensor_parallel + ) or _carry_preserved_tensor_intent( + preserved = llama_backend.layer_preserves_tensor_intent, + same_model = _same_model_loaded, + explicit_drop = _explicit_tensor_drop, + ) + # Run a single load attempt with the given tensor flag + extras. async def _attempt_gguf_load( tensor_parallel: bool, attempt_extra_args: Optional[list[str]] @@ -2823,6 +2902,12 @@ async def load_model( **_source_load_kwargs, **attempt_kwargs, tensor_parallel = tensor_parallel, + # True on the layer fallback retry (tensor wanted overall but not on + # this attempt): keep multi-GPU. Mirrors the fallback's key. + preserve_multi_gpu_on_layer = bool( + _tensor_intent_overall + and not _effective_tensor_parallel(attempt_extra_args, tensor_parallel) + ), ) # Tensor parallelism is arch-gated in llama.cpp and crashes some loads diff --git a/studio/backend/tests/test_tp_vision_regression.py b/studio/backend/tests/test_tp_vision_regression.py new file mode 100644 index 0000000000..09af876da6 --- /dev/null +++ b/studio/backend/tests/test_tp_vision_regression.py @@ -0,0 +1,805 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression guards for silent tensor-parallel downgrades in load_model. + +PR #6416 blanket-disabled tensor parallelism for vision models to dodge a +--split-mode tensor + --mmproj GGML_ASSERT (#6415), which silently single-GPU'd +any mmproj/MTP GGUF that fit on one card. The fix makes the skip self-healing: +tensor is tried by default and recorded per (binary, model) only on a real abort. + +load_model is too entangled to drive end-to-end, so these tests inspect the +source / drive the pure helpers. The headline test pins the set of TP-drop +conditions, so a new silent drop fails CI. No GPU; fully deterministic. +""" + +from __future__ import annotations + +import ast +import importlib.util +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 / httpx / +# loggers -- but only when the real module is missing, so a lightweight stub never +# shadows the real package (or `loggers.handlers` submodule) for tests collected +# later in the same pytest process. +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 as _httpx_real # 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 + +_GB = 1024**3 + + +def _load_inference_routes_module(): + """Load routes/inference.py directly, bypassing routes/__init__.py (which imports + every router, dragging in unrelated deps like python-multipart) (Codex #6659).""" + route_path = Path(_BACKEND_DIR) / "routes" / "inference.py" + spec = importlib.util.spec_from_file_location( + "tp_vision_regression_inference_routes", route_path + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _load_model_ast() -> ast.FunctionDef: + """Parse load_model into an AST FunctionDef (no import side effects).""" + src = textwrap.dedent(inspect.getsource(LlamaCppBackend.load_model)) + return ast.parse(src).body[0] + + +def _tensor_parallel_false_drop_guards() -> list[str]: + """Source of the guard expression for every `if ...: tensor_parallel = False` + (the LOCAL variable, not self._tensor_parallel) inside load_model.""" + fn = _load_model_ast() + + def _body_drops_tp(body) -> bool: + for n in body: + if ( + isinstance(n, ast.Assign) + and any(isinstance(t, ast.Name) and t.id == "tensor_parallel" for t in n.targets) + and isinstance(n.value, ast.Constant) + and n.value.value is False + ): + return True + return False + + return [ + ast.unparse(node.test) + for node in ast.walk(fn) + if isinstance(node, ast.If) and _body_drops_tp(node.body) + ] + + +# Every condition that may flip a requested tensor_parallel back to False. Adding +# one must be conscious: update this allowlist and keep multi-GPU where possible. +_ALLOWED_TP_DROP_GUARDS = { + # Capability: --split-mode tensor aborted for this (binary, model) (#6415). + # Self-healing -- tried by default, skipped only after a real abort (vs #6416). + "tensor_parallel and self._tensor_split_aborts(binary, model_identifier)", + # Capacity: tensor needs >= 2 GPUs clearing the compute-buffer reserve. + "tensor_parallel and len(tp_gpus) < 2", + # Capacity: pooled usable VRAM can't hold weights + MTP reserve -> layer split. + "_tp_weight_budget_mib <= _tp_required_mib", +} + + +def test_tensor_parallel_drop_sites_match_allowlist(): + """The set of reasons a requested TP can be dropped is fixed and reviewed: a new + drop site fails this set-equality until consciously allowlisted (would catch #6416).""" + found = set(_tensor_parallel_false_drop_guards()) + assert found == _ALLOWED_TP_DROP_GUARDS, ( + "tensor_parallel drop sites changed.\n" + f" unexpected (new) : {sorted(found - _ALLOWED_TP_DROP_GUARDS)}\n" + f" missing (removed): {sorted(_ALLOWED_TP_DROP_GUARDS - found)}\n" + "A new drop means a user's TP request is ignored for a new reason -- " + "review it, keep multi-GPU where possible, surface it, then update " + "_ALLOWED_TP_DROP_GUARDS." + ) + + +def test_every_tp_drop_is_logged_not_silent(): + """Each tensor_parallel downgrade must log why, so it never disappears silently.""" + fn = _load_model_ast() + + def _body_drops_tp(body): + return any( + isinstance(n, ast.Assign) + and any(isinstance(t, ast.Name) and t.id == "tensor_parallel" for t in n.targets) + and isinstance(n.value, ast.Constant) + and n.value.value is False + for n in body + ) + + def _body_logs(body) -> bool: + for n in ast.walk(ast.Module(body = list(body), type_ignores = [])): + if ( + isinstance(n, ast.Call) + and isinstance(n.func, ast.Attribute) + and isinstance(n.func.value, ast.Name) + and n.func.value.id == "logger" + ): + return True + return False + + for node in ast.walk(fn): + if isinstance(node, ast.If) and _body_drops_tp(node.body): + assert _body_logs(node.body), ( + f"TP drop under `{ast.unparse(node.test)}` has no logger call -- " + "downgrades must explain themselves." + ) + + +def test_tensor_split_gate_is_self_healing_not_blanket(): + """Skip is conditional on a recorded (binary, model) abort, not a blanket + is_vision disable (the #6416 regression).""" + src = inspect.getsource(LlamaCppBackend.load_model) + assert "self._tensor_split_aborts(binary, model_identifier)" in src + assert "if tensor_parallel and is_vision:" not in src + assert "if tensor_parallel and effective_is_vision:" not in src + + +def test_tensor_split_skip_documents_layer_split_fallback(): + """When the skip fires (known-bad binary+model), it states the fallback.""" + src = inspect.getsource(LlamaCppBackend.load_model) + gate = src.find("self._tensor_split_aborts(binary, model_identifier)") + assert gate != -1 + block = src[gate : gate + 600] + assert "layer split" in block, "the skip should state it falls back to layer split" + + +def test_tensor_split_abort_recorded_early_on_first_spawn(): + """Recorded on the first spawn showing the marker, before the flash-attn-off + retry (which can't run tensor so drops the marker) -- else it loops (oobabooga, #6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + idx = src.find("_record_tensor_split_abort(binary, model_identifier)") + assert idx != -1, "load_model must record a (binary, model) tensor-split abort" + guard = src[max(0, idx - 600) : idx] + assert "self._tensor_parallel" in guard + assert ( + "_should_record_tensor_split_abort" in guard + ), "record must be gated on the marker-plus-hard-crash decision helper" + # Recorded before the flash-attn-off retry, not after the full ladder. + fa_off = src.find("_with_flash_attn_off") + assert 0 <= idx < fa_off, "recording must latch on the first spawn, before flash-off" + + +def test_vision_downgrade_preserves_multi_gpu_intent(): + """The vision downgrade raises _layer_min_gpus and threads it into both the + _select_gpus and auto-context layer paths, so a fitting model still spreads.""" + src = inspect.getsource(LlamaCppBackend.load_model) + assert "_layer_min_gpus = max(_layer_min_gpus, len(gpus))" in src + assert src.count("min_gpus = _layer_min_gpus") >= 2 + assert "range(_auto_min_gpus, len(ranked) + 1)" in src + auto = src.find("_auto_min_gpus = max(") + assert auto != -1 and "_layer_min_gpus" in src[auto : auto + 200] + + +# ── per-binary capability cache (pure) ─────────────────────────────── + + +def test_tensor_attempted_by_default_for_unknown_binary(): + """A (binary, model) not seen to abort -> tensor is attempted (not skipped).""" + assert LlamaCppBackend._tensor_split_aborts("/never/seen/llama-server", "m") is False + assert LlamaCppBackend._tensor_split_aborts(None, "m") is False + assert LlamaCppBackend._tensor_split_aborts("/x", None) is False + + +def test_recorded_tensor_abort_is_per_model(): + """A recorded (binary, model) abort trips the gate for that model only -- a + different model on the same binary still attempts tensor (oobabooga, #6659).""" + b = f"/tmp/llama-server-{id(object())}" + try: + assert LlamaCppBackend._tensor_split_aborts(b, "model-a") is False + LlamaCppBackend._record_tensor_split_abort(b, "model-a") + assert LlamaCppBackend._tensor_split_aborts(b, "model-a") is True + # a different model on the same binary is unaffected + assert LlamaCppBackend._tensor_split_aborts(b, "model-b") is False + finally: + LlamaCppBackend._tensor_split_abort_keys.discard( + LlamaCppBackend._tensor_split_cache_key(b, "model-a") + ) + + +# ── _select_gpus: single-GPU collapse vs honored multi-GPU intent (pure) ── + + +def test_select_gpus_collapses_to_single_gpu_when_model_fits(): + """Default (min_gpus=1): a 39 GB model on four 183 GB GPUs pins ONE GPU -- the + 'single GPU' symptom once TP drops, and why the downgrade needs min_gpus.""" + gpus = [(0, 180000), (1, 180000), (2, 180000), (3, 180000)] # (idx, free MiB) + gpu_indices, _use_fit = LlamaCppBackend._select_gpus(int(39 * _GB), gpus) + assert gpu_indices is not None and len(gpu_indices) == 1 + + +def test_select_gpus_min_gpus_keeps_multi_gpu_for_fitting_model(): + """min_gpus>=2 must NOT collapse to one GPU for a model that fits on one.""" + gpus = [(0, 180000), (1, 180000), (2, 180000), (3, 180000)] + gpu_indices, _ = LlamaCppBackend._select_gpus(int(39 * _GB), gpus, min_gpus = 2) + assert gpu_indices is not None and len(gpu_indices) >= 2 + + +def test_select_gpus_min_gpus_capped_to_available(): + """min_gpus larger than the GPU count is capped, not an error.""" + gpus = [(0, 180000), (1, 180000)] + gi, _ = LlamaCppBackend._select_gpus(int(10 * _GB), gpus, min_gpus = 8) + assert gi is not None and len(gi) == 2 + + +def test_select_gpus_uses_multiple_gpus_when_model_does_not_fit(): + """Sanity: selection spreads across GPUs when one card can't hold the model.""" + gpus = [(0, 40000), (1, 40000), (2, 40000), (3, 40000)] # 40 GB free each + gpu_indices, _use_fit = LlamaCppBackend._select_gpus(int(120 * _GB), gpus) + assert gpu_indices is not None and len(gpu_indices) >= 2 + + +def test_select_gpus_min_gpus_excludes_unusable_gpu(): + """min_gpus caps to usable cards: 2 free + 1 nearly-full -> 2-GPU split, not + forcing the full card (OOM) or tripping --fit (#6659).""" + gpus = [(0, 180000), (1, 180000), (2, 500)] # GPU 2 is nearly full + total = {0: 180000, 1: 180000, 2: 180000} + gi, _ = LlamaCppBackend._select_gpus( + int(39 * _GB), + gpus, + min_gpus = 3, + total_by_idx = total, + per_device_overhead_bytes = int(1 * _GB), + ) + assert gi is not None + assert 2 not in gi, "a nearly-full GPU must not be forced in to satisfy min_gpus" + assert len(gi) == 2 + + +def test_tensor_abort_cache_invalidated_on_binary_mtime_change(tmp_path): + """Cache keys on (path, mtime, model), so a binary swapped in place (in-app + update, no restart) is re-probed instead of inheriting the old abort (#6659).""" + binp = tmp_path / "llama-server" + binp.write_text("v1") + p = str(binp) + try: + LlamaCppBackend._record_tensor_split_abort(p, "m") + assert LlamaCppBackend._tensor_split_aborts(p, "m") is True + # Simulate an in-place update bumping the binary's mtime. + st = binp.stat() + os.utime(p, (st.st_atime, st.st_mtime + 10)) + assert ( + LlamaCppBackend._tensor_split_aborts(p, "m") is False + ), "a binary swapped in place (new mtime) must be re-probed" + # A same-second replacement (sub-second mtime bump) must also re-probe: + # second-resolution mtime would inherit the stale abort (reviewer.py P2). + sec_ns = (binp.stat().st_mtime_ns // 1_000_000_000) * 1_000_000_000 + os.utime(p, ns = (sec_ns, sec_ns)) + LlamaCppBackend._record_tensor_split_abort(p, "m") + binp.write_text("v2") + os.utime(p, ns = (sec_ns, sec_ns + 1)) + assert ( + LlamaCppBackend._tensor_split_aborts(p, "m") is False + ), "a same-second in-place swap (ns mtime bump) must be re-probed" + finally: + for key in list(LlamaCppBackend._tensor_split_abort_keys): + if key and key[0] == p: + LlamaCppBackend._tensor_split_abort_keys.discard(key) + + +def test_tensor_split_abort_raises_early_to_layer_fallback(): + """The first-spawn abort raises to the route's layer fallback (not the text-only + mmproj strip), before the flash-attn-off retry, preserving the projector (#6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + raise_idx = src.find("(split-axis geometry); retrying with layer split") + assert raise_idx != -1, "the split-axis abort must raise to trigger a layer retry" + # raises before both the flash-attn-off retry and the text-only mmproj strip + assert raise_idx < src.find("_with_flash_attn_off") + assert raise_idx < src.find("_strip_mmproj_args(_last_spawn_cmd)") + # gated on the marker-plus-crash helper, which also drives the record just above + guard = src[max(0, raise_idx - 600) : raise_idx] + assert "_should_record_tensor_split_abort" in guard + rec_idx = src.find("_record_tensor_split_abort(binary, model_identifier)") + assert rec_idx != -1 and rec_idx < raise_idx + + +def test_budget_downgrade_preserves_multi_gpu_intent(): + """The pooled-VRAM downgrade raises _layer_min_gpus from the usable tensor GPUs + too, symmetric with the vision downgrade (reviewer.py asymmetric fix, #6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + budget = src.find("_tp_weight_budget_mib <= _tp_required_mib") + assert budget != -1 + block = src[budget : budget + 1000] + assert "tensor_parallel = False" in block + assert ( + "_layer_min_gpus = max(_layer_min_gpus, len(tp_gpus))" in block + ), "the budget downgrade must preserve multi-GPU intent like the vision gate" + + +def test_compute_buffer_downgrade_preserves_multi_gpu_intent(): + """The len(tp_gpus) < 2 compute-buffer downgrade raises _layer_min_gpus from the + full GPU set too, so it is symmetric with the budget/geometry downgrades and + doesn't collapse a multi-GPU layer load to one card (reviewer.py P1 on #6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + gate = src.find("tensor_parallel and len(tp_gpus) < 2") + assert gate != -1 + # Bound to exactly this block: from its gate to the next (budget) downgrade. + nxt = src.find("_tp_weight_budget_mib <= _tp_required_mib", gate) + assert nxt != -1 + block = src[gate:nxt] + assert "tensor_parallel = False" in block + assert ( + "_layer_min_gpus = max(_layer_min_gpus, len(gpus))" in block + ), "the compute-buffer downgrade must preserve multi-GPU intent like the others" + + +def test_tensor_split_layer_min_gpus_bump_requires_tensor_request(): + """Every guard that bumps _layer_min_gpus off the abort cache also tests + tensor_parallel, so a non-tensor load on a known-bad binary doesn't grab every + GPU for a fitting model (#6659).""" + fn = _load_model_ast() + checked = 0 + for node in ast.walk(fn): + if isinstance(node, ast.If): + test_src = ast.unparse(node.test) + if "self._tensor_split_aborts(binary, model_identifier)" not in test_src: + continue + body = "\n".join(ast.unparse(n) for n in node.body) + if "_layer_min_gpus" in body: + checked += 1 + assert "tensor_parallel" in test_src, ( + "the cached _layer_min_gpus bump must require a current tensor " + f"request, but fires under `{test_src}`" + ) + assert checked >= 1, "expected an abort-cache guard that bumps _layer_min_gpus" + + +# ── round-2 follow-up: route-fallback retry + auto-context cap + assert marker ── + + +def test_layer_fallback_retry_preserves_multi_gpu_intent(): + """load_model takes a preserve_multi_gpu_on_layer hint and raises _layer_min_gpus + for it, so the tensor-off fallback retry still spreads a fitting model (#6659).""" + sig = inspect.signature(LlamaCppBackend.load_model) + assert "preserve_multi_gpu_on_layer" in sig.parameters + assert sig.parameters["preserve_multi_gpu_on_layer"].default is False + fn = _load_model_ast() + found = any( + isinstance(n, ast.If) + and "preserve_multi_gpu_on_layer" in ast.unparse(n.test) + and "_layer_min_gpus" in "\n".join(ast.unparse(b) for b in n.body) + for n in ast.walk(fn) + ) + assert found, "preserve_multi_gpu_on_layer must raise _layer_min_gpus" + + +def test_auto_context_layer_loops_capped_to_usable_gpus(): + """The auto-context loops bypass _select_gpus, so they apply its cap: a card + counts only if usable VRAM clears the per-device layer overhead (#6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + assert ( + "range(max(1, _layer_min_gpus), len(ranked) + 1)" not in src + ), "auto-context loops must cap _layer_min_gpus to usable GPUs, not use it raw" + assert "_auto_min_gpus" in src + assert "range(_auto_min_gpus, len(ranked) + 1)" in src + # the eligibility threshold is the per-device layer overhead, not bare > 0 + auto = src.find("_auto_min_gpus = max(") + assert auto != -1 + block = src[auto : auto + 400] + assert "_pipeline_overhead_mib" in block, ( + "a card must clear the per-device layer overhead to count, mirroring " + "_select_gpus, so a nearly-full GPU is not exposed and OOMs" + ) + + +def test_fallback_hint_uses_effective_tensor_request_not_just_toggle(): + """Tensor intent keys off _effective_tensor_parallel (toggle + extras + env), not + just the toggle, so extra/env-driven tensor users keep multi-GPU (#6659).""" + route = Path(_BACKEND_DIR) / "routes" / "inference.py" + src = route.read_text() + idx = src.find("_tensor_intent_overall = _effective_tensor_parallel(") + assert idx != -1, "the GGUF load closure must compute tensor intent" + block = src[idx : idx + 300] + assert "extra_llama_args, request.tensor_parallel" in block + pres = src.find("preserve_multi_gpu_on_layer = bool(") + assert ( + "_effective_tensor_parallel(attempt_extra_args, tensor_parallel)" in src[pres : pres + 200] + ) + # not the toggle-only form this replaced + assert ( + "bool(\n request.tensor_parallel and not tensor_parallel" not in src + ) + + +def test_carry_preserved_tensor_intent_truth_table(): + """Behavioral check of the carry-forward decision: carried only for the SAME + model, preserved, and not an explicit drop. Catches a `not` inversion (ctx-only + collapse) and a missing same-model guard (cross-model leak) (#6659).""" + inference_routes = _load_inference_routes_module() + f = inference_routes._carry_preserved_tensor_intent + assert f(preserved = True, same_model = True, explicit_drop = False) is True + assert f(preserved = True, same_model = True, explicit_drop = True) is False # explicit drop + assert f(preserved = True, same_model = False, explicit_drop = False) is False # model switch + assert f(preserved = False, same_model = True, explicit_drop = False) is False # not a fallback + + +def test_preserved_fallback_carried_across_non_drop_reload(): + """The hint carries the preserved fallback via _carry_preserved_tensor_intent, + gated on the same model loaded, so a ctx-only reload keeps multi-GPU but a model + switch / explicit drop doesn't inherit it (#6659).""" + route = Path(_BACKEND_DIR) / "routes" / "inference.py" + src = route.read_text() + idx = src.find("_tensor_intent_overall = _effective_tensor_parallel(") + assert idx != -1 + block = src[idx : idx + 400] + assert "_carry_preserved_tensor_intent(" in block + assert "preserved = llama_backend.layer_preserves_tensor_intent" in block + assert "same_model = _same_model_loaded" in block + assert "explicit_drop = _explicit_tensor_drop" in block + + +def test_same_model_guard_checks_path_and_variant(): + """The same-model guard matches the resolved config.identifier (what load_model + stores, after from_identifier normalizes shorthands) -- not the raw request id -- + and also matches the loaded quant by path (local multi-variant dir) else variant (HF + repo), so a reload keeps the carry-forward and a different variant doesn't inherit + the prior one's preserved tensor intent (#6659).""" + route = Path(_BACKEND_DIR) / "routes" / "inference.py" + src = route.read_text() + idx = src.find("_same_model_loaded = (") + assert idx != -1 + block = src[idx : idx + 1300] + # Identity compares the normalized config.identifier, not the raw model_identifier. + head = src[idx : idx + 200] + assert "config.identifier" in head and "== (model_identifier" not in head + assert "llama_backend.gguf_path" in block and "config.gguf_file" in block + assert "llama_backend.hf_variant" in block and "config.gguf_variant" in block + + +def test_diffusion_load_clears_preserved_tensor_flag(): + """The diffusion early-return path (skips the command builder) clears the + preserved-fallback flag, so a prior tensor fallback doesn't churn it (#6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + diff = src.find("if self._is_diffusion:") + assert diff != -1 + start = src.find("return self._start_diffusion_server", diff) + assert start != -1 + assert "self._layer_preserves_tensor_intent = False" in src[diff:start] + + +def test_is_tensor_split_assert_marker(): + """Matches the specific #6415 split-axis assert, not any ggml assert/abort, so + an unrelated invariant a corrupt GGUF/projector trips isn't cached (#6659).""" + f = LlamaCppBackend._is_tensor_split_assert + # the real #6415 warmup assert (split-axis enum, in ggml-backend-meta) + assert ( + f( + "ggml-backend-meta.cpp:541: GGML_ASSERT(src_ss[0].axis != " + "GGML_BACKEND_SPLIT_AXIS_0) failed" + ) + is True + ) + # the split-axis token alone (file path elided / reworded) still matches + assert f("GGML_ASSERT(x.axis != GGML_BACKEND_SPLIT_AXIS_1) failed") is True + # UNRELATED asserts must NOT match -- including a different invariant from the + # same multi-assert source file (matched on the token, not the file name). + assert f("ggml-backend-meta.cpp:99: GGML_ASSERT(buf != NULL) failed") is False + assert f("/x/ggml.c:1234: GGML_ASSERT(ne == 1) failed") is False + assert f("ggml_abort: something else entirely") is False + assert f("Segmentation fault (core dumped)") is False + assert f("") is False + assert f(None) is False + + +def test_layer_preserve_hint_replayed_on_respawn(): + """The preserve hint is in the replay snapshot (_pending_load_kwargs), so a + respawn keeps the downgraded model multi-GPU (Codex review on #6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + pend = src.find("_pending_load_kwargs = {") + assert pend != -1 + block = src[pend : src.find("}", pend) + 1] + assert '"preserve_multi_gpu_on_layer": preserve_multi_gpu_on_layer' in block, ( + "the layer-preserve hint must be in the replay snapshot so _respawn_if_dead " + "keeps the multi-GPU placement" + ) + + +def test_should_record_tensor_split_abort_decision(): + """Behavioral check of marker AND (signal crash OR Windows abort), so an + or->and typo or caching a generic crash fails here, not just the source pins.""" + f = LlamaCppBackend._should_record_tensor_split_abort + marker = "ggml-backend-meta.cpp:541: GGML_ASSERT(x.axis != GGML_BACKEND_SPLIT_AXIS_0) failed" + # marker + a hard crash records, across every platform's abort encoding + assert f(-6, marker) is True # POSIX SIGABRT + assert f(-11, marker) is True # POSIX SIGSEGV + assert f(3, marker) is True # Windows CRT abort() exit (not a signal) + assert f(0xC0000005, marker) is True # Windows NTSTATUS access violation + # marker present but no hard crash -> not recorded + assert f(0, marker) is False # clean exit + assert f(-9, marker) is False # SIGKILL (OOM / unload), not a fault + assert f(None, marker) is False # still running + # hard crash but not the split-axis marker -> not recorded (no over-caching) + assert f(3, "some other failure") is False + assert f(-6, "GGML_ASSERT(buf != NULL) failed") is False + assert f(0xC0000005, "") is False + + +def test_fit_off_retry_skipped_on_split_axis_abort(): + """The fit-independent --fit off retry is skipped on the split-axis marker, else + the model crashes a second time before the latch records it (reviewer.py, #6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + retry = src.find('run_cmd = [*run_cmd, "--fit", "off"]') + assert retry != -1 + guard = src[max(0, retry - 1000) : retry] + assert "_fit_retry_allowed" in guard and "_startup_crashed" in guard + assert ( + "not _split_axis_crash" in guard + ), "the fit-off retry must be skipped when the crash is a split-axis abort" + + +def test_is_abort_exit_recognizes_windows_crt_abort(): + """exit code 3 (MSVC abort()) counts as a crash; signals / clean exits do not.""" + f = LlamaCppBackend._is_abort_exit + assert f(3) is True + assert f(0) is False + assert f(-6) is False # POSIX SIGABRT is handled by _is_signal_crash, not here + assert f(None) is False + + +# ── tensor-off after a multi-GPU fallback forces a reload (route dedup) ─ + + +class _NoopProcess: + """Stand-in for Popen so is_loaded is True and atexit cleanup doesn't crash.""" + + def terminate(self): + pass + + def wait(self, timeout = None): + return 0 + + def kill(self): + pass + + def poll(self): + return 0 + + +def _fallback_loaded_backend(layer_preserves_tensor_intent: bool) -> LlamaCppBackend: + """A loaded backend in the tensor->layer fallback state (tensor off, --split-mode + layer stored), differing only in the preserved-multi-GPU flag.""" + b = LlamaCppBackend() + b._model_identifier = "owner/repo" + b._requested_n_ctx = 0 + b._cache_type_kv = None + b._tensor_parallel = False + b._layer_preserves_tensor_intent = layer_preserves_tensor_intent + b._extra_args = ["--split-mode", "layer"] + b._requested_spec_mode = "auto" + b._chat_template_override = None + b._gguf_path = None + return b + + +def test_tensor_off_echo_preserves_multi_gpu_fallback(): + """The Studio UI always sends tensor_parallel and echoes the /load response's + resolved value, so after a fallback a ctx/settings reload carries tensor_parallel= + false even though the user never changed it. That echo must NOT collapse the + preserved multi-GPU placement -- it dedupes (Codex #6659).""" + from models.inference import LoadRequest + + inference_routes = _load_inference_routes_module() + + req = LoadRequest(model_path = "owner/repo", tensor_parallel = False) + assert "tensor_parallel" in req.model_fields_set, "the UI always sends the field" + + # Preserved fallback + bare tensor=false echo: dedupe, keep multi-GPU (no collapse). + assert ( + inference_routes._request_matches_loaded_settings( + req, _fallback_loaded_backend(layer_preserves_tensor_intent = True) + ) + is True + ) + # A genuine layer load (no preserved intent): tensor-off also dedupes, no churn. + assert ( + inference_routes._request_matches_loaded_settings( + req, _fallback_loaded_backend(layer_preserves_tensor_intent = False) + ) + is True + ) + + +def test_explicit_split_mode_layer_extras_reloads_after_multi_gpu_fallback(): + """Tensor intent can be dropped via extras too: an explicit --split-mode layer + matches the stored fallback extras but must still reload (reviewer.py P1, #6659).""" + from models.inference import LoadRequest + + inference_routes = _load_inference_routes_module() + + req = LoadRequest(model_path = "owner/repo", llama_extra_args = ["--split-mode", "layer"]) + assert "llama_extra_args" in req.model_fields_set + assert ( + inference_routes._request_matches_loaded_settings( + req, _fallback_loaded_backend(layer_preserves_tensor_intent = True) + ) + is False + ) + + +def test_tensor_off_reload_requires_explicit_toggle(): + """An Apply that doesn't touch the toggle (e.g. a context change) isn't churned + by the preserved-fallback reload -- the working server is kept (Codex #6659).""" + from models.inference import LoadRequest + + inference_routes = _load_inference_routes_module() + + req = LoadRequest(model_path = "owner/repo") # tensor_parallel left unset + assert "tensor_parallel" not in req.model_fields_set + assert ( + inference_routes._request_matches_loaded_settings( + req, _fallback_loaded_backend(layer_preserves_tensor_intent = True) + ) + is True + ) + + +def test_tensor_off_under_env_tensor_does_not_reload_loop(monkeypatch): + """With LLAMA_ARG_SPLIT_MODE=tensor set, a tensor-off request can't drop tensor + intent, so the env-aware guard dedupes instead of reload-looping (Codex #6659).""" + from models.inference import LoadRequest + + inference_routes = _load_inference_routes_module() + monkeypatch.setenv("LLAMA_ARG_SPLIT_MODE", "tensor") + + req = LoadRequest(model_path = "owner/repo", tensor_parallel = False) + assert "tensor_parallel" in req.model_fields_set + # env still forces tensor -> not a real drop -> dedupe (no reload loop). + assert ( + inference_routes._request_matches_loaded_settings( + req, _fallback_loaded_backend(layer_preserves_tensor_intent = True) + ) + is True + ) + + +def test_is_explicit_tensor_drop_truth_table(): + """Only an explicit non-tensor --split-mode override is a drop. A bare + tensor_parallel field (the UI always sends it and echoes the fallback's false), an + empty clear, an unrelated extra (--top-k), or inherit (None) must NOT collapse a + preserved fallback; --split-mode tensor / tensor_parallel=true re-engage (Codex + #6659).""" + from models.inference import LoadRequest + + f = _load_inference_routes_module()._is_explicit_tensor_drop + # A non-tensor split-mode override is the one deliberate departure -> drop. + assert ( + f(LoadRequest(model_path = "owner/repo", llama_extra_args = ["--split-mode", "layer"])) is True + ) + # tensor / retry re-engages, never a drop. + assert ( + f(LoadRequest(model_path = "owner/repo", llama_extra_args = ["--split-mode", "tensor"])) + is False + ) + # A bare tensor_parallel field is the UI echo, not a drop (would collapse on reload). + assert f(LoadRequest(model_path = "owner/repo", tensor_parallel = False)) is False + assert f(LoadRequest(model_path = "owner/repo", tensor_parallel = True)) is False + # Unrelated extra / empty clear / inherit all keep the preserved placement. + assert f(LoadRequest(model_path = "owner/repo", llama_extra_args = ["--top-k", "20"])) is False + assert f(LoadRequest(model_path = "owner/repo", llama_extra_args = [])) is False + assert f(LoadRequest(model_path = "owner/repo")) is False + + +def test_explicit_tensor_drop_uses_shared_helper_in_both_readers(): + """Both the already-loaded dedup and the load carry-forward derive the drop from + _is_explicit_tensor_drop, so they agree on what counts as a drop -- a reload for + an unrelated extra still carries the preserved intent rather than collapsing to one + GPU (Codex #6659).""" + src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text() + # Dedup reader (the preserved-fallback reload guard). + assert "layer_preserves_tensor_intent and _is_explicit_tensor_drop(request)" in src + # Load carry-forward reader feeds the same decision into the carry-forward. + assert "_explicit_tensor_drop = _is_explicit_tensor_drop(request)" in src + + +def test_layer_preserves_tensor_intent_set_only_on_preserved_downgrade(): + """load_model latches the flag from _layer_min_gpus (raised only when a tensor + request is downgraded but kept multi-GPU), and clears it when tensor stays on.""" + src = inspect.getsource(LlamaCppBackend.load_model) + on = src.find("self._tensor_parallel = True") + off = src.find("self._tensor_parallel = False") + assert 0 <= on and 0 <= off + assert "self._layer_preserves_tensor_intent = False" in src[on : on + 120] + assert "self._layer_preserves_tensor_intent = _layer_min_gpus > 1" in src[off : off + 400] + + +def test_layer_min_gpus_bound_before_gpu_selection_try(): + """_layer_min_gpus is bound before the GPU-selection try, so the --fit-on except + path can't UnboundLocalError when the command builder reads it (Codex #6659).""" + src = inspect.getsource(LlamaCppBackend.load_model) + assert src.count("_layer_min_gpus = 1") == 1, "exactly one init, before the try" + init = src.find("_layer_min_gpus = 1") + try_body = src.find("gguf_size = self._get_gguf_size_bytes") + fit_except = src.find("GPU selection failed") + use_after = src.find("self._layer_preserves_tensor_intent = _layer_min_gpus > 1") + assert ( + -1 < init < try_body < fit_except < use_after + ), "the init must precede the try body, the except, and the command-builder use" + + +def test_already_in_target_state_reloads_on_tensor_off_after_fallback(): + """The backend fast path mirrors the route dedup: a preserved fallback reloads on + an EXPLICIT tensor-off request, but an implicit same-settings reload (carry-forward + preserve_multi_gpu_on_layer=True) still dedupes (Codex #6659).""" + + def _backend(layer_preserves: bool) -> LlamaCppBackend: + b = _fallback_loaded_backend(layer_preserves_tensor_intent = layer_preserves) + b._process = _NoopProcess() + b._healthy = True + return b + + kwargs = dict( + gguf_path = None, + mtp_draft_path = None, + model_identifier = "owner/repo", + hf_variant = None, + n_ctx = 0, + cache_type_kv = None, + speculative_type = None, + spec_draft_n_max = None, + tensor_parallel = False, + chat_template_override = None, + extra_args = ["--split-mode", "layer"], + is_vision = False, + ) + # Preserved fallback + EXPLICIT tensor drop -> reload (not already in target state). + assert _backend(True)._already_in_target_state(**kwargs) is False + # Same preserved fallback but an implicit reload that carries the intent forward + # (HF auto-pick / local-dir flows skip the route guard and reach here) -> dedupe. + assert ( + _backend(True)._already_in_target_state(**kwargs, preserve_multi_gpu_on_layer = True) is True + ) + # A genuine layer load (no preserved intent) -> dedupe, no churn. + assert _backend(False)._already_in_target_state(**kwargs) is True