diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index e54b5269c1..397cba8842 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -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, ): diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index cf9fde7118..92089d26ec 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -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): diff --git a/studio/backend/tests/test_tensor_parallel.py b/studio/backend/tests/test_tensor_parallel.py index 00c7aeac69..23c70f8499 100644 --- a/studio/backend/tests/test_tensor_parallel.py +++ b/studio/backend/tests/test_tensor_parallel.py @@ -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.