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

src/llm/pool.py

@@ -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:
+        # A host that passed check() can still vanish mid-run (mDNS name
-            began = time.monotonic()
+        # stops resolving, box sleeps): drop it and retry the batch on the
-            resp = client.post(
+        # survivors instead of failing the whole run (#796). A host that
-                f"{host}/api/embed", json={"model": model, "input": batch}
+        # answers with an HTTP error is a model/config problem and still
-            )
+        # raises through raise_for_status.
-            resp.raise_for_status()
+        while True:
-            results[start] = resp.json()["embeddings"]
+            with pool.acquire() as host, httpx.Client(timeout=_TIMEOUT) as client:
-            metrics.embedded(host, len(batch), time.monotonic() - began)
+                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:

tests/llm/test_pool.py

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