ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Runtime MTP fallback for tensor parallelism (try MTP, recover if it crashes) (#6324)

Browse source 
* Disable MTP speculative decoding under tensor parallelism

Follow-up to #6040 (Studio tensor-parallel support).

MTP-draft speculative decoding plus --split-mode tensor crashes the CUDA
flash-attn kernel at decode time. The startup /health probe only checks that
llama-server comes up, so the existing MTP-drop fallback (keyed on startup
health) never fires and the server dies on the first generation instead.

Gate MTP off when a tensor attempt actually engages: this runs before the
VRAM planner (so no drafter memory is reserved) and before the speculative
flag build (so no --model-draft / --spec-type is emitted). Ngram modes use no
draft model and are kept, and mtp+ngram degrades to ngram rather than off. The
layer-split fallback re-runs with tensor_parallel False and restores MTP.

The reason is surfaced as spec_fallback_reason "tensor_parallel" so the
settings sheet explains why MTP is off instead of prompting a llama.cpp update.

Verified on unsloth/gemma-4-26B-A4B-it-GGUF:UD-Q4_K_XL across 4x B200: the
load now emits --split-mode tensor with no MTP flags and generation completes
without the prior decode crash.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Make tensor-parallel MTP gate test format-independent

The assertion pinned the multi-line `speculative_type = (` form, but ruff
collapses it onto one line, so match `speculative_type =` instead.

* Recover from MTP+tensor-parallel crashes at runtime instead of banning MTP

MTP-draft speculative decoding under --split-mode tensor usually works, but can
crash llama-server's CUDA flash-attn kernel at decode time (the prompt-cache
checkpoint-restore path). The earlier fix statically disabled MTP whenever
tensor parallelism was on, which is not future-proof and gives up the MTP
speedup even though it normally works.

Replace the static ban with a try/recover, mirroring the existing load-time
MTP-drop fallback:

- Load-time decode probe: after the server passes /health under tensor +
  MTP, run one tiny /completion to exercise the draft path. A failure flips
  the load unhealthy so the existing fallback respawns with --spec-default.
  Catches a hard incompatibility that crashes on the first decode.

- Generation-time recovery: snapshot the load kwargs after a healthy load,
  and if llama-server exits mid-generation while MTP + tensor parallelism were
  active, quietly reload the same model with speculative decoding off (one
  single-flight background reload) and surface spec_fallback_reason=runtime_error.
  Catches the rare mid-generation crash the probe and load-time fallback miss.

No persistent ban: a later fresh load re-tries MTP, so this self-heals if a
future llama.cpp supports the combo. Verified on gemma-4-26B-A4B + 4x B200:
MTP runs normally, and killing llama-server mid-generation reloads it without
MTP and serves the next request cleanly.

* Address review feedback on the MTP runtime fallback

- Authenticate the decode probe: direct-stream mode runs llama-server with
  --api-key, so the unauthenticated /completion probe got a 401 and falsely
  dropped MTP. Attach the same bearer auth the other internal requests use.
- Re-check the cancel flag inside the recovery thread after the death poll,
  so an /unload that races the reload can't resurrect the dropped model.
- Schedule the no-MTP recovery on the connection-error paths it was missing:
  generate_chat_completion's ConnectError branch, the OpenAI passthrough
  typed (RemoteProtocolError/ReadError/CloseError) stream catch, and the
  Anthropic passthrough generic stream catch. Previously a server that died
  before reconnect, or a typed mid-stream error, skipped the reload.

* Cover every request path with the MTP+tensor crash recovery via a watchdog

The runtime MTP-crash recovery only fired from request handlers that
observed the failure, so the direct llama-server proxy endpoints
(/v1/completions, /v1/responses, the OpenAI/Anthropic passthrough
transports) -- and a crash with no request in flight -- could leave a
dead server. Add a single background watchdog, armed only on a healthy
MTP + tensor-parallel load, that polls the subprocess and routes an
unexpected death into the existing single-flight no-MTP reload. It is
stopped inside _kill_process (the one deliberate-termination chokepoint)
so a planned reload/unload is never mistaken for a crash, and re-checks
the stop flag after a detected exit to close the kill-vs-poll race. The
reload turns MTP off, so the replacement server arms no watchdog and the
fallback cannot loop; a later fresh load still re-tries MTP.

* Harden MTP+tensor crash recovery: stale-load race, pass-through MTP, requested mode

Address review findings on the runtime MTP-crash recovery:

- Stale-load race: the recovery thread snapshotted the crashed load, waited up
  to 5s for the process to confirm dead, then only checked the cancel flag
  before replaying load_model. A concurrent user load clears that flag, so the
  stale snapshot could reload the old model over the user's new one. Make the
  load lock re-entrant and run the staleness check (cancel + same process +
  unchanged snapshot) under it, atomically with the reload.

- Pass-through MTP: MTP can also be requested via a user --spec-type in
  extra_args or LLAMA_ARG_SPEC_TYPE, where Studio emits no spec flags and
  _speculative_type stays unset, so the probe/watchdog/recovery never engaged.
  Track _mtp_runtime_fallback_active from the actual launched config and gate on
  it; on the no-MTP reload, append a last-wins --spec-default so the replay drops
  MTP regardless of source (and the load-time fallback does the same).

- Requested mode: the off-reload reset _requested_spec_mode to off, so after a
  status refresh the UI showed a bare Off with the runtime-error note suppressed
  and would not retry MTP. Restore the original requested mode after the reload,
  matching the startup MTP fallback.

- Snapshot the extra_args list by value so a caller mutating it cannot corrupt
  the recovery snapshot.

Tests: test_tensor_parallel.py + test_llama_server_args.py green (303 passed).

* Trim verbose comments in the MTP+tensor crash recovery

Tighten the docstrings and inline comments added for the runtime MTP recovery
(watchdog, probe, reload, gating) to succinct one/two-line forms; no code
change (verified comment-only).

---------

Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-06-18 05:42:21 -07:00 • committed by GitHub
commit 3bfc83781d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 627 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

228
studio/backend/core/inference/llama_cpp.py
View file

@ -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 ──────────────────────────────

11
studio/backend/routes/inference.py
View file

@ -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),

394
studio/backend/tests/test_tensor_parallel.py
View file

@ -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 ─────────────────