Merge pull request #39 from nickroci/fix/priming-funnel-index-build

fix(priming): funnel index build through inference thread + incremental indexing

Author: nickroci

daemon/agent_mem_daemon/__main__.py

@@ -324,13 +324,12 @@ def _prewarm_indexes(knowledge_root: Path, log: logging.Logger) -> None:
     try:
         from embeddings import load_or_build as emb_load  # noqa: PLC0415
 
-        idx = emb_load(knowledge_root)
-        # Dummy search to JIT-compile the MPS kernels for the query path.
-        # ``load_or_build`` only loads weights and computes doc vectors — it
-        # never runs the query-side forward pass, so the first real priming
-        # request would pay a ~2-3s graph-compile cost without this. Run it on
-        # the single inference thread so the kernels compile on the thread that
-        # actually serves requests (else the first real request recompiles).
+        # Build AND query both run model.encode() — run both on the single
+        # inference thread so the doc-encode and query-encode kernels compile on
+        # the thread that actually serves requests (else the first real request
+        # recompiles), and so warmup never encodes on a different thread than
+        # serving (the cross-thread case is what deadlocks MPS).
+        idx = _inference.run(emb_load, knowledge_root)
         _inference.run(idx.search, "warmup query for kernel compile", k=1)
         log.info("startup: embedding index ready (%d docs)", len(idx.docs))
     except Exception:

daemon/agent_mem_daemon/library_tools.py

@@ -152,15 +152,16 @@ def _run_embedding_search(args: Dict[str, Any], root: Path) -> Dict[str, Any]:
     if not root.exists():
         return _text_response("(library is empty — no entries to search)")
     try:
-        index = embeddings_load_or_build(root)
+        # Build + query both run model.encode() → both on the single inference
+        # thread. The Scholar/Librarian share the MPS model with priming, and a
+        # stale-index rebuild here re-encodes the whole corpus; doing that on a
+        # worker thread concurrently with priming deadlocks MPS (see _inference).
+        raw = _inference.run(lambda: embeddings_load_or_build(root).search(query, k=k))
     except FileNotFoundError:
         return _text_response("(library has no entries yet)")
     except Exception as e:
-        log.exception("embedding index load/build failed")
+        log.exception("embedding load/search failed")
         return _text_response(f"(embedding backend error: {e})")
-    # Embedding inference on the daemon's single inference thread — the
-    # Scholar/Librarian workers share the MPS model with priming (see _inference).
-    raw = _inference.run(index.search, query, k=k)
     as_tuples: list[Tuple[Path, float, str]] = [(h.path, h.score, h.snippet) for h in raw]
     return _format_hit_lines(as_tuples, root, query=query)
 

daemon/agent_mem_daemon/priming.py

@@ -386,19 +386,20 @@ def _embedding_search(
         return []
 
     try:
-        index = embeddings_load_or_build(knowledge_dir)
+        # BOTH the index build and the query encode run model.encode(), so BOTH
+        # must run on the single inference thread. ``load_or_build`` re-encodes
+        # the whole corpus whenever the Scholar has written since the last build
+        # (constant during a session) — running that on the handler pool means
+        # concurrent priming requests each re-encode all docs on their own
+        # thread → MPS deadlock. (#37 funnelled the query but not the build,
+        # which is what actually hangs; confirmed via py-spy.)
+        hits = _inference.run(
+            lambda: embeddings_load_or_build(knowledge_dir).search(query, k=max(1, k))
+        )
     except FileNotFoundError:
         return []
     except Exception:
-        log.exception("priming: embedding index load/build failed")
-        return []
-
-    try:
-        # MUST run on the single inference thread — concurrent encode() calls
-        # from multiple threads deadlock MPS and corrupt tensors (see _inference).
-        hits = _inference.run(index.search, query, k=max(1, k))
-    except Exception:
-        log.exception("priming: embedding search raised")
+        log.exception("priming: embedding load/search raised")
         return []
 
     # Filter out very weak semantic matches (see EMBEDDING_NOISE_FLOOR).

daemon/tests/test_priming_internals.py

@@ -7,6 +7,7 @@
 
 from __future__ import annotations
 
+import threading
 from pathlib import Path
 
 import pytest
@@ -235,6 +236,40 @@ def _missing(*args, **kwargs):
     assert priming._embedding_search(k, "x", k=5) == []
 
 
+def test_embedding_search_build_and_query_run_on_inference_thread(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """The index BUILD (``embeddings_load_or_build``) and the query must BOTH run
+    on the single inference thread, not the calling (handler-pool) thread.
+
+    Regression guard: #37 funnelled only ``index.search``, leaving the build —
+    which re-encodes the whole corpus when the index is stale — on the handler
+    thread. Concurrent priming requests then each rebuild on their own thread,
+    and concurrent ``model.encode()`` deadlocks MPS (py-spy-confirmed). On the
+    pre-fix code ``seen["build"]`` is the caller thread and this assertion fails.
+    """
+    kdir = tmp_path / "knowledge"
+    kdir.mkdir()
+    seen: dict[str, str] = {}
+
+    class _Idx:
+        def search(self, _query: str, k: int) -> list[object]:  # noqa: ARG002
+            seen["query"] = threading.current_thread().name
+            return []
+
+    def _load(*_args: object, **_kwargs: object) -> _Idx:
+        seen["build"] = threading.current_thread().name
+        return _Idx()
+
+    monkeypatch.setattr(priming, "embeddings_load_or_build", _load)
+    caller = threading.current_thread().name
+    priming._embedding_search(kdir, "x", k=5)
+
+    assert seen["build"].startswith("mps-inference"), seen
+    assert seen["query"].startswith("mps-inference"), seen
+    assert seen["build"] != caller
+
+
 def test_embedding_search_handles_search_failure(tmp_path: Path, monkeypatch) -> None:
     k = tmp_path / "knowledge"
     k.mkdir()

tools/search/embeddings.py

@@ -23,12 +23,16 @@
     swallows unpickle errors and rebuilds.
 
 Threading:
-  - Reads (``search``) are safe to call concurrently.
-  - ``save_index`` is concurrency-safe (see ``bm25.save_pickled``): racing
-    saves to the same path each write a unique temp file then atomically
-    replace, so they never interleave into a corrupt file. Last writer wins.
-    ``build_index`` has no shared state, so parallel rebuilds are wasteful
-    (duplicated encode work) but correct.
+  - ``search``, ``build_index`` and ``_incremental_build`` all call
+    ``model.encode()``, which is NOT thread-safe on MPS: concurrent encodes from
+    different threads DEADLOCK (and corrupt each other's tensors). The daemon
+    serialises ALL inference — search and (incremental) rebuild — onto a single
+    thread via ``agent_mem_daemon/_inference.py``; callers outside that funnel
+    must not run these concurrently. (The earlier note here — "parallel rebuilds
+    are wasteful but correct" — was wrong: they hang.)
+  - ``save_index`` is concurrency-safe (see ``bm25.save_pickled``): racing saves
+    to the same path each write a unique temp file then atomically replace, so
+    they never interleave into a corrupt file. Last writer wins.
 """
 
 from __future__ import annotations
@@ -325,6 +329,90 @@ def build_index(
     )
 
 
+def _incremental_build(
+    cached: EmbeddingIndex,
+    knowledge_dir: Path,
+    model_name: str,
+) -> EmbeddingIndex:
+    """Rebuild re-encoding ONLY the entries that changed since ``cached``.
+
+    Reuses ``cached``'s embedding row for every entry whose path is unchanged and
+    whose mtime hasn't advanced (matching :func:`is_stale`); encodes only new or
+    modified entries; drops deleted ones. A full :func:`build_index` re-encodes
+    the whole corpus — expensive on the priming hot path, which rebuilds after
+    every Scholar write. The ``model.encode`` here is funnelled onto the daemon's
+    single inference thread by the callers (``_inference.run``), so it stays
+    MPS-safe. Falls back to a full build if the cached rows can't be reused.
+    """
+    knowledge_dir = knowledge_dir.expanduser().resolve()
+    n_cached = cached.embeddings.shape[0] if cached.embeddings.ndim == 2 else 0
+    by_path = {rec.path: i for i, rec in enumerate(cached.docs)}
+
+    docs: list[_DocRecord] = []
+    rows: list[_FloatArray] = []
+    encode_texts: list[str] = []
+    encode_slots: list[int] = []
+    for md in _iter_markdown(knowledge_dir):
+        try:
+            mtime = md.stat().st_mtime
+        except OSError:
+            continue
+        path = str(md)
+        ci = by_path.get(path)
+        if ci is not None and ci < n_cached and cached.docs[ci].mtime + 1e-6 >= mtime:
+            # Unchanged → reuse the cached embedding row (refresh stored mtime so
+            # the next staleness check is exact).
+            docs.append(_DocRecord(path=path, mtime=mtime, raw_text=cached.docs[ci].raw_text))
+            rows.append(cached.embeddings[ci])
+            continue
+        try:
+            raw = md.read_text(encoding="utf-8")
+        except (OSError, UnicodeDecodeError):
+            continue
+        text = _embedding_text(raw).strip()
+        if not text:
+            continue
+        docs.append(_DocRecord(path=path, mtime=mtime, raw_text=raw))
+        rows.append(np.zeros((0,), dtype=np.float32))  # placeholder; overwritten below
+        encode_texts.append(text)
+        encode_slots.append(len(docs) - 1)
+
+    if not docs:
+        return EmbeddingIndex(
+            knowledge_dir=knowledge_dir,
+            model_name=model_name,
+            docs=[],
+            embeddings=np.zeros((0, 0), dtype=np.float32),
+            built_at=time.time(),
+        )
+
+    if encode_texts:
+        model = _load_model(model_name)
+        fresh = model.encode(
+            [_format_doc(t, model_name) for t in encode_texts],
+            convert_to_numpy=True,
+            normalize_embeddings=True,
+            show_progress_bar=False,
+            batch_size=8,
+        ).astype(np.float32)
+        for i, slot in enumerate(encode_slots):
+            rows[slot] = fresh[i]
+
+    try:
+        embeddings: _FloatArray = np.vstack(rows).astype(np.float32)
+    except ValueError:
+        # Ragged rows (corrupt / dimension-mismatched cache) — full rebuild.
+        return build_index(knowledge_dir, model_name=model_name)
+
+    return EmbeddingIndex(
+        knowledge_dir=knowledge_dir,
+        model_name=model_name,
+        docs=docs,
+        embeddings=embeddings,
+        built_at=time.time(),
+    )
+
+
 def save_index(index: EmbeddingIndex, path: Path | None = None) -> Path:
     """Pickle the index atomically. Returns the path written.
 
@@ -343,11 +431,15 @@ def load_or_build(
     force_rebuild: bool = False,
     index_path: Path | None = None,
 ) -> EmbeddingIndex:
-    """Load the persisted index if fresh, otherwise rebuild and save.
+    """Load the persisted index if fresh; otherwise rebuild (incrementally when
+    possible) and save.
 
     "Fresh" means the index was built against the same ``knowledge_dir`` and
     ``model_name``, every file we tracked is still present with the same mtime,
-    and no new ``.md`` files have appeared.
+    and no new ``.md`` files have appeared. When the cache is compatible but
+    stale we re-encode ONLY the changed/new entries (see :func:`_incremental_build`)
+    rather than the whole corpus — the corpus rebuild is what made priming slow
+    on the hot path after every Scholar write.
     """
     knowledge_dir = knowledge_dir.expanduser().resolve()
     target = index_path or _default_index_path(knowledge_dir)
@@ -358,12 +450,17 @@ def load_or_build(
             cached is not None
             and cached.knowledge_dir == knowledge_dir
             and cached.model_name == model_name
-            and not is_stale(
-                cached.docs,
-                {str(p): p.stat().st_mtime for p in _iter_markdown(knowledge_dir)},
-            )
         ):
-            return cached
+            current = {str(p): p.stat().st_mtime for p in _iter_markdown(knowledge_dir)}
+            if not is_stale(cached.docs, current):
+                return cached
+            # Compatible but stale → re-encode only the entries that changed.
+            index = _incremental_build(cached, knowledge_dir, model_name)
+            try:
+                save_index(index, target)
+            except OSError:
+                pass
+            return index
 
     index = build_index(knowledge_dir, model_name=model_name)
     try:

tools/search/test_embeddings.py

@@ -15,8 +15,10 @@
 import time
 from pathlib import Path
 
+import numpy as np
 import pytest
 
+import embeddings
 from embeddings import (
     EmbeddingHit,
     EmbeddingIndex,
@@ -191,6 +193,53 @@ def test_load_or_build_caches_and_rebuilds(tmp_path: Path) -> None:
     assert third.embeddings.shape[0] == len(third.docs) == 5
 
 
+def test_load_or_build_incremental_reencodes_only_changed(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """A stale-but-compatible index re-encodes ONLY the changed entries, not the
+    whole corpus — the fix for the priming-hot-path full rebuild after every
+    Scholar write. On the old code this re-encoded everything (batches [3, 3]).
+    """
+    knowledge = tmp_path / "knowledge"
+    knowledge.mkdir()
+    for name in ("a", "b", "c"):
+        (knowledge / f"{name}.md").write_text(f"# {name}\nbody about {name}\n", encoding="utf-8")
+
+    batches: list[int] = []
+
+    class _FakeModel:
+        def encode(self, texts: list[str], **_kwargs: object) -> object:
+            batches.append(len(texts))
+            arr = np.zeros((len(texts), 4), dtype=np.float32)
+            for i, text in enumerate(texts):
+                arr[i, len(text) % 4] = 1.0
+            return arr
+
+    fake = _FakeModel()
+    monkeypatch.setattr(embeddings, "_load_model", lambda *_a, **_k: fake)
+
+    first = load_or_build(knowledge)
+    assert batches == [3]  # full build encoded all 3 entries
+    assert len(first.docs) == 3
+
+    # Modify ONE entry and push its mtime into the future so the index is stale.
+    (knowledge / "b.md").write_text("# b\ntotally different body now\n", encoding="utf-8")
+    future = time.time() + 5
+    os.utime(knowledge / "b.md", (future, future))
+
+    second = load_or_build(knowledge)
+    assert batches == [3, 1]  # incremental encoded ONLY the changed entry
+    assert len(second.docs) == 3
+
+    # Unchanged entries kept their exact cached embedding rows (reused, not re-encoded).
+    def _row(idx: EmbeddingIndex, name: str) -> object:
+        pos = next(i for i, d in enumerate(idx.docs) if d.path.endswith(name))
+        return idx.embeddings[pos]
+
+    assert np.array_equal(_row(first, "a.md"), _row(second, "a.md"))
+    assert np.array_equal(_row(first, "c.md"), _row(second, "c.md"))
+
+
 # ── Edge cases that don't need the model loaded ──────────────────────────────