test(embed): exercise loss guards through public paths

Replace redundant private-helper checks with public runtime, embedder, and writer coverage. Clarify the dense-only and None-conversion behavior in the shipped documentation.

Author: rhdong

docs/docs/extraction/troubleshoot.md

@@ -49,28 +49,33 @@ and embedding. It does not automatically raise for:
   an embedding failure to `None`, which can still be dropped silently as described
   below.
 
-  Row-level embedding failures are unaffected and still populate the error
-  column. A plain CUDA out-of-memory is not treated as an engine failure, because
-  on the HuggingFace embedding backend a smaller next batch can succeed. That is
-  about the embed stage, not the run: if the retry succeeds nothing changes. If
-  the batch is lost, whether the ingest then fails depends on the shape those
-  rows carry. The writer refuses `[]` but not `None`, as described below. If you
-  see `Refusing to build an incomplete index` with no engine-startup message in
-  the embed actor logs, look for an out-of-memory instead.
-
-  `LanceDB(on_bad_vectors=...)` does not suppress that error. A row that arrives
-  with an empty list or tuple embedding used to be counted as a wrong-length
-  vector and silently excluded; it now fails the run, and no policy value
-  restores the old behavior. Fix the embed stage.
+  Alternate handlers that convert row-level failures to `None` still populate
+  the error column. The default local runtime is stricter: if its embedder
+  returns no vector for any submitted input, the embedder raises and aborts the
+  ingest. A plain CUDA out-of-memory is not classified as an engine failure,
+  because on the HuggingFace embedding backend a smaller next batch can succeed.
+  That is about the embed stage, not the run: if the retry succeeds nothing
+  changes. If the batch is lost, whether the ingest then fails depends on the
+  shape those rows carry. The dense writer refuses `[]` but not `None`, as
+  described below. If you see `Refusing to build an incomplete index` with no
+  engine-startup message in the embed actor logs, look for an out-of-memory
+  instead.
+
+  `LanceDB(on_bad_vectors=...)` does not suppress that error. On a dense write, a
+  row that arrives with an empty list or tuple embedding used to be counted as a
+  wrong-length vector and silently excluded; it now fails the run, and no policy
+  value restores the old behavior. Fix the embed stage.
 
   This covers empty list and tuple values only. A failed embedding that arrives
-  as `None` keeps its pre-existing silent drop, counted as
-  `dropped_no_embedding`. Some embed
-  operators produce `None` for a lost batch as well as for a row they chose not
-  to embed. The two carry different `error` payloads upstream, but the writer
-  does not see that key, so it drops both alike. If recall is low and the run
-  reported success, check `dropped_no_embedding` in the ingest logs, not just
-  the guard.
+  as `None` keeps its pre-existing handling. For direct nested LanceDB records,
+  the writer silently drops the row and counts it as `dropped_no_embedding`.
+  Graph rows are converted before the writer: a mixed batch filters out `None`
+  rows without incrementing that writer counter, while a batch with no
+  uploadable rows raises `VdbUploadError`. Some embed operators produce `None`
+  for a lost batch as well as for a row they chose not to embed. The two carry
+  different `error` payloads upstream, but the writer does not see that key. If
+  recall is low and the run reported success, inspect the embed-stage row errors
+  and writer logs; do not rely on the empty-vector guard alone.
 
 - Caption or remote VLM stages. Missing credentials fail at actor setup;
   inference failures can abort the entire ingest.

nemo_retriever/src/nemo_retriever/common/vdb/README.md

@@ -93,9 +93,9 @@ When `vdb_op="lancedb"` (or `vdb=LanceDB(...)` is passed explicitly), `_construc
 1. **`create_index`** — connects with `lancedb.connect(self.uri)`, transforms ingestion batches into Arrow rows (`vector`, `text`, `metadata`, `source`), and **`db.create_table(...)`** with schema and `on_bad_vectors` policy.
 2. **`write_to_index`** — builds the **vector index** (e.g. IVF/HNSW) and optionally an **FTS/BM25** index over the ingested `text` column when `hybrid=True`.
 
-During step 1, a row that arrives with an **empty list or tuple** embedding raises `RuntimeError`, and no table rows are written. The embed failure path writes `[]` when it produces no vector for a row. Such a row used to be counted as a wrong-length vector and silently excluded, letting a run publish a short index and still report `success: true`.
+During a dense-vector step 1, a row that arrives with an **empty list or tuple** embedding raises `RuntimeError`, and no table rows are written. The embed failure path writes `[]` when it produces no vector for a row. Such a row used to be counted as a wrong-length vector and silently excluded, letting a run publish a short index and still report `success: true`.
 
-The new guard treats only an empty list or tuple as fatal. A row whose embedding is absent or `None` keeps its pre-existing silent drop. For this policy, `None` has two relevant meanings: a deliberate skip and a genuine embedding failure. For example, `operators/embed/text_embed.py` writes `error: None` alongside a blank-text skip, while its failure handler writes a populated `error` dict with stage, type, message, and traceback. Other embedding paths can produce the same two meanings. `common/vdb/records.py` already reads the payload one layer above this writer, so the discriminator is available. Threading it down is follow-up work, and making `None` fatal without it would fail ingests that work today. Until then, a batch lost through that path is still dropped silently. `on_bad_vectors` is unaffected: it governs malformed vectors, and an empty embedding is the absence of a vector rather than a short one. The returned `counts` dict gains an `empty_embedding` key.
+The new dense-writer guard treats only an empty list or tuple as fatal. A row whose embedding is absent or `None` keeps its pre-existing handling. For direct nested LanceDB records, the writer silently drops that row and counts it as `dropped_no_embedding`. Graph rows are converted first: a mixed batch can omit failed `None` rows while indexing healthy rows, but an all-rejected batch raises `VdbUploadError`. For this policy, `None` has two relevant meanings: a deliberate skip and a genuine embedding failure. For example, `operators/embed/text_embed.py` writes `error: None` alongside a blank-text skip, while its failure handler writes a populated `error` dict with stage, type, message, and traceback. Other embedding paths can produce the same two meanings. `common/vdb/records.py` reads that payload during conversion, so the discriminator is available there, but it is not forwarded to the writer. Making every `None` fatal would fail deliberate skips; source-aware handling is follow-up work. `on_bad_vectors` is unaffected: it governs malformed vectors, and an empty embedding is the absence of a vector rather than a short one. The returned `counts` dict gains an `empty_embedding` key.
 
 The other half of the guard is upstream: a local embedder that fails to embed some rows raises `LocalEmbedderRowsLostError` from `_finalize_vectors` instead of zero-padding them, which this writer could not otherwise detect.
 

nemo_retriever/tests/test_embed_engine_failure_propagation.py

@@ -8,14 +8,15 @@
 
 from typing import Any
 
+import httpx
 import pandas as pd
 import pytest
 
 from nemo_retriever.models.embed_errors import (
     LocalEmbedderReturnedNothingError,
     LocalEmbedderRowsLostError,
 )
-from nemo_retriever.models.inference import main_text_embed, runtime
+from nemo_retriever.models.inference import runtime
 from nemo_retriever.models.nim.error_reporter import drain_errors
 
 
@@ -62,11 +63,16 @@ def _batch(rows: int = 4) -> pd.DataFrame:
     )
 
 
-def _raise(exc: BaseException):
-    def _fail(*_args: Any, **_kwargs: Any) -> pd.DataFrame:
-        raise exc
+class _TextModel:
+    """Model-boundary double for the local runtime."""
 
-    return _fail
+    def __init__(self, result: Any) -> None:
+        self.result = result
+
+    def embed(self, _texts: Any, *, batch_size: int) -> Any:
+        if isinstance(self.result, BaseException):
+            raise self.result
+        return self.result
 
 
 def _wrapped(outer: BaseException, cause: BaseException) -> BaseException:
@@ -103,60 +109,56 @@ def _wrapped(outer: BaseException, cause: BaseException) -> BaseException:
         ),
     ],
 )
-def test_local_engine_failure_propagates(monkeypatch: pytest.MonkeyPatch, exc: BaseException) -> None:
-    monkeypatch.setattr(runtime, "_embed_group", _raise(exc))
-
+def test_local_engine_failure_propagates(exc: BaseException) -> None:
     with pytest.raises(type(exc)):
-        runtime.embed_text_main_text_embed(_batch(), model=object(), inference_batch_size=2)
+        runtime.embed_text_main_text_embed(_batch(), model=_TextModel(exc), inference_batch_size=2)
 
 
 def test_local_embedder_returning_nothing_is_raised_as_a_classified_failure() -> None:
     with pytest.raises(LocalEmbedderReturnedNothingError):
-        main_text_embed._callable_runner([["page one", "page two"]], embedder=lambda _texts: [], batch_size=2)
+        runtime.embed_text_main_text_embed(_batch(2), model=_TextModel([]), inference_batch_size=2)
 
-    assert runtime._is_engine_lifecycle_failure(LocalEmbedderReturnedNothingError("no vectors"))
 
+def test_partial_local_result_keeps_the_pre_existing_fallback() -> None:
+    out_df = runtime.embed_text_main_text_embed(_batch(2), model=_TextModel([[1.0, 2.0]]), inference_batch_size=2)
 
-def test_partial_local_result_is_a_plain_value_error() -> None:
-    with pytest.raises(ValueError) as excinfo:
-        main_text_embed._callable_runner([["page one", "page two"]], embedder=lambda _texts: [[1.0, 2.0]], batch_size=2)
-
-    assert not isinstance(excinfo.value, LocalEmbedderReturnedNothingError)
-    assert not runtime._is_engine_lifecycle_failure(excinfo.value)
+    assert list(out_df["text_embeddings_1b_v2_has_embedding"]) == [False, False]
+    assert [payload["embedding"] for payload in out_df["text_embeddings_1b_v2"]] == [[], []]
 
 
 @pytest.mark.parametrize(
     "model",
     [pytest.param(None, id="model-nulled-by-the-operator"), pytest.param(object(), id="model-also-set")],
 )
 def test_endpoint_mode_absorbs_engine_lifecycle_failure(monkeypatch: pytest.MonkeyPatch, model: object) -> None:
-    monkeypatch.setattr(runtime, "_embed_group", _raise(RuntimeError(ENGINE_INIT_FAILED)))
+    original_client = httpx.Client
+
+    def client_factory(*_args: Any, **_kwargs: Any) -> httpx.Client:
+        transport = httpx.MockTransport(lambda _request: httpx.Response(400, text=ENGINE_INIT_FAILED))
+        return original_client(transport=transport)
+
+    monkeypatch.setattr(httpx, "Client", client_factory)
 
     out_df = runtime.embed_text_main_text_embed(
         _batch(),
         model=model,
         embedding_endpoint="http://embed.example/v1",
-        inference_batch_size=2,
+        inference_batch_size=4,
     )
 
     assert list(out_df["text_embeddings_1b_v2_has_embedding"]) == [False] * 4
     assert [payload["embedding"] for payload in out_df["text_embeddings_1b_v2"]] == [[]] * 4
 
 
-def test_an_oom_alone_is_not_classified_as_an_engine_failure(monkeypatch: pytest.MonkeyPatch) -> None:
-    monkeypatch.setattr(runtime, "_embed_group", _raise(OutOfMemoryError(OOM_IN_GELU)))
-
-    out_df = runtime.embed_text_main_text_embed(_batch(), model=object(), inference_batch_size=2)
+@pytest.mark.parametrize(
+    "exc",
+    [
+        pytest.param(OutOfMemoryError(OOM_IN_GELU), id="bare-oom"),
+        pytest.param(EngineGenerateError("generate() failed"), id="recoverable-generate-error"),
+    ],
+)
+def test_recoverable_local_failure_keeps_the_pre_existing_fallback(exc: BaseException) -> None:
+    out_df = runtime.embed_text_main_text_embed(_batch(), model=_TextModel(exc), inference_batch_size=2)
 
     assert list(out_df["text_embeddings_1b_v2_has_embedding"]) == [False] * 4
     assert [payload["embedding"] for payload in out_df["text_embeddings_1b_v2"]] == [[]] * 4
-
-
-def test_engine_lifecycle_classifier_rejects_unrelated_failures() -> None:
-    assert not runtime._is_engine_lifecycle_failure(ValueError("could not decode image payload"))
-    assert not runtime._is_engine_lifecycle_failure(TimeoutError("read timed out"))
-    # Recoverable on the HuggingFace backend; see the module docstring.
-    assert not runtime._is_engine_lifecycle_failure(OutOfMemoryError(OOM_IN_GELU))
-    # vLLM documents this one as recoverable in its own source.
-    assert not runtime._is_engine_lifecycle_failure(EngineGenerateError("generate() failed"))
-    assert runtime._is_engine_lifecycle_failure(RuntimeError(ENGINE_INIT_FAILED))

nemo_retriever/tests/test_lancedb_collections.py

@@ -1130,19 +1130,6 @@ def delete_target():
     assert delete_errors == []
 
 
-def _collection_context() -> CollectionWriteContext:
-    return CollectionWriteContext(
-        scope="workspace-a",
-        collection_name="collection-a",
-        document_id="document-a",
-        document_version="v1",
-        content_sha256="sha-v1",
-        filename="source.pdf",
-        job_id="job-a",
-        operation="append",
-    )
-
-
 def _collection_record(embedding: Any, *, text: str = "first chunk") -> dict:
     return {
         "document_type": "text",
@@ -1156,56 +1143,25 @@ def _collection_record(embedding: Any, *, text: str = "first chunk") -> dict:
 
 
 @pytest.mark.parametrize("empty_embedding", [pytest.param([], id="list"), pytest.param((), id="tuple")])
-def test_an_empty_embedding_on_the_collection_path_fails_the_write(empty_embedding: Any) -> None:
+def test_an_empty_embedding_on_the_collection_path_fails_the_write(tmp_path, empty_embedding: Any) -> None:
+    backend = _backend_with_collection(tmp_path)
     records = [[_collection_record(empty_embedding), _collection_record([1.0, 0.0])]]
 
     with pytest.raises(RuntimeError) as excinfo:
-        _collection_rows(records, context=_collection_context())
+        backend.write_collection(records, context=_context())
 
     message = str(excinfo.value)
     assert "incomplete document" in message
     assert "empty_embedding=1" in message
+    with pytest.raises(VDBResourceNotFound):
+        backend.get_document(scope="workspace-a", collection_name="collection-a", document_id="document-a")
 
 
-@pytest.mark.parametrize(
-    "skipped_record",
-    [
-        pytest.param(None, id="non-dict-record"),
-        pytest.param({"metadata": None}, id="non-dict-metadata"),
-        pytest.param(_collection_record(None), id="malformed-embedding"),
-        pytest.param(_collection_record([1.0, 0.0], text="   "), id="text-free-non-image"),
-    ],
-)
-def test_collection_failure_counts_each_pre_existing_silent_skip(skipped_record: Any) -> None:
-    records = [
-        [
-            skipped_record,
-            _collection_record([]),
-            _collection_record([1.0, 0.0], text="good chunk"),
-        ]
-    ]
-
-    with pytest.raises(RuntimeError) as excinfo:
-        _collection_rows(records, context=_collection_context())
-
-    message = str(excinfo.value)
-    assert "empty_embedding=1" in message
-    assert "skipped_other=1" in message
-    assert "accepted=1" in message
-
-
-def test_a_numpy_collection_embedding_does_not_raise_on_truthiness() -> None:
+def test_collection_write_does_not_apply_sequence_truthiness_to_numpy(tmp_path) -> None:
     numpy = pytest.importorskip("numpy")
+    backend = _backend_with_collection(tmp_path)
+    records = [[_collection_record(numpy.array([1.0, 0.0])), _collection_record([1.0, 0.0])]]
 
-    records = [
-        [
-            _collection_record(numpy.array([1.0, 0.0])),
-            _collection_record(numpy.array([])),
-            _collection_record([1.0, 0.0], text="good chunk"),
-        ]
-    ]
-
-    rows = _collection_rows(records, context=_collection_context())
+    result = backend.write_collection(records, context=_context())
 
-    assert len(rows) == 1
-    assert rows[0]["text"] == "good chunk"
+    assert result.written == 1

nemo_retriever/tests/test_lancedb_write_policy.py

@@ -13,7 +13,7 @@
 
 lancedb = pytest.importorskip("lancedb")
 
-from nemo_retriever.common.vdb.lancedb import LanceDB, _create_lancedb_results
+from nemo_retriever.common.vdb.lancedb import LanceDB
 
 
 def _records(text: str = "hello", vector: list[float] | None = None) -> list[list[dict]]:
@@ -298,7 +298,10 @@ def test_dense_write_fails_on_image_only_row_with_an_empty_embedding(tmp_path: P
 
 
 def test_dense_write_drops_a_none_embedding_end_to_end(tmp_path: Path) -> None:
-    records = [[_record(None), _record([1.0, 2.0], text="good row")]]
+    records = _records(text="good row", vector=[1.0, 2.0])
+    missing = _records(text="lost row", vector=[1.0, 2.0])[0][0]
+    missing["metadata"]["embedding"] = None
+    records[0].insert(0, missing)
 
     table_rows = _write_rows(tmp_path, records)
 
@@ -336,6 +339,18 @@ def test_dense_write_keeps_on_bad_vectors_fill_reachable(tmp_path: Path) -> None
     assert len(table_rows[0]["vector"]) == 2
 
 
+def test_dense_write_does_not_apply_sequence_truthiness_to_numpy(tmp_path: Path) -> None:
+    numpy = pytest.importorskip("numpy")
+    records = _records(text="good row", vector=[1.0, 2.0])
+    unsupported = _records(text="unsupported row", vector=[1.0, 2.0])[0][0]
+    unsupported["metadata"]["embedding"] = numpy.array([1.0, 2.0])
+    records[0].insert(0, unsupported)
+
+    table_rows = _write_rows(tmp_path, records)
+
+    assert [row["text"] for row in table_rows] == ["good row"]
+
+
 def test_sparse_write_drops_image_only_row_without_text(tmp_path: Path) -> None:
     table_rows = _write_rows(tmp_path, _image_only_records([1.0, 0.0]), sparse=True)
 
@@ -346,55 +361,3 @@ def test_sparse_write_drops_whitespace_only_text(tmp_path: Path) -> None:
     table_rows = _write_rows(tmp_path, _records(text=" \n\t "), sparse=True)
 
     assert table_rows == []
-
-
-def _record(embedding: Any, *, text: str = "page text") -> dict:
-    return {
-        "document_type": "text",
-        "metadata": {
-            "embedding": embedding,
-            "content": text,
-            "content_metadata": {"page_number": 1, "id": "row-1"},
-            "source_metadata": {"source_name": "doc.pdf"},
-        },
-    }
-
-
-def test_create_lancedb_results_rejects_empty_embeddings_without_length_check() -> None:
-    with pytest.raises(RuntimeError, match="empty_embedding=1"):
-        _create_lancedb_results([[_record([])]], expected_dim=None)
-
-
-def test_wrong_length_rows_are_forwarded_when_the_wrapper_check_is_disabled() -> None:
-    records = [[_record([1.0]), _record([1.0, 2.0])]]
-
-    rows, counts = _create_lancedb_results(records, expected_dim=None)
-
-    assert len(rows) == 2
-    assert counts["dropped_bad_length"] == 0
-    assert counts["dropped_no_embedding"] == 0
-    assert counts["empty_embedding"] == 0
-
-
-@pytest.mark.parametrize(
-    "embedding",
-    [
-        pytest.param([0.0, 0.0], id="all-zero-but-present"),
-        pytest.param((1.0, 2.0), id="tuple"),
-    ],
-)
-def test_present_vectors_are_accepted_whatever_their_values(embedding) -> None:
-    rows, counts = _create_lancedb_results([[_record(embedding)]], expected_dim=2)
-
-    assert len(rows) == 1
-    assert counts["accepted"] == 1
-
-
-def test_an_empty_numpy_array_remains_a_nonfatal_bad_length_row() -> None:
-    numpy = pytest.importorskip("numpy")
-
-    rows, counts = _create_lancedb_results([[_record(numpy.array([]))]], expected_dim=2)
-
-    assert rows == []
-    assert counts["empty_embedding"] == 0
-    assert counts["dropped_bad_length"] == 1