diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 062944264a..e7030b77ce 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1213,8 +1213,9 @@ class LlamaCppBackend: self._nextn_predict_layers: Optional[int] = None self._lock = threading.Lock() # Wraps load_model() end-to-end so concurrent loads serialise and never - # coexist as two llama-server processes (#5401). - self._serial_load_lock = threading.Lock() + # coexist as two llama-server processes (#5401). RLock so MTP-crash + # recovery can re-acquire it for its nested load_model. + self._serial_load_lock = threading.RLock() # Set by the in-app updater while it swaps prebuilt binaries; load_model() # rejects fast so no server starts from a half-swapped binary. self._llama_update_in_progress = False @@ -1225,6 +1226,18 @@ class LlamaCppBackend: self._extra_args: Optional[List[str]] = None self._extra_args_source: Optional[tuple[str, Optional[str]]] = None self._requested_n_ctx: int = 0 + # Raw kwargs of the last healthy load, for the MTP-crash reload. Memory-only + # (carries hf_token, never logged); single-flight via the lock below. + self._last_load_kwargs: Optional[dict] = None + self._mtp_runtime_fallback_lock = threading.Lock() + self._mtp_runtime_fallback_in_progress = False + # Background watchdog so an MTP+tensor crash recovers even when no request + # observes it (direct proxy endpoints, or nothing in flight). + self._mtp_watchdog_thread: Optional[threading.Thread] = None + self._mtp_watchdog_stop = threading.Event() + # True when the launch actually runs MTP+tensor (Studio- or user/env-driven); + # gates the probe, watchdog, and recovery so pass-through MTP is covered. + self._mtp_runtime_fallback_active = False self._stdout_lines: list[str] = [] self._stdout_thread: Optional[threading.Thread] = None # llama-server tee log (see _drain_stdout / _kill_process). @@ -4156,6 +4169,28 @@ class LlamaCppBackend: Returns True if the server started and the health check passed. """ + # Raw load inputs so the runtime MTP-crash reload can replay this model + # without MTP. Committed to _last_load_kwargs only on a healthy load. + _pending_load_kwargs = { + "gguf_path": gguf_path, + "mmproj_path": mmproj_path, + "mtp_draft_path": mtp_draft_path, + "hf_repo": hf_repo, + "hf_variant": hf_variant, + "hf_token": hf_token, + "model_identifier": model_identifier, + "is_vision": is_vision, + "n_ctx": n_ctx, + "chat_template_override": chat_template_override, + "cache_type_kv": cache_type_kv, + "speculative_type": speculative_type, + "spec_draft_n_max": spec_draft_n_max, + "tensor_parallel": tensor_parallel, + "n_threads": n_threads, + "n_gpu_layers": n_gpu_layers, + "n_parallel": n_parallel, + "extra_args": list(extra_args) if extra_args is not None else None, + } # Serialise the whole load so concurrent /load calls never leave two # llama-server processes alive (#5401 / #5161). Doesn't block /unload. with self._serial_load_lock: @@ -5481,6 +5516,37 @@ class LlamaCppBackend: ) healthy = _spawn_and_wait(cmd) + # MTP from Studio's spec flags or the user's (extra_args + # --spec-type / LLAMA_ARG_SPEC_TYPE). The env reaches the child + # only when neither emits a spec flag, so consult it only then. + _launch_spec_env: Mapping[str, str] = ( + os.environ + if (not _extra_args_set_spec_type(extra_args) and not spec_flags) + else {} + ) + _spec_requested_mtp = any( + "mtp" in str(t).lower() for t in spec_flags + ) or _extra_args_requests_mtp(extra_args, env = _launch_spec_env) + # Is the launched server actually running MTP+tensor? Gates the + # probe/watchdog/recovery; cleared if the MTP-drop fallback wins. + _mtp_active_for_launched_server = bool( + self._tensor_parallel and _spec_requested_mtp + ) + # MTP can pass /health then crash the flash-attn kernel on the + # first decode under tensor; probe one generation so the fallback + # catches that too. Tensor-only, so ordinary MTP stays probe-free. + if ( + healthy + and self._tensor_parallel + and _spec_requested_mtp + and not self._cancel_event.is_set() + and not self._probe_mtp_decode() + ): + logger.warning( + "MTP speculative decoding crashed on the first decode " + "under tensor parallelism; retrying without it." + ) + healthy = False # Any MTP request can abort the server: a separate drafter # (Gemma) on a binary that predates its arch, or an embedded # head (Qwen) the binary cannot build. Retry once with the @@ -5488,8 +5554,8 @@ class LlamaCppBackend: # loads. Gate on the spec block (not the drafter path, which # off/ngram local loads also carry) and keep # _requested_spec_mode so a duplicate /load doesn't thrash. The - # cancel check stops an /unload-killed attempt respawning. - _spec_requested_mtp = any("mtp" in str(t).lower() for t in spec_flags) + # 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(): # Blame the binary only when the output shows MTP itself # failing (unknown arch / draft or context build); an @@ -5535,9 +5601,14 @@ class LlamaCppBackend: + ["--spec-default"] + cmd[_spec_start + len(spec_flags) :] ) + # User/env MTP survives in the tail; llama.cpp takes the last + # spec flag, so a trailing --spec-default overrides it too. + if _extra_args_requests_mtp(extra_args, env = _launch_spec_env): + fallback_cmd.append("--spec-default") healthy = _spawn_and_wait(fallback_cmd, label = "-retry") if healthy: self._speculative_type = "default" + _mtp_active_for_launched_server = False # A vision GGUF launched with --mmproj can abort when the # installed llama.cpp is too old for the model's projector @@ -5587,6 +5658,11 @@ class LlamaCppBackend: self._extra_args = list(extra_args) self._extra_args_source = (model_identifier, hf_variant) self._requested_n_ctx = int(n_ctx) + # Commit the known-good snapshot + whether MTP+tensor is live, then + # watch this load for a mid-generation crash. + self._last_load_kwargs = _pending_load_kwargs + self._mtp_runtime_fallback_active = _mtp_active_for_launched_server + self._start_mtp_crash_watchdog() # Catch silent CPU fallback when GPU was intended (#5106). self._gpu_offload_active = self._classify_gpu_offload( @@ -6050,6 +6126,8 @@ class LlamaCppBackend: self._hf_repo = None self._mtp_draft_path = None self._spec_fallback_reason = None + self._last_load_kwargs = None + self._mtp_runtime_fallback_active = False self._hf_variant = None self._is_vision = False self._is_audio = False @@ -6113,6 +6191,9 @@ class LlamaCppBackend: def _kill_process(self): """Terminate the subprocess if running.""" + # Stop the watchdog before a deliberate kill so a planned reload/unload + # isn't seen as a crash; a real crash never routes through here. + self._stop_mtp_crash_watchdog() if self._process is None: return try: @@ -6342,6 +6423,139 @@ class LlamaCppBackend: return False return True + def _probe_mtp_decode(self, timeout: float = 60.0) -> bool: + """One tiny /completion to confirm MTP survives the first decode. + + MTP-draft can pass /health yet crash the flash-attn kernel only once + tokens generate (e.g. under --split-mode tensor). False on any error so + the caller can drop MTP and retry. + """ + url = f"http://127.0.0.1:{self._port}/completion" + payload = {"prompt": "Hi", "n_predict": 4, "temperature": 0.0, "stream": False} + # Match the --api-key auth direct-stream mode uses, else a spurious 401. + headers = {"Authorization": f"Bearer {self._api_key}"} if self._api_key else None + try: + resp = httpx.post(url, json = payload, timeout = timeout, headers = headers) + except Exception as e: + logger.debug(f"MTP decode probe failed: {e}") + return False + if resp.status_code != 200: + logger.debug(f"MTP decode probe returned HTTP {resp.status_code}") + return False + # A crash can drop the connection or kill the process right after a reply. + if self._process is not None and self._process.poll() is not None: + return False + return True + + def _maybe_recover_from_mtp_crash(self, exc: Optional[BaseException] = None) -> bool: + """Schedule one background reload without MTP after a mid-generation death. + + MTP+tensor can crash the flash-attn kernel on a later request, after + load_model returned, past the load-time fallback and decode probe. Not a + persistent ban: a fresh load re-tries MTP. Returns True if scheduled. + """ + # Cheap async-safe gate: only our live MTP+tensor launch, not cancelled, + # with a snapshot to replay. + if self._cancel_event.is_set(): + return False + if not self._mtp_runtime_fallback_active: + return False + if not self._last_load_kwargs or self._process is None: + return False + # Single-flight: the first failure claims the reload. + with self._mtp_runtime_fallback_lock: + if self._mtp_runtime_fallback_in_progress: + return False + self._mtp_runtime_fallback_in_progress = True + snapshot = dict(self._last_load_kwargs) + proc = self._process + + def _recover(): + try: + # Confirm the process really exited (the error can arrive a beat + # early) so a transient stream error can't disable MTP. + deadline = time.monotonic() + 5.0 + while proc.poll() is None and time.monotonic() < deadline: + time.sleep(0.1) + if proc.poll() is None: + logger.debug("Generation error but llama-server is alive; keeping MTP.") + return + logger.warning( + "llama-server exited mid-generation with MTP under tensor " + "parallelism (%s); reloading without speculative decoding.", + type(exc).__name__ if exc is not None else "server exited", + ) + # Re-check under the load lock (RLock allows the nested + # load_model) so a newer load isn't clobbered by this stale replay. + requested_mode = snapshot.get("speculative_type") + with self._serial_load_lock: + if self._cancel_event.is_set(): + logger.info("MTP-crash reload skipped: load was cancelled/unloaded.") + return + if self._process is not proc: + logger.info("MTP-crash reload skipped: a newer load is already active.") + return + if self._last_load_kwargs != snapshot: + logger.info("MTP-crash reload skipped: load settings changed.") + return + snapshot["speculative_type"] = "off" + # Drop user/env MTP too: append a last-wins --spec-default. + _ea = list(snapshot.get("extra_args") or []) + if _extra_args_requests_mtp(_ea, env = os.environ): + _ea.append("--spec-default") + snapshot["extra_args"] = _ea + self.load_model(**snapshot) + # Restore the requested mode + reason load_model("off") cleared, + # so /status shows the user's mode + note (like the startup fallback). + self._requested_spec_mode = _canonicalize_spec_mode(requested_mode) + self._spec_fallback_reason = "runtime_error" + logger.info("Reloaded without MTP after the tensor-parallel crash.") + except Exception as e: + logger.error(f"Reload without MTP failed: {e}") + finally: + with self._mtp_runtime_fallback_lock: + self._mtp_runtime_fallback_in_progress = False + + threading.Thread(target = _recover, daemon = True, name = "mtp-crash-reload").start() + return True + + def _start_mtp_crash_watchdog(self) -> None: + """Background poll that recovers on an MTP+tensor crash even when no + request observes it (direct proxy endpoints, or nothing in flight). + + Armed only for a live MTP+tensor launch; the no-MTP reload disarms it, so + it can't loop. + """ + if not self._mtp_runtime_fallback_active: + return + proc = self._process + if proc is None: + return + # Replace any prior watchdog (loads are serialised, so at most one). + self._stop_mtp_crash_watchdog() + stop = threading.Event() + self._mtp_watchdog_stop = stop + + def _watch(): + # Exit on stop or process death. _kill_process sets stop before + # terminating, so re-check it: only a real crash (stop unset) recovers. + while not stop.wait(1.0): + if proc.poll() is not None: + if not stop.is_set(): + self._maybe_recover_from_mtp_crash() + return + + t = threading.Thread(target = _watch, daemon = True, name = "mtp-crash-watchdog") + self._mtp_watchdog_thread = t + t.start() + + def _stop_mtp_crash_watchdog(self) -> None: + """Signal the crash watchdog to exit; called before any deliberate kill.""" + stop = getattr(self, "_mtp_watchdog_stop", None) + if stop is not None: + stop.set() + self._mtp_watchdog_thread = None + def _wait_for_health( self, timeout: float = 120.0, @@ -6821,11 +7035,15 @@ class LlamaCppBackend: "finish_reason": _metadata_finish_reason, } - except httpx.ConnectError: + except httpx.ConnectError as e: + # Server already down (e.g. crashed on a prior request): recover MTP. + self._maybe_recover_from_mtp_crash(e) raise RuntimeError("Lost connection to llama-server") except Exception as e: if cancel_event is not None and cancel_event.is_set(): return + # Died mid-generation: recover MTP, re-raise unchanged for this request. + self._maybe_recover_from_mtp_crash(e) raise # ── Tool-calling agentic loop ────────────────────────────── diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index bfceb9b904..be8bf77eb6 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -4869,6 +4869,8 @@ async def openai_chat_completions( tb = traceback.format_exc() logger.error(f"Error during GGUF tool streaming: {e}\n{tb}") api_monitor.fail(monitor_id, _friendly_error(e)) + # Recover if an MTP+tensor crash killed the server mid-stream. + get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) error_chunk = _openai_stream_error_chunk(e) yield f"data: {json.dumps(error_chunk)}\n\n" finally: @@ -5115,6 +5117,8 @@ async def openai_chat_completions( except Exception as e: logger.error(f"Error during GGUF completion: {e}", exc_info = True) api_monitor.fail(monitor_id, _friendly_error(e)) + # Recover if an MTP+tensor crash killed the server. + get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) # An over-context prompt makes llama-server return 400; map any # upstream 4xx to a 400 client error rather than leaking a 500. _cls = _classify_llama_generation_error(e) @@ -8641,6 +8645,7 @@ async def _anthropic_passthrough_stream( except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e: if not cancel_event.is_set(): logger.error("anthropic_messages passthrough stream error: %s", e) + get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) event = _anthropic_stream_error_event( e, force = True, @@ -8651,6 +8656,7 @@ async def _anthropic_passthrough_stream( except Exception as e: if not cancel_event.is_set(): logger.error("anthropic_messages passthrough stream error: %s", e) + get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) event = _anthropic_stream_error_event( e, force = True, @@ -9282,11 +9288,12 @@ async def _openai_passthrough_stream( except asyncio.CancelledError: api_monitor.finish(monitor_id, "cancelled") raise - except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError): + except (httpx.RemoteProtocolError, httpx.ReadError, httpx.CloseError) as e: # Watcher closed resp on cancel. Emit nothing extra; the client # initiated the cancel or already disconnected. if not cancel_event.is_set(): api_monitor.fail(monitor_id, "Stream interrupted") + get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) raise api_monitor.finish(monitor_id, "cancelled") except Exception as e: @@ -9296,6 +9303,7 @@ async def _openai_passthrough_stream( # 200 headers already flushed; errors must go in the SSE body. logger.error("openai passthrough stream error: %s", e) api_monitor.fail(monitor_id, _friendly_error(e)) + get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) err = _openai_stream_error_chunk(e) yield f"data: {json.dumps(err)}\n\n" finally: @@ -9374,6 +9382,7 @@ async def _openai_passthrough_non_streaming( # a bare 500 with no diagnostic. logger.error("openai passthrough non-streaming: upstream unreachable: %s", e) api_monitor.fail(monitor_id, _friendly_error(e)) + get_llama_cpp_backend()._maybe_recover_from_mtp_crash(e) raise HTTPException( status_code = 502, detail = _friendly_error(e), diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py index 8ec629bd2c..4ddd0224a2 100644 --- a/studio/backend/tests/test_tensor_parallel.py +++ b/studio/backend/tests/test_tensor_parallel.py @@ -20,6 +20,8 @@ from __future__ import annotations import asyncio import inspect import sys +import threading +import time import types as _types from pathlib import Path @@ -265,6 +267,398 @@ def test_proportional_tensor_split_is_emitted_in_tensor_mode(): assert 0 <= gate < ts < nxt_else, "--tensor-split must be emitted under `if tensor_parallel:`" +def test_mtp_decode_probe_wired_under_tensor_parallel(): + # MTP-draft can pass /health and crash the CUDA FA kernel only on the first + # decode under --split-mode tensor. Rather than statically banning MTP+TP + # (which a future llama.cpp may support), load_model probes a decode and + # routes a failure into the existing MTP-drop fallback. + src = _load_model_source() + probe = src.find("_probe_mtp_decode()") + assert probe != -1, "load_model must decode-probe MTP under tensor parallelism" + # Gated on tensor mode AND an MTP request (ordinary MTP loads stay unprobed). + guard = src[max(0, probe - 400) : probe] + assert "self._tensor_parallel" in guard and "_spec_requested_mtp" in guard + # A failed probe flips healthy so the shared MTP-drop fallback fires. + assert "healthy = False" in src[probe : probe + 400] + fallback = src.find("if not healthy and _spec_requested_mtp") + assert 0 <= probe < fallback, "the probe must precede the MTP-drop fallback" + + +def test_probe_mtp_decode_returns_false_on_crash(monkeypatch): + # The probe is the decode-time health gate: True only on a clean 200 from a + # live server; any error (dropped connection, non-200, dead process) is a + # failed probe so the caller drops MTP and retries. + backend = LlamaCppBackend() + backend._port = 0 + + class _Resp: + def __init__(self, code): + self.status_code = code + + backend._process = None # liveness check skipped; exercise the HTTP result + monkeypatch.setattr(llama_cpp_module.httpx, "post", lambda *a, **k: _Resp(200), raising = False) + assert backend._probe_mtp_decode(timeout = 1.0) is True + + def _drop(*a, **k): + raise llama_cpp_module.httpx.RemoteProtocolError("peer closed connection") + + monkeypatch.setattr(llama_cpp_module.httpx, "post", _drop, raising = False) + assert backend._probe_mtp_decode(timeout = 1.0) is False + + monkeypatch.setattr(llama_cpp_module.httpx, "post", lambda *a, **k: _Resp(500), raising = False) + assert backend._probe_mtp_decode(timeout = 1.0) is False + + # 200 but the server aborted right after (poll() returns an exit code). + backend._process = _FakeProcess() + monkeypatch.setattr(llama_cpp_module.httpx, "post", lambda *a, **k: _Resp(200), raising = False) + assert backend._probe_mtp_decode(timeout = 1.0) is False + + +# ── generation-time MTP recovery (mid-stream crash) ────────────────── + + +def _recovery_backend() -> LlamaCppBackend: + # A backend that loaded MTP under tensor parallelism and whose server has + # since exited (the _FakeProcess poll() returns 0 -> a dead subprocess). + b = LlamaCppBackend() + b._tensor_parallel = True + b._speculative_type = "draft-mtp" + b._mtp_runtime_fallback_active = True + b._process = _FakeProcess() + b._last_load_kwargs = { + "model_identifier": "owner/repo", + "tensor_parallel": True, + "speculative_type": "auto", + "n_parallel": 4, + } + return b + + +def test_generate_chat_completion_wires_runtime_recovery(): + # The non-tool generation path must route a mid-stream server death into the + # recovery helper (the tool + passthrough paths do so from the routes). + src = inspect.getsource(LlamaCppBackend.generate_chat_completion) + assert "_maybe_recover_from_mtp_crash" in src + + +def test_runtime_recovery_reloads_without_mtp(monkeypatch): + # One background reload with speculative_type="off" (rest of snapshot kept), + # then spec_fallback_reason="runtime_error" and single-flight released. + b = _recovery_backend() + done = threading.Event() + captured = {} + + def _fake_load_model(**kwargs): + captured.update(kwargs) + done.set() + return True + + monkeypatch.setattr(b, "load_model", _fake_load_model) + assert b._maybe_recover_from_mtp_crash(RuntimeError("peer closed")) is True + assert done.wait(timeout = 5) + assert captured["speculative_type"] == "off" + assert captured["model_identifier"] == "owner/repo" + assert captured["n_parallel"] == 4 # snapshot replayed faithfully + deadline = time.monotonic() + 2 + while b._spec_fallback_reason != "runtime_error" and time.monotonic() < deadline: + time.sleep(0.02) + assert b._spec_fallback_reason == "runtime_error" + assert b._mtp_runtime_fallback_in_progress is False + + +@pytest.mark.parametrize( + "mutate", + [ + lambda b: setattr(b, "_mtp_runtime_fallback_active", False), + lambda b: setattr(b, "_last_load_kwargs", None), + lambda b: setattr(b, "_process", None), + lambda b: b._cancel_event.set(), + ], +) +def test_runtime_recovery_skips_when_not_applicable(monkeypatch, mutate): + # No reload when this launch is not running MTP+tensor, there is no snapshot, + # the process handle is gone, or the request was cancelled. + b = _recovery_backend() + mutate(b) + calls = [] + monkeypatch.setattr(b, "load_model", lambda **k: calls.append(k)) + assert b._maybe_recover_from_mtp_crash(RuntimeError()) is False + assert calls == [] + + +class _BlockingDeadProc: + # Reports alive until released, then dead -- lets a test mutate backend state + # while the recovery thread is still in its death-confirm poll. + def __init__(self): + self._dead = threading.Event() + + def poll(self): + return 0 if self._dead.is_set() else None + + def terminate(self): + self._dead.set() + + def kill(self): + self._dead.set() + + def wait(self, timeout = None): + self._dead.set() + return 0 + + def release(self): + self._dead.set() + + +def test_runtime_recovery_fires_for_user_env_mtp(monkeypatch): + # MTP driven by user extra_args / LLAMA_ARG_SPEC_TYPE leaves _speculative_type + # unset, but the launch flag still gates recovery on (pass-through MTP). + b = _recovery_backend() + b._speculative_type = None # Studio stepped back; user/env owns the spec + done = threading.Event() + captured = {} + + def _fake_load_model(**kwargs): + captured.update(kwargs) + done.set() + return True + + monkeypatch.setattr(b, "load_model", _fake_load_model) + assert b._maybe_recover_from_mtp_crash(RuntimeError()) is True + assert done.wait(timeout = 5) + assert captured["speculative_type"] == "off" + + +def test_runtime_recovery_strips_user_mtp_extra_args(monkeypatch): + # A user --spec-type draft-mtp in extra_args must be neutralised on the reload + # (append a last-wins --spec-default) so MTP can't re-engage and loop. + b = _recovery_backend() + b._last_load_kwargs = dict(b._last_load_kwargs, extra_args = ["--spec-type", "draft-mtp"]) + done = threading.Event() + captured = {} + + def _fake_load_model(**kwargs): + captured.update(kwargs) + done.set() + return True + + monkeypatch.setattr(b, "load_model", _fake_load_model) + assert b._maybe_recover_from_mtp_crash(RuntimeError()) is True + assert done.wait(timeout = 5) + assert captured["speculative_type"] == "off" + assert captured["extra_args"][-1] == "--spec-default" + + +def test_runtime_recovery_restores_requested_mode(monkeypatch): + # After the off-reload, /status must show the user's requested mode + the + # runtime-error note, not a bare "off" (matches the startup MTP fallback). + b = _recovery_backend() + b._last_load_kwargs = dict(b._last_load_kwargs, speculative_type = "mtp") + done = threading.Event() + + def _fake_load_model(**kwargs): + b._requested_spec_mode = "off" # what a real off-reload would leave behind + done.set() + return True + + monkeypatch.setattr(b, "load_model", _fake_load_model) + assert b._maybe_recover_from_mtp_crash(RuntimeError()) is True + assert done.wait(timeout = 5) + deadline = time.monotonic() + 2 + while b._requested_spec_mode != "mtp" and time.monotonic() < deadline: + time.sleep(0.02) + assert b._requested_spec_mode == "mtp" + assert b._spec_fallback_reason == "runtime_error" + + +def test_runtime_recovery_skips_when_process_replaced(monkeypatch): + # A newer user load that replaces the process during the death-confirm poll + # must not be clobbered by the stale recovery replay. + b = _recovery_backend() + p1 = _BlockingDeadProc() + b._process = p1 + calls = [] + monkeypatch.setattr(b, "load_model", lambda **k: calls.append(k)) + assert b._maybe_recover_from_mtp_crash(RuntimeError()) is True # captures p1 + b._process = _FakeProcess() # a newer load swapped the live process + p1.release() # p1 now reports dead -> recovery runs its staleness check + time.sleep(0.6) + assert calls == [], "stale recovery replayed over a newer load" + + +def test_runtime_recovery_skips_when_snapshot_changed(monkeypatch): + # If the recorded load changed during the poll, the stale snapshot is dropped. + b = _recovery_backend() + p1 = _BlockingDeadProc() + b._process = p1 + calls = [] + monkeypatch.setattr(b, "load_model", lambda **k: calls.append(k)) + assert b._maybe_recover_from_mtp_crash(RuntimeError()) is True + b._last_load_kwargs = dict(b._last_load_kwargs, model_identifier = "other/model") + p1.release() + time.sleep(0.6) + assert calls == [] + + +def test_runtime_recovery_is_single_flight(monkeypatch): + # Concurrent failures schedule only one reload. + b = _recovery_backend() + started = threading.Event() + release = threading.Event() + + def _slow_load(**kwargs): + started.set() + release.wait(timeout = 5) + return True + + monkeypatch.setattr(b, "load_model", _slow_load) + assert b._maybe_recover_from_mtp_crash(RuntimeError()) is True + assert started.wait(timeout = 5) + # Second failure while the first reload is in flight is a no-op. + assert b._maybe_recover_from_mtp_crash(RuntimeError()) is False + release.set() + + +def test_runtime_recovery_rechecks_cancel_before_reload(): + # recover() must re-check the cancel flag after the death poll (load_model + # clears it), so a reload scheduled just before /unload can't resurrect it. + src = inspect.getsource(LlamaCppBackend._maybe_recover_from_mtp_crash) + cancel = src.rfind("self._cancel_event.is_set()") + load = src.find("self.load_model(") + assert 0 <= cancel < load, "recovery must re-check cancel before reloading" + + +def test_probe_mtp_decode_uses_api_key_auth(monkeypatch): + # Direct-stream mode runs llama-server with --api-key; the probe must send + # the same bearer auth or it gets a spurious 401 and falsely drops MTP. + backend = LlamaCppBackend() + backend._port = 0 + backend._process = None + captured = {} + + class _Resp: + status_code = 200 + + def _capture(*a, **k): + captured.clear() + captured.update(k) + return _Resp() + + monkeypatch.setattr(llama_cpp_module.httpx, "post", _capture, raising = False) + backend._api_key = "secret" + backend._probe_mtp_decode(timeout = 1.0) + assert captured["headers"] == {"Authorization": "Bearer secret"} + backend._api_key = None + backend._probe_mtp_decode(timeout = 1.0) + assert captured["headers"] is None + + +class _ToggleProcess: + """A subprocess stand-in whose liveness can be flipped at runtime.""" + + def __init__(self): + self._alive = True + + def poll(self): + return None if self._alive else 0 + + def terminate(self): + self._alive = False + + def kill(self): + self._alive = False + + def wait(self, timeout = None): + self._alive = False + return 0 + + def die(self): + self._alive = False + + +def test_crash_watchdog_triggers_recovery_on_death(monkeypatch): + # The watchdog must notice the process exit and recover even when no request + # handler observed it (e.g. the direct proxy endpoints). + b = _recovery_backend() + proc = _ToggleProcess() + b._process = proc + fired = threading.Event() + monkeypatch.setattr(b, "_maybe_recover_from_mtp_crash", lambda *a, **k: fired.set()) + b._start_mtp_crash_watchdog() + assert b._mtp_watchdog_thread is not None + proc.die() + assert fired.wait(timeout = 3) + + +def test_crash_watchdog_ignores_intentional_termination(monkeypatch): + # A planned reload/unload stops the watchdog before killing the process, so + # the resulting death must not be mistaken for a crash. + b = _recovery_backend() + proc = _ToggleProcess() + b._process = proc + fired = threading.Event() + monkeypatch.setattr(b, "_maybe_recover_from_mtp_crash", lambda *a, **k: fired.set()) + b._start_mtp_crash_watchdog() + b._stop_mtp_crash_watchdog() # what _kill_process does first + proc.die() + assert not fired.wait(timeout = 2) + assert b._mtp_watchdog_thread is None + + +@pytest.mark.parametrize( + "mutate", + [ + lambda b: setattr(b, "_mtp_runtime_fallback_active", False), + lambda b: setattr(b, "_process", None), + ], +) +def test_crash_watchdog_not_armed_when_inapplicable(mutate): + # Only a launch actually running MTP+tensor with a live process arms it. + b = _recovery_backend() + b._process = _ToggleProcess() + mutate(b) + b._start_mtp_crash_watchdog() + assert b._mtp_watchdog_thread is None + + +def test_kill_process_stops_crash_watchdog(monkeypatch): + # _kill_process is the single deliberate-termination chokepoint; it must + # stop the watchdog so the planned kill isn't seen as a crash. + b = _recovery_backend() + proc = _ToggleProcess() + b._process = proc + fired = threading.Event() + monkeypatch.setattr(b, "_maybe_recover_from_mtp_crash", lambda *a, **k: fired.set()) + b._start_mtp_crash_watchdog() + b._kill_process() + assert b._mtp_watchdog_thread is None + assert b._process is None + assert not fired.wait(timeout = 2) + + +def test_kill_process_stops_watchdog_before_terminate(): + # Ordering matters: stop the watchdog before terminating so the watchdog's + # post-death stop re-check reliably sees a planned kill. + src = inspect.getsource(LlamaCppBackend._kill_process) + stop = src.find("_stop_mtp_crash_watchdog()") + term = src.find(".terminate(") + assert 0 <= stop < term, "must stop the watchdog before terminating" + + +def test_crash_watchdog_rechecks_stop_before_recovery(): + # After a detected exit the watchdog re-checks the stop flag so a kill that + # raced in between the poll-wait and the poll-read can't fire recovery. + src = inspect.getsource(LlamaCppBackend._start_mtp_crash_watchdog) + check = src.find("stop.is_set()") + recover = src.find("_maybe_recover_from_mtp_crash") + assert 0 <= check < recover, "must re-check stop before recovering" + + +def test_load_model_arms_crash_watchdog(): + # The healthy-load commit arms the watchdog for this load. + src = inspect.getsource(LlamaCppBackend.load_model) + assert "_start_mtp_crash_watchdog" in src + + # ── tensor-mode allocation: conservative VRAM budget ─────────────────