Studio: recover tool-enabled GGUF chats after llama-server exits (#7424)

* Fix GGUF tool chat server recovery

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

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

* Cover MTP precedence and loosen the replay assertion for PR #7424

Add a regression test for the MTP branch of the tool-loop respawn retry: the
file-wide _make_backend stub forces _maybe_recover_from_mtp_crash to False, so
nothing exercised the case where an MTP crash reload is already claimed and an
ordinary same-config respawn must not run on top of it. Cover both the next
tool-loop request and the final synthesis pass.

Replace the whole-payload equality assertions with a field-wise check. Comparing
the full dict pins max_tokens to the value derived from the dead server's
effective context, so a later fix that rebuilds server-derived defaults after a
respawn would read as a test failure rather than an improvement.

Document that the one-retry budget is per model request, not per chat turn.

* Recover from prefill-time deaths and stop respawn racing the MTP reload

Two gaps in the tool-loop respawn retry, both reproduced before fixing.

A child that exits during prefill has already accepted the socket, so httpx
raises ReadError, WriteError or RemoteProtocolError rather than ConnectError.
Those all arrive before the response opens, which is exactly the window where a
replay is safe, but the helper only caught ConnectError and gave up. Widen the
catch to NetworkError plus RemoteProtocolError. Timeouts stay excluded on
purpose: they mean the server is slow, not dead, and retrying one would spend
the 20 minute first-token budget twice. Windows resets connections where Linux
refuses them, so this also covers the common Windows presentation.

_maybe_recover_from_mtp_crash returns False both when the crash is not an MTP
crash and when an MTP-free reload is already in flight. Callers read that as
permission to respawn, so _respawn_if_dead replayed the crashing MTP kwargs and,
by replacing the process, made the in-flight reload abort on its own newer-load
check. Skip the respawn while that reload owns the corpse. The guard lives in
_respawn_if_dead so the plain chat path gets it too.

Regression tests for both, including a guard against retrying prefill timeouts.

* Release the MTP single-flight claim when the reload never starts

_mtp_runtime_fallback_in_progress is claimed before the reload thread exists, and
only that thread's finally clears it. Two statements ran in between with no unwind
path: re-reading _last_load_kwargs, which an unload can null underneath us, and
Thread.start(), which raises under the thread exhaustion that is exactly the
pressure killing llama-server in the first place. Nothing else ever resets the
flag, so a failure there latched it for the life of the process.

That was survivable before, since respawn ignored the flag. It is not now: the
guard added in db78184be keys off the flag alone, so a latch would silently
disable auto-respawn for every later model, including plain non-MTP ones. Read
the kwargs and process once before claiming, and release the claim if the thread
cannot start.

Restore the whole-payload equality assertions. Comparing field-wise was meant to
leave room for rebuilding server-derived defaults on replay, but the payload is
built once before the retry and re-sent unchanged, so the looser check only
dropped seven real keys and added a vacuous seed comparison.

Also correct the docstring: llama-server flushes its 200 at slot start, so a
death during decode arrives with the response already open. The pre-header window
this covers is an upload still in flight or a request waiting behind busy slots.

* Confirm the child exited before spending the retry

A closing llama-server can beat its own exit status: the socket error arrives while
poll() still reports the process running. _respawn_if_dead then took the alive
branch, handed back the stale _healthy, and the caller read that as a successful
respawn and spent its single retry on the same corpse. When that retry failed,
attempt was no longer 0, so no respawn ever happened and the turn died, with a log
line claiming a respawn that had not occurred. The window matters most for the
pre-header ReadError and RemoteProtocolError shutdowns the retry now covers.

Wait a bounded second for the exit status before calling the child alive. The same
race is already conceded in _maybe_recover_from_mtp_crash, whose recovery thread
polls for 5s because the error can arrive a beat early; 1s here because this runs
on the request path, and a genuinely live server, including one a concurrent caller
has just respawned, still returns promptly.

* Tighten the recovery comments

* Harden the respawn path around concurrent unloads and replacements

Two problems with the reap grace loop, both found by review.

Skip the grace when the server was already replaced. A caller queued on
_respawn_lock behind someone else's respawn woke holding the healthy replacement,
could not tell it from the child its own request had used, and waited out the full
grace. That sleep is under the lock, so the waits serialised: four concurrent
generations cost roughly three grace periods before any retry began. Capture the
process before taking the lock and return early once it has been swapped.

Do not respawn a server that is being torn down on purpose. unload_model() sets
_cancel_event and only clears _last_load_kwargs after the kill, so a request losing
its connection mid-unload could watch that deliberate exit through the grace loop,
read the stale kwargs and load the model straight back; a model switch landing
during the wait was reverted the same way. Re-check the cancel flag and the process
identity under _serial_load_lock before capturing the replay kwargs, matching what
the MTP-crash reload already does.

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

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

* Tighten the respawn comments

* Do not charge the reap grace to a server that is still serving

The grace loop added for the not-yet-reaped race waits on poll(), which for a
live child never returns, so every transient transport error paid the full
_RESPAWN_REAP_GRACE_S. That sleep is held under _respawn_lock, so the cost
serialised: measured 1002 ms for one caller and 8.02 s for eight concurrent ones,
against 0 ms on main. A working install pays this, not a broken one.

A llama-server's listening socket dies with the process, so a loopback connect
separates the two cases in microseconds. Probe it first and return immediately
when the port still accepts; fall through to the grace only when the port is
gone, which is the case the grace exists for. Back to 0.7 ms for one caller and
0.00 s for eight.

Cross-checked on real hardware over Qwen3.5-2B, Llama-3.2-1B, Gemma-3-4B with
mmproj and Qwen3-30B-A3B: decode throughput within noise of main (-0.06%, -3.71%,
+2.57%, +0.29%, against a 54-232% spread between rounds of a single run), output
byte-identical on every round, tool-path recovery restored on the three families
whose model calls the tool, and plain-chat recovery still working on all four.

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

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

* Make the respawn lose to a deliberate unload in every window

Two follow-ups on the respawn path, both reproduced first.

Check _cancel_event before the socket fast path. unload_model sets the flag before
it kills, so the child is still accepting when the probe runs; returning the stale
_healthy there aims the retry at a server that is deliberately going away.

Close the unload TOCTOU. The old cancel check sat under _serial_load_lock, which
unload_model never takes, so an unload could land entirely between that check and
load_model and the captured kwargs would restart a model the user had stopped.
Snapshot the kwargs, the flag and a new _unload_epoch together under _lock, the
lock unload does hold, so a teardown is either wholly before the snapshot or
wholly after it. load_model clears _cancel_event on the way in, so the epoch is
the only evidence that survives; when it moves during the reload the replacement
is unloaded again rather than left running.

_lock stays uncontended across load_model, which would deadlock a plain Lock and
block /status for the length of a load. Error-path latency is unchanged: 0.6 ms
for a live server and 0.00 s for eight concurrent callers.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <unslothai@gmail.com>
Co-authored-by: danielhanchen <danielhanchen@gmail.com>
This commit is contained in:
oobabooga 2026-07-26 08:53:45 -03:00 committed by GitHub
commit aefeb5821d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 794 additions and 24 deletions

View file

@ -307,6 +307,9 @@ def _native_linux_system_rocm_lib_dirs(binary_dir: str = "") -> "list[str]":
# enough for reasoning-heavy GGUFs and max_tokens-omitting API clients.
_DEFAULT_MAX_TOKENS_FLOOR = 32768
_DEFAULT_FIRST_TOKEN_TIMEOUT_S = 1200.0 # 20 min
# A transport error can arrive before the child is reapable; a request path cannot
# afford the 5s the background MTP reload spends on the same race.
_RESPAWN_REAP_GRACE_S = 1.0
def _finalize_reasoning_only_cumulative(
@ -2099,6 +2102,9 @@ class LlamaCppBackend:
# Serialises mid-session respawns so many generations hitting a killed
# server trigger at most one reload (see _respawn_if_dead).
self._respawn_lock = threading.Lock()
# Bumped by every unload. load_model clears _cancel_event, so a respawn that
# raced an unload needs a signal that survives the clear (see _respawn_if_dead).
self._unload_epoch = 0
# 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
@ -9308,6 +9314,7 @@ class LlamaCppBackend:
"""Terminate the subprocess and cancel any in-flight download."""
self._cancel_event.set()
with self._lock:
self._unload_epoch += 1
self._kill_process()
logger.info(f"Unloaded GGUF model: {self._model_identifier}")
self._model_identifier = None
@ -10107,15 +10114,18 @@ class LlamaCppBackend:
return False
if not self._mtp_runtime_fallback_active:
return False
if not self._last_load_kwargs or self._process is None:
# Read before claiming: a raise after the claim strands the flag, and nothing
# else clears it, blocking every later respawn.
kwargs = self._last_load_kwargs
proc = self._process
if not kwargs or proc 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
snapshot = dict(kwargs)
def _recover():
try:
@ -10163,7 +10173,14 @@ class LlamaCppBackend:
with self._mtp_runtime_fallback_lock:
self._mtp_runtime_fallback_in_progress = False
threading.Thread(target = _recover, daemon = True, name = "mtp-crash-reload").start()
try:
threading.Thread(target = _recover, daemon = True, name = "mtp-crash-reload").start()
except RuntimeError as exc:
# Release the claim: a reload that never started would block respawn forever.
with self._mtp_runtime_fallback_lock:
self._mtp_runtime_fallback_in_progress = False
logger.error(f"Could not start the MTP-crash reload: {exc}")
return False
return True
def _start_mtp_crash_watchdog(self) -> None:
@ -10635,6 +10652,21 @@ class LlamaCppBackend:
finally:
_cancel_closed.set()
def _server_socket_is_open(self, timeout_s: float = 0.15) -> bool:
"""True if anything still accepts on the server port.
The listening socket dies with the process, so this tells a live server
from a dead one without waiting for the child to become reapable.
"""
port = self._port
if not port:
return False
try:
with socket.create_connection(("127.0.0.1", port), timeout = timeout_s):
return True
except OSError:
return False
def _respawn_if_dead(self) -> bool:
"""Relaunch the llama-server if its process has exited.
@ -10644,28 +10676,114 @@ class LlamaCppBackend:
recover, returning True once healthy. Serialised on ``_respawn_lock`` so
many generations hitting the dead server trigger at most one reload.
"""
# Read outside the lock so a queued caller can tell the replacement from the child
# its own error came from; otherwise each burns the grace wait below, and that
# sleep is held under the lock, so the waits serialise.
served_by = self._process
with self._respawn_lock:
proc = self._process
if proc is None:
return False
if proc.poll() is None:
# Process is alive: either a concurrent caller already respawned
# it (healthy), or this connection error wasn't a dead server.
if self._cancel_event.is_set():
# unload_model sets this before it kills, so the child can still be
# accepting. Reporting it healthy would aim the retry at a server
# that is deliberately going away.
return False
if proc is not served_by:
# Replaced while we queued: this child never served our request.
return self._healthy
kwargs = self._last_load_kwargs
if not kwargs:
return False
logger.warning(
f"llama-server for '{self._model_identifier}' exited "
f"(code {proc.returncode}); respawning to recover the session"
)
with self._lock:
self._healthy = False
if proc.poll() is None:
# Still serving, so the error was transient. Charging it the grace below
# would cost a second per caller, serialised under this lock.
if self._server_socket_is_open():
return self._healthy
# A closing server can beat its own exit status: calling it alive returns
# the stale _healthy and spends the retry on the corpse.
deadline = time.monotonic() + _RESPAWN_REAP_GRACE_S
while proc.poll() is None and time.monotonic() < deadline:
time.sleep(0.05)
if proc.poll() is None:
# Alive: either a concurrent caller already respawned it (healthy), or
# this connection error wasn't a dead server.
return self._healthy
with self._mtp_runtime_fallback_lock:
if self._mtp_runtime_fallback_in_progress:
# An MTP-free reload owns this corpse; replaying the old kwargs
# restarts the crashing config and aborts that reload.
logger.info("Respawn skipped: an MTP-free reload is already recovering.")
return False
# The RLock lets the load_model below re-enter it.
with self._serial_load_lock:
if self._process is not proc:
logger.info("Respawn skipped: a newer load is already active.")
return self._healthy
# Snapshot under _lock, the one unload_model holds, so a teardown is
# either wholly before us (flag set) or wholly after (epoch bumped).
# _serial_load_lock alone would not exclude it: unload never takes it.
with self._lock:
if self._cancel_event.is_set():
logger.info("Respawn skipped: the model was unloaded.")
return False
kwargs = dict(self._last_load_kwargs or {})
if not kwargs:
return False
epoch = self._unload_epoch
self._healthy = False
logger.warning(
f"llama-server for '{self._model_identifier}' exited "
f"(code {proc.returncode}); respawning to recover the session"
)
try:
started = bool(self.load_model(**kwargs))
except Exception as exc:
logger.error(f"Failed to respawn llama-server: {exc}")
return False
if started and self._unload_epoch != epoch:
# An unload landed mid-reload. load_model cleared _cancel_event on
# the way in, so the epoch is the only surviving evidence; undo the
# replacement rather than leave a model the user stopped running.
logger.info("Respawn undone: the model was unloaded during the reload.")
self.unload_model()
return False
return started
@contextlib.contextmanager
def _open_chat_stream_with_respawn_retry(self, payload: dict, cancel_event):
"""Open a chat stream, respawning a dead llama-server once before streaming.
Retry only when opening the response fails: once it is open a consumer may
already have emitted content or tool events, so a replay could duplicate
output and side effects. ``base_url`` is resolved per attempt because a
respawn may pick a new port. The budget is one retry per model request, not
per chat turn, so a long tool loop never discards a completed tool.
A child dying after the accept but before the headers surfaces as
ReadError/WriteError/RemoteProtocolError rather than ConnectError, and which
one differs per OS. llama-server flushes its 200 at slot start, so that window
is an upload still in flight or a request behind busy slots; a death during
decode arrives with the response open and is not replayed. Timeouts are
excluded: the server is slow, not dead, and a replay would spend the
first-token budget twice.
"""
for attempt in range(2):
response_opened = False
try:
return bool(self.load_model(**kwargs))
except Exception as exc:
logger.error(f"Failed to respawn llama-server: {exc}")
return False
url = f"{self.base_url}/v1/chat/completions"
with self._open_stream(url, payload, cancel_event) as opened:
response_opened = True
yield opened
return
except (httpx.NetworkError, httpx.RemoteProtocolError) as exc:
if response_opened:
raise
if self._maybe_recover_from_mtp_crash(exc):
raise RuntimeError("Lost connection to llama-server") from exc
if attempt == 0 and self._respawn_if_dead():
logger.warning(
"llama-server was unreachable; respawned it and retrying the generation"
)
continue
raise
def generate_chat_completion(
self,
@ -10963,7 +11081,6 @@ class LlamaCppBackend:
yield _ev
conversation.extend(_auto["messages"])
url = f"{self.base_url}/v1/chat/completions"
_accumulated_completion_tokens = 0
_accumulated_predicted_ms = 0.0
_accumulated_predicted_n = 0
@ -11223,7 +11340,7 @@ class LlamaCppBackend:
_text_args_name = ""
_confirm_gated_iteration = bool(confirm_tool_calls) and not bypass_permissions
with self._open_stream(url, payload, cancel_event) as (
with self._open_chat_stream_with_respawn_retry(payload, cancel_event) as (
response,
first_token_deadline,
):
@ -12260,7 +12377,7 @@ class LlamaCppBackend:
_stream_done = False
try:
with self._open_stream(url, stream_payload, cancel_event) as (
with self._open_chat_stream_with_respawn_retry(stream_payload, cancel_event) as (
response,
first_token_deadline,
):

View file

@ -14,6 +14,7 @@ import contextlib
import copy
import json
import sys
import threading
from pathlib import Path
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
@ -55,7 +56,12 @@ def _finish(reason: str) -> str:
)
def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]):
def _make_backend(
monkeypatch,
streams: list[object],
payloads: list[dict],
urls: list[str] | None = None,
):
backend = LlamaCppBackend.__new__(LlamaCppBackend)
backend._process = object()
backend._healthy = True
@ -77,7 +83,12 @@ def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]):
first_token_deadline = None,
):
payloads.append(copy.deepcopy(payload))
yield type("FakeResponse", (), {"status_code": 200, "chunks": streams.pop(0)})()
if urls is not None:
urls.append(_url)
stream = streams.pop(0)
if isinstance(stream, BaseException):
raise stream
yield type("FakeResponse", (), {"status_code": 200, "chunks": stream})()
def fake_iter_text_cancellable(
response,
@ -88,9 +99,27 @@ def _make_backend(monkeypatch, streams: list[list[str]], payloads: list[dict]):
monkeypatch.setattr(backend, "_stream_with_retry", fake_stream_with_retry)
monkeypatch.setattr(backend, "_iter_text_cancellable", fake_iter_text_cancellable)
monkeypatch.setattr(backend, "_maybe_recover_from_mtp_crash", lambda *_a, **_k: False)
return backend
def _patch_successful_respawn(
monkeypatch,
backend,
port: int | None = None,
) -> list[bool]:
calls: list[bool] = []
def fake_respawn():
calls.append(True)
if port is not None:
backend._port = port
return True
monkeypatch.setattr(backend, "_respawn_if_dead", fake_respawn)
return calls
def _tool_names(payload: dict) -> list[str]:
return [
(tool.get("function") or {}).get("name")
@ -2239,7 +2268,13 @@ def test_connect_error_during_tool_call_closes_provisional_card(monkeypatch):
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [raising_stream()], payloads)
respawn_calls: list[bool] = []
monkeypatch.setattr(
backend,
"_respawn_if_dead",
lambda: respawn_calls.append(True) or True,
)
monkeypatch.setattr("core.inference.tools.execute_tool", lambda *_a, **_k: "OK")
collected: list[dict] = []
@ -2270,6 +2305,271 @@ def test_connect_error_during_tool_call_closes_provisional_card(monkeypatch):
# The closing card is marked as an error, not an empty success, so the UI
# renders it as failed.
assert "Error" in (closing[0].get("result") or "")
assert respawn_calls == []
def test_connect_error_before_tool_stream_respawns_and_retries(monkeypatch):
"""A dead server before the first tool-loop response is opened is safe to retry."""
import httpx
payloads: list[dict] = []
urls: list[str] = []
backend = _make_backend(
monkeypatch,
[
httpx.ConnectError("server is down"),
[_sse({"content": "Recovered."}), _done()],
],
payloads,
urls,
)
respawn_calls = _patch_successful_respawn(monkeypatch, backend, port = 49999)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "hello"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
)
assert respawn_calls == [True]
assert len(payloads) == 2
assert payloads[0] == payloads[1]
assert urls == [
"http://127.0.0.1:48847/v1/chat/completions",
"http://127.0.0.1:49999/v1/chat/completions",
]
assert any(e.get("type") == "content" and e.get("text") == "Recovered." for e in events)
def test_connect_error_after_tool_result_recovers_both_generation_paths(monkeypatch):
"""Recover either post-tool generation path without rerunning the tool."""
import httpx
for max_tool_iterations, final_text in (
(2, "The result is 1."),
(1, "Final answer."),
):
payloads: list[dict] = []
backend = _make_backend(
monkeypatch,
[
_structured_tool_call("python", {"code": "print(1)"}, "call_once"),
httpx.ConnectError("server died between turns"),
[_sse({"content": final_text}), _done()],
],
payloads,
)
respawn_calls = _patch_successful_respawn(monkeypatch, backend)
tool_calls: list[tuple[str, dict]] = []
def fake_execute_tool(name, arguments, **_kwargs):
tool_calls.append((name, arguments))
return "1"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "print one"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = max_tool_iterations,
)
)
assert respawn_calls == [True]
assert tool_calls == [("python", {"code": "print(1)"})]
assert len(payloads) == 3
assert payloads[1] == payloads[2]
assert any(e.get("type") == "content" and e.get("text") == final_text for e in events)
def test_connect_error_retry_is_bounded(monkeypatch):
"""A failed retry surfaces the error without another respawn attempt."""
import httpx
payloads: list[dict] = []
backend = _make_backend(
monkeypatch,
[
httpx.ConnectError("server is down"),
httpx.ConnectError("replacement is also down"),
],
payloads,
)
respawn_calls = _patch_successful_respawn(monkeypatch, backend)
raised = False
try:
list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "hello"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
)
except RuntimeError as exc:
raised = True
assert "Lost connection" in str(exc)
assert raised
assert respawn_calls == [True]
assert len(payloads) == 2
def test_pre_header_transport_errors_also_respawn(monkeypatch):
"""A child that dies during prefill already accepted the socket, so it does
not surface as ConnectError. Nothing has streamed yet, so replay is safe."""
import httpx
for exc in (
httpx.RemoteProtocolError("server disconnected without sending a response"),
httpx.ReadError("connection reset by peer"),
httpx.WriteError("broken pipe"),
):
payloads: list[dict] = []
backend = _make_backend(
monkeypatch, [exc, [_sse({"content": "Recovered."}), _done()]], payloads
)
respawn_calls = _patch_successful_respawn(monkeypatch, backend)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "hello"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
)
assert respawn_calls == [True], type(exc).__name__
assert len(payloads) == 2, type(exc).__name__
assert any(e.get("type") == "content" and e.get("text") == "Recovered." for e in events)
def test_a_not_yet_reaped_child_does_not_burn_the_retry(monkeypatch):
"""A closing server can beat its own exit status, so poll() briefly reports it
alive. Without a grace wait _respawn_if_dead hands back the stale _healthy and the
single retry is spent on the corpse rather than on a replacement."""
import httpx
class _Dying:
# reapable only from the 4th poll, mimicking teardown lagging the socket close
def __init__(self):
self.polls = 0
self.returncode = None
def poll(self):
self.polls += 1
if self.polls > 3:
self.returncode = -9
return -9
return None
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [], payloads)
backend._process = _Dying()
backend._healthy = True
backend._respawn_lock = threading.RLock()
backend._lock = threading.RLock()
backend._mtp_runtime_fallback_lock = threading.Lock()
backend._serial_load_lock = threading.RLock()
backend._cancel_event = threading.Event()
backend._unload_epoch = 0
backend._mtp_runtime_fallback_in_progress = False
backend._mtp_runtime_fallback_active = False
backend._last_load_kwargs = {"gguf_path": "/m.gguf"}
backend._model_identifier = "m"
dying = backend._process
loads: list[dict] = []
@contextlib.contextmanager
def dead_until_respawned(
_c,
_url,
payload,
_ce,
headers = None,
first_token_deadline = None,
):
payloads.append(copy.deepcopy(payload))
if backend._process is dying:
raise httpx.ReadError("connection reset while shutting down")
yield type(
"FakeResponse",
(),
{"status_code": 200, "chunks": [_sse({"content": "Recovered."}), _done()]},
)()
def fake_load(**kwargs):
loads.append(kwargs)
backend._process = type("Live", (), {"poll": lambda self: None, "returncode": None})()
backend._healthy = True
return True
monkeypatch.setattr(backend, "_stream_with_retry", dead_until_respawned)
monkeypatch.setattr(backend, "load_model", fake_load)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "hello"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
)
assert len(loads) == 1
assert any(e.get("type") == "content" and e.get("text") == "Recovered." for e in events)
def test_prefill_timeout_is_not_retried(monkeypatch):
"""A slow-but-alive server must not have its first-token budget spent twice."""
import httpx
for exc in (httpx.ReadTimeout("no first token"), httpx.PoolTimeout("pool")):
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [exc], payloads)
respawn_calls = _patch_successful_respawn(monkeypatch, backend)
raised = False
try:
list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "hello"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = 1,
)
)
except httpx.TimeoutException:
raised = True
assert raised, type(exc).__name__
assert respawn_calls == [], type(exc).__name__
assert len(payloads) == 1, type(exc).__name__
def test_mtp_crash_recovery_wins_over_respawn(monkeypatch):
"""An MTP crash reloads without MTP, so never respawn the same config on top."""
import httpx
for max_tool_iterations in (2, 1):
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [httpx.ConnectError("mtp crash")], payloads)
monkeypatch.setattr(backend, "_maybe_recover_from_mtp_crash", lambda *_a, **_k: True)
respawn_calls = _patch_successful_respawn(monkeypatch, backend)
raised = False
try:
list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "hello"}],
tools = [{"type": "function", "function": {"name": "python"}}],
max_tool_iterations = max_tool_iterations,
)
)
except RuntimeError as exc:
raised = True
assert "Lost connection" in str(exc)
assert raised
assert respawn_calls == []
assert len(payloads) == 1
def test_empty_tool_call_id_does_not_emit_provisional_card(monkeypatch):

View file

@ -19,6 +19,7 @@ from __future__ import annotations
import asyncio
import inspect
import socket
import sys
import threading
import time
@ -528,6 +529,358 @@ def test_runtime_recovery_is_single_flight(monkeypatch):
release.set()
def test_single_flight_claim_is_released_when_the_reload_cannot_start(monkeypatch):
# Only the reload thread's finally clears the claim, so if starting it raises the
# claim must not latch: nothing else resets it, and _respawn_if_dead then refuses
# forever, for every later model.
b = _recovery_backend()
class _NoThread:
def __init__(self, *args, **kwargs):
pass
def start(self):
raise RuntimeError("can't start new thread")
monkeypatch.setattr(llama_cpp_module.threading, "Thread", _NoThread)
assert b._maybe_recover_from_mtp_crash(RuntimeError()) is False
assert b._mtp_runtime_fallback_in_progress is False
def test_load_kwargs_are_read_once_before_the_claim(monkeypatch):
# Gate and snapshot must share one read: reading twice lets an unload null
# _last_load_kwargs in between, so dict(None) raises after the claim and strands
# the flag with no thread alive to clear it.
b = _recovery_backend()
class _CountingKwargs: # data descriptor, so it wins over the instance dict
def __init__(self, value):
self.value = value
self.reads = 0
def __get__(self, obj, owner):
if obj is None:
return self
self.reads += 1
return self.value
def __set__(self, obj, value):
self.value = value
counter = _CountingKwargs({"model_identifier": "owner/repo"})
monkeypatch.setattr(type(b), "_last_load_kwargs", counter, raising = False)
class _UnstartedThread: # keep the reload off-thread so only sync reads count
def __init__(self, *args, **kwargs):
pass
def start(self):
pass
monkeypatch.setattr(llama_cpp_module.threading, "Thread", _UnstartedThread)
assert b._maybe_recover_from_mtp_crash(RuntimeError()) is True
assert counter.reads == 1, f"read {counter.reads} times; an unload can race the claim"
def test_respawn_defers_to_an_inflight_mtp_reload(monkeypatch):
# "Already recovering" must not read as "not an MTP crash": respawning replays the
# crashing MTP kwargs and aborts the in-flight no-MTP reload on its "newer load" check.
b = _recovery_backend()
b._mtp_runtime_fallback_in_progress = True
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
assert b._respawn_if_dead() is False
assert loads == []
# Once that reload finishes, an ordinary respawn works again.
b._mtp_runtime_fallback_in_progress = False
b._process.returncode = -9 # only the respawn path logs it
assert b._respawn_if_dead() is True
assert [kw.get("speculative_type") for kw in loads] == ["auto"]
def test_respawn_does_not_wait_out_the_grace_on_a_replacement(monkeypatch):
# Callers losing the same child queue on _respawn_lock and wake holding the healthy
# REPLACEMENT. Unable to tell it from their own child, each burns the reap grace, and
# that sleep is held under the lock, so N callers cost N grace periods.
class _LiveProcess(_FakeProcess):
returncode = None
def __init__(self):
self.polls = 0
def poll(self): # never reapable, so the grace loop runs to its deadline
self.polls += 1
return None
workers = 4
b = _recovery_backend()
b._healthy = True
b._process.returncode = -9 # only the respawn path logs it
live = _LiveProcess()
loads: list[dict] = []
guard = threading.Lock()
all_in_flight = threading.Event()
# Subclass this instance, not the class: a descriptor on LlamaCppBackend would
# redirect _process for every other live backend, including atexit-registered ones.
state = {"proc": b._process, "readers": set()}
class _Tracked(type(b)):
@property
def _process(self):
"""Reports when every worker has taken its pre-lock look at the child."""
with guard:
state["readers"].add(threading.get_ident())
everyone = len(state["readers"]) >= workers
if everyone:
all_in_flight.set()
return state["proc"]
@_process.setter
def _process(self, value):
state["proc"] = value
b.__class__ = _Tracked
def _load(**kwargs):
# A real load_model takes seconds, so every caller that lost this child is in
# flight before the replacement appears; waiting reproduces that ordering. The
# timeout keeps the pre-fix build, where losers cannot read until the lock is
# free, from hanging instead of failing.
all_in_flight.wait(timeout = 2)
with guard:
loads.append(kwargs)
b._process = live
b._healthy = True # the real load_model marks the new server healthy
return True
monkeypatch.setattr(b, "load_model", _load)
results: list[bool] = []
def _respawn():
outcome = b._respawn_if_dead()
with guard:
results.append(outcome)
threads = [threading.Thread(target = _respawn) for _ in range(workers)]
started = time.monotonic()
for thread in threads:
thread.start()
for thread in threads:
thread.join(timeout = 30)
elapsed = time.monotonic() - started
assert results == [True] * workers, results
assert len(loads) == 1, f"{len(loads)} reloads, expected one"
# The grace loop is the only poll() of a live process, so any count means a queued
# caller charged the wait to a server that never failed.
assert live.polls == 0, "queued caller waited out the grace on a healthy server"
assert elapsed < llama_cpp_module._RESPAWN_REAP_GRACE_S * (workers - 1)
class _DyingChild(_FakeProcess):
"""Alive for the first polls, then reapable: what a terminate() looks like."""
def __init__(
self,
code = -15,
alive_polls = 2,
on_death = None,
):
self.polls = 0
self.returncode = None
self._code = code
self._alive_polls = alive_polls
self._on_death = on_death
def poll(self):
self.polls += 1
if self.polls <= self._alive_polls:
return None
if self.returncode is None:
self.returncode = self._code
if self._on_death is not None:
self._on_death()
return self._code
def test_respawn_does_not_resurrect_a_deliberate_unload(monkeypatch):
# unload_model() sets _cancel_event before killing, so a request that loses the
# connection can watch that deliberate exit through the grace loop and call it a
# crash, with _last_load_kwargs still populated (unload clears it after the kill).
b = _recovery_backend()
b._healthy = True
b._process = _DyingChild()
b._cancel_event.set()
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
assert b._respawn_if_dead() is False
assert loads == [], "resurrected a model the user unloaded"
def test_respawn_rechecks_the_cancel_flag_after_the_grace_wait(monkeypatch):
# The unload can also begin while we are already sleeping in the grace loop.
b = _recovery_backend()
b._healthy = True
b._process = _DyingChild(on_death = b._cancel_event.set)
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
assert b._respawn_if_dead() is False
assert loads == [], "checked the cancel flag only before the wait"
def test_respawn_does_not_revert_a_newer_load(monkeypatch):
# A model switch landing while we wait must win; replaying the old kwargs would
# swap the user's new model back out.
b = _recovery_backend()
b._healthy = True
replacement = _DyingChild(alive_polls = 10**6)
b._process = _DyingChild(on_death = lambda: setattr(b, "_process", replacement))
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
b._respawn_if_dead()
assert loads == [], "replayed stale kwargs over a newer load"
assert b._process is replacement
def test_respawn_still_recovers_an_ordinary_crash(monkeypatch):
# Guard rail: none of the above may disable the recovery this path exists for.
b = _recovery_backend()
b._healthy = True
b._process = _DyingChild(code = -9)
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
assert b._respawn_if_dead() is True
assert len(loads) == 1
class _NeverReapable(_FakeProcess):
"""A child that stays unreapable, so only the port can tell alive from dead."""
returncode = None
def poll(self):
return None
def test_a_transient_error_against_a_live_server_costs_nothing(monkeypatch):
# The reap grace must not be charged to a server that never died: the sleep is
# held under _respawn_lock, so a full grace per caller serialises into N seconds
# of added latency on an install that is working fine.
listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
listener.bind(("127.0.0.1", 0))
listener.listen(16)
try:
b = _recovery_backend()
b._healthy = True
b._process = _NeverReapable()
b._port = listener.getsockname()[1]
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
started = time.monotonic()
assert b._respawn_if_dead() is True
elapsed = time.monotonic() - started
assert loads == [], "a live server must not be reloaded"
assert (
elapsed < llama_cpp_module._RESPAWN_REAP_GRACE_S / 2
), f"waited {elapsed:.2f}s on a server that is still accepting"
finally:
listener.close()
def test_a_closed_port_still_waits_for_the_child_to_be_reapable(monkeypatch):
# The other half: no listener means the server really is gone, so the grace
# still runs and the reap-race fix is preserved.
probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
probe.bind(("127.0.0.1", 0))
dead_port = probe.getsockname()[1]
probe.close()
b = _recovery_backend()
b._healthy = True
b._process = _DyingChild(code = -9)
b._port = dead_port
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
assert b._respawn_if_dead() is True
assert len(loads) == 1
def test_socket_fast_path_honours_a_pending_unload(monkeypatch):
# unload_model() sets _cancel_event before it kills, so the child is still
# accepting when the probe runs. Reporting it healthy aims the retry at a server
# that is deliberately going away.
listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
listener.bind(("127.0.0.1", 0))
listener.listen(8)
try:
b = _recovery_backend()
b._healthy = True
b._process = _NeverReapable()
b._port = listener.getsockname()[1]
b._cancel_event.set()
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
assert b._respawn_if_dead() is False
assert loads == []
finally:
listener.close()
def test_an_unload_landing_during_the_reload_is_undone(monkeypatch):
# The cancel check cannot live under _serial_load_lock alone: unload_model never
# takes that lock, so it can land entirely between the check and load_model and
# the captured kwargs then restart a model the user stopped. load_model clears
# _cancel_event on the way in, so _unload_epoch is the surviving evidence.
b = _recovery_backend()
b._healthy = True
b._process = _FakeProcess()
b._process.returncode = -9
loads: list[dict] = []
monkeypatch.setattr(b, "load_model", lambda **kwargs: loads.append(kwargs) or True)
unloads: list[int] = []
real_unload = b.unload_model
monkeypatch.setattr(b, "unload_model", lambda: unloads.append(1) or real_unload())
# The warning marks the window: after the snapshot, before the reload.
real_warning = llama_cpp_module.logger.warning
fired: list[int] = []
def racing_warning(*args, **kwargs):
if not fired:
fired.append(1)
real_unload()
return real_warning(*args, **kwargs)
monkeypatch.setattr(llama_cpp_module.logger, "warning", racing_warning)
assert b._respawn_if_dead() is False
assert unloads, "the racing unload was not honoured"
def test_socket_probe_is_false_without_a_port():
# Unloaded backends have no port; the probe must not raise, and the caller
# then falls back to the poll-based grace.
b = _recovery_backend()
b._port = None
assert b._server_socket_is_open() is False
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.