fix: stop the re-embed migration from destroying the index it rebuilds

Five defects from a review of the embeddings work, four of them silent.

- Write local files atomically. `LocalStorage.save_file` streamed straight onto
  the destination, so an interrupted write left a truncated file. `reembed`
  rewrites every index it touches, and a half-written `index.faiss` loads at
  neither the old width nor the new one -- the source was unrecoverable, with
  no backup and no temp file left behind. Bytes now land beside the destination
  and move into place with `os.replace`. S3 was already safe (single PUT).

- Read pgvector chunks a page at a time. `reembed_pgvector` materialised every
  `(id, text)` row for a source before embedding -- ~1.6 GB at 200k chunks and
  several times that for non-Latin scripts, with the `PGresult` held alongside
  until the cursor closed. Inside the shipped 4Gi limit, while also holding the
  model, that is an OOMKill -- which is exactly the SIGKILL the point above
  turned into a destroyed index. It now walks the source by keyset.

- Bound the first wave of delegated embeds. The failure cooldown is only latched
  once the first `get()` returns, so every request already in flight paid the
  full EMBEDDINGS_DELEGATE_TIMEOUT: measured 64 threads all timing out together,
  and at the shipped 60s across a 96-thread WSGI pool that is an API serving
  nothing at all, health checks included. One caller now probes while the rest
  fail fast; after a single success the gate leaves the path entirely.

- Ship EMBEDDINGS_NAME commented in .env-template. The comment directly above it
  says to leave it commented when upgrading, and the line shipped set. Any value
  reaching `.env` lands in `model_fields_set`, which makes `resolve_embeddings_pin`
  bail -- so a template-derived `.env` disabled the legacy pin outright and
  repointed a populated index at a different 768-dim model, where no width check
  fires. The pin already picks granite for a fresh install and mpnet for an
  existing one, so nothing needs to be set by hand.

- Stamp `sources.model` on wiki sources. They were created with the column NULL
  and then embedded like any other source, and the boot check reads NULL as
  "pre-dates the column, therefore the legacy model" -- reporting a correctly
  embedded source as stale on every startup of every process. Stamped at
  creation, and again on each page re-embed so existing rows heal.

The two docs that promised the FAISS index survives a failed run said so of the
embed only; both now describe the write, and upgrading.mdx says to stop ingest
for the duration.

application/api/user/sources/routes.py

@@ -800,6 +800,12 @@ class CreateWikiSource(Resource):
                     config={"kind": "wiki"},
                     directory_structure={},
                     tokens=0,
+                    # Wiki pages are embedded like any other source, so record
+                    # which model did it. Left unset, the column reads as NULL,
+                    # which the boot mismatch check takes to mean "pre-dates the
+                    # column, therefore the legacy model" -- and reports a
+                    # correctly-embedded source as stale on every startup.
+                    model=settings.EMBEDDINGS_NAME,
                 )
                 if initial_content:
                     WikiPagesRepository(conn).upsert(

application/scripts/reembed.py

@@ -197,15 +197,89 @@ def reembed_graph_nodes(store, conn, source_id: str, batch_size: int, dry_run: b
     return written
 
 
+def _count_source_chunks(conn, table: str, source_id: str) -> int:
+    """How many chunk rows one source holds."""
+    cursor = conn.cursor()
+    try:
+        cursor.execute(
+            sql.SQL("SELECT count(*) FROM {table} WHERE source_id = %s").format(
+                table=sql.Identifier(table)
+            ),
+            (source_id,),
+        )
+        return int(cursor.fetchone()[0])
+    finally:
+        cursor.close()
+
+
+def _fetch_chunk_page(
+    conn, table: str, source_id: str, after_id: Optional[Any], limit: int
+) -> List[Tuple[Any, str]]:
+    """One page of ``(id, text)`` ordered by id, starting after ``after_id``.
+
+    Keyed off ``id`` rather than ``OFFSET`` so each page costs an index seek,
+    and the page boundaries stay stable across the updates the caller makes
+    between pages -- only the vector column changes, never ``id``.
+    """
+    cursor = conn.cursor()
+    try:
+        if after_id is None:
+            cursor.execute(
+                sql.SQL(
+                    "SELECT id, text FROM {table} WHERE source_id = %s "
+                    "ORDER BY id LIMIT %s"
+                ).format(table=sql.Identifier(table)),
+                (source_id, limit),
+            )
+        else:
+            cursor.execute(
+                sql.SQL(
+                    "SELECT id, text FROM {table} WHERE source_id = %s AND id > %s "
+                    "ORDER BY id LIMIT %s"
+                ).format(table=sql.Identifier(table)),
+                (source_id, after_id, limit),
+            )
+        return cursor.fetchall()
+    finally:
+        cursor.close()
+
+
+def _write_chunk_page(conn, table: str, vector_column: str, store, page) -> None:
+    """Embed one page and commit its vectors."""
+    vectors = store._embedding.embed_documents([row[1] or "" for row in page])
+    cursor = conn.cursor()
+    try:
+        cursor.executemany(
+            sql.SQL("UPDATE {table} SET {column} = %s::vector WHERE id = %s").format(
+                table=sql.Identifier(table),
+                column=sql.Identifier(vector_column),
+            ),
+            [(str(list(vector)), row[0]) for vector, row in zip(vectors, page)],
+        )
+        conn.commit()
+    except Exception:
+        conn.rollback()
+        raise
+    finally:
+        cursor.close()
+
+
 def reembed_pgvector(source_id: str, batch_size: int, dry_run: bool) -> Tuple[int, int]:
     """Rewrite one source's vectors in place.
 
     Rows are updated, never deleted and re-inserted, so an interrupted run
     leaves every chunk present and simply re-does its batch next time.
 
+    Chunks are read a page at a time rather than all at once. A source's text is
+    the whole corpus: at the default chunk size 200k chunks materialise to about
+    1.6 GB of Python strings, several times that for non-Latin scripts, and the
+    result set is held alongside them until the cursor closes. That is enough to
+    be OOM-killed inside a 4 GiB container while also holding the model -- and a
+    kill here is what a torn index costs the most.
+
     Args:
         source_id: Source whose chunks to re-embed.
-        batch_size: Chunks per embed call and per transaction.
+        batch_size: Chunks per page, per embed call and per transaction.
         dry_run: When true, count the work and change nothing.
 
     Returns:
@@ -215,44 +289,22 @@ def reembed_pgvector(source_id: str, batch_size: int, dry_run: bool) -> Tuple[in
     table, vector_column = store._table_name, store._vector_column
     conn = store._get_connection()
 
-    cursor = conn.cursor()
+    seen = written = 0
     try:
-        cursor.execute(
-            sql.SQL("SELECT id, text FROM {table} WHERE source_id = %s ORDER BY id").format(
-                table=sql.Identifier(table)
-            ),
-            (source_id,),
-        )
-        rows = cursor.fetchall()
-    finally:
-        cursor.close()
-
-    written = 0
-    try:
-        if not dry_run:
-            for batch in _batched(rows, batch_size):
-                vectors = store._embedding.embed_documents([row[1] or "" for row in batch])
-                cursor = conn.cursor()
-                try:
-                    cursor.executemany(
-                        sql.SQL(
-                            "UPDATE {table} SET {column} = %s::vector WHERE id = %s"
-                        ).format(
-                            table=sql.Identifier(table),
-                            column=sql.Identifier(vector_column),
-                        ),
-                        [
-                            (str(list(vector)), row[0])
-                            for vector, row in zip(vectors, batch)
-                        ],
-                    )
-                    conn.commit()
-                except Exception:
-                    conn.rollback()
-                    raise
-                finally:
-                    cursor.close()
-                written += len(batch)
+        if dry_run:
+            seen = _count_source_chunks(conn, table, source_id)
+        else:
+            after_id: Optional[Any] = None
+            while True:
+                page = _fetch_chunk_page(
+                    conn, table, source_id, after_id, batch_size
+                )
+                if not page:
+                    break
+                _write_chunk_page(conn, table, vector_column, store, page)
+                after_id = page[-1][0]
+                seen += len(page)
+                written += len(page)
 
         # The graph seeds every traversal from its own vectors, so leaving them
         # in the old model's space is the same silent mismatch this script
@@ -268,7 +320,7 @@ def reembed_pgvector(source_id: str, batch_size: int, dry_run: bool) -> Tuple[in
                 )
     finally:
         store.close()
-    return len(rows), written
+    return seen, written
 
 
 def reembed_faiss(source_id: str, batch_size: int, dry_run: bool) -> Tuple[int, int]:
@@ -276,7 +328,9 @@ def reembed_faiss(source_id: str, batch_size: int, dry_run: bool) -> Tuple[int,
 
     A FAISS index is a flat array, so this rebuilds rather than updates. Every
     new vector is computed *before* the old index is touched, so an interrupt
-    during embedding leaves the existing index intact.
+    during embedding leaves the existing index intact -- and ``save_local``
+    writes each file to a temporary path and moves it into place, so an
+    interrupt during the write does too.
 
     Args:
         source_id: Source whose index to rebuild.

application/storage/local.py

@@ -1,6 +1,7 @@
 """Local file system implementation."""
 import os
 import shutil
+import tempfile
 from typing import BinaryIO, List, Callable
 
 from application.storage.base import BaseStorage
@@ -36,16 +37,49 @@ class LocalStorage(BaseStorage):
         return resolved
 
     def save_file(self, file_data: BinaryIO, path: str, **kwargs) -> dict:
-        """Save a file to local storage."""
+        """Save a file, replacing any existing one atomically.
+
+        The bytes land on a temporary file beside the destination and are moved
+        into place with ``os.replace``, so an interrupted write leaves the
+        previous file intact instead of a truncated one. Streaming straight onto
+        the destination is unrecoverable for a file that is rewritten in place:
+        a half-written ``index.faiss`` loads at neither the old width nor the
+        new one, and ``application.scripts.reembed`` rewrites every index it
+        touches.
+        """
         full_path = self._get_full_path(path)
+        directory = os.path.dirname(full_path)
+        os.makedirs(directory, exist_ok=True)
 
-        os.makedirs(os.path.dirname(full_path), exist_ok=True)
-
-        if hasattr(file_data, 'save'):
-            file_data.save(full_path)
-        else:
-            with open(full_path, 'wb') as f:
-                shutil.copyfileobj(file_data, f)
+        # Same directory as the destination, so the replace below is a rename
+        # within one filesystem rather than a copy.
+        fd, temp_path = tempfile.mkstemp(
+            dir=directory, prefix=f".{os.path.basename(full_path)}.", suffix=".tmp"
+        )
+        try:
+            with os.fdopen(fd, "wb") as f:
+                if hasattr(file_data, "save"):
+                    file_data.save(f)
+                else:
+                    shutil.copyfileobj(file_data, f)
+                f.flush()
+                os.fsync(f.fileno())
+            # mkstemp is 0600; keep whatever the file already had, or fall back
+            # to what a plain open() would have produced.
+            try:
+                mode = os.stat(full_path).st_mode & 0o777
+            except FileNotFoundError:
+                mode = 0o644
+            os.chmod(temp_path, mode)
+            os.replace(temp_path, full_path)
+        except BaseException:
+            # A successful replace consumes the temp file; every other exit
+            # leaves it behind.
+            try:
+                os.unlink(temp_path)
+            except OSError:
+                pass
+            raise
 
         return {
             'storage_type': 'local'

application/vectorstore/embeddings_delegated.py

@@ -22,6 +22,7 @@ network hop rather than a broker round trip.
 from __future__ import annotations
 
 import logging
+import threading
 import time
 from typing import Any, List, Optional
 
@@ -40,6 +41,13 @@ EMBED_TASK = "application.vectorstore.embeddings_tasks.embed_texts"
 #: single retrieval into one timeout rather than one per source.
 _FAILURE_COOLDOWN = 30.0
 
+#: How long a caller waits for the outcome of the dispatch already in flight
+#: before giving up on its own. Only applies while the worker is unproven --
+#: once one dispatch has succeeded, every caller goes straight to the broker.
+#: Comfortably above a healthy round trip (~60 ms on a prefork worker) and far
+#: below ``EMBEDDINGS_DELEGATE_TIMEOUT``, which is the point.
+_PROBE_WAIT = 2.0
+
 _NO_WORKER_HINT = (
     "Start a worker consuming it, point EMBEDDINGS_BASE_URL at an embedding "
     "service, or set EMBEDDINGS_DELEGATE_TO_WORKER=false to load the model in "
@@ -89,12 +97,21 @@ class DelegatedEmbeddings:
         self._local: Any = None
         self._dimension: Optional[int] = dimension_for(embeddings_name)
         self._failed_at: Optional[float] = None
+        # A dispatch has completed successfully, so the worker is known to be
+        # consuming the queue and callers need not take turns proving it.
+        self._verified = False
+        self._probing = False
+        self._state_lock = threading.Lock()
+        self._probe_done = threading.Event()
 
     def _cooldown_remaining(self) -> float:
         """Seconds left of the fail-fast window after a failed dispatch."""
-        if self._failed_at is None:
+        # One load: a concurrent success clearing the latch between two reads
+        # would otherwise subtract from None.
+        failed_at = self._failed_at
+        if failed_at is None:
             return 0.0
-        return max(0.0, _FAILURE_COOLDOWN - (time.monotonic() - self._failed_at))
+        return max(0.0, _FAILURE_COOLDOWN - (time.monotonic() - failed_at))
 
     def _local_embeddings(self):
         """The in-process model, built once, for use inside a worker task."""
@@ -104,26 +121,10 @@ class DelegatedEmbeddings:
             self._local = build_local_embeddings(self.embeddings_name, self.embeddings_key)
         return self._local
 
-    def _dispatch(self, texts: List[str]) -> List[List[float]]:
-        """Run the embed task on the worker and wait for its vectors."""
+    def _send(self, texts: List[str], queue: str, timeout: int) -> List[List[float]]:
+        """Publish the embed task and wait for its vectors."""
         from application.celery_init import celery
 
-        queue = getattr(settings, "EMBEDDINGS_QUEUE", "embeddings")
-        timeout = getattr(settings, "EMBEDDINGS_DELEGATE_TIMEOUT", 60)
-
-        # A missing worker is a property of the deployment, not of this call,
-        # so once one dispatch has timed out the next is not worth another full
-        # timeout. Without this latch a single retrieval pays the timeout twice
-        # -- once in ``fanout.embed_questions``, then again per source when it
-        # falls back to letting each store embed its own query.
-        remaining = self._cooldown_remaining()
-        if remaining > 0:
-            raise RuntimeError(
-                f"Skipping the embed dispatch: a previous request to the {queue!r} "
-                f"queue failed and the {_FAILURE_COOLDOWN}s cooldown has "
-                f"{remaining:.0f}s left. {_NO_WORKER_HINT}"
-            )
-
         result = celery.send_task(EMBED_TASK, args=[texts, self.embeddings_name], queue=queue)
         try:
             vectors = result.get(timeout=timeout)
@@ -137,8 +138,76 @@ class DelegatedEmbeddings:
         finally:
             _forget(result)
         self._failed_at = None
+        self._verified = True
         return vectors
 
+    def _cooldown_error(self, queue: str, remaining: float) -> RuntimeError:
+        return RuntimeError(
+            f"Skipping the embed dispatch: a previous request to the {queue!r} "
+            f"queue failed and the {_FAILURE_COOLDOWN}s cooldown has "
+            f"{remaining:.0f}s left. {_NO_WORKER_HINT}"
+        )
+
+    def _dispatch(self, texts: List[str]) -> List[List[float]]:
+        """Run the embed task on the worker and wait for its vectors.
+
+        A missing worker is a property of the deployment, not of this call, so
+        at most one caller waits out ``EMBEDDINGS_DELEGATE_TIMEOUT`` to discover
+        it. Two guards do that:
+
+        The cooldown latch covers requests arriving *after* a failure -- without
+        it a single retrieval pays the timeout twice, once in
+        ``fanout.embed_questions`` and again per source when it falls back to
+        letting each store embed its own query.
+
+        The probe covers requests already in flight *alongside* the first one,
+        which the latch cannot: nothing is latched until that first ``get()``
+        returns, so every thread in the opening wave would otherwise block for
+        the full timeout at once -- at the shipped 60s and 96 WSGI threads, an
+        API that serves nothing at all, health checks included.
+        """
+        queue = getattr(settings, "EMBEDDINGS_QUEUE", "embeddings")
+        timeout = getattr(settings, "EMBEDDINGS_DELEGATE_TIMEOUT", 60)
+
+        remaining = self._cooldown_remaining()
+        if remaining > 0:
+            raise self._cooldown_error(queue, remaining)
+
+        if self._verified:
+            return self._send(texts, queue, timeout)
+
+        with self._state_lock:
+            # "send" -- proven while we waited for the lock, just go.
+            # "wait" -- another caller is already finding out; don't pay a
+            #           second full timeout to learn the same thing.
+            # "probe" -- nobody is; this call is the one that finds out.
+            role = "send" if self._verified else "wait" if self._probing else "probe"
+            if role == "probe":
+                self._probing = True
+                self._probe_done.clear()
+
+        if role == "probe":
+            try:
+                return self._send(texts, queue, timeout)
+            finally:
+                with self._state_lock:
+                    self._probing = False
+                self._probe_done.set()
+
+        if role == "wait":
+            self._probe_done.wait(_PROBE_WAIT)
+            remaining = self._cooldown_remaining()
+            if remaining > 0:
+                raise self._cooldown_error(queue, remaining)
+            if not self._verified:
+                raise RuntimeError(
+                    f"Skipping the embed dispatch: an earlier request to the "
+                    f"{queue!r} queue is still unanswered after {_PROBE_WAIT}s, so "
+                    f"no worker appears to be consuming it. {_NO_WORKER_HINT}"
+                )
+
+        return self._send(texts, queue, timeout)
+
     def embed_documents(self, documents: List[str]) -> List[List[float]]:
         """Embed a list of texts, preserving order."""
         if not documents:

application/vectorstore/faiss.py

@@ -325,7 +325,13 @@ class FaissStore(BaseVectorStore):
             f.write(dump_pickle_sidecar(self.documents, self.index_to_docstore_id))
 
     def _save_to_storage(self) -> bool:
-        """Persist the index through the configured storage backend."""
+        """Persist the index through the configured storage backend.
+
+        Each file is replaced atomically by the backend, so an interrupted save
+        leaves the previous one readable. The three are still written in
+        sequence: a crash between them pairs a new index with an older sidecar,
+        which stays consistent because the ids and row order are preserved.
+        """
         with tempfile.TemporaryDirectory() as temp_dir:
             self._write_index_files(temp_dir)
             storage_path = get_vectorstore(self.source_id)

application/worker.py

@@ -2523,6 +2523,14 @@ def reembed_wiki_page_worker(self, source_id, path, content_hash, user):
 
         with db_session() as conn:
             WikiPagesRepository(conn).set_embed_status(source_id, path, "embedded")
+            # These chunks were just embedded with the configured model, so the
+            # source now names it. Wiki sources created before this was recorded
+            # carry NULL, which the boot mismatch check reads as the legacy
+            # model and reports as stale; stamping here heals them on the next
+            # page edit.
+            SourcesRepository(conn).update(
+                source_id, user, {"model": settings.EMBEDDINGS_NAME}
+            )
     except Exception:
         with db_session() as conn:
             WikiPagesRepository(conn).set_embed_status(source_id, path, "failed")