cbc55c95d4
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.
237 lines
8.6 KiB
Python
237 lines
8.6 KiB
Python
"""Least-loaded fan-out across Ollama hosts + largest-GPU generation.
|
|
|
|
Remote GPUs are plain Ollama endpoints (rig 4090 / laptop 5070 Ti run
|
|
``ollama serve``); this pool spreads embed calls across them while all
|
|
bib/pgvector I/O stays on the server. Uses the raw ``/api/embed`` HTTP API
|
|
rather than langchain-ollama so one pool can juggle N base URLs;
|
|
``PoolEmbeddings`` adapts it back to the langchain interface for PGVector.
|
|
|
|
Hosts carry a declared VRAM size (``LLM_OLLAMA_HOSTS`` ``url@gb``);
|
|
``acquire_generation`` hands chat to the largest live host, while embeds
|
|
stay least-loaded (``acquire``).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import threading
|
|
import time
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from contextlib import contextmanager
|
|
from typing import Iterator, Mapping, Sequence
|
|
|
|
import httpx
|
|
from langchain_core.embeddings import Embeddings
|
|
|
|
from llm import metrics
|
|
|
|
_TIMEOUT = httpx.Timeout(120.0, connect=5.0)
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class HostPool:
|
|
"""Tracks in-flight requests per host; hands out the idlest one for
|
|
embeds and the largest live one for generation."""
|
|
|
|
def __init__(
|
|
self, hosts: Sequence[str], *, vram_gb: Mapping[str, float] | None = None
|
|
) -> None:
|
|
self._lock = threading.Lock()
|
|
self._in_flight: dict[str, int] = {h.rstrip("/"): 0 for h in hosts}
|
|
vram = vram_gb or {}
|
|
self._vram: dict[str, float] = {
|
|
h: float(vram.get(h, vram.get(h + "/", 0.0))) for h in self._in_flight
|
|
}
|
|
# Model names (full tag + bare prefix) seen on each host at the last
|
|
# check(); empty until a check has run.
|
|
self._models: dict[str, set[str]] = {}
|
|
|
|
@classmethod
|
|
def from_config(cls, cfg) -> "HostPool":
|
|
return cls(cfg.ollama_hosts, vram_gb=getattr(cfg, "host_vram", None))
|
|
|
|
@property
|
|
def hosts(self) -> list[str]:
|
|
return list(self._in_flight)
|
|
|
|
def vram(self, host: str) -> float:
|
|
return self._vram.get(host.rstrip("/"), 0.0)
|
|
|
|
def reachable(self, host: str) -> bool:
|
|
"""True when the last ``check()`` got a model list from ``host``."""
|
|
return host.rstrip("/") in self._models
|
|
|
|
def serves(self, host: str, model: str) -> bool:
|
|
names = self._models.get(host.rstrip("/"), set())
|
|
return model in names or model.split(":")[0] in names
|
|
|
|
def status(self) -> list[dict]:
|
|
return [
|
|
{
|
|
"host": h,
|
|
"vram_gb": self._vram.get(h, 0.0),
|
|
"models": sorted(n for n in self._models.get(h, set()) if ":" in n),
|
|
}
|
|
for h in self._in_flight
|
|
]
|
|
|
|
def check(self, model: str) -> list[str]:
|
|
"""Keep only hosts that are up and serve ``model``.
|
|
|
|
Ollama tags models ``name:latest``; match on the bare prefix.
|
|
Remembers every model each live host serves (see ``serves``).
|
|
"""
|
|
alive: list[str] = []
|
|
with httpx.Client(timeout=_TIMEOUT) as client:
|
|
for host in self.hosts:
|
|
try:
|
|
resp = client.get(f"{host}/api/tags")
|
|
full = {m["name"] for m in resp.json().get("models", [])}
|
|
names = full | {n.split(":")[0] for n in full}
|
|
self._models[host] = names
|
|
if model.split(":")[0] in names:
|
|
alive.append(host)
|
|
except (httpx.HTTPError, ValueError, KeyError, TypeError):
|
|
# Down, or up but returning a malformed/non-JSON body —
|
|
# either way, drop it rather than aborting the whole run.
|
|
self._models.pop(host, None)
|
|
continue
|
|
with self._lock:
|
|
self._in_flight = {h: self._in_flight.get(h, 0) for h in alive}
|
|
if not alive:
|
|
raise RuntimeError(
|
|
f"no Ollama host in pool serves {model!r} — "
|
|
f"pull it or fix LLM_OLLAMA_HOSTS"
|
|
)
|
|
return alive
|
|
|
|
def _release(self, host: str) -> None:
|
|
with self._lock:
|
|
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:
|
|
with metrics.dispatch(host, "embed"):
|
|
yield host
|
|
finally:
|
|
self._release(host)
|
|
|
|
@contextmanager
|
|
def acquire_generation(self) -> Iterator[str]:
|
|
"""The largest live host (declared VRAM), ties → least in-flight.
|
|
|
|
Call ``check(model)`` first so the pool holds only live hosts that
|
|
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]),
|
|
)
|
|
self._in_flight[host] += 1
|
|
try:
|
|
with metrics.dispatch(host, "generate"):
|
|
yield host
|
|
finally:
|
|
self._release(host)
|
|
|
|
|
|
def pick_model(cfg, pool: HostPool, host: str) -> str:
|
|
"""``instruct_model_large`` when ``host`` declares enough VRAM and
|
|
serves it; else ``instruct_model``."""
|
|
large = getattr(cfg, "instruct_model_large", "")
|
|
if (
|
|
large
|
|
and pool.vram(host) >= float(getattr(cfg, "large_min_vram_gb", 20.0))
|
|
and pool.serves(host, large)
|
|
):
|
|
return large
|
|
return cfg.instruct_model
|
|
|
|
|
|
def embed_texts(
|
|
pool: HostPool,
|
|
model: str,
|
|
texts: Sequence[str],
|
|
*,
|
|
batch_size: int = 64,
|
|
) -> list[list[float]]:
|
|
"""Embed ``texts`` in order, fanning batches across the pool."""
|
|
batches = [
|
|
(i, list(texts[i : i + batch_size])) for i in range(0, len(texts), batch_size)
|
|
]
|
|
results: dict[int, list[list[float]]] = {}
|
|
|
|
def run(start: int, batch: list[str]) -> None:
|
|
# 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:
|
|
for future in [ex.submit(run, s, b) for s, b in batches]:
|
|
future.result()
|
|
return [vec for start in sorted(results) for vec in results[start]]
|
|
|
|
|
|
class PoolEmbeddings(Embeddings):
|
|
"""langchain adapter so PGVector can query through the pool."""
|
|
|
|
def __init__(self, pool: HostPool, model: str) -> None:
|
|
self._pool = pool
|
|
self._model = model
|
|
|
|
def embed_documents(self, texts: list[str]) -> list[list[float]]:
|
|
return embed_texts(self._pool, self._model, texts)
|
|
|
|
def embed_query(self, text: str) -> list[float]:
|
|
return self.embed_documents([text])[0]
|