feat(llm): VRAM-aware generation routing — acquire_generation + pick_model (refs #654)

src/llm/pool.py

@@ -1,10 +1,14 @@
-"""Least-loaded fan-out across Ollama hosts.
+"""Least-loaded fan-out across Ollama hosts + largest-GPU generation.
 
-Remote GPUs are plain Ollama endpoints (rig 4090 / laptop 5080 run
+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
@@ -12,7 +16,7 @@ from __future__ import annotations
 import threading
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
-from typing import Iterator, Sequence
+from typing import Iterator, Mapping, Sequence
 
 import httpx
 from langchain_core.embeddings import Embeddings
@@ -21,38 +25,67 @@ _TIMEOUT = httpx.Timeout(120.0, connect=5.0)
 
 
 class HostPool:
-    """Tracks in-flight requests per host; hands out the idlest one."""
+    """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]) -> None:
+    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)
+        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 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")
-                    names = {
-                        m["name"].split(":")[0] for m in resp.json().get("models", [])
-                    }
+                    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}
@@ -63,6 +96,11 @@ class HostPool:
             )
         return alive
 
+    def _release(self, host: str) -> None:
+        with self._lock:
+            if host in self._in_flight:
+                self._in_flight[host] -= 1
+
     @contextmanager
     def acquire(self) -> Iterator[str]:
         with self._lock:
@@ -71,9 +109,38 @@ class HostPool:
         try:
             yield host
         finally:
-            with self._lock:
-                if host in self._in_flight:
-                    self._in_flight[host] -= 1
+            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:
+            host = max(
+                self._in_flight,
+                key=lambda h: (self._vram.get(h, 0.0), -self._in_flight[h]),
+            )
+            self._in_flight[host] += 1
+        try:
+            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(

tests/llm/test_pool.py

@@ -155,3 +155,108 @@ class TestPoolEmbeddings:
         emb = PoolEmbeddings(pool, "m")
         assert emb.embed_documents(["x"]) == [[0.1]]
         assert emb.embed_query("x") == [0.1]
+
+
+class TestGeneration:
+    def _pool(self):
+        return HostPool([H1, H2], vram_gb={H1: 12, H2: 24})
+
+    @patch("llm.pool.httpx.Client")
+    def test_acquire_generation_prefers_largest_live_host(self, MockClient):
+        client = MockClient.return_value.__enter__.return_value
+        client.get.side_effect = [
+            _resp({"models": [{"name": "chat:latest"}]}),
+            _resp({"models": [{"name": "chat:latest"}, {"name": "big:latest"}]}),
+        ]
+        pool = self._pool()
+        pool.check("chat")
+        with pool.acquire_generation() as host:
+            assert host == H2
+            assert pool._in_flight[H2] == 1
+        assert pool._in_flight[H2] == 0
+
+    @patch("llm.pool.httpx.Client")
+    def test_falls_back_when_largest_is_down(self, MockClient):
+        import httpx
+
+        client = MockClient.return_value.__enter__.return_value
+        client.get.side_effect = [
+            _resp({"models": [{"name": "chat:latest"}]}),
+            httpx.ConnectError("rig down"),
+        ]
+        pool = self._pool()
+        pool.check("chat")
+        with pool.acquire_generation() as host:
+            assert host == H1
+
+    def test_tie_breaks_on_least_in_flight(self):
+        pool = HostPool([H1, H2], vram_gb={H1: 24, H2: 24})
+        pool._in_flight[H1] = 3
+        with pool.acquire_generation() as host:
+            assert host == H2
+
+    def test_vram_defaults_to_zero(self):
+        assert HostPool([H1]).vram(H1) == 0.0
+
+    @patch("llm.pool.httpx.Client")
+    def test_serves_reflects_last_check(self, MockClient):
+        client = MockClient.return_value.__enter__.return_value
+        client.get.side_effect = [
+            _resp({"models": [{"name": "chat:latest"}, {"name": "big:latest"}]}),
+        ]
+        pool = HostPool([H1])
+        pool.check("chat")
+        assert pool.serves(H1, "big") is True
+        assert pool.serves(H1, "big:latest") is True
+        assert pool.serves(H1, "nope") is False
+        assert pool.status() == [
+            {"host": H1, "vram_gb": 0.0, "models": ["big:latest", "chat:latest"]}
+        ]
+
+
+class TestPickModel:
+    def _cfg(self, **kw):
+        base = dict(
+            ollama_hosts=(H1, H2),
+            embed_model="e",
+            instruct_model="chat",
+            instruct_model_large="big",
+            large_min_vram_gb=20,
+            embed_dim=768,
+            build_ann_index=False,
+            pg_host="x",
+            pg_port=5432,
+            pg_db="llm",
+            pg_user="llm",
+        )
+        base.update(kw)
+        return LlmConfig(**base)
+
+    def _pool_serving(self, models_by_host):
+        pool = HostPool(list(models_by_host), vram_gb={H1: 12, H2: 24})
+        pool._models = {h: set(ms) for h, ms in models_by_host.items()}
+        return pool
+
+    def test_large_on_big_host_that_serves_it(self):
+        from llm.pool import pick_model
+
+        pool = self._pool_serving({H1: {"chat"}, H2: {"chat", "big"}})
+        assert pick_model(self._cfg(), pool, H2) == "big"
+
+    def test_baseline_on_small_host(self):
+        from llm.pool import pick_model
+
+        pool = self._pool_serving({H1: {"chat", "big"}, H2: {"chat", "big"}})
+        assert pick_model(self._cfg(), pool, H1) == "chat"
+
+    def test_baseline_when_big_host_lacks_large_model(self):
+        from llm.pool import pick_model
+
+        pool = self._pool_serving({H1: {"chat"}, H2: {"chat"}})
+        assert pick_model(self._cfg(), pool, H2) == "chat"
+
+    def test_baseline_when_no_large_configured(self):
+        from llm.pool import pick_model
+
+        pool = self._pool_serving({H2: {"chat", "big"}})
+        assert pick_model(self._cfg(instruct_model_large=""), pool, H2) == "chat"