fix(llm): drop a pool host that dies mid-run and retry the batch on the survivors (refs #796)

check() filters the pool at start, but a host can vanish afterwards
(rig.local stopped resolving during the 2026-09-22 CMS-2026-2377
re-index; same on 2026-09-04) and embed_texts let the first transport
error abort the whole run. Now an httpx.TransportError drops that host
(HostPool.drop) and the batch is retried elsewhere; only when no host
is left does a RuntimeError name the last failure. acquire() and
acquire_generation() on an empty pool raise the same clear error
instead of min()/max()'s ValueError. HTTP error responses still raise
through raise_for_status — that is a model/config problem, not a dead
host.
This commit is contained in:
kert
2026-09-22 15:52:07 -04:00
parent 4bd86b091f
commit cbc55c95d4
2 changed files with 88 additions and 8 deletions

View File

@@ -13,6 +13,7 @@ stay least-loaded (``acquire``).
from __future__ import annotations
import logging
import threading
import time
from concurrent.futures import ThreadPoolExecutor
@@ -25,6 +26,7 @@ from langchain_core.embeddings import Embeddings
from llm import metrics
_TIMEOUT = httpx.Timeout(120.0, connect=5.0)
log = logging.getLogger(__name__)
class HostPool:
@@ -108,9 +110,20 @@ class HostPool:
if host in self._in_flight:
self._in_flight[host] -= 1
def drop(self, host: str) -> None:
"""Forget ``host`` for the rest of the run — a transport failure
after ``check()`` passed means it went away mid-run (#796)."""
with self._lock:
self._in_flight.pop(host.rstrip("/"), None)
self._models.pop(host.rstrip("/"), None)
@contextmanager
def acquire(self) -> Iterator[str]:
with self._lock:
if not self._in_flight:
raise RuntimeError(
"no live Ollama host in pool — check() or drop() left none"
)
host = min(self._in_flight, key=self._in_flight.__getitem__)
self._in_flight[host] += 1
try:
@@ -127,6 +140,10 @@ class HostPool:
serve the model; this then realises "largest GPU available now".
"""
with self._lock:
if not self._in_flight:
raise RuntimeError(
"no live Ollama host in pool — check() or drop() left none"
)
host = max(
self._in_flight,
key=lambda h: (self._vram.get(h, 0.0), -self._in_flight[h]),
@@ -166,14 +183,37 @@ def embed_texts(
results: dict[int, list[list[float]]] = {}
def run(start: int, batch: list[str]) -> None:
with pool.acquire() as host, httpx.Client(timeout=_TIMEOUT) as client:
began = time.monotonic()
resp = client.post(
f"{host}/api/embed", json={"model": model, "input": batch}
)
resp.raise_for_status()
results[start] = resp.json()["embeddings"]
metrics.embedded(host, len(batch), time.monotonic() - began)
# A host that passed check() can still vanish mid-run (mDNS name
# stops resolving, box sleeps): drop it and retry the batch on the
# survivors instead of failing the whole run (#796). A host that
# answers with an HTTP error is a model/config problem and still
# raises through raise_for_status.
while True:
with pool.acquire() as host, httpx.Client(timeout=_TIMEOUT) as client:
began = time.monotonic()
try:
resp = client.post(
f"{host}/api/embed", json={"model": model, "input": batch}
)
except httpx.TransportError as exc:
pool.drop(host)
left = pool.hosts
log.warning(
"dropping %s after %s: %s; %d host(s) left",
host,
type(exc).__name__,
exc,
len(left),
)
if not left:
raise RuntimeError(
f"every Ollama host failed; last {host}: {exc}"
) from exc
continue
resp.raise_for_status()
results[start] = resp.json()["embeddings"]
metrics.embedded(host, len(batch), time.monotonic() - began)
return
workers = max(1, min(len(pool.hosts) * 2, len(batches)))
with ThreadPoolExecutor(max_workers=workers) as ex:

View File

@@ -2,6 +2,7 @@
from unittest.mock import MagicMock, patch
import httpx
import pytest
from llm.config import LlmConfig
@@ -116,6 +117,45 @@ class TestEmbedTexts:
out = embed_texts(pool, "m", ["a", "bb", "ccc"], batch_size=1)
assert out == [[1.0], [2.0], [3.0]]
@patch("llm.pool.httpx.Client")
def test_dying_host_is_dropped_and_batch_retried_on_survivor(self, MockClient):
"""#796: a transport error from one host must not abort the run —
the host leaves the pool and the batch completes elsewhere."""
client = MockClient.return_value.__enter__.return_value
calls: list[str] = []
def fake_post(url, json=None, **kw):
calls.append(url)
if url.startswith(H1):
raise httpx.ConnectError("[Errno -2] Name or service not known")
return _resp({"embeddings": [[float(len(t))] for t in json["input"]]})
client.post.side_effect = fake_post
pool = HostPool([H1, H2])
out = embed_texts(pool, "m", ["a", "bb", "ccc"], batch_size=1)
assert out == [[1.0], [2.0], [3.0]]
assert pool.hosts == [H2]
assert any(u.startswith(H1) for u in calls) # H1 was tried, then dropped
assert pool._in_flight == {H2: 0} # dropped host never left a count behind
@patch("llm.pool.httpx.Client")
def test_every_host_dying_raises_runtime_error(self, MockClient):
client = MockClient.return_value.__enter__.return_value
client.post.side_effect = httpx.ReadTimeout("timed out")
pool = HostPool([H1, H2])
with pytest.raises(RuntimeError, match="every Ollama host failed"):
embed_texts(pool, "m", ["a"])
assert pool.hosts == []
def test_acquire_on_empty_pool_is_a_clear_error(self):
pool = HostPool([])
with pytest.raises(RuntimeError, match="no live Ollama host"):
with pool.acquire():
pass
with pytest.raises(RuntimeError, match="no live Ollama host"):
with pool.acquire_generation():
pass
@patch("llm.pool.httpx.Client")
def test_http_error_raises(self, MockClient):
client = MockClient.return_value.__enter__.return_value