From 128e93661c5d2229b7d8b60d850af3d6e2740d0a Mon Sep 17 00:00:00 2001 From: Haidong Rong Date: Tue, 18 Aug 2026 08:46:46 -0700 Subject: [PATCH 1/4] fix(embed): fail the run when embed rows arrive with no embedding When several embed workers share one GPU, vLLM can refuse to start an engine because free memory is below `gpu_memory_utilization x total`, or an engine can die after it has started. Both failures were caught by a bare `except` in the embed runtime and turned into `embedding: []`. Those empty rows reached LanceDB, where the length check counted them as wrong-length vectors and dropped them, so the index came out short and the run still reported `success: true` with `exit_code 0`. Be precise about the defect, because the obvious guess is wrong. Empty vectors were *not* written into the index. The measured unpatched run recorded `accepted=41181 dropped_no_embedding=0 dropped_bad_length=38053 dropped_no_text=0`. The length check was on - `dropped_bad_length` can only increment when `enforce_length` is true, and the sibling text+image leg logged `expected_dim=2048`, the constructor default - so the 38,053 rows whose embedding had failed were counted as wrong-length and silently excluded. The failure is silent exclusion, not silent insertion. Measured on a single L40S with the shipped `embed_workers: 3`, on the unpatched tree: two of three engines were refused for the first 42 minutes of a 100-minute run, 38,053 of 79,234 rows lost their embedding, the published index held 41,181 rows and was missing 48% of the corpus, the run exited 0, and the benchmark published nDCG@10 0.2598 against an expected 0.75. The admission gate scales with card size while the extraction stage holds an amount of memory that does not, so smaller cards lose engines while 80 GB and larger ones do not. The guarantee comes from the writer, and it is general: a row must not reach the index without an embedding, and a run that loses rows must fail rather than publish a short index. It is enforced at both writers. The other two changes are narrower and sit on top of it. - `common/vdb/lancedb.py` refuses to build an index when rows arrive with an empty embedding. `embedding: []` had no category of its own: it is not `None`, so the absent-embedding branch missed it, and it fell through to the length check, which counted it as a wrong-length vector and dropped it. The row was therefore excluded from the index rather than written to it, and nothing reported that the run had lost it. This layer carries the correctness guarantee: it sees rows, not causes, so it holds for any backend, any GPU and any failure mode, including the endpoint path. `create_index` calls the row builder twice when `vector_dim` is `None`: pass 1 with `expected_dim=None` to infer the dimension, pass 2 with the inferred value to filter. The new check ignores `expected_dim`, so it raises in pass 1 and pass 2 never runs. That is why the error reports `expected_dim=None` while the unpatched run's only summary came from pass 2 and reported `expected_dim=2048` - same configuration, different pass. The derivation of `expected_dim` is untouched by this change. - `common/vdb/lancedb_collections.py` carries the same guarantee on the collection-managed write path, which is a second, independent writer. `_collection_rows` skipped an empty embedding with `not vector`, folding it into the same silent skip as a malformed value: no counter, no log, no failure, so a collection document was published short while the ingest reported success. It is shipped - `POST /v1/ingest/job/{job_id}/document` reaches it through `IngestVdbOperator` and the `/internal/vectordb/write` route - so this is a live path, not defensive cover. It now counts `empty_embedding` and raises before any row reaches LanceDB, using the same counter name and the same explicit `isinstance(v, (list, tuple)) and len(v) == 0` spelling. The check is inserted ahead of the existing skip rather than replacing it, so every value other than `[]` keeps its current route; the other skips gain only a `skipped_other` count. - `models/inference/main_text_embed.py` `_multimodal_callable_runner` fails when the engine answers a multimodal batch short. It reassembles rows by walking an iterator of returned vectors, so a short answer silently became `None` for the shortfall - and `None` is the one shape both writers deliberately ignore, so those rows were dropped and the run still succeeded. The same shape on the text path was already fatal; only the multimodal path was silent, which is the path this defect was reported on. The check counts rows *submitted with an image*, not rows in the chunk. The embedders filter empty entries before inference and a row with no image is owed `None` by contract, so comparing against the chunk size would fail runs on image-free chunks. Both subsets of `text_image` are checked the same way: the paired subset against `mm_images`, the text-only fallback against `fb_texts`. - `models/local/*_embedder.py` `_finalize_vectors` stops discarding the loss it already knows about. It is the only place that holds both the batch that was sent and the vectors that came back, so `len(vectors) - len(valid)` is exact there and nowhere else. It used to compute that number and throw it away, zero-padding the failed rows: a padded row has the right width and a non-zero length, so every shape check downstream accepts it and `has_embedding` reports `True` for a row carrying nothing. It now raises `LocalEmbedderRowsLostError` with the count. The writer cannot detect a padded row, so this is the one loss the general guard cannot see. - `models/inference/runtime.py` no longer swallows engine-lifecycle failures. An engine refused at startup, or dead during inference, aborts the run. Failures are identified by class name and message, walking the `__cause__` chain, so the module still imports where vLLM and torch are absent. The endpoint path is unaffected: the re-raise requires a local model and no endpoint, and row-level failures stay non-fatal. This layer is a fast-fail optimisation, not the correctness guarantee. It turns a failure at the terminal write into one within minutes. Because it only buys latency, its fatal set is deliberately narrow: classifying an unmeasured exception would ship false failures for no correctness gain. The asymmetry is intentional - general writer, narrow classifier. The four layers are complementary, not redundant. Each catches a shape the others cannot see, so deleting one because another looks sufficient loses a real case. This was nearly done twice during review, on both counts wrongly. - An engine that raises - admission refused at startup, or dead mid-run - is caught by the `runtime.py` classifier. - An engine that returns nothing at all raises nowhere. When vLLM yields no outputs for a batch, `embed_with_vllm_llm` appends nothing and returns `[]`; `_finalize_vectors` then counts no loss, because there are no rows to count, and its `if not valid` early return hands back a 0-row tensor - the same answer as an empty input. `main_text_embed._callable_runner` is the only layer that sees that shape, which is why its `LocalEmbedderReturnedNothingError` guard is not speculative generality for a hypothetical custom callable. - An engine that loses only some rows is caught by `_finalize_vectors`, the one place holding both the batch that was sent and the vectors that came back. The writer cannot see those rows: they were zero-padded to the right width. - Anything all three miss reaches the `lancedb.py` writer guard, which sees rows rather than causes and so holds for any backend and any failure mode. Only the new `empty_embedding` counter is fatal, because `[]` is the only value with no legitimate producer: the sole places that write it for an embedding are the whole-batch failure path in `models/inference/runtime.py` and `models/inference/vllm.py`, where it means "this output carried no embedding". Every other counter here can be reached by a legitimate row and none of them is folded in: - `dropped_bad_length` is the category `on_bad_vectors` governs, so it stays a counted drop under the user's configured `drop`/`fill`/`null`/`error`. - `dropped_no_text` is a content filter, with a deliberate carve-out for canonical image rows. - `dropped_no_embedding` (absent key or `None`) keeps its pre-existing silent drop. `operators/embed/text_embed.py` writes `{"embedding": None}` on purpose for a blank-text row it chose not to embed, so making it fatal would fail ingests that work today. Folding any of them in would turn a tolerance into a hard failure on upgrade. The counter is named `empty_embedding`, without the `dropped_` prefix the other three carry, and the key is externally visible in the returned `counts` dict. `dropped_no_embedding`, `dropped_bad_length` and `dropped_no_text` all mean "row excluded, run continues". This one never does: whenever it is non-zero the run fails and no table is written at all, and the `continue` in the loop exists only so the total can be counted before raising. `dropped_empty_embedding=38053` in a log would tell an operator the opposite of what happened. The log line and the error message state the consequence rather than the disposal, for the same reason. One user-visible behaviour change follows, and it is stated rather than buried. A row whose embedding failed used to be counted as a wrong-length vector and silently excluded, letting a run publish a short index and still report success. It now fails the run. That silent exclusion was the defect, so the new failure is the fix working, not a regression - but a deployment that today absorbs failed embed batches will see the ingest stop instead. There is no per-policy story worth telling. `validate_vector_length` and `on_bad_vectors` have no production writers: only the constructor defaults (`True` and `"drop"`), a `policy.py` allowlist for service-mode clients, and tests. No environment variable, config file, harness preset or runfile key sets either one. Every source change is additive. Against the merge base the eight changed source files are +N -0: no pre-existing branch, counter, log line or docstring contract was rewritten, so no existing behaviour could change as a side effect. The new writer check is deliberately narrow - `isinstance(embedding, (list, tuple)) and len(embedding) == 0` - so `None` cannot reach it and numpy or tensor values keep their existing route. The obvious spelling, `if not embedding`, would have raised `ValueError` on a multi-element array and converted a counted drop into a crash. A bare `OutOfMemoryError` is deliberately NOT classified as fatal. It was never in the repo; this change chooses not to add it. The HuggingFace local backend is reachable and is the base default at `graph/retriever.py:116`, and it raises the same exception from an ordinary forward pass, where it is per-batch and recoverable. The two backends are not distinguishable at the point of classification. The cost is time-to-failure, not safety: a vLLM engine that OOMs mid-inference raises `EngineDeadError` from the next batch onward, which is fatal, and the writer guard refuses the first batch's rows in any case. Do not re-add it for faster failure without weighing that trade. Be precise about what that carve-out buys, because an earlier revision of this change overstated it. "Recoverable" is a claim about the batch, not the run. A batch absorbed after an OOM becomes `{"embedding": []}` for every row, which is exactly the shape the writer now treats as fatal, so an HF-backend run that loses a batch to OOM still fails - at the terminal write, after extraction and embedding are already paid for. The carve-out delivers two things: the runtime does not abort on a per-batch condition, and a run whose retry succeeded and lost nothing is not failed at all. It does not deliver end-to-end survival for a run that lost rows. The test that asserted otherwise was renamed from `test_torch_oom_alone_stays_recoverable_for_the_huggingface_backend` to `test_torch_oom_alone_does_not_abort_the_runtime_but_the_run_still_fails_at_the_writer`, and now also asserts the `[]` payload that makes the writer fail. The matching sentences in `docs/docs/extraction/troubleshoot.md` were corrected the same way. Known gap, documented rather than fixed: two embed entry points swallow `LocalEmbedderRowsLostError`. It subclasses `RuntimeError`, and both `operators/embed/text_embed.py::embed_text_1b_v2` and `common/modality/pipeline/embedding.py::embed_text_main_text_embed` catch `BaseException` around the embedder call and rewrite the batch as `{"embedding": None}` - the one shape the writer deliberately ignores, because `text_embed.py` also writes `None` legitimately for a blank-text row. Widening the fatal set to cover `None` would fail ingests that work today, which is the false-failure class this change exists to avoid, so it was not done. Neither route is shipped: all three embed actors reached from `graph/retriever.py`'s `_BatchEmbedActor` import `embed_text_main_text_embed` from `models/inference/runtime.py`, the patched one; `text_embed.py`'s actors have no importer in `src` at all, only in tests; and nothing outside `common/modality/pipeline/` imports that package. The residual risk is naming - the modality-pipeline duplicate exports an identical function name and is re-exported from its package `__init__`, so a future import-site flip would remove all three guard layers with no visible diff at the call site. Three tests pin this: one asserting the shipped actors resolve to the guarded function, and one per swallowing route. Closing the gap properly means adding a re-raise at both of `text_embed.py`'s nested `BaseException` handlers, which changes that operator's documented never-raise contract, or deleting the near-duplicate module; both are out of scope here. The `on_bad_vectors="fill"` guard test asserts the row survives at full schema width rather than the exact filled composition. How LanceDB spreads `fill_value` over a short vector is its own detail and differs by version - 0.34 replaces the whole vector, 0.37 pads and keeps the produced component - and `lancedb` is unpinned in `nemo_retriever/pyproject.toml`, so the resolved version varies by environment. What the guard owns is that `fill` still reaches the writer instead of being pre-empted, and that is what the test now asserts. Tests cover an engine refused at startup, an engine that dies after starting, an engine that answers with nothing, an embedder that loses part of a batch, and the writer rejecting rows with no embedding; each fails on the unpatched tree. There is also one test per case the fatal condition must NOT fire on: the endpoint path, row-level failures, wrong-length vectors under each `on_bad_vectors` value, text-free rows, canonical image rows, `None` and absent embeddings, all-zero vectors, numpy values, and a fully healthy batch. Each is labelled in-file as a guard rather than a regression test, and each says whether it can run on the unpatched tree - several cannot, because they assert on the new counter key, and their docstrings say so rather than claiming to pass both ways. The three tests added for the known gap above are pins on untouched pre-existing behaviour in two modules this change does not modify, so they pass on the unpatched tree by design; their job is to make a future import-site flip fail loudly rather than to prove a defect. Signed-off-by: Haidong Rong --- docs/docs/extraction/troubleshoot.md | 25 +- .../src/nemo_retriever/common/vdb/README.md | 6 + .../src/nemo_retriever/common/vdb/lancedb.py | 39 ++ .../common/vdb/lancedb_collections.py | 51 +- .../src/nemo_retriever/models/embed_errors.py | 82 +++ .../models/inference/main_text_embed.py | 39 ++ .../models/inference/runtime.py | 62 +++ .../llama_nemotron_embed_1b_v2_embedder.py | 4 + .../llama_nemotron_embed_vl_1b_v2_embedder.py | 6 + .../test_embed_engine_failure_propagation.py | 516 ++++++++++++++++++ .../test_lancedb_incomplete_index_guard.py | 510 +++++++++++++++++ .../tests/test_lancedb_write_policy.py | 67 ++- nemo_retriever/tests/test_vllm_embed.py | 176 +++++- 13 files changed, 1566 insertions(+), 17 deletions(-) create mode 100644 nemo_retriever/src/nemo_retriever/models/embed_errors.py create mode 100644 nemo_retriever/tests/test_embed_engine_failure_propagation.py create mode 100644 nemo_retriever/tests/test_lancedb_incomplete_index_guard.py diff --git a/docs/docs/extraction/troubleshoot.md b/docs/docs/extraction/troubleshoot.md index 72d88d195d..5e6343b6da 100644 --- a/docs/docs/extraction/troubleshoot.md +++ b/docs/docs/extraction/troubleshoot.md @@ -38,7 +38,28 @@ configured invoke URL: Page Elements, OCR, Table Structure, Nemotron Parse, and embedding. It does not automatically raise for: - Local-only pipelines (`pdfium` without remote URLs), even when rows contain - `metadata.error` or column-level error payloads. + `metadata.error` or column-level error payloads. One exception: an in-process + vLLM embedding engine that has stopped serving - refused at startup for lack of + free GPU memory, dead after a crash, or returning no vectors - aborts the + ingest under every error policy. Such an engine produces no embedding for every + batch it is handed, and those rows are then excluded at the writer, so + continuing would publish an index covering only part of the corpus while the + run reported success. The error message names the knobs to change. + + 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, but + if the batch is lost its rows reach the writer with no embedding and the ingest + fails there. 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 whose + embedding failed used to be counted as a wrong-length vector and silently + excluded; it now fails the run. Fix the embed stage - no policy value restores + the old behaviour. + - Caption or remote VLM stages. Missing credentials fail at actor setup; inference failures can abort the entire ingest. - Audio or video ASR over gRPC or HTTP. Failures can drop individual rows and @@ -81,7 +102,7 @@ troubleshooting path. A single document can pass through several stages. | `ExtractParams(method="ocr")` | Page rendering, Page Elements, and the local or remote OCR backend | Missing local model dependencies, invalid image payload, authentication/transport status, or OCR row-level failure | | `ExtractParams(method="nemotron_parse")` | PDF rendering and local Nemotron Parse model or configured Nemotron Parse NIM | Missing `open_clip`, missing local model configuration, unsupported image input, or Nemotron Parse row-level/HTTP failure | | `.caption(...)` | Local caption model or remote VLM endpoint | `ValueError` at setup when credentials or endpoint/protocol are invalid; remote inference failures can abort the whole ingest rather than populate a row error column | -| `.embed(...)` | Local embedding model or configured embedding NIM | Model/dependency error, input-size or schema rejection, authentication/transport status, or embedding row-level failure; `GraphIngestionError` when a remote embed URL is configured | +| `.embed(...)` | Local embedding model or configured embedding NIM | Model/dependency error, input-size or schema rejection, authentication/transport status, or embedding row-level failure; `GraphIngestionError` when a remote embed URL is configured. An in-process vLLM embedding engine that has stopped serving always aborts the ingest, whatever the error policy | | Audio or video extraction | `ffmpeg`/`ffprobe`, media decoding, frame/chunk creation, and local or remote ASR | Missing executable, malformed media, codec failure, gRPC status, or credential error; ASR failures may omit rows and log warnings instead of raising, so verify logs when output is unexpectedly empty | `pdfium` itself is primarily a local parser, so a Page Elements, Table diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/README.md b/nemo_retriever/src/nemo_retriever/common/vdb/README.md index fbd543abf0..5372e842af 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/README.md +++ b/nemo_retriever/src/nemo_retriever/common/vdb/README.md @@ -93,6 +93,12 @@ 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** embedding, `[]`, raises `RuntimeError` and no table is written. `[]` is written only on embed failure paths, so it means no vector was produced for that 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`. + +Only `[]` is fatal. A row whose embedding is absent or `None` keeps its pre-existing silent drop, because `operators/embed/text_embed.py` writes `None` on purpose for a blank-text row. `on_bad_vectors` is unaffected: it governs malformed vectors, and an empty list 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. + Common constructor arguments include: | Parameter | Purpose | diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py index 81b78d6d38..d10ee2a571 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py @@ -368,10 +368,21 @@ def _create_lancedb_results( and ``counts`` is a dict containing ``accepted``, ``dropped_no_embedding``, ``dropped_bad_length``, and ``dropped_no_text`` keys. + Also ``empty_embedding``, added by the incomplete-index guard. + + An empty embedding, ``[]``, means the embed stage produced no vector for that + row. Such rows are counted as ``empty_embedding`` and make this function + raise :class:`RuntimeError` after the loop, so no table is written. + + :meth:`LanceDB.create_index` calls this function twice when ``vector_dim`` is + ``None``: pass 1 infers the dimension and its rows are discarded, pass 2 + filters. The check ignores ``expected_dim``, so it raises in pass 1 and pass 2 + never runs - which is why the error reports ``expected_dim=None``. """ lancedb_rows: list = [] accepted = 0 dropped_no_embedding = 0 + empty_embedding = 0 dropped_bad_length = 0 dropped_no_text = 0 @@ -388,6 +399,13 @@ def _create_lancedb_results( dropped_no_embedding += 1 continue + # ``[]`` is not ``None``, so the branch above misses it. Explicit + # isinstance/len, never ``not embedding``, which raises on an array. + if isinstance(embedding, (list, tuple)) and len(embedding) == 0: + empty_embedding += 1 + logger.debug("Dropping row with an empty embedding (doc_type=%s)", doc_type) + continue + if enforce_length and (not isinstance(embedding, (list, tuple)) or len(embedding) != expected_dim_int): dropped_bad_length += 1 got_len: Any = len(embedding) if hasattr(embedding, "__len__") else "n/a" @@ -434,6 +452,7 @@ def _create_lancedb_results( counts: dict[str, int] = { "accepted": accepted, "dropped_no_embedding": dropped_no_embedding, + "empty_embedding": empty_embedding, "dropped_bad_length": dropped_bad_length, "dropped_no_text": dropped_no_text, } @@ -450,6 +469,26 @@ def _create_lancedb_results( expected_dim_repr, ) + if empty_embedding: + logger.warning( + "_create_lancedb_results: empty_embedding=%d", + empty_embedding, + ) + total = accepted + dropped_no_embedding + empty_embedding + dropped_bad_length + dropped_no_text + raise RuntimeError( + "Refusing to build an incomplete index: " + f"{empty_embedding} of {total} rows had no embedding. No table is written and " + "this run fails; the alternative is an index that is silently short by those " + "rows while the run reports success. " + f"Counters: empty_embedding={empty_embedding}, no_embedding={dropped_no_embedding}, " + f"bad_length={dropped_bad_length}, no_text={dropped_no_text}, " + f"expected_dim={expected_dim_int if enforce_length else 'None'}. " + "Only empty_embedding caused this failure - the other three are rows filtered " + "under the configured policy. This normally means the embed stage failed for " + "whole batches: check the embed actor logs for engine initialization or " + "out-of-memory errors." + ) + return lancedb_rows, counts diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py index ddfa0ec533..e020b9a37d 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py @@ -168,22 +168,52 @@ def _collection_rows( *, context: CollectionWriteContext, ) -> list[dict[str, Any]]: - """Convert canonical NRL record batches into collection-managed LanceDB rows.""" + """Convert canonical NRL record batches into collection-managed LanceDB rows. + + An empty embedding, ``[]``, means the embed stage produced no vector for that + row. Such rows are counted as ``empty_embedding`` and make this function raise + :class:`RuntimeError` after the loop, so no rows are written for the document. + This is the same guarantee :func:`nemo_retriever.common.vdb.lancedb._create_lancedb_results` + carries on the pipeline write path: a row must not reach the index without an + embedding, and a write that loses rows must fail rather than publish a short + document while reporting success. + + Every other skip keeps its pre-existing silent behaviour and is counted only + as ``skipped_other``: a malformed batch, record, metadata or embedding value, + and a text-free non-image row, are all still dropped without failing the + write. Only ``[]`` is fatal, because it is the only value with no legitimate + producer - the embed stage writes it to mean "this output carried no + embedding". + """ rows: list[dict[str, Any]] = [] created_at = _now() row_index = 0 + empty_embedding = 0 + skipped_other = 0 for batch in records or []: if not isinstance(batch, list): continue for record in batch: if not isinstance(record, dict): + skipped_other += 1 continue metadata = record.get("metadata") if not isinstance(metadata, dict): + skipped_other += 1 continue vector = metadata.get("embedding") + # ``[]`` is folded into the skip below by ``not vector``, which drops + # it silently. Check it first, with an explicit isinstance/len and + # never ``not vector``: the latter raises ValueError on a + # multi-element array. The skip below is left exactly as it was, so + # every value other than ``[]`` keeps its current route. + if isinstance(vector, (list, tuple)) and len(vector) == 0: + empty_embedding += 1 + logger.debug("Row has an empty embedding (document_id=%s)", context.document_id) + continue if not isinstance(vector, (list, tuple)) or not vector: + skipped_other += 1 continue content_metadata = metadata.get("content_metadata") if not isinstance(content_metadata, dict): @@ -200,6 +230,7 @@ def _collection_rows( content_metadata["type"] = content_type content_metadata["_content_type"] = content_type if not text.strip() and content_type != "image": + skipped_other += 1 continue source_id = str( @@ -238,6 +269,24 @@ def _collection_rows( } ) row_index += 1 + + if empty_embedding: + total = len(rows) + empty_embedding + skipped_other + logger.warning("_collection_rows: empty_embedding=%d", empty_embedding) + raise RuntimeError( + "Refusing to write an incomplete document: " + f"{empty_embedding} of {total} rows had no embedding. No rows are written and " + "this write fails; the alternative is a document that is silently short by those " + "rows while the ingest reports success. " + f"Counters: empty_embedding={empty_embedding}, skipped_other={skipped_other}, " + f"accepted={len(rows)}, document_id={context.document_id}, " + f"document_version={context.document_version}. " + "Only empty_embedding caused this failure - skipped_other counts rows filtered " + "under the pre-existing rules. This normally means the embed stage failed for " + "whole batches: check the embed actor logs for engine initialization or " + "out-of-memory errors." + ) + return rows diff --git a/nemo_retriever/src/nemo_retriever/models/embed_errors.py b/nemo_retriever/src/nemo_retriever/models/embed_errors.py new file mode 100644 index 0000000000..c32ce224ae --- /dev/null +++ b/nemo_retriever/src/nemo_retriever/models/embed_errors.py @@ -0,0 +1,82 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Embed failures that mean rows were lost. + +Kept in their own module, with no third-party imports, so the local embedders +and the embed runtime can both import them on the endpoint-only path. +""" + +from __future__ import annotations + +import logging +from typing import Sequence + +from nemo_retriever.models.nim.error_reporter import report_error + +logger = logging.getLogger(__name__) + + +class LocalEmbedderReturnedNothingError(RuntimeError): + """An in-process embedder returned no vectors for a non-empty batch.""" + + +class LocalEmbedderRowsLostError(RuntimeError): + """An in-process embedder produced no vector for some rows of a batch. + + Raised by :func:`report_lost_rows`. ``models/inference/runtime.py`` + classifies both errors in this module as fatal by name. + + Args: + lost: Number of rows with no vector. + total: Size of the batch. + embedder: Class name of the embedder that lost them. + """ + + def __init__(self, *, lost: int, total: int, embedder: str) -> None: + self.lost = int(lost) + self.total = int(total) + self.embedder = str(embedder) + super().__init__( + f"{embedder} returned no vector for {lost} of {total} row(s) in this batch. " + "Those rows would be zero-padded to the right width, so nothing downstream could " + "tell them apart from real embeddings by shape alone, and indexing them would " + "publish rows that match nothing. This normally means the in-process engine " + "failed for the batch - check the embed actor logs for engine initialization or " + "out-of-memory errors." + ) + + +def _has_no_vector(vector: object) -> bool: + """Return whether ``vector`` is empty or absent. + + Explicit rather than truthy: ``not vector`` raises on a multi-element numpy + array. + """ + if vector is None: + return True + if hasattr(vector, "__len__"): + return len(vector) == 0 + return False + + +def report_lost_rows(vectors: Sequence[Sequence[float]], *, embedder: str) -> int: + """Raise :class:`LocalEmbedderRowsLostError` if any row came back empty. + + Returns ``0`` when nothing was lost; never returns a non-zero count. + + Called from the local embedders' ``_finalize_vectors``, the only place that + holds both the batch that was sent and the vectors that came back. It used to + zero-pad the missing rows, which made them indistinguishable downstream: a + padded row has the right width, so ``has_embedding`` reported ``True`` for a + row carrying nothing. + """ + lost = sum(1 for vector in vectors if _has_no_vector(vector)) + if not lost: + return 0 + + exc = LocalEmbedderRowsLostError(lost=lost, total=len(vectors), embedder=embedder) + logger.error("%s", exc) + report_error("embed", exc) + raise exc diff --git a/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py b/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py index 3d14bbc42e..61c49ad9e2 100644 --- a/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py +++ b/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py @@ -52,6 +52,7 @@ def local_embedder(texts): ) from nemo_retriever.common.params.models import IMAGE_MODALITIES from nemo_retriever.models import _DEFAULT_EMBED_MODEL +from nemo_retriever.models.embed_errors import LocalEmbedderReturnedNothingError, LocalEmbedderRowsLostError logger = logging.getLogger(__name__) @@ -213,6 +214,18 @@ def _multimodal_callable_runner( tolist = getattr(vecs, "tolist", None) vecs_list = tolist() if callable(tolist) else list(vecs) + # Only rows submitted with an image are owed a vector: the + # embedders drop empty entries before inference, and a row with no + # image gets ``None`` by contract. Comparing against ``size`` would + # false-fire on an image-free chunk. + submitted = sum(1 for b64 in images_b64 if b64) + if len(vecs_list) != submitted: + raise LocalEmbedderRowsLostError( + lost=max(submitted - len(vecs_list), 0), + total=submitted, + embedder=type(embedder).__name__, + ) + if len(vecs_list) == size: flat_embeddings.extend(vecs_list) else: @@ -233,6 +246,16 @@ def _multimodal_callable_runner( vecs = embedder.embed_text_image(mm_texts, mm_images, batch_size=bs) tolist = getattr(vecs, "tolist", None) mm_vecs_list = tolist() if callable(tolist) else list(vecs) + # ``mm_images`` is already image-only, so one vector per entry + # is the contract. A short answer means the engine lost rows; + # padding them with ``None`` here would hide that, because the + # writers ignore ``None``. + if len(mm_vecs_list) != len(mm_images): + raise LocalEmbedderRowsLostError( + lost=max(len(mm_images) - len(mm_vecs_list), 0), + total=len(mm_images), + embedder=type(embedder).__name__, + ) # text-only fallback subset fb_texts = [t for t, h in zip(texts, has_image) if not h and t.strip()] @@ -241,6 +264,12 @@ def _multimodal_callable_runner( vecs = embedder.embed(fb_texts, batch_size=bs) tolist = getattr(vecs, "tolist", None) fb_vecs_list = tolist() if callable(tolist) else list(vecs) + if len(fb_vecs_list) != len(fb_texts): + raise LocalEmbedderRowsLostError( + lost=max(len(fb_texts) - len(fb_vecs_list), 0), + total=len(fb_texts), + embedder=type(embedder).__name__, + ) # reassemble in original order mm_iter = iter(mm_vecs_list) @@ -567,6 +596,16 @@ def _callable_runner( chunk = prompt_batch[i : i + max(1, int(batch_size))] vecs = embedder(chunk) vecs_list = list(vecs) + if not vecs_list: + # One row per input is the contract, so no rows at all means + # the engine is not serving. Classified as fatal by + # ``runtime.embed_text_main_text_embed``. The shipped local + # embedders fail earlier, in ``_finalize_vectors``; this is the + # backstop for an arbitrary callable. + raise LocalEmbedderReturnedNothingError( + f"Local embedder returned no embeddings for a batch of {len(chunk)} input(s). " + "The in-process engine produced nothing, so it is not serving." + ) if len(vecs_list) != len(chunk): raise ValueError( "Local embedder returned a mismatched number of embeddings " diff --git a/nemo_retriever/src/nemo_retriever/models/inference/runtime.py b/nemo_retriever/src/nemo_retriever/models/inference/runtime.py index 2edd8b9cd4..f4de9f4638 100644 --- a/nemo_retriever/src/nemo_retriever/models/inference/runtime.py +++ b/nemo_retriever/src/nemo_retriever/models/inference/runtime.py @@ -19,6 +19,53 @@ from nemo_retriever.models.inference.main_text_embed import TextEmbeddingConfig, create_text_embeddings_for_df +# Exception classes and message fragments that mean "the in-process embed engine +# is not serving", as opposed to "this batch failed". Matched by name and text +# rather than by import so this module keeps working without vLLM or torch, +# which is the case on the endpoint-only path. +# +# ``OutOfMemoryError`` is deliberately absent. The HuggingFace backend raises it +# from an ordinary forward pass, where a smaller next batch can succeed, and the +# two backends are not distinguishable here - this function receives only the +# embedder object. Classifying it would abort on a per-batch condition. The cost +# is time-to-failure only: a vLLM engine that OOMs raises ``EngineDeadError`` +# from the next batch onward. +_ENGINE_LIFECYCLE_EXC_NAMES: frozenset[str] = frozenset( + { + "EngineDeadError", + "LocalEmbedderReturnedNothingError", + "LocalEmbedderRowsLostError", + } +) +_ENGINE_LIFECYCLE_MESSAGES: tuple[str, ...] = ( + "Engine core initialization failed", + "less than desired GPU memory utilization", +) +_ENGINE_LIFECYCLE_CAUSE_DEPTH: int = 5 + + +def _is_engine_lifecycle_failure(exc: BaseException) -> bool: + """Return whether ``exc`` means the local embed engine stopped serving. + + Walks the ``__cause__``/``__context__`` chain because vLLM re-raises the + original failure wrapped in its own error, so the OOM that killed the engine + is often not the outermost exception. + """ + seen: set[int] = set() + current: BaseException | None = exc + for _depth in range(_ENGINE_LIFECYCLE_CAUSE_DEPTH): + if current is None or id(current) in seen: + return False + seen.add(id(current)) + if any(klass.__name__ in _ENGINE_LIFECYCLE_EXC_NAMES for klass in type(current).__mro__): + return True + message = str(current) + if any(fragment in message for fragment in _ENGINE_LIFECYCLE_MESSAGES): + return True + current = current.__cause__ or current.__context__ + return False + + def _embed_group( group_df: pd.DataFrame, *, @@ -190,6 +237,21 @@ def embed_text_main_text_embed( logger.debug("torch.cuda.empty_cache() failed during error cleanup: %s", _cache_exc) logger.error("Embedding failed: %s: %s", type(exc).__name__, exc, exc_info=True) report_error("embed", exc) + if endpoint is None and model is not None and _is_engine_lifecycle_failure(exc): + logger.error( + "Local embed engine failure is fatal: aborting instead of returning %d empty " + "embedding(s) for this batch (embedder=%s, failure=%s: %s). A failed in-process " + "engine keeps accepting batches and returning empty rows for as long as the stage " + "runs, so continuing would write a partial index. Most often the engine was " + "refused admission because several embed workers share one GPU: lower " + "`embed_workers`, lower `gpu_memory_utilization`, give the run a larger card, or " + "point `embedding_endpoint` at a separate embedding service.", + len(batch_df), + type(model).__name__, + type(exc).__name__, + exc, + ) + raise out_df = batch_df.copy() out_df[output_column] = [{"embedding": [], "error": str(exc)}] * len(out_df) out_df[embedding_dim_column] = 0 diff --git a/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_1b_v2_embedder.py b/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_1b_v2_embedder.py index 4ef170705c..b896e36111 100644 --- a/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_1b_v2_embedder.py +++ b/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_1b_v2_embedder.py @@ -11,6 +11,7 @@ import torch from nemo_retriever.models.hf_cache import configure_global_hf_cache_base +from nemo_retriever.models.embed_errors import report_lost_rows from nemo_retriever.models.embed_model_spec import resolve_embed_model_revision @@ -84,6 +85,9 @@ def is_remote(self) -> bool: return False def _finalize_vectors(self, vectors: List[List[float]]) -> torch.Tensor: + # See the mirror of this function in + # ``models/local/llama_nemotron_embed_vl_1b_v2_embedder.py``. + report_lost_rows(vectors, embedder=type(self).__name__) valid = [v for v in vectors if v] if not valid: return torch.empty((0, 0), dtype=torch.float32) diff --git a/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_vl_1b_v2_embedder.py b/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_vl_1b_v2_embedder.py index c9c682210d..6e1be0a94c 100644 --- a/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_vl_1b_v2_embedder.py +++ b/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_vl_1b_v2_embedder.py @@ -11,6 +11,7 @@ import torch from nemo_retriever.models.hf_cache import configure_global_hf_cache_base +from nemo_retriever.models.embed_errors import report_lost_rows from nemo_retriever.models.embed_model_spec import resolve_embed_model_revision from nemo_retriever.common.nvtx import gpu_inference_range @@ -250,6 +251,11 @@ def is_remote(self) -> bool: def _finalize_vectors(self, vectors: Sequence[Sequence[float]]) -> torch.Tensor: """Zero-pad rows vLLM failed to embed, then optionally normalize.""" + # Fail here when vLLM did not embed every row. This is the only place + # that knows the count; it used to zero-pad the missing rows, which made + # them indistinguishable from real embeddings downstream. Mirrored in + # ``models/local/llama_nemotron_embed_1b_v2_embedder.py``. + report_lost_rows(vectors, embedder=type(self).__name__) valid = [v for v in vectors if v] if not valid: return torch.empty((0, self.output_dimension), dtype=torch.float32) diff --git a/nemo_retriever/tests/test_embed_engine_failure_propagation.py b/nemo_retriever/tests/test_embed_engine_failure_propagation.py new file mode 100644 index 0000000000..09c431fac6 --- /dev/null +++ b/nemo_retriever/tests/test_embed_engine_failure_propagation.py @@ -0,0 +1,516 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""A run that loses embedding rows must fail, not publish a short index. + +Correctness comes from the LanceDB writer guard, which is general: any row, any +cause, any backend. It refuses to build an index when rows arrive with no +embedding. See ``test_lancedb_incomplete_index_guard.py``. + +This module tests the layer above it, which is narrow on purpose. +``embed_text_main_text_embed`` catches every exception from ``_embed_group`` and +returns ``{"embedding": [], "error": ...}`` for the whole batch. That per-batch +resilience is right for the endpoint path, where one HTTP call can fail alone. +It is wrong for an in-process engine that has stopped serving: every later batch +fails too, and a failed engine returns instantly, so it drains the queue far +faster than a healthy one. Aborting there is a fast-fail optimisation on top of +the writer guard - it turns a failure at the terminal write into one within +minutes. It is not what makes the result correct. + +Because it only buys latency, the fatal set stays limited to signals that mean +"the engine is not serving" and that a recoverable backend cannot produce. +Classifying an unmeasured exception would ship false failures. That is why a +bare ``OutOfMemoryError`` is excluded: the HuggingFace backend raises it from an +ordinary forward pass where a smaller next batch can succeed, and the two +backends are not distinguishable at the point of classification. A run that +loses rows to an OOM still fails - at the writer. + +``LocalEmbedderRowsLostError`` is the signal for *partial* loss and comes from +where the loss happens. The local embedders' ``_finalize_vectors`` is the only +function that knows how many rows failed to embed; it used to zero-pad them and +discard the count, which is why ``has_embedding`` could report ``True`` for a +row carrying nothing. The writer cannot detect a padded row, so neither layer +subsumes the other. + +The tests use stand-in exception classes rather than importing vLLM or torch, +matching the module under test, which identifies failures by class name and +message so it keeps working on the endpoint-only path. +""" + +from __future__ import annotations + +from typing import Any + +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 + + +class EngineDeadError(Exception): + """Stand-in for ``vllm.v1.engine.exceptions.EngineDeadError``.""" + + +class OutOfMemoryError(Exception): + """Stand-in for ``torch.OutOfMemoryError``.""" + + +class EngineGenerateError(Exception): + """Stand-in for ``vllm.v1.engine.exceptions.EngineGenerateError``.""" + + +ADMISSION_REFUSAL = ( + "Free memory on device cuda:0 (14.73/44.39 GiB) on startup is less than desired " + "GPU memory utilization (0.45, 19.98 GiB). Decrease GPU memory utilization or " + "reduce GPU memory used by other processes." +) +ENGINE_INIT_FAILED = "Engine core initialization failed. See root cause above. Failed core proc(s): {}" +OOM_IN_GELU = ( + "CUDA out of memory. Tried to allocate 380.00 MiB. GPU 0 has a total capacity of " + "44.39 GiB of which 274.69 MiB is free." +) + + +def _batch(rows: int = 4) -> pd.DataFrame: + return pd.DataFrame( + { + "text": [f"page {index}" for index in range(rows)], + "metadata": [ + {"content": f"page {index}", "content_metadata": {"page_number": index}, "source_metadata": {}} + for index in range(rows) + ], + } + ) + + +def _raise(exc: BaseException): + def _fail(*_args: Any, **_kwargs: Any) -> pd.DataFrame: + raise exc + + return _fail + + +def _wrapped(outer: BaseException, cause: BaseException) -> BaseException: + """Build ``outer`` raised from ``cause``, as vLLM re-raises engine errors.""" + try: + try: + raise cause + except BaseException as inner: # noqa: BLE001 - constructing a chain on purpose + raise outer from inner + except BaseException as chained: # noqa: BLE001 + return chained + + +@pytest.mark.parametrize( + ("exc", "label"), + [ + pytest.param(RuntimeError(ENGINE_INIT_FAILED), "engine-core-init-failed", id="refused-at-startup"), + pytest.param(ValueError(ADMISSION_REFUSAL), "free-memory-gate", id="admission-gate-valueerror"), + pytest.param(EngineDeadError("EngineCore encountered an issue"), "engine-dead", id="dead-after-admission"), + pytest.param( + _wrapped(RuntimeError("Worker proc died unexpectedly"), EngineDeadError("EngineCore is dead")), + "wrapped-engine-dead", + id="engine-dead-wrapped-in-runtimeerror", + ), + pytest.param( + LocalEmbedderReturnedNothingError( + "Local embedder returned no embeddings for a batch of 2 input(s). " + "The in-process engine produced nothing, so it is not serving." + ), + "returned-nothing", + id="engine-answered-with-nothing", + ), + pytest.param( + LocalEmbedderRowsLostError(lost=7, total=64, embedder="LlamaNemotronEmbedVL1BV2VLLMEmbedder"), + "rows-lost", + id="engine-lost-part-of-the-batch", + ), + ], +) +def test_local_engine_failure_propagates(monkeypatch: pytest.MonkeyPatch, exc: BaseException, label: str) -> None: + """Known-bad: returns a full batch of ``embedding: []`` and the run continues.""" + monkeypatch.setattr(runtime, "_embed_group", _raise(exc)) + + with pytest.raises(type(exc)): + runtime.embed_text_main_text_embed(_batch(), model=object(), inference_batch_size=2) + + +def test_local_embedder_returning_nothing_is_raised_as_a_classified_failure() -> None: + """An all-empty local batch must not escape as a plain ``ValueError``. + + An embedder that answers a non-empty batch with nothing is not serving. + ``_callable_runner`` used to report that as a bare count mismatch, which + ``_is_engine_lifecycle_failure`` does not match, so the run continued and + emptied the index anyway. + + This is not a backstop for a hypothetical callable. It is the only layer + that sees a shipped failure: when vLLM yields no outputs for a batch, + ``embed_with_vllm_llm`` returns ``[]``, and ``_finalize_vectors`` counts no + loss because there are no rows to count - see + ``test_vllm_embed.py::test_finalize_vectors_cannot_see_a_zero_output_batch``. + That path raises nothing anywhere else. + + Known-bad: raises ``ValueError`` with "mismatched number of embeddings", + which the classifier rejects. + """ + with pytest.raises(LocalEmbedderReturnedNothingError): + main_text_embed._callable_runner([["page one", "page two"]], embedder=lambda _texts: [], batch_size=2) + + assert runtime._is_engine_lifecycle_failure(LocalEmbedderReturnedNothingError("no vectors")) + + +def test_partial_local_result_is_still_a_plain_value_error() -> None: + """Guard: a count mismatch that is not "nothing at all" stays unclassified. + + The embedder did produce vectors, so this is a data-shape problem rather + than a dead engine, and it must keep its per-batch handling. Passes before + and after. + """ + 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) + + +@pytest.mark.parametrize( + "model", + [pytest.param(None, id="model-nulled-by-the-operator"), pytest.param(object(), id="model-also-set")], +) +def test_endpoint_mode_still_absorbs_the_same_failure(monkeypatch: pytest.MonkeyPatch, model: object) -> None: + """Guard: endpoint mode has no in-process engine, so nothing changes there. + + Passes before and after. A service-mode user must keep per-batch resilience. + + Both parameters matter. In production the model is always ``None`` when an + endpoint is configured - ``operators/embed/cpu_operator.py:36`` and + ``operators/embed/gpu_operator.py:37`` null it - so the ``model is None`` + case alone would still pass if the ``endpoint is None`` term were dropped + from the re-raise condition. The second parameter pins that term directly, + as defence against a future change to those two actors. + """ + monkeypatch.setattr(runtime, "_embed_group", _raise(RuntimeError(ENGINE_INIT_FAILED))) + + out_df = runtime.embed_text_main_text_embed( + _batch(), + model=model, + embedding_endpoint="http://embed.example/v1", + 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_an_oom_alone_is_not_classified_as_an_engine_failure(monkeypatch: pytest.MonkeyPatch) -> None: + """A bare OOM is absorbed here; a run that loses rows still fails at the writer. + + The HuggingFace backend raises ``OutOfMemoryError`` from an ordinary forward + pass, where a smaller next batch can succeed. The backend is not visible at + the point of classification, so treating the exception as fatal would abort + on a per-batch condition. + + Absorbing it is not the same as the run surviving. The batch becomes + ``{"embedding": []}`` for every row, which the writer guard treats as fatal, + so a run that really lost those rows fails at the write. The second + assertion pins that shape, so this test cannot be read as a claim that the + run recovers. + """ + 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) + + # Claim 1: the runtime absorbed it rather than re-raising. + assert list(out_df["text_embeddings_1b_v2_has_embedding"]) == [False] * 4 + + # Claim 2: what it absorbed into is the writer-fatal shape, so the run ends + # at the write, not here. + assert [payload["embedding"] for payload in out_df["text_embeddings_1b_v2"]] == [[]] * 4 + + +def test_engine_death_after_an_oom_is_fatal_from_the_next_batch(monkeypatch: pytest.MonkeyPatch) -> None: + """Excluding the OOM costs one batch, not the run. + + An engine that dies mid-inference raises ``EngineDeadError`` from the next + batch onward, and that is fatal. This pins the cost of the OOM carve-out as + bounded rather than open-ended. + """ + calls: list[int] = [] + + def _fail_then_die(*_args: Any, **_kwargs: Any) -> pd.DataFrame: + calls.append(1) + if len(calls) == 1: + raise OutOfMemoryError(OOM_IN_GELU) + raise EngineDeadError("EngineCore encountered an issue") + + monkeypatch.setattr(runtime, "_embed_group", _fail_then_die) + + absorbed = runtime.embed_text_main_text_embed(_batch(), model=object(), inference_batch_size=2) + assert list(absorbed["text_embeddings_1b_v2_has_embedding"]) == [False] * 4 + + with pytest.raises(EngineDeadError): + runtime.embed_text_main_text_embed(_batch(), model=object(), inference_batch_size=2) + + +def test_local_batch_level_failure_is_still_absorbed(monkeypatch: pytest.MonkeyPatch) -> None: + """Guard: a failure that is not an engine-lifecycle failure stays non-fatal. + + Passes before and after. The change is deliberately narrow - only failures + that mean the engine stopped serving are fatal, because only those were + measured to affect every subsequent batch. + """ + monkeypatch.setattr(runtime, "_embed_group", _raise(ValueError("could not decode image payload for row 3"))) + + out_df = runtime.embed_text_main_text_embed(_batch(), model=object(), inference_batch_size=2) + + assert list(out_df["text_embeddings_1b_v2_has_embedding"]) == [False] * 4 + + +def test_engine_lifecycle_classifier_rejects_unrelated_failures() -> None: + """Pins the classifier's specificity - it must not fire on ordinary failures. + + Cannot run on unmodified HEAD: the helper does not exist there, so this + fails with ``AttributeError`` rather than with a wrong answer. + """ + 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)) + + +def test_the_fatal_set_stays_narrow() -> None: + """Pins the exact fatal set, so widening it is a deliberate edit. + + Each name here has a cited reason in ``runtime.py``. ``EngineDeadError`` is + "Unrecoverable" in vLLM's own definition; ``LocalEmbedderReturnedNothingError`` + and ``LocalEmbedderRowsLostError`` are raised by this repo, each in one place + with one meaning. ``OutOfMemoryError`` and ``EngineGenerateError`` were + removed because a recoverable backend raises them. + + Cannot run on unmodified HEAD: the set does not exist there. + """ + assert set(runtime._ENGINE_LIFECYCLE_EXC_NAMES) == { + "EngineDeadError", + "LocalEmbedderReturnedNothingError", + "LocalEmbedderRowsLostError", + } + + +def test_engine_lifecycle_classifier_visits_a_cyclic_chain_once_per_link() -> None: + """The ``seen`` set, not the depth cap, is what stops a cyclic chain. + + Termination alone does not test ``seen``: the depth-5 cap ends the walk + either way, so an assertion that the call returns would still pass with + ``seen`` deleted. Visit count does distinguish them. A two-link cycle costs + two visits with ``seen`` and five without, because the cap then does the + stopping. + + Same status as the test above: it cannot run on unmodified HEAD, where the + helper does not exist. + """ + visits: list[str] = [] + + class _CountingError(RuntimeError): + def __str__(self) -> str: + visits.append(self.args[0]) + return str(self.args[0]) + + first = _CountingError("first") + second = _CountingError("second") + first.__context__ = second + second.__context__ = first + + assert not runtime._is_engine_lifecycle_failure(first) + assert visits == ["first", "second"] + + +# --------------------------------------------------------------------------- +# Known gap: two embed entry points swallow ``LocalEmbedderRowsLostError``. +# +# Neither is on the shipped route, so these tests pin the gap rather than close +# it. They exist so a future import-site change that flips onto one of these +# functions shows up as a failing, clearly named test instead of as a silently +# short index in production. See ``models/embed_errors.py`` +# ``LocalEmbedderRowsLostError`` for the full analysis and for why the writer's +# fatal set must NOT be widened to cover ``None``. +# --------------------------------------------------------------------------- + + +def test_the_shipped_embed_actors_use_the_guarded_runtime_function() -> None: + """All three shipped embed actors resolve to the guarded implementation. + + This is the reachability claim the two gap tests below depend on. If it ever + fails, the gap stops being theoretical and the swallowing paths must be + fixed. + """ + from nemo_retriever.operators.embed import cpu_operator, gpu_operator, operators + + for module in (operators, gpu_operator, cpu_operator): + assert module.embed_text_main_text_embed is runtime.embed_text_main_text_embed + + +def test_text_embed_operator_route_swallows_the_rows_lost_error(monkeypatch: pytest.MonkeyPatch) -> None: + """Pins the gap in ``operators/embed/text_embed.py``. + + Its ``except BaseException`` around ``model.embed(...)`` converts the raise + into ``{"embedding": None}``, which the LanceDB writer treats as a + legitimate silent drop. Unprotected today, and off the shipped route. + + Passes before and after this change: it asserts pre-existing behaviour. + """ + from nemo_retriever.operators.embed import text_embed + + class _LosingEmbedder: + def embed(self, texts: list[str], **_kwargs: Any) -> list: + raise LocalEmbedderRowsLostError(lost=1, total=len(texts), embedder="_LosingEmbedder") + + out_df = text_embed.embed_text_1b_v2(_batch(), model=_LosingEmbedder(), inference_batch_size=2) + + payloads = list(out_df["text_embeddings_1b_v2"]) + assert [p["embedding"] for p in payloads] == [None] * 4 + assert all(p["error"]["type"] == "LocalEmbedderRowsLostError" for p in payloads) + # The shape the writer deliberately does not treat as fatal, because + # ``text_embed.py`` also writes ``None`` for a legitimate blank-text row. + assert not any(p["embedding"] == [] for p in payloads) + + +def test_modality_pipeline_duplicate_swallows_the_rows_lost_error(monkeypatch: pytest.MonkeyPatch) -> None: + """Pins the gap in ``common/modality/pipeline/embedding.py``. + + This module exports a function with the *identical* name as the patched one + and is re-exported from its package ``__init__``, so it is the higher-risk + of the two: a future import-site change could flip to it and remove all + three guard layers with no visible diff at the call site. + + Passes before and after this change: it asserts pre-existing behaviour. + """ + from nemo_retriever.common.modality.pipeline import embedding as duplicate + + assert duplicate.embed_text_main_text_embed is not runtime.embed_text_main_text_embed + + monkeypatch.setattr( + duplicate, + "_embed_group", + _raise(LocalEmbedderRowsLostError(lost=2, total=4, embedder="LlamaNemotronEmbed1BV2Embedder")), + ) + + out_df = duplicate.embed_text_main_text_embed(_batch(), model=object(), inference_batch_size=2) + + payloads = list(out_df["text_embeddings_1b_v2"]) + assert [p["embedding"] for p in payloads] == [None] * 4 + assert all(p["error"]["type"] == "LocalEmbedderRowsLostError" for p in payloads) + + +class _StubVLEmbedder: + """VL embedder stub whose per-call return length is scripted by the test.""" + + def __init__(self, *, images=None, text_image=None, text=None): + self._images = images + self._text_image = text_image + self._text = text + + def embed_images(self, images_b64, *, batch_size=64): + return self._images + + def embed_text_image(self, texts, images_b64, *, batch_size=64): + return self._text_image + + def embed(self, texts, *, batch_size=64): + return self._text + + +def _image_frame(image_values): + return pd.DataFrame({"_image_b64": list(image_values), "text": [""] * len(image_values)}) + + +# --- the guard must fire: the engine answered short, rows would be lost --- + + +def test_image_mode_short_answer_is_fatal(): + """A row submitted with an image and answered with nothing must fail the run. + + Without the guard the shortfall is padded with ``None``, which both writers + ignore, so the row is dropped and the run still succeeds. + """ + df = _image_frame(["b64-a", "b64-b", "b64-c"]) + embedder = _StubVLEmbedder(images=[[0.1, 0.2]]) # 1 vector for 3 images + + with pytest.raises(LocalEmbedderRowsLostError) as excinfo: + main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") + assert excinfo.value.lost == 2 + assert excinfo.value.total == 3 + + +def test_image_mode_empty_answer_is_fatal(): + """Zero vectors for a chunk that did submit images is the whole-batch failure.""" + df = _image_frame(["b64-a", "b64-b"]) + embedder = _StubVLEmbedder(images=[]) + + with pytest.raises(LocalEmbedderRowsLostError): + main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") + + +def test_text_image_mode_short_multimodal_answer_is_fatal(): + df = pd.DataFrame({"_image_b64": ["b64-a", "b64-b"], "text": ["alpha", "beta"]}) + embedder = _StubVLEmbedder(text_image=[[0.1, 0.2]]) # 1 vector for 2 paired rows + + with pytest.raises(LocalEmbedderRowsLostError): + main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="text_image") + + +def test_text_image_mode_short_text_fallback_answer_is_fatal(): + """The text-only fallback subset is owed one vector per row as well.""" + df = pd.DataFrame({"_image_b64": ["", ""], "text": ["alpha", "beta"]}) + embedder = _StubVLEmbedder(text=[[0.1, 0.2]]) # 1 vector for 2 fallback rows + + with pytest.raises(LocalEmbedderRowsLostError): + main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="text_image") + + +# --- the guard must stay silent: these are legitimate short answers --- + + +def test_image_mode_rows_without_images_do_not_fire_the_guard(): + """Rows with no image are owed nothing; they get ``None`` by contract. + + This is the false-failure the naive ``if not vecs_list: raise`` would cause. + """ + df = _image_frame(["b64-a", "", "b64-c"]) + embedder = _StubVLEmbedder(images=[[0.1, 0.2], [0.3, 0.4]]) # 2 vectors for 2 images + + out = main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") + assert out["embeddings"] == [[0.1, 0.2], None, [0.3, 0.4]] + + +def test_image_mode_chunk_with_no_images_at_all_does_not_fire_the_guard(): + """An image-free chunk submits nothing, so zero vectors back is correct.""" + df = _image_frame(["", ""]) + embedder = _StubVLEmbedder(images=[]) + + out = main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") + assert out["embeddings"] == [None, None] + + +def test_text_image_mode_mixed_rows_do_not_fire_the_guard(): + """Paired, text-only and empty rows each get their own correct treatment.""" + df = pd.DataFrame({"_image_b64": ["b64-a", "", ""], "text": ["alpha", "beta", " "]}) + embedder = _StubVLEmbedder(text_image=[[0.1, 0.2]], text=[[0.3, 0.4]]) + + out = main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="text_image") + assert out["embeddings"] == [[0.1, 0.2], [0.3, 0.4], None] + + +def test_image_mode_healthy_batch_does_not_fire_the_guard(): + df = _image_frame(["b64-a", "b64-b"]) + embedder = _StubVLEmbedder(images=[[0.1, 0.2], [0.3, 0.4]]) + + out = main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") + assert out["embeddings"] == [[0.1, 0.2], [0.3, 0.4]] diff --git a/nemo_retriever/tests/test_lancedb_incomplete_index_guard.py b/nemo_retriever/tests/test_lancedb_incomplete_index_guard.py new file mode 100644 index 0000000000..c0d65dab77 --- /dev/null +++ b/nemo_retriever/tests/test_lancedb_incomplete_index_guard.py @@ -0,0 +1,510 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""The LanceDB writer must refuse to build an index that is missing rows. + +The property under test: a row must not reach the index without an embedding, +and a run that loses rows must fail rather than publish a short index and report +success. It holds for any cause, any backend, and any GPU - the writer sees only +the rows, not what produced them. + +Concretely, a row whose embedding is ``[]`` fails the write. The embed stage +emits that shape for a whole batch it could not embed. It used to be counted as +a wrong-length vector and dropped with a warning, so the run completed with a +short index. + +Three things are deliberately NOT covered here. + +* A wrong-length vector. That is a real vector against the wrong schema, which + is what ``on_bad_vectors`` exists for; folding it in would turn a configured + tolerance into a hard failure on upgrade. +* An absent or ``None`` embedding, which keeps its pre-existing silent drop + because ``operators/embed/text_embed.py`` writes ``None`` on purpose for a + blank-text row. +* A row the embedder zero-padded. It has the correct width and a non-zero + length, so the writer cannot see it. That half is marked at the source; see + ``test_vllm_embed.py``. + +Scope: this guard covers the LanceDB writer. Other paths in the repo still drop +or zero-fill rows of their own accord. + +Every test here is deterministic and needs no GPU. They assert on the write +contract rather than on end-to-end accuracy, which would drift with corpus size. +""" + +from __future__ import annotations + +from typing import Any + +import pandas as pd +import pytest + +from nemo_retriever.common.vdb.adt_vdb import CollectionWriteContext +from nemo_retriever.common.vdb.lancedb import _create_lancedb_results +from nemo_retriever.common.vdb.lancedb_collections import _collection_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_an_empty_embedding_fails_the_run() -> None: + """Known-bad: returned ``(rows, counts)`` and only logged a WARNING. + + ``embedding: []`` is what ``embed_text_main_text_embed`` writes for every + row of a batch it could not embed. It was not counted before this change: + it is not ``None``, and it passed the length check as a wrong-length vector. + """ + records = [[_record([]), _record([1.0, 2.0])]] + + with pytest.raises(RuntimeError) as excinfo: + _create_lancedb_results(records, expected_dim=2) + + message = str(excinfo.value) + assert "Refusing to build an incomplete index" in message + assert "1 of 2 rows" in message + assert "empty_embedding=1" in message + + +def test_create_lancedb_results_rejects_empty_embeddings_without_length_check() -> None: + """The non-enforcing path must not write empty vectors into the index. + + Known-bad: with ``expected_dim=None`` an ``embedding: []`` row was accepted + and counted in ``accepted``. + """ + with pytest.raises(RuntimeError, match="empty_embedding=1"): + _create_lancedb_results([[_record([])]], expected_dim=None) + + +@pytest.mark.parametrize("on_bad_vectors", ["drop", "fill", "error"]) +def test_wrong_length_rows_stay_under_the_on_bad_vectors_policy(on_bad_vectors: str) -> None: + """A short vector is a schema mismatch, not a missing embedding. + + ``on_bad_vectors`` is a documented, user-configured tolerance + (``common/vdb/lancedb.py`` ``create_index``). The incomplete-index guard must + not reach into it, or a user who deliberately configured ``drop`` or ``fill`` + would go from silent dropping to a hard run failure on upgrade. + + Known-bad for the first revision of this fix, which folded + ``dropped_bad_length`` into the fatal condition and raised for all three + values. It pins a contract rather than a code change, so it is a guard test: + the drop it asserts is identical before and after. It cannot run on the + unpatched tree, because the final assertion reads the new + ``empty_embedding`` key. + """ + records = [[_record([1.0]), _record([1.0, 2.0])]] + + # ``expected_dim=None`` is the shape ``create_index`` uses when the caller + # asked LanceDB to own the policy (``on_bad_vectors="error"``) or turned the + # wrapper's length check off. + rows, counts = _create_lancedb_results(records, expected_dim=None if on_bad_vectors == "error" else 2) + + if on_bad_vectors == "error": + # The wrapper forwards both rows so LanceDB itself raises, per the + # documented strict-fail semantics of that policy. + assert len(rows) == 2 + assert counts["dropped_bad_length"] == 0 + else: + assert len(rows) == 1 + assert counts["dropped_bad_length"] == 1 + assert counts["dropped_no_embedding"] == 0 + assert counts["empty_embedding"] == 0 + + +def test_empty_embeddings_from_the_endpoint_path_do_not_reach_the_index_silently( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The writer is the backstop for producers that legitimately keep going. + + A local engine failure now propagates from + ``models/inference/runtime.py`` and never reaches here. The endpoint path + keeps its per-batch resilience by design - a single failed HTTP call should + not kill a service-mode run - so it can still emit ``embedding: []`` rows. + This is the case that makes the writer-side check load-bearing rather than + redundant. + + Known-bad: the writer dropped those rows and returned normally. + """ + from nemo_retriever.models.inference import runtime + + def _refuse(*_args: Any, **_kwargs: Any) -> pd.DataFrame: + raise TimeoutError("read timed out waiting for the embedding endpoint") + + monkeypatch.setattr(runtime, "_embed_group", _refuse) + + batch_df = pd.DataFrame( + { + "text": ["page one", "page two"], + "metadata": [ + {"content": "page one", "content_metadata": {"page_number": 1}, "source_metadata": {}}, + {"content": "page two", "content_metadata": {"page_number": 2}, "source_metadata": {}}, + ], + } + ) + + out_df = runtime.embed_text_main_text_embed( + batch_df, + embedding_endpoint="http://embed.example/v1", + inference_batch_size=2, + ) + + # The stage returns successfully - deliberate for the endpoint path - but + # every row is empty. + assert list(out_df["text_embeddings_1b_v2_has_embedding"]) == [False, False] + assert [payload["embedding"] for payload in out_df["text_embeddings_1b_v2"]] == [[], []] + + records = [ + [ + { + "document_type": "text", + "metadata": {**row["metadata"], "embedding": row["text_embeddings_1b_v2"]["embedding"]}, + } + for _index, row in out_df.iterrows() + ] + ] + + with pytest.raises(RuntimeError, match="Refusing to build an incomplete index"): + _create_lancedb_results(records, expected_dim=2048) + + +# -------------------------------------------------------------------------- +# No false positives: everything the fatal condition must NOT fire on. +# +# Catching the defect is half the evidence. The other half is that a normal run +# still finishes, so there is one case here per legitimate drop and per value +# type the new check could have swallowed by accident. +# -------------------------------------------------------------------------- + + +def test_wrong_length_is_counted_and_never_fatal() -> None: + """``dropped_bad_length`` stays out of the fatal condition. + + It is the category ``on_bad_vectors`` governs: a real vector against the + wrong schema. Folding it in is exactly the regression that would turn a + user's configured ``drop``/``fill``/``null`` tolerance into a hard failure + on upgrade. + + Cannot run on the unpatched tree: it asserts on the new + ``empty_embedding`` key, which does not exist there. What it pins is + a contract, not a behaviour change - the wrong-length row is dropped and not + raised on, before and after. + """ + rows, counts = _create_lancedb_results([[_record([1.0, 2.0, 3.0]), _record([1.0, 2.0])]], expected_dim=2) + + assert len(rows) == 1 + assert counts["dropped_bad_length"] == 1 + assert counts["empty_embedding"] == 0 + + +def test_canonical_image_row_without_text_is_accepted_not_dropped() -> None: + """The text carve-out for canonical image rows still works. + + An image row legitimately carries ``text=""``. It has a real embedding, so + nothing here may touch it. + """ + record = { + "document_type": "image", + "metadata": { + "embedding": [1.0, 2.0], + "content": "", + "content_metadata": {"page_number": 3, "type": "image"}, + "source_metadata": {"source_name": "scan.pdf"}, + }, + } + rows, counts = _create_lancedb_results([[record]], expected_dim=2) + + assert len(rows) == 1 + assert counts["accepted"] == 1 + assert counts["dropped_no_text"] == 0 + + +def test_text_free_non_image_row_is_dropped_and_never_fatal() -> None: + """``dropped_no_text`` stays out of the fatal condition. + + It is a content filter, not a loss: the row embedded successfully and was + excluded for having nothing to search on. + + Cannot run on the unpatched tree: it asserts on the new + ``empty_embedding`` key. The drop itself is unchanged. + """ + rows, counts = _create_lancedb_results([[_record([1.0, 2.0], text=" ")]], expected_dim=2) + + assert rows == [] + assert counts["dropped_no_text"] == 1 + 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: + """The check tests presence, never values. + + Deciding from the numbers would mean guessing which embeddings are "real", + and an all-zero vector is a legal thing for a model to emit. Only the + absence of a vector is fatal. + """ + rows, counts = _create_lancedb_results([[_record(embedding)]], expected_dim=2) + + assert len(rows) == 1 + assert counts["accepted"] == 1 + + +def test_numpy_embeddings_keep_their_existing_handling() -> None: + """The new check must not change how non-list types are treated. + + ``not embedding`` raises ``ValueError`` on a multi-element numpy array, so + the check is written as ``isinstance(embedding, (list, tuple)) and + len(embedding) == 0``. A numpy row therefore falls through to the + pre-existing length check and behaves exactly as it does today: counted as + ``dropped_bad_length``, not raised on, not crashed on. + + Not a known-bad test: it pins a hazard the obvious spelling of this check + would have introduced. It cannot run on the unpatched tree, which has no + ``empty_embedding`` key. + """ + numpy = pytest.importorskip("numpy") + + rows, counts = _create_lancedb_results( + [[_record(numpy.array([1.0, 2.0])), _record([1.0, 2.0])]], + expected_dim=2, + ) + + assert len(rows) == 1 + assert counts["dropped_bad_length"] == 1 + assert counts["empty_embedding"] == 0 + + +def test_an_empty_numpy_array_is_not_swallowed_by_the_new_check() -> None: + """An empty ndarray is not a list, so it keeps its pre-existing route. + + This is the narrowness of the check made explicit: it fires on ``[]`` and + ``()`` and nothing else. + """ + 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 + + +def test_a_fully_healthy_batch_raises_nothing() -> None: + """The whole point: a normal run is untouched. + + Cannot run on the unpatched tree, which has no ``empty_embedding`` + key; the acceptance of all 50 rows is identical there. + """ + records = [[_record([1.0, 2.0]) for _ in range(50)]] + + rows, counts = _create_lancedb_results(records, expected_dim=2) + + assert len(rows) == 50 + assert counts["accepted"] == 50 + assert counts["empty_embedding"] == 0 + assert counts["dropped_no_embedding"] == 0 + + +def test_a_none_embedding_keeps_its_pre_existing_silent_drop() -> None: + """``None`` is NOT fatal, because it has a legitimate producer. + + ``operators/embed/text_embed.py`` writes ``{"embedding": None}`` on purpose + for a row whose text was blank and which it therefore chose not to embed. + Making ``dropped_no_embedding`` fatal would fail ingests containing such + rows, which work today. Only ``[]`` - written solely on failure paths, at + ``models/inference/runtime.py`` and ``models/inference/vllm.py`` - is fatal. + + Cannot run on the unpatched tree, which has no ``empty_embedding`` + key. The drop it asserts is behaviour this change deliberately does not + touch. + """ + rows, counts = _create_lancedb_results([[_record(None), _record([1.0, 2.0])]], expected_dim=2) + + assert len(rows) == 1 + assert counts["dropped_no_embedding"] == 1 + assert counts["empty_embedding"] == 0 + + +# --------------------------------------------------------------------------- +# The collection-managed write path. +# +# ``common/vdb/lancedb_collections.py::_collection_rows`` is the second writer +# in the repo. It builds the rows for a collection-managed document ingest and +# had the same defect as ``_create_lancedb_results``: ``not vector`` is true for +# ``[]``, so a row the embed stage failed to embed was skipped with no counter, +# no log and no failure, and the document was published short while the ingest +# reported success. Same property, same fatal condition, same explicit +# ``isinstance``/``len`` spelling. +# --------------------------------------------------------------------------- + + +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", + "metadata": { + "embedding": embedding, + "content": text, + "content_metadata": {"type": "text", "page_number": 2}, + "source_metadata": {"source_id": "/inputs/source.pdf", "source_name": "source.pdf"}, + }, + } + + +def test_an_empty_embedding_on_the_collection_path_fails_the_write() -> None: + """Known-bad: returned the surviving rows and skipped the empty one silently. + + On the unpatched tree ``not vector`` swallowed ``[]`` into the same branch as + a malformed value, so this returned one row and no error, and the collection + document was written short. It now raises before any row reaches LanceDB. + + Fails on the unpatched tree: no exception is raised. + """ + records = [[_collection_record([]), _collection_record([1.0, 0.0])]] + + with pytest.raises(RuntimeError) as excinfo: + _collection_rows(records, context=_collection_context()) + + message = str(excinfo.value) + assert "incomplete document" in message + assert "empty_embedding=1" in message + + +def test_the_collection_path_reports_the_same_counter_name_as_the_pipeline_path() -> None: + """Both writers name the condition ``empty_embedding`` so one grep finds both. + + Fails on the unpatched tree: no exception, so nothing to read the name from. + """ + with pytest.raises(RuntimeError) as excinfo: + _collection_rows([[_collection_record([])]], context=_collection_context()) + + assert "empty_embedding" in str(excinfo.value) + + +def test_a_healthy_collection_document_still_writes_every_row() -> None: + """False-failure guard: the fatal branch must not fire on good input. + + Runs on the unpatched tree and passes there too, which is the point. + """ + records = [[_collection_record([1.0, 0.0]), _collection_record([0.0, 1.0], text="second chunk")]] + + rows = _collection_rows(records, context=_collection_context()) + + assert len(rows) == 2 + assert [row["text"] for row in rows] == ["first chunk", "second chunk"] + + +@pytest.mark.parametrize( + "embedding", + [ + None, + "not-a-vector", + 123, + {"dense": [1.0, 0.0]}, + ], + ids=["none", "string", "int", "dict"], +) +def test_a_malformed_collection_embedding_keeps_its_pre_existing_silent_skip(embedding: Any) -> None: + """False-failure guard, and the one that matters most. + + ``not isinstance(vector, (list, tuple))`` covers malformed values. It was + NOT made fatal: only ``[]`` was, because ``[]`` is the sole value written + exclusively on a failure path. Widening the fatal set to these would fail + ingests that work today. The healthy sibling row must still be written. + + Runs on the unpatched tree and passes there too. + """ + records = [[_collection_record(embedding), _collection_record([1.0, 0.0], text="good chunk")]] + + rows = _collection_rows(records, context=_collection_context()) + + assert len(rows) == 1 + assert rows[0]["text"] == "good chunk" + + +def test_a_malformed_collection_batch_or_record_keeps_its_pre_existing_silent_skip() -> None: + """False-failure guard for the batch- and record-shaped skips. + + ``not isinstance(batch, list)``, ``not isinstance(record, dict)`` and + ``not isinstance(metadata, dict)`` are untouched: they still skip silently + and never fail the write. + + Runs on the unpatched tree and passes there too. + """ + records = [ + "not-a-batch", + ["not-a-record"], + [{"document_type": "text", "metadata": "not-a-dict"}], + [_collection_record([1.0, 0.0], text="good chunk")], + ] + + rows = _collection_rows(records, context=_collection_context()) + + assert len(rows) == 1 + assert rows[0]["text"] == "good chunk" + + +def test_a_text_free_non_image_collection_row_is_skipped_and_never_fatal() -> None: + """False-failure guard: the content filter keeps its silent drop. + + Runs on the unpatched tree and passes there too. + """ + records = [[_collection_record([1.0, 0.0], text=" "), _collection_record([0.0, 1.0], text="good chunk")]] + + rows = _collection_rows(records, context=_collection_context()) + + assert len(rows) == 1 + assert rows[0]["text"] == "good chunk" + + +def test_a_numpy_collection_embedding_does_not_raise_on_truthiness() -> None: + """The new branch must never evaluate ``not vector`` on an array. + + A multi-element numpy array raises ``ValueError`` on ``bool()``. The check + is written as ``isinstance(vector, (list, tuple)) and len(vector) == 0``, so + an array short-circuits out of it and keeps its pre-existing skip rather + than crashing the ingest. + + Runs on the unpatched tree, where the existing ``or`` short-circuits for the + same reason, and passes there too. + """ + numpy = pytest.importorskip("numpy") + + 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()) + + assert len(rows) == 1 + assert rows[0]["text"] == "good chunk" diff --git a/nemo_retriever/tests/test_lancedb_write_policy.py b/nemo_retriever/tests/test_lancedb_write_policy.py index d4c3d80f62..c96b47d4ff 100644 --- a/nemo_retriever/tests/test_lancedb_write_policy.py +++ b/nemo_retriever/tests/test_lancedb_write_policy.py @@ -284,14 +284,75 @@ def test_dense_write_requires_both_canonical_image_fields(tmp_path: Path, missin @pytest.mark.parametrize( "vector", - [pytest.param(None, id="missing"), pytest.param([1.0], id="wrong-length")], + [ + pytest.param([], id="empty-embed-failure"), + ], ) -def test_dense_write_drops_image_only_row_with_invalid_embedding(tmp_path: Path, vector: list[float] | None) -> None: - table_rows = _write_rows(tmp_path, _image_only_records(vector)) +def test_dense_write_fails_on_image_only_row_with_no_usable_embedding( + tmp_path: Path, vector: list[float] | None +) -> None: + """A row without a usable embedding must fail the write, not be dropped. + + The embed stage writes ``embedding: []`` for every row of a batch whose + engine failed. On the unpatched tree that row was accepted here, because + ``[]`` is not ``None`` and this write infers the dimension, so the length + check does not run - which is how a run publishes a short index and still + reports success. + + A ``None`` embedding is deliberately not covered: it has a legitimate + producer and keeps its pre-existing drop. See + ``test_lancedb_incomplete_index_guard.py``, + ``test_a_none_embedding_keeps_its_pre_existing_silent_drop``. + """ + with pytest.raises(RuntimeError, match="Refusing to build an incomplete index"): + _write_rows(tmp_path, _image_only_records(vector)) + + +def test_dense_write_still_drops_wrong_length_row_under_the_default_policy(tmp_path: Path) -> None: + """``on_bad_vectors="drop"`` keeps working: a short vector is dropped, not fatal. + + Known-bad for the first revision of this fix, which folded + ``dropped_bad_length`` into the fatal condition and made the shipped default + unreachable. Passes on unmodified HEAD too - it pins the documented contract + rather than a code change, so it is a guard test. + """ + table_rows = _write_rows(tmp_path, _image_only_records([1.0])) assert table_rows == [] +def test_dense_write_keeps_on_bad_vectors_fill_reachable(tmp_path: Path) -> None: + """``on_bad_vectors="fill"`` with the wrapper check off still reaches LanceDB. + + With ``validate_vector_length=False`` the short row is forwarded and LanceDB + fills it, which is what a user who configured ``fill`` asked for. The guard + must not pre-empt that. + + Known-bad for the first revision of this fix, which raised before LanceDB + ever saw the row. Passes on unmodified HEAD; guard test. + + Asserts the row survives at full schema width, not the exact filled + composition: how LanceDB distributes ``fill_value`` over a short vector is + its own detail and differs by version (0.34 replaces the whole vector, 0.37 + pads and keeps the produced component), and ``lancedb`` is unpinned here. + What this guard owns is that the row reached the writer at all. + """ + op = LanceDB( + uri=str(tmp_path), + table_name="t", + vector_dim=2, + create_index=False, + on_bad_vectors="fill", + fill_value=0.5, + validate_vector_length=False, + ) + op.run(_image_only_records([1.0])) + + table_rows = lancedb.connect(str(tmp_path)).open_table("t").to_arrow().to_pylist() + assert len(table_rows) == 1 + assert len(table_rows[0]["vector"]) == 2 + + 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) diff --git a/nemo_retriever/tests/test_vllm_embed.py b/nemo_retriever/tests/test_vllm_embed.py index ddf5548139..e14cbd2795 100644 --- a/nemo_retriever/tests/test_vllm_embed.py +++ b/nemo_retriever/tests/test_vllm_embed.py @@ -22,6 +22,8 @@ embed_with_vllm_llm, ) from nemo_retriever.models.local.llama_nemotron_embed_1b_v2_embedder import LlamaNemotronEmbed1BV2Embedder +from nemo_retriever.models.embed_errors import LocalEmbedderRowsLostError +from nemo_retriever.models.nim.error_reporter import drain_errors def _make_output(embedding): @@ -499,14 +501,77 @@ def test_output_is_unnormalized_when_normalize_false(self): assert mock_mm.call_args.kwargs["normalize"] is False assert result.tolist() == [[3.0, 4.0]] - def test_no_valid_embeddings_returns_empty_tensor(self): + def test_no_valid_embeddings_no_longer_returns_an_empty_tensor(self): + # Rewritten, not weakened. This test previously asserted the defect: + # a batch vLLM failed to embed came back as a 0-row tensor, which is + # the same answer as an empty input, so the loss was unobservable. b64 = _make_minimal_b64() with patch( "nemo_retriever.models.inference.vllm.embed_multimodal_with_vllm_llm", return_value=[[]], ): - result = self.embedder.embed_images([b64]) - assert result.shape[0] == 0 + with pytest.raises(LocalEmbedderRowsLostError): + self.embedder.embed_images([b64]) + + def test_a_partially_lost_batch_fails_with_the_exact_count(self): + """The count exists here and nowhere else, so it must leave the function. + + ``_finalize_vectors`` holds both the batch it sent and the vectors that + came back, so ``len(vectors) - len(valid)`` is exact. It used to compute + that and discard it: the failed row was zero-padded to the right width, + which every shape check downstream accepts, and ``has_embedding`` then + reported ``True`` for a row carrying nothing. + + The LanceDB writer guard cannot cover this row - correct width, non-zero + length - without guessing from its values, which is why the loss is + marked at the source instead. + + Known-bad: returns a ``(2, 2)`` tensor whose second row is ``[0, 0]``, + with nothing raised and nothing collected. + """ + drain_errors() + b64 = _make_minimal_b64() + with patch( + "nemo_retriever.models.inference.vllm.embed_multimodal_with_vllm_llm", + return_value=[[3.0, 4.0], []], + ): + with pytest.raises(LocalEmbedderRowsLostError) as excinfo: + self.embedder.embed_images([b64, b64]) + + assert (excinfo.value.lost, excinfo.value.total) == (1, 2) + assert excinfo.value.embedder == "LlamaNemotronEmbedVL1BV2VLLMEmbedder" + collected = drain_errors() + assert [(error.exc_type, error.stage) for error in collected] == [("LocalEmbedderRowsLostError", "embed")] + + def test_a_wholly_lost_batch_fails_naming_every_row(self): + """Total loss must name the count too, not just fail. + + Known-bad: returns shape ``(0, 2048)``, the same answer as + ``embed_images([])``, so the whole batch vanishes with no padding to + notice and nothing raised. + """ + b64 = _make_minimal_b64() + with patch( + "nemo_retriever.models.inference.vllm.embed_multimodal_with_vllm_llm", + return_value=[[], []], + ): + with pytest.raises(LocalEmbedderRowsLostError) as excinfo: + self.embedder.embed_images([b64, b64]) + + assert (excinfo.value.lost, excinfo.value.total) == (2, 2) + + def test_a_complete_batch_is_untouched(self): + """Guard: no loss, nothing raised, nothing collected. Passes both ways.""" + drain_errors() + b64 = _make_minimal_b64() + with patch( + "nemo_retriever.models.inference.vllm.embed_multimodal_with_vllm_llm", + return_value=[[3.0, 4.0], [0.0, 5.0]], + ): + result = self.embedder.embed_images([b64, b64]) + + assert drain_errors() == [] + assert result.shape == (2, 2) def _make_text_embedder(): @@ -520,15 +585,104 @@ class TestLlamaNemotronEmbed1BV2Embedder: def setup_method(self): self.embedder = _make_text_embedder() - def test_finalize_vectors_all_empty_returns_empty_tensor(self): - result = self.embedder._finalize_vectors([[], []]) - assert isinstance(result, torch.Tensor) - assert result.shape[0] == 0 + def test_finalize_vectors_all_empty_no_longer_returns_empty_tensor(self): + # Rewritten, not weakened: the old assertion pinned the defect. + with pytest.raises(LocalEmbedderRowsLostError): + self.embedder._finalize_vectors([[], []]) + + def test_finalize_vectors_cannot_see_a_zero_output_batch(self): + """The conservation check counts empty rows, so zero rows count as zero loss. + + When vLLM yields no outputs at all for a non-empty batch, + ``embed_with_vllm_llm`` returns ``[]``. ``report_lost_rows`` then sums + over nothing and reports no loss, and the ``if not valid`` early return + hands back a 0-row tensor - the same answer as an empty input. Nothing + raises here. + + This is a pin on a real gap, not a defect this change fixes, so it + passes before and after. It is why the guard in + ``main_text_embed._callable_runner`` is not redundant with + ``_finalize_vectors``: that guard is the only layer that sees this + shape, and it fires before the LanceDB writer would. + """ + assert self.embedder._finalize_vectors([]).shape == (0, 0) - def test_finalize_vectors_zero_pads_missing(self): - result = self.embedder._finalize_vectors([[1.0, 0.0], []]) - assert result.shape == (2, 2) - assert result[1].tolist() == [0.0, 0.0] + def test_finalize_vectors_no_longer_zero_pads_missing(self): + # Rewritten, not weakened: the old assertion required the zero padding + # that made a lost row indistinguishable from a real embedding. + with pytest.raises(LocalEmbedderRowsLostError): + self.embedder._finalize_vectors([[1.0, 0.0], []]) + + def test_finalize_vectors_fails_on_the_rows_it_would_pad(self): + """The text embedder's ``_finalize_vectors`` has the same contract. + + Both embedders pad, so both must refuse to. Known-bad: the padding in + the test above happens and nothing records or stops it. + """ + drain_errors() + with pytest.raises(LocalEmbedderRowsLostError) as excinfo: + self.embedder._finalize_vectors([[1.0, 0.0], []]) + + assert (excinfo.value.lost, excinfo.value.total) == (1, 2) + assert excinfo.value.embedder == "LlamaNemotronEmbed1BV2Embedder" + assert [error.exc_type for error in drain_errors()] == ["LocalEmbedderRowsLostError"] + + def test_report_lost_rows_never_fires_on_a_healthy_batch(self): + """No false positives: every batch in every run goes through this call. + + A false positive here would fail runs that work today, so the healthy + shapes are pinned explicitly: full-width vectors, a single component, a + zero component, and an all-zero vector. An all-zero vector is a legal + thing for a model to emit; the check tests presence, never values. + """ + from nemo_retriever.models.embed_errors import report_lost_rows + + drain_errors() + for batch in ( + [[0.1] * 2048, [0.2] * 2048], + [[0.0]], + [[0.0, 0.0], [0.0, 0.0]], + [[1.0]] * 256, + ): + assert report_lost_rows(batch, embedder="X") == 0 + assert drain_errors() == [] + + def test_report_lost_rows_does_not_evaluate_truthiness_of_arrays(self): + """Explicit length test, because ``not vector`` raises on an array. + + ``bool(numpy.array([1.0, 2.0]))`` raises ``ValueError: The truth value + of an array with more than one element is ambiguous``. This call runs + before any pre-existing code, so a crash here would be a new one. + """ + numpy = pytest.importorskip("numpy") + from nemo_retriever.models.embed_errors import report_lost_rows + + drain_errors() + assert report_lost_rows([numpy.array([1.0, 2.0]), numpy.array([3.0, 4.0])], embedder="X") == 0 + assert drain_errors() == [] + + def test_report_lost_rows_is_silent_when_nothing_was_lost(self): + """Guard: the added call is inert on the healthy path. + + Every batch goes through it, so a false positive here would break runs + that work today. Cannot run on unmodified HEAD: the helper does not + exist there. + """ + from nemo_retriever.models.embed_errors import report_lost_rows + + drain_errors() + assert report_lost_rows([[1.0], [0.0]], embedder="X") == 0 + assert report_lost_rows([], embedder="X") == 0 + assert drain_errors() == [] + + def test_rows_lost_error_message_names_the_consequence(self): + """The message must say why a padded row matters, not just that it exists. + + Cannot run on unmodified HEAD: the class does not exist there. + """ + exc = LocalEmbedderRowsLostError(lost=7, total=64, embedder="SomeEmbedder") + assert (exc.lost, exc.total, exc.embedder) == (7, 64, "SomeEmbedder") + assert "match nothing" in str(exc) def test_embed_uses_passage_prefix_by_default(self): with patch("nemo_retriever.models.inference.vllm.embed_with_vllm_llm", return_value=[[0.6, 0.8]]) as mock_fn: From b7275b97c24753157e290c77a18153a6461ca8c4 Mon Sep 17 00:00:00 2001 From: Haidong Rong Date: Wed, 19 Aug 2026 09:44:53 -0700 Subject: [PATCH 2/4] test(embed): consolidate guard tests into the existing suites Addresses review feedback that the change was mostly test code, spread over two new files of a little over 1,000 lines that also covered behaviour this change does not touch. Deletes `test_lancedb_incomplete_index_guard.py` and moves its cases to the suites that already own those surfaces: the pipeline-writer cases to `test_lancedb_write_policy.py`, the collection-writer cases to `test_lancedb_collections.py`, and the multimodal row-loss cases to `test_multimodal_embed.py`. `test_embed_engine_failure_propagation.py` keeps only the runtime classifier, which has no existing home. Drops 11 cases that asserted pre-existing untouched behaviour rather than this change: numpy and malformed-value handling, the collection path's existing silent skips, a duplicate wrong-length assertion already covered by the `on_bad_vectors` case, the fatal-set contents, the cause-chain cycle walk, and three import-topology pins on modules this change does not modify. What stays is the pair that matters: for each guard, a case that fails on the unpatched tree, and a case asserting the guard does not fire - the endpoint path, wrong-length vectors under the configured policy, blank-text and canonical image rows, `None` and absent embeddings, all-zero vectors, empty numpy arrays, image-free chunks, and healthy batches. Net 721 deletions against 430 insertions. Local run of the five affected files: 195 passed, 3 skipped. Signed-off-by: Haidong Rong --- .../test_embed_engine_failure_propagation.py | 210 -------- .../tests/test_lancedb_collections.py | 73 +++ .../test_lancedb_incomplete_index_guard.py | 510 ------------------ .../tests/test_lancedb_write_policy.py | 248 ++++++++- nemo_retriever/tests/test_multimodal_embed.py | 110 ++++ 5 files changed, 430 insertions(+), 721 deletions(-) delete mode 100644 nemo_retriever/tests/test_lancedb_incomplete_index_guard.py diff --git a/nemo_retriever/tests/test_embed_engine_failure_propagation.py b/nemo_retriever/tests/test_embed_engine_failure_propagation.py index 09c431fac6..f7c8727b38 100644 --- a/nemo_retriever/tests/test_embed_engine_failure_propagation.py +++ b/nemo_retriever/tests/test_embed_engine_failure_propagation.py @@ -286,52 +286,6 @@ def test_engine_lifecycle_classifier_rejects_unrelated_failures() -> None: assert runtime._is_engine_lifecycle_failure(RuntimeError(ENGINE_INIT_FAILED)) -def test_the_fatal_set_stays_narrow() -> None: - """Pins the exact fatal set, so widening it is a deliberate edit. - - Each name here has a cited reason in ``runtime.py``. ``EngineDeadError`` is - "Unrecoverable" in vLLM's own definition; ``LocalEmbedderReturnedNothingError`` - and ``LocalEmbedderRowsLostError`` are raised by this repo, each in one place - with one meaning. ``OutOfMemoryError`` and ``EngineGenerateError`` were - removed because a recoverable backend raises them. - - Cannot run on unmodified HEAD: the set does not exist there. - """ - assert set(runtime._ENGINE_LIFECYCLE_EXC_NAMES) == { - "EngineDeadError", - "LocalEmbedderReturnedNothingError", - "LocalEmbedderRowsLostError", - } - - -def test_engine_lifecycle_classifier_visits_a_cyclic_chain_once_per_link() -> None: - """The ``seen`` set, not the depth cap, is what stops a cyclic chain. - - Termination alone does not test ``seen``: the depth-5 cap ends the walk - either way, so an assertion that the call returns would still pass with - ``seen`` deleted. Visit count does distinguish them. A two-link cycle costs - two visits with ``seen`` and five without, because the cap then does the - stopping. - - Same status as the test above: it cannot run on unmodified HEAD, where the - helper does not exist. - """ - visits: list[str] = [] - - class _CountingError(RuntimeError): - def __str__(self) -> str: - visits.append(self.args[0]) - return str(self.args[0]) - - first = _CountingError("first") - second = _CountingError("second") - first.__context__ = second - second.__context__ = first - - assert not runtime._is_engine_lifecycle_failure(first) - assert visits == ["first", "second"] - - # --------------------------------------------------------------------------- # Known gap: two embed entry points swallow ``LocalEmbedderRowsLostError``. # @@ -344,173 +298,9 @@ def __str__(self) -> str: # --------------------------------------------------------------------------- -def test_the_shipped_embed_actors_use_the_guarded_runtime_function() -> None: - """All three shipped embed actors resolve to the guarded implementation. - - This is the reachability claim the two gap tests below depend on. If it ever - fails, the gap stops being theoretical and the swallowing paths must be - fixed. - """ - from nemo_retriever.operators.embed import cpu_operator, gpu_operator, operators - - for module in (operators, gpu_operator, cpu_operator): - assert module.embed_text_main_text_embed is runtime.embed_text_main_text_embed - - -def test_text_embed_operator_route_swallows_the_rows_lost_error(monkeypatch: pytest.MonkeyPatch) -> None: - """Pins the gap in ``operators/embed/text_embed.py``. - - Its ``except BaseException`` around ``model.embed(...)`` converts the raise - into ``{"embedding": None}``, which the LanceDB writer treats as a - legitimate silent drop. Unprotected today, and off the shipped route. - - Passes before and after this change: it asserts pre-existing behaviour. - """ - from nemo_retriever.operators.embed import text_embed - - class _LosingEmbedder: - def embed(self, texts: list[str], **_kwargs: Any) -> list: - raise LocalEmbedderRowsLostError(lost=1, total=len(texts), embedder="_LosingEmbedder") - - out_df = text_embed.embed_text_1b_v2(_batch(), model=_LosingEmbedder(), inference_batch_size=2) - - payloads = list(out_df["text_embeddings_1b_v2"]) - assert [p["embedding"] for p in payloads] == [None] * 4 - assert all(p["error"]["type"] == "LocalEmbedderRowsLostError" for p in payloads) - # The shape the writer deliberately does not treat as fatal, because - # ``text_embed.py`` also writes ``None`` for a legitimate blank-text row. - assert not any(p["embedding"] == [] for p in payloads) - - -def test_modality_pipeline_duplicate_swallows_the_rows_lost_error(monkeypatch: pytest.MonkeyPatch) -> None: - """Pins the gap in ``common/modality/pipeline/embedding.py``. - - This module exports a function with the *identical* name as the patched one - and is re-exported from its package ``__init__``, so it is the higher-risk - of the two: a future import-site change could flip to it and remove all - three guard layers with no visible diff at the call site. - - Passes before and after this change: it asserts pre-existing behaviour. - """ - from nemo_retriever.common.modality.pipeline import embedding as duplicate - - assert duplicate.embed_text_main_text_embed is not runtime.embed_text_main_text_embed - - monkeypatch.setattr( - duplicate, - "_embed_group", - _raise(LocalEmbedderRowsLostError(lost=2, total=4, embedder="LlamaNemotronEmbed1BV2Embedder")), - ) - - out_df = duplicate.embed_text_main_text_embed(_batch(), model=object(), inference_batch_size=2) - - payloads = list(out_df["text_embeddings_1b_v2"]) - assert [p["embedding"] for p in payloads] == [None] * 4 - assert all(p["error"]["type"] == "LocalEmbedderRowsLostError" for p in payloads) - - -class _StubVLEmbedder: - """VL embedder stub whose per-call return length is scripted by the test.""" - - def __init__(self, *, images=None, text_image=None, text=None): - self._images = images - self._text_image = text_image - self._text = text - - def embed_images(self, images_b64, *, batch_size=64): - return self._images - - def embed_text_image(self, texts, images_b64, *, batch_size=64): - return self._text_image - - def embed(self, texts, *, batch_size=64): - return self._text - - -def _image_frame(image_values): - return pd.DataFrame({"_image_b64": list(image_values), "text": [""] * len(image_values)}) - - # --- the guard must fire: the engine answered short, rows would be lost --- -def test_image_mode_short_answer_is_fatal(): - """A row submitted with an image and answered with nothing must fail the run. - - Without the guard the shortfall is padded with ``None``, which both writers - ignore, so the row is dropped and the run still succeeds. - """ - df = _image_frame(["b64-a", "b64-b", "b64-c"]) - embedder = _StubVLEmbedder(images=[[0.1, 0.2]]) # 1 vector for 3 images - - with pytest.raises(LocalEmbedderRowsLostError) as excinfo: - main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") - assert excinfo.value.lost == 2 - assert excinfo.value.total == 3 - - -def test_image_mode_empty_answer_is_fatal(): - """Zero vectors for a chunk that did submit images is the whole-batch failure.""" - df = _image_frame(["b64-a", "b64-b"]) - embedder = _StubVLEmbedder(images=[]) - - with pytest.raises(LocalEmbedderRowsLostError): - main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") - - -def test_text_image_mode_short_multimodal_answer_is_fatal(): - df = pd.DataFrame({"_image_b64": ["b64-a", "b64-b"], "text": ["alpha", "beta"]}) - embedder = _StubVLEmbedder(text_image=[[0.1, 0.2]]) # 1 vector for 2 paired rows - - with pytest.raises(LocalEmbedderRowsLostError): - main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="text_image") - - -def test_text_image_mode_short_text_fallback_answer_is_fatal(): - """The text-only fallback subset is owed one vector per row as well.""" - df = pd.DataFrame({"_image_b64": ["", ""], "text": ["alpha", "beta"]}) - embedder = _StubVLEmbedder(text=[[0.1, 0.2]]) # 1 vector for 2 fallback rows - - with pytest.raises(LocalEmbedderRowsLostError): - main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="text_image") - - # --- the guard must stay silent: these are legitimate short answers --- -def test_image_mode_rows_without_images_do_not_fire_the_guard(): - """Rows with no image are owed nothing; they get ``None`` by contract. - - This is the false-failure the naive ``if not vecs_list: raise`` would cause. - """ - df = _image_frame(["b64-a", "", "b64-c"]) - embedder = _StubVLEmbedder(images=[[0.1, 0.2], [0.3, 0.4]]) # 2 vectors for 2 images - - out = main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") - assert out["embeddings"] == [[0.1, 0.2], None, [0.3, 0.4]] - - -def test_image_mode_chunk_with_no_images_at_all_does_not_fire_the_guard(): - """An image-free chunk submits nothing, so zero vectors back is correct.""" - df = _image_frame(["", ""]) - embedder = _StubVLEmbedder(images=[]) - - out = main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") - assert out["embeddings"] == [None, None] - - -def test_text_image_mode_mixed_rows_do_not_fire_the_guard(): - """Paired, text-only and empty rows each get their own correct treatment.""" - df = pd.DataFrame({"_image_b64": ["b64-a", "", ""], "text": ["alpha", "beta", " "]}) - embedder = _StubVLEmbedder(text_image=[[0.1, 0.2]], text=[[0.3, 0.4]]) - - out = main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="text_image") - assert out["embeddings"] == [[0.1, 0.2], [0.3, 0.4], None] - - -def test_image_mode_healthy_batch_does_not_fire_the_guard(): - df = _image_frame(["b64-a", "b64-b"]) - embedder = _StubVLEmbedder(images=[[0.1, 0.2], [0.3, 0.4]]) - - out = main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") - assert out["embeddings"] == [[0.1, 0.2], [0.3, 0.4]] diff --git a/nemo_retriever/tests/test_lancedb_collections.py b/nemo_retriever/tests/test_lancedb_collections.py index 3606cbaf32..8e1826bcd6 100644 --- a/nemo_retriever/tests/test_lancedb_collections.py +++ b/nemo_retriever/tests/test_lancedb_collections.py @@ -13,6 +13,7 @@ import math import threading from dataclasses import replace +from typing import Any import lancedb import pytest @@ -1127,3 +1128,75 @@ def delete_target(): assert delete_finished.is_set() assert query_errors == [] assert delete_errors == [] + +# --- incomplete-index guard: a row must not reach the index without an embedding --- +# ``[]`` is not ``None``, so it used to fall through to the length check, be counted a +# wrong-length vector, and be dropped - a short index published with exit 0. + + +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", + "metadata": { + "embedding": embedding, + "content": text, + "content_metadata": {"type": "text", "page_number": 2}, + "source_metadata": {"source_id": "/inputs/source.pdf", "source_name": "source.pdf"}, + }, + } + + +def test_an_empty_embedding_on_the_collection_path_fails_the_write() -> None: + """Known-bad: returned the surviving rows and skipped the empty one silently. + + On the unpatched tree ``not vector`` swallowed ``[]`` into the same branch as + a malformed value, so this returned one row and no error, and the collection + document was written short. It now raises before any row reaches LanceDB. + + Fails on the unpatched tree: no exception is raised. + """ + records = [[_collection_record([]), _collection_record([1.0, 0.0])]] + + with pytest.raises(RuntimeError) as excinfo: + _collection_rows(records, context=_collection_context()) + + message = str(excinfo.value) + assert "incomplete document" in message + assert "empty_embedding=1" in message + + +def test_the_collection_path_reports_the_same_counter_name_as_the_pipeline_path() -> None: + """Both writers name the condition ``empty_embedding`` so one grep finds both. + + Fails on the unpatched tree: no exception, so nothing to read the name from. + """ + with pytest.raises(RuntimeError) as excinfo: + _collection_rows([[_collection_record([])]], context=_collection_context()) + + assert "empty_embedding" in str(excinfo.value) + + +def test_a_healthy_collection_document_still_writes_every_row() -> None: + """False-failure guard: the fatal branch must not fire on good input. + + Runs on the unpatched tree and passes there too, which is the point. + """ + records = [[_collection_record([1.0, 0.0]), _collection_record([0.0, 1.0], text="second chunk")]] + + rows = _collection_rows(records, context=_collection_context()) + + assert len(rows) == 2 + assert [row["text"] for row in rows] == ["first chunk", "second chunk"] diff --git a/nemo_retriever/tests/test_lancedb_incomplete_index_guard.py b/nemo_retriever/tests/test_lancedb_incomplete_index_guard.py deleted file mode 100644 index c0d65dab77..0000000000 --- a/nemo_retriever/tests/test_lancedb_incomplete_index_guard.py +++ /dev/null @@ -1,510 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. -# All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""The LanceDB writer must refuse to build an index that is missing rows. - -The property under test: a row must not reach the index without an embedding, -and a run that loses rows must fail rather than publish a short index and report -success. It holds for any cause, any backend, and any GPU - the writer sees only -the rows, not what produced them. - -Concretely, a row whose embedding is ``[]`` fails the write. The embed stage -emits that shape for a whole batch it could not embed. It used to be counted as -a wrong-length vector and dropped with a warning, so the run completed with a -short index. - -Three things are deliberately NOT covered here. - -* A wrong-length vector. That is a real vector against the wrong schema, which - is what ``on_bad_vectors`` exists for; folding it in would turn a configured - tolerance into a hard failure on upgrade. -* An absent or ``None`` embedding, which keeps its pre-existing silent drop - because ``operators/embed/text_embed.py`` writes ``None`` on purpose for a - blank-text row. -* A row the embedder zero-padded. It has the correct width and a non-zero - length, so the writer cannot see it. That half is marked at the source; see - ``test_vllm_embed.py``. - -Scope: this guard covers the LanceDB writer. Other paths in the repo still drop -or zero-fill rows of their own accord. - -Every test here is deterministic and needs no GPU. They assert on the write -contract rather than on end-to-end accuracy, which would drift with corpus size. -""" - -from __future__ import annotations - -from typing import Any - -import pandas as pd -import pytest - -from nemo_retriever.common.vdb.adt_vdb import CollectionWriteContext -from nemo_retriever.common.vdb.lancedb import _create_lancedb_results -from nemo_retriever.common.vdb.lancedb_collections import _collection_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_an_empty_embedding_fails_the_run() -> None: - """Known-bad: returned ``(rows, counts)`` and only logged a WARNING. - - ``embedding: []`` is what ``embed_text_main_text_embed`` writes for every - row of a batch it could not embed. It was not counted before this change: - it is not ``None``, and it passed the length check as a wrong-length vector. - """ - records = [[_record([]), _record([1.0, 2.0])]] - - with pytest.raises(RuntimeError) as excinfo: - _create_lancedb_results(records, expected_dim=2) - - message = str(excinfo.value) - assert "Refusing to build an incomplete index" in message - assert "1 of 2 rows" in message - assert "empty_embedding=1" in message - - -def test_create_lancedb_results_rejects_empty_embeddings_without_length_check() -> None: - """The non-enforcing path must not write empty vectors into the index. - - Known-bad: with ``expected_dim=None`` an ``embedding: []`` row was accepted - and counted in ``accepted``. - """ - with pytest.raises(RuntimeError, match="empty_embedding=1"): - _create_lancedb_results([[_record([])]], expected_dim=None) - - -@pytest.mark.parametrize("on_bad_vectors", ["drop", "fill", "error"]) -def test_wrong_length_rows_stay_under_the_on_bad_vectors_policy(on_bad_vectors: str) -> None: - """A short vector is a schema mismatch, not a missing embedding. - - ``on_bad_vectors`` is a documented, user-configured tolerance - (``common/vdb/lancedb.py`` ``create_index``). The incomplete-index guard must - not reach into it, or a user who deliberately configured ``drop`` or ``fill`` - would go from silent dropping to a hard run failure on upgrade. - - Known-bad for the first revision of this fix, which folded - ``dropped_bad_length`` into the fatal condition and raised for all three - values. It pins a contract rather than a code change, so it is a guard test: - the drop it asserts is identical before and after. It cannot run on the - unpatched tree, because the final assertion reads the new - ``empty_embedding`` key. - """ - records = [[_record([1.0]), _record([1.0, 2.0])]] - - # ``expected_dim=None`` is the shape ``create_index`` uses when the caller - # asked LanceDB to own the policy (``on_bad_vectors="error"``) or turned the - # wrapper's length check off. - rows, counts = _create_lancedb_results(records, expected_dim=None if on_bad_vectors == "error" else 2) - - if on_bad_vectors == "error": - # The wrapper forwards both rows so LanceDB itself raises, per the - # documented strict-fail semantics of that policy. - assert len(rows) == 2 - assert counts["dropped_bad_length"] == 0 - else: - assert len(rows) == 1 - assert counts["dropped_bad_length"] == 1 - assert counts["dropped_no_embedding"] == 0 - assert counts["empty_embedding"] == 0 - - -def test_empty_embeddings_from_the_endpoint_path_do_not_reach_the_index_silently( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """The writer is the backstop for producers that legitimately keep going. - - A local engine failure now propagates from - ``models/inference/runtime.py`` and never reaches here. The endpoint path - keeps its per-batch resilience by design - a single failed HTTP call should - not kill a service-mode run - so it can still emit ``embedding: []`` rows. - This is the case that makes the writer-side check load-bearing rather than - redundant. - - Known-bad: the writer dropped those rows and returned normally. - """ - from nemo_retriever.models.inference import runtime - - def _refuse(*_args: Any, **_kwargs: Any) -> pd.DataFrame: - raise TimeoutError("read timed out waiting for the embedding endpoint") - - monkeypatch.setattr(runtime, "_embed_group", _refuse) - - batch_df = pd.DataFrame( - { - "text": ["page one", "page two"], - "metadata": [ - {"content": "page one", "content_metadata": {"page_number": 1}, "source_metadata": {}}, - {"content": "page two", "content_metadata": {"page_number": 2}, "source_metadata": {}}, - ], - } - ) - - out_df = runtime.embed_text_main_text_embed( - batch_df, - embedding_endpoint="http://embed.example/v1", - inference_batch_size=2, - ) - - # The stage returns successfully - deliberate for the endpoint path - but - # every row is empty. - assert list(out_df["text_embeddings_1b_v2_has_embedding"]) == [False, False] - assert [payload["embedding"] for payload in out_df["text_embeddings_1b_v2"]] == [[], []] - - records = [ - [ - { - "document_type": "text", - "metadata": {**row["metadata"], "embedding": row["text_embeddings_1b_v2"]["embedding"]}, - } - for _index, row in out_df.iterrows() - ] - ] - - with pytest.raises(RuntimeError, match="Refusing to build an incomplete index"): - _create_lancedb_results(records, expected_dim=2048) - - -# -------------------------------------------------------------------------- -# No false positives: everything the fatal condition must NOT fire on. -# -# Catching the defect is half the evidence. The other half is that a normal run -# still finishes, so there is one case here per legitimate drop and per value -# type the new check could have swallowed by accident. -# -------------------------------------------------------------------------- - - -def test_wrong_length_is_counted_and_never_fatal() -> None: - """``dropped_bad_length`` stays out of the fatal condition. - - It is the category ``on_bad_vectors`` governs: a real vector against the - wrong schema. Folding it in is exactly the regression that would turn a - user's configured ``drop``/``fill``/``null`` tolerance into a hard failure - on upgrade. - - Cannot run on the unpatched tree: it asserts on the new - ``empty_embedding`` key, which does not exist there. What it pins is - a contract, not a behaviour change - the wrong-length row is dropped and not - raised on, before and after. - """ - rows, counts = _create_lancedb_results([[_record([1.0, 2.0, 3.0]), _record([1.0, 2.0])]], expected_dim=2) - - assert len(rows) == 1 - assert counts["dropped_bad_length"] == 1 - assert counts["empty_embedding"] == 0 - - -def test_canonical_image_row_without_text_is_accepted_not_dropped() -> None: - """The text carve-out for canonical image rows still works. - - An image row legitimately carries ``text=""``. It has a real embedding, so - nothing here may touch it. - """ - record = { - "document_type": "image", - "metadata": { - "embedding": [1.0, 2.0], - "content": "", - "content_metadata": {"page_number": 3, "type": "image"}, - "source_metadata": {"source_name": "scan.pdf"}, - }, - } - rows, counts = _create_lancedb_results([[record]], expected_dim=2) - - assert len(rows) == 1 - assert counts["accepted"] == 1 - assert counts["dropped_no_text"] == 0 - - -def test_text_free_non_image_row_is_dropped_and_never_fatal() -> None: - """``dropped_no_text`` stays out of the fatal condition. - - It is a content filter, not a loss: the row embedded successfully and was - excluded for having nothing to search on. - - Cannot run on the unpatched tree: it asserts on the new - ``empty_embedding`` key. The drop itself is unchanged. - """ - rows, counts = _create_lancedb_results([[_record([1.0, 2.0], text=" ")]], expected_dim=2) - - assert rows == [] - assert counts["dropped_no_text"] == 1 - 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: - """The check tests presence, never values. - - Deciding from the numbers would mean guessing which embeddings are "real", - and an all-zero vector is a legal thing for a model to emit. Only the - absence of a vector is fatal. - """ - rows, counts = _create_lancedb_results([[_record(embedding)]], expected_dim=2) - - assert len(rows) == 1 - assert counts["accepted"] == 1 - - -def test_numpy_embeddings_keep_their_existing_handling() -> None: - """The new check must not change how non-list types are treated. - - ``not embedding`` raises ``ValueError`` on a multi-element numpy array, so - the check is written as ``isinstance(embedding, (list, tuple)) and - len(embedding) == 0``. A numpy row therefore falls through to the - pre-existing length check and behaves exactly as it does today: counted as - ``dropped_bad_length``, not raised on, not crashed on. - - Not a known-bad test: it pins a hazard the obvious spelling of this check - would have introduced. It cannot run on the unpatched tree, which has no - ``empty_embedding`` key. - """ - numpy = pytest.importorskip("numpy") - - rows, counts = _create_lancedb_results( - [[_record(numpy.array([1.0, 2.0])), _record([1.0, 2.0])]], - expected_dim=2, - ) - - assert len(rows) == 1 - assert counts["dropped_bad_length"] == 1 - assert counts["empty_embedding"] == 0 - - -def test_an_empty_numpy_array_is_not_swallowed_by_the_new_check() -> None: - """An empty ndarray is not a list, so it keeps its pre-existing route. - - This is the narrowness of the check made explicit: it fires on ``[]`` and - ``()`` and nothing else. - """ - 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 - - -def test_a_fully_healthy_batch_raises_nothing() -> None: - """The whole point: a normal run is untouched. - - Cannot run on the unpatched tree, which has no ``empty_embedding`` - key; the acceptance of all 50 rows is identical there. - """ - records = [[_record([1.0, 2.0]) for _ in range(50)]] - - rows, counts = _create_lancedb_results(records, expected_dim=2) - - assert len(rows) == 50 - assert counts["accepted"] == 50 - assert counts["empty_embedding"] == 0 - assert counts["dropped_no_embedding"] == 0 - - -def test_a_none_embedding_keeps_its_pre_existing_silent_drop() -> None: - """``None`` is NOT fatal, because it has a legitimate producer. - - ``operators/embed/text_embed.py`` writes ``{"embedding": None}`` on purpose - for a row whose text was blank and which it therefore chose not to embed. - Making ``dropped_no_embedding`` fatal would fail ingests containing such - rows, which work today. Only ``[]`` - written solely on failure paths, at - ``models/inference/runtime.py`` and ``models/inference/vllm.py`` - is fatal. - - Cannot run on the unpatched tree, which has no ``empty_embedding`` - key. The drop it asserts is behaviour this change deliberately does not - touch. - """ - rows, counts = _create_lancedb_results([[_record(None), _record([1.0, 2.0])]], expected_dim=2) - - assert len(rows) == 1 - assert counts["dropped_no_embedding"] == 1 - assert counts["empty_embedding"] == 0 - - -# --------------------------------------------------------------------------- -# The collection-managed write path. -# -# ``common/vdb/lancedb_collections.py::_collection_rows`` is the second writer -# in the repo. It builds the rows for a collection-managed document ingest and -# had the same defect as ``_create_lancedb_results``: ``not vector`` is true for -# ``[]``, so a row the embed stage failed to embed was skipped with no counter, -# no log and no failure, and the document was published short while the ingest -# reported success. Same property, same fatal condition, same explicit -# ``isinstance``/``len`` spelling. -# --------------------------------------------------------------------------- - - -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", - "metadata": { - "embedding": embedding, - "content": text, - "content_metadata": {"type": "text", "page_number": 2}, - "source_metadata": {"source_id": "/inputs/source.pdf", "source_name": "source.pdf"}, - }, - } - - -def test_an_empty_embedding_on_the_collection_path_fails_the_write() -> None: - """Known-bad: returned the surviving rows and skipped the empty one silently. - - On the unpatched tree ``not vector`` swallowed ``[]`` into the same branch as - a malformed value, so this returned one row and no error, and the collection - document was written short. It now raises before any row reaches LanceDB. - - Fails on the unpatched tree: no exception is raised. - """ - records = [[_collection_record([]), _collection_record([1.0, 0.0])]] - - with pytest.raises(RuntimeError) as excinfo: - _collection_rows(records, context=_collection_context()) - - message = str(excinfo.value) - assert "incomplete document" in message - assert "empty_embedding=1" in message - - -def test_the_collection_path_reports_the_same_counter_name_as_the_pipeline_path() -> None: - """Both writers name the condition ``empty_embedding`` so one grep finds both. - - Fails on the unpatched tree: no exception, so nothing to read the name from. - """ - with pytest.raises(RuntimeError) as excinfo: - _collection_rows([[_collection_record([])]], context=_collection_context()) - - assert "empty_embedding" in str(excinfo.value) - - -def test_a_healthy_collection_document_still_writes_every_row() -> None: - """False-failure guard: the fatal branch must not fire on good input. - - Runs on the unpatched tree and passes there too, which is the point. - """ - records = [[_collection_record([1.0, 0.0]), _collection_record([0.0, 1.0], text="second chunk")]] - - rows = _collection_rows(records, context=_collection_context()) - - assert len(rows) == 2 - assert [row["text"] for row in rows] == ["first chunk", "second chunk"] - - -@pytest.mark.parametrize( - "embedding", - [ - None, - "not-a-vector", - 123, - {"dense": [1.0, 0.0]}, - ], - ids=["none", "string", "int", "dict"], -) -def test_a_malformed_collection_embedding_keeps_its_pre_existing_silent_skip(embedding: Any) -> None: - """False-failure guard, and the one that matters most. - - ``not isinstance(vector, (list, tuple))`` covers malformed values. It was - NOT made fatal: only ``[]`` was, because ``[]`` is the sole value written - exclusively on a failure path. Widening the fatal set to these would fail - ingests that work today. The healthy sibling row must still be written. - - Runs on the unpatched tree and passes there too. - """ - records = [[_collection_record(embedding), _collection_record([1.0, 0.0], text="good chunk")]] - - rows = _collection_rows(records, context=_collection_context()) - - assert len(rows) == 1 - assert rows[0]["text"] == "good chunk" - - -def test_a_malformed_collection_batch_or_record_keeps_its_pre_existing_silent_skip() -> None: - """False-failure guard for the batch- and record-shaped skips. - - ``not isinstance(batch, list)``, ``not isinstance(record, dict)`` and - ``not isinstance(metadata, dict)`` are untouched: they still skip silently - and never fail the write. - - Runs on the unpatched tree and passes there too. - """ - records = [ - "not-a-batch", - ["not-a-record"], - [{"document_type": "text", "metadata": "not-a-dict"}], - [_collection_record([1.0, 0.0], text="good chunk")], - ] - - rows = _collection_rows(records, context=_collection_context()) - - assert len(rows) == 1 - assert rows[0]["text"] == "good chunk" - - -def test_a_text_free_non_image_collection_row_is_skipped_and_never_fatal() -> None: - """False-failure guard: the content filter keeps its silent drop. - - Runs on the unpatched tree and passes there too. - """ - records = [[_collection_record([1.0, 0.0], text=" "), _collection_record([0.0, 1.0], text="good chunk")]] - - rows = _collection_rows(records, context=_collection_context()) - - assert len(rows) == 1 - assert rows[0]["text"] == "good chunk" - - -def test_a_numpy_collection_embedding_does_not_raise_on_truthiness() -> None: - """The new branch must never evaluate ``not vector`` on an array. - - A multi-element numpy array raises ``ValueError`` on ``bool()``. The check - is written as ``isinstance(vector, (list, tuple)) and len(vector) == 0``, so - an array short-circuits out of it and keeps its pre-existing skip rather - than crashing the ingest. - - Runs on the unpatched tree, where the existing ``or`` short-circuits for the - same reason, and passes there too. - """ - numpy = pytest.importorskip("numpy") - - 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()) - - assert len(rows) == 1 - assert rows[0]["text"] == "good chunk" diff --git a/nemo_retriever/tests/test_lancedb_write_policy.py b/nemo_retriever/tests/test_lancedb_write_policy.py index c96b47d4ff..afef136570 100644 --- a/nemo_retriever/tests/test_lancedb_write_policy.py +++ b/nemo_retriever/tests/test_lancedb_write_policy.py @@ -7,12 +7,14 @@ import json import logging from pathlib import Path +from typing import Any +import pandas as pd import pytest lancedb = pytest.importorskip("lancedb") -from nemo_retriever.common.vdb.lancedb import LanceDB +from nemo_retriever.common.vdb.lancedb import LanceDB, _create_lancedb_results def _records(text: str = "hello", vector: list[float] | None = None) -> list[list[dict]]: @@ -363,3 +365,247 @@ 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 == [] + +# --- incomplete-index guard: a row must not reach the index without an embedding --- +# ``[]`` is not ``None``, so it used to fall through to the length check, be counted a +# wrong-length vector, and be dropped - a short index published with exit 0. + + +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_an_empty_embedding_fails_the_run() -> None: + """Known-bad: returned ``(rows, counts)`` and only logged a WARNING. + + ``embedding: []`` is what ``embed_text_main_text_embed`` writes for every + row of a batch it could not embed. It was not counted before this change: + it is not ``None``, and it passed the length check as a wrong-length vector. + """ + records = [[_record([]), _record([1.0, 2.0])]] + + with pytest.raises(RuntimeError) as excinfo: + _create_lancedb_results(records, expected_dim=2) + + message = str(excinfo.value) + assert "Refusing to build an incomplete index" in message + assert "1 of 2 rows" in message + assert "empty_embedding=1" in message + + +def test_create_lancedb_results_rejects_empty_embeddings_without_length_check() -> None: + """The non-enforcing path must not write empty vectors into the index. + + Known-bad: with ``expected_dim=None`` an ``embedding: []`` row was accepted + and counted in ``accepted``. + """ + with pytest.raises(RuntimeError, match="empty_embedding=1"): + _create_lancedb_results([[_record([])]], expected_dim=None) + + +@pytest.mark.parametrize("on_bad_vectors", ["drop", "fill", "error"]) +def test_wrong_length_rows_stay_under_the_on_bad_vectors_policy(on_bad_vectors: str) -> None: + """A short vector is a schema mismatch, not a missing embedding. + + ``on_bad_vectors`` is a documented, user-configured tolerance + (``common/vdb/lancedb.py`` ``create_index``). The incomplete-index guard must + not reach into it, or a user who deliberately configured ``drop`` or ``fill`` + would go from silent dropping to a hard run failure on upgrade. + + Known-bad for the first revision of this fix, which folded + ``dropped_bad_length`` into the fatal condition and raised for all three + values. It pins a contract rather than a code change, so it is a guard test: + the drop it asserts is identical before and after. It cannot run on the + unpatched tree, because the final assertion reads the new + ``empty_embedding`` key. + """ + records = [[_record([1.0]), _record([1.0, 2.0])]] + + # ``expected_dim=None`` is the shape ``create_index`` uses when the caller + # asked LanceDB to own the policy (``on_bad_vectors="error"``) or turned the + # wrapper's length check off. + rows, counts = _create_lancedb_results(records, expected_dim=None if on_bad_vectors == "error" else 2) + + if on_bad_vectors == "error": + # The wrapper forwards both rows so LanceDB itself raises, per the + # documented strict-fail semantics of that policy. + assert len(rows) == 2 + assert counts["dropped_bad_length"] == 0 + else: + assert len(rows) == 1 + assert counts["dropped_bad_length"] == 1 + assert counts["dropped_no_embedding"] == 0 + assert counts["empty_embedding"] == 0 + + +def test_empty_embeddings_from_the_endpoint_path_do_not_reach_the_index_silently( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The writer is the backstop for producers that legitimately keep going. + + A local engine failure now propagates from + ``models/inference/runtime.py`` and never reaches here. The endpoint path + keeps its per-batch resilience by design - a single failed HTTP call should + not kill a service-mode run - so it can still emit ``embedding: []`` rows. + This is the case that makes the writer-side check load-bearing rather than + redundant. + + Known-bad: the writer dropped those rows and returned normally. + """ + from nemo_retriever.models.inference import runtime + + def _refuse(*_args: Any, **_kwargs: Any) -> pd.DataFrame: + raise TimeoutError("read timed out waiting for the embedding endpoint") + + monkeypatch.setattr(runtime, "_embed_group", _refuse) + + batch_df = pd.DataFrame( + { + "text": ["page one", "page two"], + "metadata": [ + {"content": "page one", "content_metadata": {"page_number": 1}, "source_metadata": {}}, + {"content": "page two", "content_metadata": {"page_number": 2}, "source_metadata": {}}, + ], + } + ) + + out_df = runtime.embed_text_main_text_embed( + batch_df, + embedding_endpoint="http://embed.example/v1", + inference_batch_size=2, + ) + + # The stage returns successfully - deliberate for the endpoint path - but + # every row is empty. + assert list(out_df["text_embeddings_1b_v2_has_embedding"]) == [False, False] + assert [payload["embedding"] for payload in out_df["text_embeddings_1b_v2"]] == [[], []] + + records = [ + [ + { + "document_type": "text", + "metadata": {**row["metadata"], "embedding": row["text_embeddings_1b_v2"]["embedding"]}, + } + for _index, row in out_df.iterrows() + ] + ] + + with pytest.raises(RuntimeError, match="Refusing to build an incomplete index"): + _create_lancedb_results(records, expected_dim=2048) + + +def test_canonical_image_row_without_text_is_accepted_not_dropped() -> None: + """The text carve-out for canonical image rows still works. + + An image row legitimately carries ``text=""``. It has a real embedding, so + nothing here may touch it. + """ + record = { + "document_type": "image", + "metadata": { + "embedding": [1.0, 2.0], + "content": "", + "content_metadata": {"page_number": 3, "type": "image"}, + "source_metadata": {"source_name": "scan.pdf"}, + }, + } + rows, counts = _create_lancedb_results([[record]], expected_dim=2) + + assert len(rows) == 1 + assert counts["accepted"] == 1 + assert counts["dropped_no_text"] == 0 + + +def test_text_free_non_image_row_is_dropped_and_never_fatal() -> None: + """``dropped_no_text`` stays out of the fatal condition. + + It is a content filter, not a loss: the row embedded successfully and was + excluded for having nothing to search on. + + Cannot run on the unpatched tree: it asserts on the new + ``empty_embedding`` key. The drop itself is unchanged. + """ + rows, counts = _create_lancedb_results([[_record([1.0, 2.0], text=" ")]], expected_dim=2) + + assert rows == [] + assert counts["dropped_no_text"] == 1 + 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: + """The check tests presence, never values. + + Deciding from the numbers would mean guessing which embeddings are "real", + and an all-zero vector is a legal thing for a model to emit. Only the + absence of a vector is fatal. + """ + rows, counts = _create_lancedb_results([[_record(embedding)]], expected_dim=2) + + assert len(rows) == 1 + assert counts["accepted"] == 1 + + +def test_an_empty_numpy_array_is_not_swallowed_by_the_new_check() -> None: + """An empty ndarray is not a list, so it keeps its pre-existing route. + + This is the narrowness of the check made explicit: it fires on ``[]`` and + ``()`` and nothing else. + """ + 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 + + +def test_a_fully_healthy_batch_raises_nothing() -> None: + """The whole point: a normal run is untouched. + + Cannot run on the unpatched tree, which has no ``empty_embedding`` + key; the acceptance of all 50 rows is identical there. + """ + records = [[_record([1.0, 2.0]) for _ in range(50)]] + + rows, counts = _create_lancedb_results(records, expected_dim=2) + + assert len(rows) == 50 + assert counts["accepted"] == 50 + assert counts["empty_embedding"] == 0 + assert counts["dropped_no_embedding"] == 0 + + +def test_a_none_embedding_keeps_its_pre_existing_silent_drop() -> None: + """``None`` is NOT fatal, because it has a legitimate producer. + + ``operators/embed/text_embed.py`` writes ``{"embedding": None}`` on purpose + for a row whose text was blank and which it therefore chose not to embed. + Making ``dropped_no_embedding`` fatal would fail ingests containing such + rows, which work today. Only ``[]`` - written solely on failure paths, at + ``models/inference/runtime.py`` and ``models/inference/vllm.py`` - is fatal. + + Cannot run on the unpatched tree, which has no ``empty_embedding`` + key. The drop it asserts is behaviour this change deliberately does not + touch. + """ + rows, counts = _create_lancedb_results([[_record(None), _record([1.0, 2.0])]], expected_dim=2) + + assert len(rows) == 1 + assert counts["dropped_no_embedding"] == 1 + assert counts["empty_embedding"] == 0 diff --git a/nemo_retriever/tests/test_multimodal_embed.py b/nemo_retriever/tests/test_multimodal_embed.py index cec96cd989..6730521a2a 100644 --- a/nemo_retriever/tests/test_multimodal_embed.py +++ b/nemo_retriever/tests/test_multimodal_embed.py @@ -15,9 +15,12 @@ import pandas as pd import pytest +from nemo_retriever.models.embed_errors import LocalEmbedderRowsLostError + # --------------------------------------------------------------------------- # Pure helpers from main_text_embed (no transitive-import issues) # --------------------------------------------------------------------------- +from nemo_retriever.models.inference import main_text_embed from nemo_retriever.models.inference.main_text_embed import ( _format_image_input_string, _format_text_image_pair_input_string, @@ -353,3 +356,110 @@ def test_non_dataframe_passthrough(self): """Non-DataFrame input is returned as-is.""" result = collapse_content_to_page_rows(None) assert result is None + +# --- multimodal row-loss guard --- +# A short answer used to be padded with ``None`` for the shortfall, and ``None`` is the +# one shape both writers ignore, so those rows were dropped and the run still succeeded. +# The guard counts rows submitted *with an image*, not rows in the chunk: embedders drop +# empty entries before inference, so comparing against chunk size would fail image-free chunks. + + +class _StubVLEmbedder: + """VL embedder stub whose per-call return length is scripted by the test.""" + + def __init__(self, *, images=None, text_image=None, text=None): + self._images = images + self._text_image = text_image + self._text = text + + def embed_images(self, images_b64, *, batch_size=64): + return self._images + + def embed_text_image(self, texts, images_b64, *, batch_size=64): + return self._text_image + + def embed(self, texts, *, batch_size=64): + return self._text + + +def _image_frame(image_values): + return pd.DataFrame({"_image_b64": list(image_values), "text": [""] * len(image_values)}) + + +def test_image_mode_short_answer_is_fatal(): + """A row submitted with an image and answered with nothing must fail the run. + + Without the guard the shortfall is padded with ``None``, which both writers + ignore, so the row is dropped and the run still succeeds. + """ + df = _image_frame(["b64-a", "b64-b", "b64-c"]) + embedder = _StubVLEmbedder(images=[[0.1, 0.2]]) # 1 vector for 3 images + + with pytest.raises(LocalEmbedderRowsLostError) as excinfo: + main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") + assert excinfo.value.lost == 2 + assert excinfo.value.total == 3 + + +def test_image_mode_empty_answer_is_fatal(): + """Zero vectors for a chunk that did submit images is the whole-batch failure.""" + df = _image_frame(["b64-a", "b64-b"]) + embedder = _StubVLEmbedder(images=[]) + + with pytest.raises(LocalEmbedderRowsLostError): + main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") + + +def test_text_image_mode_short_multimodal_answer_is_fatal(): + df = pd.DataFrame({"_image_b64": ["b64-a", "b64-b"], "text": ["alpha", "beta"]}) + embedder = _StubVLEmbedder(text_image=[[0.1, 0.2]]) # 1 vector for 2 paired rows + + with pytest.raises(LocalEmbedderRowsLostError): + main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="text_image") + + +def test_text_image_mode_short_text_fallback_answer_is_fatal(): + """The text-only fallback subset is owed one vector per row as well.""" + df = pd.DataFrame({"_image_b64": ["", ""], "text": ["alpha", "beta"]}) + embedder = _StubVLEmbedder(text=[[0.1, 0.2]]) # 1 vector for 2 fallback rows + + with pytest.raises(LocalEmbedderRowsLostError): + main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="text_image") + + +def test_image_mode_rows_without_images_do_not_fire_the_guard(): + """Rows with no image are owed nothing; they get ``None`` by contract. + + This is the false-failure the naive ``if not vecs_list: raise`` would cause. + """ + df = _image_frame(["b64-a", "", "b64-c"]) + embedder = _StubVLEmbedder(images=[[0.1, 0.2], [0.3, 0.4]]) # 2 vectors for 2 images + + out = main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") + assert out["embeddings"] == [[0.1, 0.2], None, [0.3, 0.4]] + + +def test_image_mode_chunk_with_no_images_at_all_does_not_fire_the_guard(): + """An image-free chunk submits nothing, so zero vectors back is correct.""" + df = _image_frame(["", ""]) + embedder = _StubVLEmbedder(images=[]) + + out = main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") + assert out["embeddings"] == [None, None] + + +def test_text_image_mode_mixed_rows_do_not_fire_the_guard(): + """Paired, text-only and empty rows each get their own correct treatment.""" + df = pd.DataFrame({"_image_b64": ["b64-a", "", ""], "text": ["alpha", "beta", " "]}) + embedder = _StubVLEmbedder(text_image=[[0.1, 0.2]], text=[[0.3, 0.4]]) + + out = main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="text_image") + assert out["embeddings"] == [[0.1, 0.2], [0.3, 0.4], None] + + +def test_image_mode_healthy_batch_does_not_fire_the_guard(): + df = _image_frame(["b64-a", "b64-b"]) + embedder = _StubVLEmbedder(images=[[0.1, 0.2], [0.3, 0.4]]) + + out = main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") + assert out["embeddings"] == [[0.1, 0.2], [0.3, 0.4]] From 698de561e0b24896351d3fc11e6aeaab2b35ea53 Mon Sep 17 00:00:00 2001 From: hrong Date: Wed, 19 Aug 2026 19:01:20 -0700 Subject: [PATCH 3/4] fix(embed): align guard scope and cardinality handling --- docs/docs/extraction/troubleshoot.md | 48 +-- .../src/nemo_retriever/common/vdb/README.md | 4 +- .../src/nemo_retriever/common/vdb/lancedb.py | 13 +- .../common/vdb/lancedb_collections.py | 43 +-- .../src/nemo_retriever/models/embed_errors.py | 18 +- .../models/inference/main_text_embed.py | 35 ++- .../llama_nemotron_embed_1b_v2_embedder.py | 5 +- .../llama_nemotron_embed_vl_1b_v2_embedder.py | 10 +- .../test_embed_engine_failure_propagation.py | 176 +---------- .../tests/test_lancedb_collections.py | 63 ++-- .../tests/test_lancedb_write_policy.py | 281 +++--------------- nemo_retriever/tests/test_multimodal_embed.py | 86 +++--- nemo_retriever/tests/test_vllm_embed.py | 135 +-------- 13 files changed, 238 insertions(+), 679 deletions(-) diff --git a/docs/docs/extraction/troubleshoot.md b/docs/docs/extraction/troubleshoot.md index 5e6343b6da..0df5ab85d1 100644 --- a/docs/docs/extraction/troubleshoot.md +++ b/docs/docs/extraction/troubleshoot.md @@ -38,27 +38,39 @@ configured invoke URL: Page Elements, OCR, Table Structure, Nemotron Parse, and embedding. It does not automatically raise for: - Local-only pipelines (`pdfium` without remote URLs), even when rows contain - `metadata.error` or column-level error payloads. One exception: an in-process - vLLM embedding engine that has stopped serving - refused at startup for lack of - free GPU memory, dead after a crash, or returning no vectors - aborts the - ingest under every error policy. Such an engine produces no embedding for every - batch it is handed, and those rows are then excluded at the writer, so - continuing would publish an index covering only part of the corpus while the - run reported success. The error message names the knobs to change. + `metadata.error` or column-level error payloads. One exception applies to the + default local embedding runtime: an in-process vLLM embedding engine that is + refused at startup, dies after a crash, or returns no vectors aborts the ingest + under every error policy. Continuing on that route would publish an index that + covers only part of the corpus while the run reports success. The error message + names the settings to change. + + The Designer **Text Embedder** component uses a different handler. It converts + 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, but - if the batch is lost its rows reach the writer with no embedding and the ingest - fails there. 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 whose - embedding failed used to be counted as a wrong-length vector and silently - excluded; it now fails the run. Fix the embed stage - no policy value restores - the old behaviour. + 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. + + 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. - Caption or remote VLM stages. Missing credentials fail at actor setup; inference failures can abort the entire ingest. @@ -102,7 +114,7 @@ troubleshooting path. A single document can pass through several stages. | `ExtractParams(method="ocr")` | Page rendering, Page Elements, and the local or remote OCR backend | Missing local model dependencies, invalid image payload, authentication/transport status, or OCR row-level failure | | `ExtractParams(method="nemotron_parse")` | PDF rendering and local Nemotron Parse model or configured Nemotron Parse NIM | Missing `open_clip`, missing local model configuration, unsupported image input, or Nemotron Parse row-level/HTTP failure | | `.caption(...)` | Local caption model or remote VLM endpoint | `ValueError` at setup when credentials or endpoint/protocol are invalid; remote inference failures can abort the whole ingest rather than populate a row error column | -| `.embed(...)` | Local embedding model or configured embedding NIM | Model/dependency error, input-size or schema rejection, authentication/transport status, or embedding row-level failure; `GraphIngestionError` when a remote embed URL is configured. An in-process vLLM embedding engine that has stopped serving always aborts the ingest, whatever the error policy | +| `.embed(...)` | Local embedding model or configured embedding NIM | Model/dependency error, input-size or schema rejection, authentication/transport status, or embedding row-level failure; `GraphIngestionError` when a remote embed URL is configured. On the default local embedding runtime, an in-process vLLM engine that has stopped serving aborts the ingest regardless of the error policy. The Designer **Text Embedder** component retains the `None` gap described above | | Audio or video extraction | `ffmpeg`/`ffprobe`, media decoding, frame/chunk creation, and local or remote ASR | Missing executable, malformed media, codec failure, gRPC status, or credential error; ASR failures may omit rows and log warnings instead of raising, so verify logs when output is unexpectedly empty | `pdfium` itself is primarily a local parser, so a Page Elements, Table diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/README.md b/nemo_retriever/src/nemo_retriever/common/vdb/README.md index 5372e842af..bea3c9d4d5 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/README.md +++ b/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** embedding, `[]`, raises `RuntimeError` and no table is written. `[]` is written only on embed failure paths, so it means no vector was produced for that 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 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`. -Only `[]` is fatal. A row whose embedding is absent or `None` keeps its pre-existing silent drop, because `operators/embed/text_embed.py` writes `None` on purpose for a blank-text row. `on_bad_vectors` is unaffected: it governs malformed vectors, and an empty list is the absence of a vector rather than a short one. The returned `counts` dict gains an `empty_embedding` key. +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 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. diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py index d10ee2a571..e519365685 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py @@ -366,13 +366,12 @@ def _create_lancedb_results( ``(rows, counts)`` where ``rows`` is the list of dicts shaped for LanceDB ingestion (``vector``, ``text``, ``metadata``, ``source``) and ``counts`` is a dict containing ``accepted``, - ``dropped_no_embedding``, ``dropped_bad_length``, and - ``dropped_no_text`` keys. - Also ``empty_embedding``, added by the incomplete-index guard. + ``dropped_no_embedding``, ``empty_embedding``, + ``dropped_bad_length``, and ``dropped_no_text`` keys. - An empty embedding, ``[]``, means the embed stage produced no vector for that - row. Such rows are counted as ``empty_embedding`` and make this function - raise :class:`RuntimeError` after the loop, so no table is written. + An empty list or tuple embedding means the embed stage produced no vector for + that row. Such rows are counted as ``empty_embedding`` and make this function + raise :class:`RuntimeError` after the loop, so no table rows are written. :meth:`LanceDB.create_index` calls this function twice when ``vector_dim`` is ``None``: pass 1 infers the dimension and its rows are discarded, pass 2 @@ -477,7 +476,7 @@ def _create_lancedb_results( total = accepted + dropped_no_embedding + empty_embedding + dropped_bad_length + dropped_no_text raise RuntimeError( "Refusing to build an incomplete index: " - f"{empty_embedding} of {total} rows had no embedding. No table is written and " + f"{empty_embedding} of {total} rows had no embedding. No table rows are written and " "this run fails; the alternative is an index that is silently short by those " "rows while the run reports success. " f"Counters: empty_embedding={empty_embedding}, no_embedding={dropped_no_embedding}, " diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py index e020b9a37d..5003c5fdb3 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py +++ b/nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py @@ -170,20 +170,27 @@ def _collection_rows( ) -> list[dict[str, Any]]: """Convert canonical NRL record batches into collection-managed LanceDB rows. - An empty embedding, ``[]``, means the embed stage produced no vector for that - row. Such rows are counted as ``empty_embedding`` and make this function raise - :class:`RuntimeError` after the loop, so no rows are written for the document. + An empty list or tuple embedding means the embed stage produced no vector for + that row. Such rows are counted as ``empty_embedding`` and make this function + raise :class:`RuntimeError` after the loop, so no rows are written for the + document. This is the same guarantee :func:`nemo_retriever.common.vdb.lancedb._create_lancedb_results` - carries on the pipeline write path: a row must not reach the index without an - embedding, and a write that loses rows must fail rather than publish a short - document while reporting success. - - Every other skip keeps its pre-existing silent behaviour and is counted only - as ``skipped_other``: a malformed batch, record, metadata or embedding value, - and a text-free non-image row, are all still dropped without failing the - write. Only ``[]`` is fatal, because it is the only value with no legitimate - producer - the embed stage writes it to mean "this output carried no - embedding". + carries on the pipeline write path for this representation: an empty list or + tuple must not reach the index, and its presence fails the write rather than + publishing a short document while reporting success. + + Every other row-level skip remains non-fatal and is + counted only as ``skipped_other``: a malformed record, metadata or embedding + value, and a text-free non-image row, are all still dropped without failing + the write. A malformed outer batch is ignored but is not included in the row + counters. Only an empty list or tuple + is fatal. The embed stage writes ``[]`` to mean "this output carried no + embedding". ``None`` can represent either a row the embed stage chose not to + embed or a chunk whose embedding failed. Those meanings are distinguishable + by the ``error`` key that travels with the payload - + ``None`` for the deliberate skip, a populated dict for the failure - but that + key is not visible at this layer, so ``None`` keeps its silent drop until the + discriminator is threaded down. That is follow-up work, not an impossibility. """ rows: list[dict[str, Any]] = [] created_at = _now() @@ -203,11 +210,11 @@ def _collection_rows( skipped_other += 1 continue vector = metadata.get("embedding") - # ``[]`` is folded into the skip below by ``not vector``, which drops - # it silently. Check it first, with an explicit isinstance/len and + # Empty lists and tuples are folded into the skip below by ``not vector``, + # which drops them silently. Check them first with ``isinstance``/``len``; # never ``not vector``: the latter raises ValueError on a - # multi-element array. The skip below is left exactly as it was, so - # every value other than ``[]`` keeps its current route. + # multi-element array. The skip below is left exactly as it was, so all + # other values keep their current route. if isinstance(vector, (list, tuple)) and len(vector) == 0: empty_embedding += 1 logger.debug("Row has an empty embedding (document_id=%s)", context.document_id) @@ -282,7 +289,7 @@ def _collection_rows( f"accepted={len(rows)}, document_id={context.document_id}, " f"document_version={context.document_version}. " "Only empty_embedding caused this failure - skipped_other counts rows filtered " - "under the pre-existing rules. This normally means the embed stage failed for " + "under the non-fatal row rules. This normally means the embed stage failed for " "whole batches: check the embed actor logs for engine initialization or " "out-of-memory errors." ) diff --git a/nemo_retriever/src/nemo_retriever/models/embed_errors.py b/nemo_retriever/src/nemo_retriever/models/embed_errors.py index c32ce224ae..005724514c 100644 --- a/nemo_retriever/src/nemo_retriever/models/embed_errors.py +++ b/nemo_retriever/src/nemo_retriever/models/embed_errors.py @@ -40,11 +40,10 @@ def __init__(self, *, lost: int, total: int, embedder: str) -> None: self.embedder = str(embedder) super().__init__( f"{embedder} returned no vector for {lost} of {total} row(s) in this batch. " - "Those rows would be zero-padded to the right width, so nothing downstream could " - "tell them apart from real embeddings by shape alone, and indexing them would " - "publish rows that match nothing. This normally means the in-process engine " - "failed for the batch - check the embed actor logs for engine initialization or " - "out-of-memory errors." + "Continuing would pad or drop those rows, hide the loss, and allow an invalid or " + "incomplete index to be published. This normally means the in-process engine failed " + "for the batch - check the embed actor logs for engine initialization or out-of-memory " + "errors." ) @@ -66,11 +65,10 @@ def report_lost_rows(vectors: Sequence[Sequence[float]], *, embedder: str) -> in Returns ``0`` when nothing was lost; never returns a non-zero count. - Called from the local embedders' ``_finalize_vectors``, the only place that - holds both the batch that was sent and the vectors that came back. It used to - zero-pad the missing rows, which made them indistinguishable downstream: a - padded row has the right width, so ``has_embedding`` reported ``True`` for a - row carrying nothing. + Called from the local embedders' ``_finalize_vectors`` before they discard + empty placeholders. Raising before padding is required because a padded row + has the right width, so ``has_embedding`` would report ``True`` for a row + carrying nothing. """ lost = sum(1 for vector in vectors if _has_no_vector(vector)) if not lost: diff --git a/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py b/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py index 61c49ad9e2..6da99b31a2 100644 --- a/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py +++ b/nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py @@ -176,6 +176,16 @@ def _format_text_image_pair_input_string(text: str, image_b64: str, mime: str = return f"{text}\n{data_url}" +def _validate_vector_count(vectors: Sequence[Any], *, expected: int, embedder: str) -> None: + """Raise the appropriate error when a local embedder violates cardinality.""" + actual = len(vectors) + if actual == expected: + return + if actual < expected: + raise LocalEmbedderRowsLostError(lost=expected - actual, total=expected, embedder=embedder) + raise ValueError(f"{embedder} returned {actual} vectors for {expected} submitted input(s)") + + def _multimodal_callable_runner( df_slice: pd.DataFrame, *, @@ -219,13 +229,12 @@ def _multimodal_callable_runner( # image gets ``None`` by contract. Comparing against ``size`` would # false-fire on an image-free chunk. submitted = sum(1 for b64 in images_b64 if b64) - if len(vecs_list) != submitted: - raise LocalEmbedderRowsLostError( - lost=max(submitted - len(vecs_list), 0), - total=submitted, - embedder=type(embedder).__name__, - ) + _validate_vector_count(vecs_list, expected=submitted, embedder=type(embedder).__name__) + # Retained only for the no-blanks case: when the chunk has blank + # images ``submitted < size`` and the guard above already ensured + # one vector per submitted row, so the else-branch below is the + # only reachable route and produces the same list. if len(vecs_list) == size: flat_embeddings.extend(vecs_list) else: @@ -250,12 +259,7 @@ def _multimodal_callable_runner( # is the contract. A short answer means the engine lost rows; # padding them with ``None`` here would hide that, because the # writers ignore ``None``. - if len(mm_vecs_list) != len(mm_images): - raise LocalEmbedderRowsLostError( - lost=max(len(mm_images) - len(mm_vecs_list), 0), - total=len(mm_images), - embedder=type(embedder).__name__, - ) + _validate_vector_count(mm_vecs_list, expected=len(mm_images), embedder=type(embedder).__name__) # text-only fallback subset fb_texts = [t for t, h in zip(texts, has_image) if not h and t.strip()] @@ -264,12 +268,7 @@ def _multimodal_callable_runner( vecs = embedder.embed(fb_texts, batch_size=bs) tolist = getattr(vecs, "tolist", None) fb_vecs_list = tolist() if callable(tolist) else list(vecs) - if len(fb_vecs_list) != len(fb_texts): - raise LocalEmbedderRowsLostError( - lost=max(len(fb_texts) - len(fb_vecs_list), 0), - total=len(fb_texts), - embedder=type(embedder).__name__, - ) + _validate_vector_count(fb_vecs_list, expected=len(fb_texts), embedder=type(embedder).__name__) # reassemble in original order mm_iter = iter(mm_vecs_list) diff --git a/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_1b_v2_embedder.py b/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_1b_v2_embedder.py index b896e36111..3d9f86acfc 100644 --- a/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_1b_v2_embedder.py +++ b/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_1b_v2_embedder.py @@ -10,9 +10,9 @@ import torch -from nemo_retriever.models.hf_cache import configure_global_hf_cache_base from nemo_retriever.models.embed_errors import report_lost_rows from nemo_retriever.models.embed_model_spec import resolve_embed_model_revision +from nemo_retriever.models.hf_cache import configure_global_hf_cache_base def _l2_normalize(x: torch.Tensor, eps: float = 1e-12) -> torch.Tensor: @@ -85,8 +85,7 @@ def is_remote(self) -> bool: return False def _finalize_vectors(self, vectors: List[List[float]]) -> torch.Tensor: - # See the mirror of this function in - # ``models/local/llama_nemotron_embed_vl_1b_v2_embedder.py``. + """Reject empty rows before tensor conversion and normalization.""" report_lost_rows(vectors, embedder=type(self).__name__) valid = [v for v in vectors if v] if not valid: diff --git a/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_vl_1b_v2_embedder.py b/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_vl_1b_v2_embedder.py index 6e1be0a94c..86c803f01d 100644 --- a/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_vl_1b_v2_embedder.py +++ b/nemo_retriever/src/nemo_retriever/models/local/llama_nemotron_embed_vl_1b_v2_embedder.py @@ -10,10 +10,10 @@ import torch -from nemo_retriever.models.hf_cache import configure_global_hf_cache_base +from nemo_retriever.common.nvtx import gpu_inference_range from nemo_retriever.models.embed_errors import report_lost_rows from nemo_retriever.models.embed_model_spec import resolve_embed_model_revision -from nemo_retriever.common.nvtx import gpu_inference_range +from nemo_retriever.models.hf_cache import configure_global_hf_cache_base def _l2_normalize(x: torch.Tensor, eps: float = 1e-12) -> torch.Tensor: @@ -250,11 +250,7 @@ def is_remote(self) -> bool: return False def _finalize_vectors(self, vectors: Sequence[Sequence[float]]) -> torch.Tensor: - """Zero-pad rows vLLM failed to embed, then optionally normalize.""" - # Fail here when vLLM did not embed every row. This is the only place - # that knows the count; it used to zero-pad the missing rows, which made - # them indistinguishable from real embeddings downstream. Mirrored in - # ``models/local/llama_nemotron_embed_1b_v2_embedder.py``. + """Reject empty rows before tensor conversion and normalization.""" report_lost_rows(vectors, embedder=type(self).__name__) valid = [v for v in vectors if v] if not valid: diff --git a/nemo_retriever/tests/test_embed_engine_failure_propagation.py b/nemo_retriever/tests/test_embed_engine_failure_propagation.py index f7c8727b38..6b6b9a9856 100644 --- a/nemo_retriever/tests/test_embed_engine_failure_propagation.py +++ b/nemo_retriever/tests/test_embed_engine_failure_propagation.py @@ -2,41 +2,7 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""A run that loses embedding rows must fail, not publish a short index. - -Correctness comes from the LanceDB writer guard, which is general: any row, any -cause, any backend. It refuses to build an index when rows arrive with no -embedding. See ``test_lancedb_incomplete_index_guard.py``. - -This module tests the layer above it, which is narrow on purpose. -``embed_text_main_text_embed`` catches every exception from ``_embed_group`` and -returns ``{"embedding": [], "error": ...}`` for the whole batch. That per-batch -resilience is right for the endpoint path, where one HTTP call can fail alone. -It is wrong for an in-process engine that has stopped serving: every later batch -fails too, and a failed engine returns instantly, so it drains the queue far -faster than a healthy one. Aborting there is a fast-fail optimisation on top of -the writer guard - it turns a failure at the terminal write into one within -minutes. It is not what makes the result correct. - -Because it only buys latency, the fatal set stays limited to signals that mean -"the engine is not serving" and that a recoverable backend cannot produce. -Classifying an unmeasured exception would ship false failures. That is why a -bare ``OutOfMemoryError`` is excluded: the HuggingFace backend raises it from an -ordinary forward pass where a smaller next batch can succeed, and the two -backends are not distinguishable at the point of classification. A run that -loses rows to an OOM still fails - at the writer. - -``LocalEmbedderRowsLostError`` is the signal for *partial* loss and comes from -where the loss happens. The local embedders' ``_finalize_vectors`` is the only -function that knows how many rows failed to embed; it used to zero-pad them and -discard the count, which is why ``has_embedding`` could report ``True`` for a -row carrying nothing. The writer cannot detect a padded row, so neither layer -subsumes the other. - -The tests use stand-in exception classes rather than importing vLLM or torch, -matching the module under test, which identifies failures by class name and -message so it keeps working on the endpoint-only path. -""" +"""Runtime propagation tests for local embed-engine lifecycle failures.""" from __future__ import annotations @@ -50,6 +16,7 @@ LocalEmbedderRowsLostError, ) from nemo_retriever.models.inference import main_text_embed, runtime +from nemo_retriever.models.nim.error_reporter import drain_errors class EngineDeadError(Exception): @@ -76,6 +43,13 @@ class EngineGenerateError(Exception): ) +@pytest.fixture(autouse=True) +def _clear_reported_errors(): + drain_errors() + yield + drain_errors() + + def _batch(rows: int = 4) -> pd.DataFrame: return pd.DataFrame( { @@ -107,14 +81,13 @@ def _wrapped(outer: BaseException, cause: BaseException) -> BaseException: @pytest.mark.parametrize( - ("exc", "label"), + "exc", [ - pytest.param(RuntimeError(ENGINE_INIT_FAILED), "engine-core-init-failed", id="refused-at-startup"), - pytest.param(ValueError(ADMISSION_REFUSAL), "free-memory-gate", id="admission-gate-valueerror"), - pytest.param(EngineDeadError("EngineCore encountered an issue"), "engine-dead", id="dead-after-admission"), + pytest.param(RuntimeError(ENGINE_INIT_FAILED), id="refused-at-startup"), + pytest.param(ValueError(ADMISSION_REFUSAL), id="admission-gate-valueerror"), + pytest.param(EngineDeadError("EngineCore encountered an issue"), id="dead-after-admission"), pytest.param( _wrapped(RuntimeError("Worker proc died unexpectedly"), EngineDeadError("EngineCore is dead")), - "wrapped-engine-dead", id="engine-dead-wrapped-in-runtimeerror", ), pytest.param( @@ -122,18 +95,15 @@ def _wrapped(outer: BaseException, cause: BaseException) -> BaseException: "Local embedder returned no embeddings for a batch of 2 input(s). " "The in-process engine produced nothing, so it is not serving." ), - "returned-nothing", id="engine-answered-with-nothing", ), pytest.param( LocalEmbedderRowsLostError(lost=7, total=64, embedder="LlamaNemotronEmbedVL1BV2VLLMEmbedder"), - "rows-lost", id="engine-lost-part-of-the-batch", ), ], ) -def test_local_engine_failure_propagates(monkeypatch: pytest.MonkeyPatch, exc: BaseException, label: str) -> None: - """Known-bad: returns a full batch of ``embedding: []`` and the run continues.""" +def test_local_engine_failure_propagates(monkeypatch: pytest.MonkeyPatch, exc: BaseException) -> None: monkeypatch.setattr(runtime, "_embed_group", _raise(exc)) with pytest.raises(type(exc)): @@ -141,36 +111,13 @@ def test_local_engine_failure_propagates(monkeypatch: pytest.MonkeyPatch, exc: B def test_local_embedder_returning_nothing_is_raised_as_a_classified_failure() -> None: - """An all-empty local batch must not escape as a plain ``ValueError``. - - An embedder that answers a non-empty batch with nothing is not serving. - ``_callable_runner`` used to report that as a bare count mismatch, which - ``_is_engine_lifecycle_failure`` does not match, so the run continued and - emptied the index anyway. - - This is not a backstop for a hypothetical callable. It is the only layer - that sees a shipped failure: when vLLM yields no outputs for a batch, - ``embed_with_vllm_llm`` returns ``[]``, and ``_finalize_vectors`` counts no - loss because there are no rows to count - see - ``test_vllm_embed.py::test_finalize_vectors_cannot_see_a_zero_output_batch``. - That path raises nothing anywhere else. - - Known-bad: raises ``ValueError`` with "mismatched number of embeddings", - which the classifier rejects. - """ with pytest.raises(LocalEmbedderReturnedNothingError): main_text_embed._callable_runner([["page one", "page two"]], embedder=lambda _texts: [], batch_size=2) assert runtime._is_engine_lifecycle_failure(LocalEmbedderReturnedNothingError("no vectors")) -def test_partial_local_result_is_still_a_plain_value_error() -> None: - """Guard: a count mismatch that is not "nothing at all" stays unclassified. - - The embedder did produce vectors, so this is a data-shape problem rather - than a dead engine, and it must keep its per-batch handling. Passes before - and after. - """ +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) @@ -182,18 +129,7 @@ def test_partial_local_result_is_still_a_plain_value_error() -> None: "model", [pytest.param(None, id="model-nulled-by-the-operator"), pytest.param(object(), id="model-also-set")], ) -def test_endpoint_mode_still_absorbs_the_same_failure(monkeypatch: pytest.MonkeyPatch, model: object) -> None: - """Guard: endpoint mode has no in-process engine, so nothing changes there. - - Passes before and after. A service-mode user must keep per-batch resilience. - - Both parameters matter. In production the model is always ``None`` when an - endpoint is configured - ``operators/embed/cpu_operator.py:36`` and - ``operators/embed/gpu_operator.py:37`` null it - so the ``model is None`` - case alone would still pass if the ``endpoint is None`` term were dropped - from the re-raise condition. The second parameter pins that term directly, - as defence against a future change to those two actors. - """ +def test_endpoint_mode_absorbs_engine_lifecycle_failure(monkeypatch: pytest.MonkeyPatch, model: object) -> None: monkeypatch.setattr(runtime, "_embed_group", _raise(RuntimeError(ENGINE_INIT_FAILED))) out_df = runtime.embed_text_main_text_embed( @@ -208,75 +144,15 @@ def test_endpoint_mode_still_absorbs_the_same_failure(monkeypatch: pytest.Monkey def test_an_oom_alone_is_not_classified_as_an_engine_failure(monkeypatch: pytest.MonkeyPatch) -> None: - """A bare OOM is absorbed here; a run that loses rows still fails at the writer. - - The HuggingFace backend raises ``OutOfMemoryError`` from an ordinary forward - pass, where a smaller next batch can succeed. The backend is not visible at - the point of classification, so treating the exception as fatal would abort - on a per-batch condition. - - Absorbing it is not the same as the run surviving. The batch becomes - ``{"embedding": []}`` for every row, which the writer guard treats as fatal, - so a run that really lost those rows fails at the write. The second - assertion pins that shape, so this test cannot be read as a claim that the - run recovers. - """ 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) - # Claim 1: the runtime absorbed it rather than re-raising. assert list(out_df["text_embeddings_1b_v2_has_embedding"]) == [False] * 4 - - # Claim 2: what it absorbed into is the writer-fatal shape, so the run ends - # at the write, not here. assert [payload["embedding"] for payload in out_df["text_embeddings_1b_v2"]] == [[]] * 4 -def test_engine_death_after_an_oom_is_fatal_from_the_next_batch(monkeypatch: pytest.MonkeyPatch) -> None: - """Excluding the OOM costs one batch, not the run. - - An engine that dies mid-inference raises ``EngineDeadError`` from the next - batch onward, and that is fatal. This pins the cost of the OOM carve-out as - bounded rather than open-ended. - """ - calls: list[int] = [] - - def _fail_then_die(*_args: Any, **_kwargs: Any) -> pd.DataFrame: - calls.append(1) - if len(calls) == 1: - raise OutOfMemoryError(OOM_IN_GELU) - raise EngineDeadError("EngineCore encountered an issue") - - monkeypatch.setattr(runtime, "_embed_group", _fail_then_die) - - absorbed = runtime.embed_text_main_text_embed(_batch(), model=object(), inference_batch_size=2) - assert list(absorbed["text_embeddings_1b_v2_has_embedding"]) == [False] * 4 - - with pytest.raises(EngineDeadError): - runtime.embed_text_main_text_embed(_batch(), model=object(), inference_batch_size=2) - - -def test_local_batch_level_failure_is_still_absorbed(monkeypatch: pytest.MonkeyPatch) -> None: - """Guard: a failure that is not an engine-lifecycle failure stays non-fatal. - - Passes before and after. The change is deliberately narrow - only failures - that mean the engine stopped serving are fatal, because only those were - measured to affect every subsequent batch. - """ - monkeypatch.setattr(runtime, "_embed_group", _raise(ValueError("could not decode image payload for row 3"))) - - out_df = runtime.embed_text_main_text_embed(_batch(), model=object(), inference_batch_size=2) - - assert list(out_df["text_embeddings_1b_v2_has_embedding"]) == [False] * 4 - - def test_engine_lifecycle_classifier_rejects_unrelated_failures() -> None: - """Pins the classifier's specificity - it must not fire on ordinary failures. - - Cannot run on unmodified HEAD: the helper does not exist there, so this - fails with ``AttributeError`` rather than with a wrong answer. - """ 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. @@ -284,23 +160,3 @@ def test_engine_lifecycle_classifier_rejects_unrelated_failures() -> None: # 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)) - - -# --------------------------------------------------------------------------- -# Known gap: two embed entry points swallow ``LocalEmbedderRowsLostError``. -# -# Neither is on the shipped route, so these tests pin the gap rather than close -# it. They exist so a future import-site change that flips onto one of these -# functions shows up as a failing, clearly named test instead of as a silently -# short index in production. See ``models/embed_errors.py`` -# ``LocalEmbedderRowsLostError`` for the full analysis and for why the writer's -# fatal set must NOT be widened to cover ``None``. -# --------------------------------------------------------------------------- - - -# --- the guard must fire: the engine answered short, rows would be lost --- - - -# --- the guard must stay silent: these are legitimate short answers --- - - diff --git a/nemo_retriever/tests/test_lancedb_collections.py b/nemo_retriever/tests/test_lancedb_collections.py index 8e1826bcd6..91dd75a02b 100644 --- a/nemo_retriever/tests/test_lancedb_collections.py +++ b/nemo_retriever/tests/test_lancedb_collections.py @@ -1129,10 +1129,6 @@ def delete_target(): assert query_errors == [] assert delete_errors == [] -# --- incomplete-index guard: a row must not reach the index without an embedding --- -# ``[]`` is not ``None``, so it used to fall through to the length check, be counted a -# wrong-length vector, and be dropped - a short index published with exit 0. - def _collection_context() -> CollectionWriteContext: return CollectionWriteContext( @@ -1159,16 +1155,9 @@ def _collection_record(embedding: Any, *, text: str = "first chunk") -> dict: } -def test_an_empty_embedding_on_the_collection_path_fails_the_write() -> None: - """Known-bad: returned the surviving rows and skipped the empty one silently. - - On the unpatched tree ``not vector`` swallowed ``[]`` into the same branch as - a malformed value, so this returned one row and no error, and the collection - document was written short. It now raises before any row reaches LanceDB. - - Fails on the unpatched tree: no exception is raised. - """ - records = [[_collection_record([]), _collection_record([1.0, 0.0])]] +@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: + records = [[_collection_record(empty_embedding), _collection_record([1.0, 0.0])]] with pytest.raises(RuntimeError) as excinfo: _collection_rows(records, context=_collection_context()) @@ -1178,25 +1167,45 @@ def test_an_empty_embedding_on_the_collection_path_fails_the_write() -> None: assert "empty_embedding=1" in message -def test_the_collection_path_reports_the_same_counter_name_as_the_pipeline_path() -> None: - """Both writers name the condition ``empty_embedding`` so one grep finds both. +@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"), + ] + ] - Fails on the unpatched tree: no exception, so nothing to read the name from. - """ with pytest.raises(RuntimeError) as excinfo: - _collection_rows([[_collection_record([])]], context=_collection_context()) + _collection_rows(records, context=_collection_context()) - assert "empty_embedding" in str(excinfo.value) + message = str(excinfo.value) + assert "empty_embedding=1" in message + assert "skipped_other=1" in message + assert "accepted=1" in message -def test_a_healthy_collection_document_still_writes_every_row() -> None: - """False-failure guard: the fatal branch must not fire on good input. +def test_a_numpy_collection_embedding_does_not_raise_on_truthiness() -> None: + numpy = pytest.importorskip("numpy") - Runs on the unpatched tree and passes there too, which is the point. - """ - records = [[_collection_record([1.0, 0.0]), _collection_record([0.0, 1.0], text="second chunk")]] + 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()) - assert len(rows) == 2 - assert [row["text"] for row in rows] == ["first chunk", "second chunk"] + assert len(rows) == 1 + assert rows[0]["text"] == "good chunk" diff --git a/nemo_retriever/tests/test_lancedb_write_policy.py b/nemo_retriever/tests/test_lancedb_write_policy.py index afef136570..b828b22368 100644 --- a/nemo_retriever/tests/test_lancedb_write_policy.py +++ b/nemo_retriever/tests/test_lancedb_write_policy.py @@ -9,7 +9,6 @@ from pathlib import Path from typing import Any -import pandas as pd import pytest lancedb = pytest.importorskip("lancedb") @@ -284,61 +283,43 @@ def test_dense_write_requires_both_canonical_image_fields(tmp_path: Path, missin assert table_rows == [] -@pytest.mark.parametrize( - "vector", - [ - pytest.param([], id="empty-embed-failure"), - ], -) -def test_dense_write_fails_on_image_only_row_with_no_usable_embedding( - tmp_path: Path, vector: list[float] | None -) -> None: - """A row without a usable embedding must fail the write, not be dropped. - - The embed stage writes ``embedding: []`` for every row of a batch whose - engine failed. On the unpatched tree that row was accepted here, because - ``[]`` is not ``None`` and this write infers the dimension, so the length - check does not run - which is how a run publishes a short index and still - reports success. - - A ``None`` embedding is deliberately not covered: it has a legitimate - producer and keeps its pre-existing drop. See - ``test_lancedb_incomplete_index_guard.py``, - ``test_a_none_embedding_keeps_its_pre_existing_silent_drop``. - """ - with pytest.raises(RuntimeError, match="Refusing to build an incomplete index"): - _write_rows(tmp_path, _image_only_records(vector)) - - -def test_dense_write_still_drops_wrong_length_row_under_the_default_policy(tmp_path: Path) -> None: - """``on_bad_vectors="drop"`` keeps working: a short vector is dropped, not fatal. - - Known-bad for the first revision of this fix, which folded - ``dropped_bad_length`` into the fatal condition and made the shipped default - unreachable. Passes on unmodified HEAD too - it pins the documented contract - rather than a code change, so it is a guard test. - """ - table_rows = _write_rows(tmp_path, _image_only_records([1.0])) +@pytest.mark.parametrize("empty_embedding", [pytest.param([], id="list"), pytest.param((), id="tuple")]) +def test_dense_write_fails_on_image_only_row_with_an_empty_embedding(tmp_path: Path, empty_embedding: Any) -> None: + records = _image_only_records() + records[0][0]["metadata"]["embedding"] = empty_embedding + + with pytest.raises(RuntimeError) as excinfo: + _write_rows(tmp_path, records) + + message = str(excinfo.value) + assert "Refusing to build an incomplete index" in message + assert "empty_embedding=1" in message + assert "No table rows are written" in message + + +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")]] + + table_rows = _write_rows(tmp_path, records) + + assert len(table_rows) == 1 + assert table_rows[0]["text"] == "good row" + + +def test_dense_write_drops_an_image_only_row_with_no_embedding_key(tmp_path: Path) -> None: + table_rows = _write_rows(tmp_path, _image_only_records()) assert table_rows == [] -def test_dense_write_keeps_on_bad_vectors_fill_reachable(tmp_path: Path) -> None: - """``on_bad_vectors="fill"`` with the wrapper check off still reaches LanceDB. +def test_dense_write_drops_wrong_length_row_under_the_default_policy(tmp_path: Path) -> None: + table_rows = _write_rows(tmp_path, _image_only_records([1.0])) - With ``validate_vector_length=False`` the short row is forwarded and LanceDB - fills it, which is what a user who configured ``fill`` asked for. The guard - must not pre-empt that. + assert table_rows == [] - Known-bad for the first revision of this fix, which raised before LanceDB - ever saw the row. Passes on unmodified HEAD; guard test. - Asserts the row survives at full schema width, not the exact filled - composition: how LanceDB distributes ``fill_value`` over a short vector is - its own detail and differs by version (0.34 replaces the whole vector, 0.37 - pads and keeps the produced component), and ``lancedb`` is unpinned here. - What this guard owns is that the row reached the writer at all. - """ +def test_dense_write_keeps_on_bad_vectors_fill_reachable(tmp_path: Path) -> None: + """The empty-vector guard must not pre-empt LanceDB's fill policy.""" op = LanceDB( uri=str(tmp_path), table_name="t", @@ -366,10 +347,6 @@ def test_sparse_write_drops_whitespace_only_text(tmp_path: Path) -> None: assert table_rows == [] -# --- incomplete-index guard: a row must not reach the index without an embedding --- -# ``[]`` is not ``None``, so it used to fall through to the length check, be counted a -# wrong-length vector, and be dropped - a short index published with exit 0. - def _record(embedding: Any, *, text: str = "page text") -> dict: return { @@ -383,163 +360,22 @@ def _record(embedding: Any, *, text: str = "page text") -> dict: } -def test_an_empty_embedding_fails_the_run() -> None: - """Known-bad: returned ``(rows, counts)`` and only logged a WARNING. - - ``embedding: []`` is what ``embed_text_main_text_embed`` writes for every - row of a batch it could not embed. It was not counted before this change: - it is not ``None``, and it passed the length check as a wrong-length vector. - """ - records = [[_record([]), _record([1.0, 2.0])]] - - with pytest.raises(RuntimeError) as excinfo: - _create_lancedb_results(records, expected_dim=2) - - message = str(excinfo.value) - assert "Refusing to build an incomplete index" in message - assert "1 of 2 rows" in message - assert "empty_embedding=1" in message - - def test_create_lancedb_results_rejects_empty_embeddings_without_length_check() -> None: - """The non-enforcing path must not write empty vectors into the index. - - Known-bad: with ``expected_dim=None`` an ``embedding: []`` row was accepted - and counted in ``accepted``. - """ with pytest.raises(RuntimeError, match="empty_embedding=1"): _create_lancedb_results([[_record([])]], expected_dim=None) -@pytest.mark.parametrize("on_bad_vectors", ["drop", "fill", "error"]) -def test_wrong_length_rows_stay_under_the_on_bad_vectors_policy(on_bad_vectors: str) -> None: - """A short vector is a schema mismatch, not a missing embedding. - - ``on_bad_vectors`` is a documented, user-configured tolerance - (``common/vdb/lancedb.py`` ``create_index``). The incomplete-index guard must - not reach into it, or a user who deliberately configured ``drop`` or ``fill`` - would go from silent dropping to a hard run failure on upgrade. - - Known-bad for the first revision of this fix, which folded - ``dropped_bad_length`` into the fatal condition and raised for all three - values. It pins a contract rather than a code change, so it is a guard test: - the drop it asserts is identical before and after. It cannot run on the - unpatched tree, because the final assertion reads the new - ``empty_embedding`` key. - """ +def test_wrong_length_rows_are_forwarded_when_the_wrapper_check_is_disabled() -> None: records = [[_record([1.0]), _record([1.0, 2.0])]] - # ``expected_dim=None`` is the shape ``create_index`` uses when the caller - # asked LanceDB to own the policy (``on_bad_vectors="error"``) or turned the - # wrapper's length check off. - rows, counts = _create_lancedb_results(records, expected_dim=None if on_bad_vectors == "error" else 2) + rows, counts = _create_lancedb_results(records, expected_dim=None) - if on_bad_vectors == "error": - # The wrapper forwards both rows so LanceDB itself raises, per the - # documented strict-fail semantics of that policy. - assert len(rows) == 2 - assert counts["dropped_bad_length"] == 0 - else: - assert len(rows) == 1 - assert counts["dropped_bad_length"] == 1 + assert len(rows) == 2 + assert counts["dropped_bad_length"] == 0 assert counts["dropped_no_embedding"] == 0 assert counts["empty_embedding"] == 0 -def test_empty_embeddings_from_the_endpoint_path_do_not_reach_the_index_silently( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """The writer is the backstop for producers that legitimately keep going. - - A local engine failure now propagates from - ``models/inference/runtime.py`` and never reaches here. The endpoint path - keeps its per-batch resilience by design - a single failed HTTP call should - not kill a service-mode run - so it can still emit ``embedding: []`` rows. - This is the case that makes the writer-side check load-bearing rather than - redundant. - - Known-bad: the writer dropped those rows and returned normally. - """ - from nemo_retriever.models.inference import runtime - - def _refuse(*_args: Any, **_kwargs: Any) -> pd.DataFrame: - raise TimeoutError("read timed out waiting for the embedding endpoint") - - monkeypatch.setattr(runtime, "_embed_group", _refuse) - - batch_df = pd.DataFrame( - { - "text": ["page one", "page two"], - "metadata": [ - {"content": "page one", "content_metadata": {"page_number": 1}, "source_metadata": {}}, - {"content": "page two", "content_metadata": {"page_number": 2}, "source_metadata": {}}, - ], - } - ) - - out_df = runtime.embed_text_main_text_embed( - batch_df, - embedding_endpoint="http://embed.example/v1", - inference_batch_size=2, - ) - - # The stage returns successfully - deliberate for the endpoint path - but - # every row is empty. - assert list(out_df["text_embeddings_1b_v2_has_embedding"]) == [False, False] - assert [payload["embedding"] for payload in out_df["text_embeddings_1b_v2"]] == [[], []] - - records = [ - [ - { - "document_type": "text", - "metadata": {**row["metadata"], "embedding": row["text_embeddings_1b_v2"]["embedding"]}, - } - for _index, row in out_df.iterrows() - ] - ] - - with pytest.raises(RuntimeError, match="Refusing to build an incomplete index"): - _create_lancedb_results(records, expected_dim=2048) - - -def test_canonical_image_row_without_text_is_accepted_not_dropped() -> None: - """The text carve-out for canonical image rows still works. - - An image row legitimately carries ``text=""``. It has a real embedding, so - nothing here may touch it. - """ - record = { - "document_type": "image", - "metadata": { - "embedding": [1.0, 2.0], - "content": "", - "content_metadata": {"page_number": 3, "type": "image"}, - "source_metadata": {"source_name": "scan.pdf"}, - }, - } - rows, counts = _create_lancedb_results([[record]], expected_dim=2) - - assert len(rows) == 1 - assert counts["accepted"] == 1 - assert counts["dropped_no_text"] == 0 - - -def test_text_free_non_image_row_is_dropped_and_never_fatal() -> None: - """``dropped_no_text`` stays out of the fatal condition. - - It is a content filter, not a loss: the row embedded successfully and was - excluded for having nothing to search on. - - Cannot run on the unpatched tree: it asserts on the new - ``empty_embedding`` key. The drop itself is unchanged. - """ - rows, counts = _create_lancedb_results([[_record([1.0, 2.0], text=" ")]], expected_dim=2) - - assert rows == [] - assert counts["dropped_no_text"] == 1 - assert counts["empty_embedding"] == 0 - - @pytest.mark.parametrize( "embedding", [ @@ -548,24 +384,13 @@ def test_text_free_non_image_row_is_dropped_and_never_fatal() -> None: ], ) def test_present_vectors_are_accepted_whatever_their_values(embedding) -> None: - """The check tests presence, never values. - - Deciding from the numbers would mean guessing which embeddings are "real", - and an all-zero vector is a legal thing for a model to emit. Only the - absence of a vector is fatal. - """ rows, counts = _create_lancedb_results([[_record(embedding)]], expected_dim=2) assert len(rows) == 1 assert counts["accepted"] == 1 -def test_an_empty_numpy_array_is_not_swallowed_by_the_new_check() -> None: - """An empty ndarray is not a list, so it keeps its pre-existing route. - - This is the narrowness of the check made explicit: it fires on ``[]`` and - ``()`` and nothing else. - """ +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) @@ -573,39 +398,3 @@ def test_an_empty_numpy_array_is_not_swallowed_by_the_new_check() -> None: assert rows == [] assert counts["empty_embedding"] == 0 assert counts["dropped_bad_length"] == 1 - - -def test_a_fully_healthy_batch_raises_nothing() -> None: - """The whole point: a normal run is untouched. - - Cannot run on the unpatched tree, which has no ``empty_embedding`` - key; the acceptance of all 50 rows is identical there. - """ - records = [[_record([1.0, 2.0]) for _ in range(50)]] - - rows, counts = _create_lancedb_results(records, expected_dim=2) - - assert len(rows) == 50 - assert counts["accepted"] == 50 - assert counts["empty_embedding"] == 0 - assert counts["dropped_no_embedding"] == 0 - - -def test_a_none_embedding_keeps_its_pre_existing_silent_drop() -> None: - """``None`` is NOT fatal, because it has a legitimate producer. - - ``operators/embed/text_embed.py`` writes ``{"embedding": None}`` on purpose - for a row whose text was blank and which it therefore chose not to embed. - Making ``dropped_no_embedding`` fatal would fail ingests containing such - rows, which work today. Only ``[]`` - written solely on failure paths, at - ``models/inference/runtime.py`` and ``models/inference/vllm.py`` - is fatal. - - Cannot run on the unpatched tree, which has no ``empty_embedding`` - key. The drop it asserts is behaviour this change deliberately does not - touch. - """ - rows, counts = _create_lancedb_results([[_record(None), _record([1.0, 2.0])]], expected_dim=2) - - assert len(rows) == 1 - assert counts["dropped_no_embedding"] == 1 - assert counts["empty_embedding"] == 0 diff --git a/nemo_retriever/tests/test_multimodal_embed.py b/nemo_retriever/tests/test_multimodal_embed.py index 6730521a2a..d87736e89c 100644 --- a/nemo_retriever/tests/test_multimodal_embed.py +++ b/nemo_retriever/tests/test_multimodal_embed.py @@ -357,12 +357,6 @@ def test_non_dataframe_passthrough(self): result = collapse_content_to_page_rows(None) assert result is None -# --- multimodal row-loss guard --- -# A short answer used to be padded with ``None`` for the shortfall, and ``None`` is the -# one shape both writers ignore, so those rows were dropped and the run still succeeded. -# The guard counts rows submitted *with an image*, not rows in the chunk: embedders drop -# empty entries before inference, so comparing against chunk size would fail image-free chunks. - class _StubVLEmbedder: """VL embedder stub whose per-call return length is scripted by the test.""" @@ -386,61 +380,74 @@ def _image_frame(image_values): return pd.DataFrame({"_image_b64": list(image_values), "text": [""] * len(image_values)}) -def test_image_mode_short_answer_is_fatal(): - """A row submitted with an image and answered with nothing must fail the run. - - Without the guard the shortfall is padded with ``None``, which both writers - ignore, so the row is dropped and the run still succeeds. - """ - df = _image_frame(["b64-a", "b64-b", "b64-c"]) - embedder = _StubVLEmbedder(images=[[0.1, 0.2]]) # 1 vector for 3 images +@pytest.mark.parametrize( + ("images", "returned", "lost", "total"), + [ + (["b64-a", "b64-b", "b64-c"], [[0.1, 0.2]], 2, 3), + (["b64-a", "b64-b"], [], 2, 2), + ], + ids=["partial-answer", "empty-answer"], +) +def test_image_mode_short_answer_is_fatal(images, returned, lost, total): + embedder = _StubVLEmbedder(images=returned) with pytest.raises(LocalEmbedderRowsLostError) as excinfo: - main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") - assert excinfo.value.lost == 2 - assert excinfo.value.total == 3 - - -def test_image_mode_empty_answer_is_fatal(): - """Zero vectors for a chunk that did submit images is the whole-batch failure.""" - df = _image_frame(["b64-a", "b64-b"]) - embedder = _StubVLEmbedder(images=[]) - - with pytest.raises(LocalEmbedderRowsLostError): - main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") + main_text_embed._multimodal_callable_runner( + _image_frame(images), embedder=embedder, batch_size=8, embed_modality="image" + ) + assert excinfo.value.lost == lost + assert excinfo.value.total == total + assert "pad or drop those rows" in str(excinfo.value) def test_text_image_mode_short_multimodal_answer_is_fatal(): df = pd.DataFrame({"_image_b64": ["b64-a", "b64-b"], "text": ["alpha", "beta"]}) - embedder = _StubVLEmbedder(text_image=[[0.1, 0.2]]) # 1 vector for 2 paired rows + embedder = _StubVLEmbedder(text_image=[[0.1, 0.2]]) with pytest.raises(LocalEmbedderRowsLostError): main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="text_image") def test_text_image_mode_short_text_fallback_answer_is_fatal(): - """The text-only fallback subset is owed one vector per row as well.""" df = pd.DataFrame({"_image_b64": ["", ""], "text": ["alpha", "beta"]}) - embedder = _StubVLEmbedder(text=[[0.1, 0.2]]) # 1 vector for 2 fallback rows + embedder = _StubVLEmbedder(text=[[0.1, 0.2]]) with pytest.raises(LocalEmbedderRowsLostError): main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="text_image") -def test_image_mode_rows_without_images_do_not_fire_the_guard(): - """Rows with no image are owed nothing; they get ``None`` by contract. +@pytest.mark.parametrize( + ("df", "embed_modality", "embedder"), + [ + pytest.param(_image_frame(["b64-a"]), "image", _StubVLEmbedder(images=[[0.1], [0.2]]), id="image"), + pytest.param( + pd.DataFrame({"_image_b64": ["b64-a"], "text": ["alpha"]}), + "text_image", + _StubVLEmbedder(text_image=[[0.1], [0.2]]), + id="text-image", + ), + pytest.param( + pd.DataFrame({"_image_b64": [""], "text": ["alpha"]}), + "text_image", + _StubVLEmbedder(text=[[0.1], [0.2]]), + id="text-fallback", + ), + ], +) +def test_multimodal_extra_answer_reports_the_cardinality(df, embed_modality, embedder): + with pytest.raises(ValueError, match=r"returned 2 vectors for 1 submitted input"): + main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality=embed_modality) + - This is the false-failure the naive ``if not vecs_list: raise`` would cause. - """ +def test_image_mode_rows_without_images_do_not_fire_the_guard(): df = _image_frame(["b64-a", "", "b64-c"]) - embedder = _StubVLEmbedder(images=[[0.1, 0.2], [0.3, 0.4]]) # 2 vectors for 2 images + embedder = _StubVLEmbedder(images=[[0.1, 0.2], [0.3, 0.4]]) out = main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") assert out["embeddings"] == [[0.1, 0.2], None, [0.3, 0.4]] def test_image_mode_chunk_with_no_images_at_all_does_not_fire_the_guard(): - """An image-free chunk submits nothing, so zero vectors back is correct.""" df = _image_frame(["", ""]) embedder = _StubVLEmbedder(images=[]) @@ -449,17 +456,8 @@ def test_image_mode_chunk_with_no_images_at_all_does_not_fire_the_guard(): def test_text_image_mode_mixed_rows_do_not_fire_the_guard(): - """Paired, text-only and empty rows each get their own correct treatment.""" df = pd.DataFrame({"_image_b64": ["b64-a", "", ""], "text": ["alpha", "beta", " "]}) embedder = _StubVLEmbedder(text_image=[[0.1, 0.2]], text=[[0.3, 0.4]]) out = main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="text_image") assert out["embeddings"] == [[0.1, 0.2], [0.3, 0.4], None] - - -def test_image_mode_healthy_batch_does_not_fire_the_guard(): - df = _image_frame(["b64-a", "b64-b"]) - embedder = _StubVLEmbedder(images=[[0.1, 0.2], [0.3, 0.4]]) - - out = main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") - assert out["embeddings"] == [[0.1, 0.2], [0.3, 0.4]] diff --git a/nemo_retriever/tests/test_vllm_embed.py b/nemo_retriever/tests/test_vllm_embed.py index e14cbd2795..f19aa96ed5 100644 --- a/nemo_retriever/tests/test_vllm_embed.py +++ b/nemo_retriever/tests/test_vllm_embed.py @@ -26,6 +26,13 @@ from nemo_retriever.models.nim.error_reporter import drain_errors +@pytest.fixture(autouse=True) +def _clear_reported_errors(): + drain_errors() + yield + drain_errors() + + def _make_output(embedding): """Build a fake vLLM EmbeddingRequestOutput with out.outputs.embedding.""" return SimpleNamespace(outputs=SimpleNamespace(embedding=embedding)) @@ -501,34 +508,19 @@ def test_output_is_unnormalized_when_normalize_false(self): assert mock_mm.call_args.kwargs["normalize"] is False assert result.tolist() == [[3.0, 4.0]] - def test_no_valid_embeddings_no_longer_returns_an_empty_tensor(self): - # Rewritten, not weakened. This test previously asserted the defect: - # a batch vLLM failed to embed came back as a 0-row tensor, which is - # the same answer as an empty input, so the loss was unobservable. + @pytest.mark.parametrize("image_count", [1, 2]) + def test_no_valid_embeddings_raise_with_the_exact_count(self, image_count: int): b64 = _make_minimal_b64() with patch( "nemo_retriever.models.inference.vllm.embed_multimodal_with_vllm_llm", - return_value=[[]], + return_value=[[]] * image_count, ): - with pytest.raises(LocalEmbedderRowsLostError): - self.embedder.embed_images([b64]) - - def test_a_partially_lost_batch_fails_with_the_exact_count(self): - """The count exists here and nowhere else, so it must leave the function. - - ``_finalize_vectors`` holds both the batch it sent and the vectors that - came back, so ``len(vectors) - len(valid)`` is exact. It used to compute - that and discard it: the failed row was zero-padded to the right width, - which every shape check downstream accepts, and ``has_embedding`` then - reported ``True`` for a row carrying nothing. + with pytest.raises(LocalEmbedderRowsLostError) as excinfo: + self.embedder.embed_images([b64] * image_count) - The LanceDB writer guard cannot cover this row - correct width, non-zero - length - without guessing from its values, which is why the loss is - marked at the source instead. + assert (excinfo.value.lost, excinfo.value.total) == (image_count, image_count) - Known-bad: returns a ``(2, 2)`` tensor whose second row is ``[0, 0]``, - with nothing raised and nothing collected. - """ + def test_a_partially_lost_batch_fails_with_the_exact_count(self): drain_errors() b64 = _make_minimal_b64() with patch( @@ -543,36 +535,6 @@ def test_a_partially_lost_batch_fails_with_the_exact_count(self): collected = drain_errors() assert [(error.exc_type, error.stage) for error in collected] == [("LocalEmbedderRowsLostError", "embed")] - def test_a_wholly_lost_batch_fails_naming_every_row(self): - """Total loss must name the count too, not just fail. - - Known-bad: returns shape ``(0, 2048)``, the same answer as - ``embed_images([])``, so the whole batch vanishes with no padding to - notice and nothing raised. - """ - b64 = _make_minimal_b64() - with patch( - "nemo_retriever.models.inference.vllm.embed_multimodal_with_vllm_llm", - return_value=[[], []], - ): - with pytest.raises(LocalEmbedderRowsLostError) as excinfo: - self.embedder.embed_images([b64, b64]) - - assert (excinfo.value.lost, excinfo.value.total) == (2, 2) - - def test_a_complete_batch_is_untouched(self): - """Guard: no loss, nothing raised, nothing collected. Passes both ways.""" - drain_errors() - b64 = _make_minimal_b64() - with patch( - "nemo_retriever.models.inference.vllm.embed_multimodal_with_vllm_llm", - return_value=[[3.0, 4.0], [0.0, 5.0]], - ): - result = self.embedder.embed_images([b64, b64]) - - assert drain_errors() == [] - assert result.shape == (2, 2) - def _make_text_embedder(): with patch.object(LlamaNemotronEmbed1BV2Embedder, "__post_init__", lambda self: None): @@ -585,40 +547,11 @@ class TestLlamaNemotronEmbed1BV2Embedder: def setup_method(self): self.embedder = _make_text_embedder() - def test_finalize_vectors_all_empty_no_longer_returns_empty_tensor(self): - # Rewritten, not weakened: the old assertion pinned the defect. + def test_finalize_vectors_all_empty_raises_rows_lost(self): with pytest.raises(LocalEmbedderRowsLostError): self.embedder._finalize_vectors([[], []]) - def test_finalize_vectors_cannot_see_a_zero_output_batch(self): - """The conservation check counts empty rows, so zero rows count as zero loss. - - When vLLM yields no outputs at all for a non-empty batch, - ``embed_with_vllm_llm`` returns ``[]``. ``report_lost_rows`` then sums - over nothing and reports no loss, and the ``if not valid`` early return - hands back a 0-row tensor - the same answer as an empty input. Nothing - raises here. - - This is a pin on a real gap, not a defect this change fixes, so it - passes before and after. It is why the guard in - ``main_text_embed._callable_runner`` is not redundant with - ``_finalize_vectors``: that guard is the only layer that sees this - shape, and it fires before the LanceDB writer would. - """ - assert self.embedder._finalize_vectors([]).shape == (0, 0) - - def test_finalize_vectors_no_longer_zero_pads_missing(self): - # Rewritten, not weakened: the old assertion required the zero padding - # that made a lost row indistinguishable from a real embedding. - with pytest.raises(LocalEmbedderRowsLostError): - self.embedder._finalize_vectors([[1.0, 0.0], []]) - - def test_finalize_vectors_fails_on_the_rows_it_would_pad(self): - """The text embedder's ``_finalize_vectors`` has the same contract. - - Both embedders pad, so both must refuse to. Known-bad: the padding in - the test above happens and nothing records or stops it. - """ + def test_finalize_vectors_rejects_a_missing_row(self): drain_errors() with pytest.raises(LocalEmbedderRowsLostError) as excinfo: self.embedder._finalize_vectors([[1.0, 0.0], []]) @@ -628,13 +561,6 @@ def test_finalize_vectors_fails_on_the_rows_it_would_pad(self): assert [error.exc_type for error in drain_errors()] == ["LocalEmbedderRowsLostError"] def test_report_lost_rows_never_fires_on_a_healthy_batch(self): - """No false positives: every batch in every run goes through this call. - - A false positive here would fail runs that work today, so the healthy - shapes are pinned explicitly: full-width vectors, a single component, a - zero component, and an all-zero vector. An all-zero vector is a legal - thing for a model to emit; the check tests presence, never values. - """ from nemo_retriever.models.embed_errors import report_lost_rows drain_errors() @@ -648,12 +574,6 @@ def test_report_lost_rows_never_fires_on_a_healthy_batch(self): assert drain_errors() == [] def test_report_lost_rows_does_not_evaluate_truthiness_of_arrays(self): - """Explicit length test, because ``not vector`` raises on an array. - - ``bool(numpy.array([1.0, 2.0]))`` raises ``ValueError: The truth value - of an array with more than one element is ambiguous``. This call runs - before any pre-existing code, so a crash here would be a new one. - """ numpy = pytest.importorskip("numpy") from nemo_retriever.models.embed_errors import report_lost_rows @@ -661,29 +581,6 @@ def test_report_lost_rows_does_not_evaluate_truthiness_of_arrays(self): assert report_lost_rows([numpy.array([1.0, 2.0]), numpy.array([3.0, 4.0])], embedder="X") == 0 assert drain_errors() == [] - def test_report_lost_rows_is_silent_when_nothing_was_lost(self): - """Guard: the added call is inert on the healthy path. - - Every batch goes through it, so a false positive here would break runs - that work today. Cannot run on unmodified HEAD: the helper does not - exist there. - """ - from nemo_retriever.models.embed_errors import report_lost_rows - - drain_errors() - assert report_lost_rows([[1.0], [0.0]], embedder="X") == 0 - assert report_lost_rows([], embedder="X") == 0 - assert drain_errors() == [] - - def test_rows_lost_error_message_names_the_consequence(self): - """The message must say why a padded row matters, not just that it exists. - - Cannot run on unmodified HEAD: the class does not exist there. - """ - exc = LocalEmbedderRowsLostError(lost=7, total=64, embedder="SomeEmbedder") - assert (exc.lost, exc.total, exc.embedder) == (7, 64, "SomeEmbedder") - assert "match nothing" in str(exc) - def test_embed_uses_passage_prefix_by_default(self): with patch("nemo_retriever.models.inference.vllm.embed_with_vllm_llm", return_value=[[0.6, 0.8]]) as mock_fn: self.embedder.embed(["hello"]) From 59989662d2245f1fb4064a7fe37902a2377ae9f1 Mon Sep 17 00:00:00 2001 From: hrong Date: Thu, 20 Aug 2026 11:12:16 -0700 Subject: [PATCH 4/4] 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. --- docs/docs/extraction/troubleshoot.md | 45 ++++++------ .../src/nemo_retriever/common/vdb/README.md | 4 +- .../test_embed_engine_failure_propagation.py | 68 +++++++++--------- .../tests/test_lancedb_collections.py | 64 +++-------------- .../tests/test_lancedb_write_policy.py | 71 +++++-------------- nemo_retriever/tests/test_multimodal_embed.py | 36 ++++++---- nemo_retriever/tests/test_vllm_embed.py | 35 +++------ 7 files changed, 120 insertions(+), 203 deletions(-) diff --git a/docs/docs/extraction/troubleshoot.md b/docs/docs/extraction/troubleshoot.md index 0df5ab85d1..e19f3decef 100644 --- a/docs/docs/extraction/troubleshoot.md +++ b/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. diff --git a/nemo_retriever/src/nemo_retriever/common/vdb/README.md b/nemo_retriever/src/nemo_retriever/common/vdb/README.md index bea3c9d4d5..4757d765c1 100644 --- a/nemo_retriever/src/nemo_retriever/common/vdb/README.md +++ b/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. diff --git a/nemo_retriever/tests/test_embed_engine_failure_propagation.py b/nemo_retriever/tests/test_embed_engine_failure_propagation.py index 6b6b9a9856..b3bfa1e907 100644 --- a/nemo_retriever/tests/test_embed_engine_failure_propagation.py +++ b/nemo_retriever/tests/test_embed_engine_failure_propagation.py @@ -8,6 +8,7 @@ from typing import Any +import httpx import pandas as pd import pytest @@ -15,7 +16,7 @@ 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,26 +109,21 @@ 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( @@ -130,33 +131,34 @@ def test_partial_local_result_is_a_plain_value_error() -> None: [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)) diff --git a/nemo_retriever/tests/test_lancedb_collections.py b/nemo_retriever/tests/test_lancedb_collections.py index 91dd75a02b..d6b55aeea0 100644 --- a/nemo_retriever/tests/test_lancedb_collections.py +++ b/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 diff --git a/nemo_retriever/tests/test_lancedb_write_policy.py b/nemo_retriever/tests/test_lancedb_write_policy.py index b828b22368..2b7e30059d 100644 --- a/nemo_retriever/tests/test_lancedb_write_policy.py +++ b/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 diff --git a/nemo_retriever/tests/test_multimodal_embed.py b/nemo_retriever/tests/test_multimodal_embed.py index d87736e89c..d0470a59ec 100644 --- a/nemo_retriever/tests/test_multimodal_embed.py +++ b/nemo_retriever/tests/test_multimodal_embed.py @@ -20,12 +20,13 @@ # --------------------------------------------------------------------------- # Pure helpers from main_text_embed (no transitive-import issues) # --------------------------------------------------------------------------- -from nemo_retriever.models.inference import main_text_embed from nemo_retriever.models.inference.main_text_embed import ( + TextEmbeddingConfig, _format_image_input_string, _format_text_image_pair_input_string, _image_from_row, _multimodal_callable_runner, + create_text_embeddings_for_df, ) # --------------------------------------------------------------------------- @@ -380,6 +381,16 @@ def _image_frame(image_values): return pd.DataFrame({"_image_b64": list(image_values), "text": [""] * len(image_values)}) +def _run_multimodal(df, *, embedder, embed_modality): + config = TextEmbeddingConfig(embed_modality=embed_modality, output_payload_column="embedding_result") + out, _ = create_text_embeddings_for_df( + df, + task_config={"endpoint_url": None, "multimodal_embedder": embedder, "local_batch_size": 8}, + transform_config=config, + ) + return [payload["embedding"] for payload in out["embedding_result"]] + + @pytest.mark.parametrize( ("images", "returned", "lost", "total"), [ @@ -392,9 +403,7 @@ def test_image_mode_short_answer_is_fatal(images, returned, lost, total): embedder = _StubVLEmbedder(images=returned) with pytest.raises(LocalEmbedderRowsLostError) as excinfo: - main_text_embed._multimodal_callable_runner( - _image_frame(images), embedder=embedder, batch_size=8, embed_modality="image" - ) + _run_multimodal(_image_frame(images), embedder=embedder, embed_modality="image") assert excinfo.value.lost == lost assert excinfo.value.total == total assert "pad or drop those rows" in str(excinfo.value) @@ -405,7 +414,7 @@ def test_text_image_mode_short_multimodal_answer_is_fatal(): embedder = _StubVLEmbedder(text_image=[[0.1, 0.2]]) with pytest.raises(LocalEmbedderRowsLostError): - main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="text_image") + _run_multimodal(df, embedder=embedder, embed_modality="text_image") def test_text_image_mode_short_text_fallback_answer_is_fatal(): @@ -413,7 +422,7 @@ def test_text_image_mode_short_text_fallback_answer_is_fatal(): embedder = _StubVLEmbedder(text=[[0.1, 0.2]]) with pytest.raises(LocalEmbedderRowsLostError): - main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="text_image") + _run_multimodal(df, embedder=embedder, embed_modality="text_image") @pytest.mark.parametrize( @@ -436,28 +445,29 @@ def test_text_image_mode_short_text_fallback_answer_is_fatal(): ) def test_multimodal_extra_answer_reports_the_cardinality(df, embed_modality, embedder): with pytest.raises(ValueError, match=r"returned 2 vectors for 1 submitted input"): - main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality=embed_modality) + _run_multimodal(df, embedder=embedder, embed_modality=embed_modality) def test_image_mode_rows_without_images_do_not_fire_the_guard(): df = _image_frame(["b64-a", "", "b64-c"]) embedder = _StubVLEmbedder(images=[[0.1, 0.2], [0.3, 0.4]]) - out = main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") - assert out["embeddings"] == [[0.1, 0.2], None, [0.3, 0.4]] + assert _run_multimodal(df, embedder=embedder, embed_modality="image") == [[0.1, 0.2], None, [0.3, 0.4]] def test_image_mode_chunk_with_no_images_at_all_does_not_fire_the_guard(): df = _image_frame(["", ""]) embedder = _StubVLEmbedder(images=[]) - out = main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="image") - assert out["embeddings"] == [None, None] + assert _run_multimodal(df, embedder=embedder, embed_modality="image") == [None, None] def test_text_image_mode_mixed_rows_do_not_fire_the_guard(): df = pd.DataFrame({"_image_b64": ["b64-a", "", ""], "text": ["alpha", "beta", " "]}) embedder = _StubVLEmbedder(text_image=[[0.1, 0.2]], text=[[0.3, 0.4]]) - out = main_text_embed._multimodal_callable_runner(df, embedder=embedder, batch_size=8, embed_modality="text_image") - assert out["embeddings"] == [[0.1, 0.2], [0.3, 0.4], None] + assert _run_multimodal(df, embedder=embedder, embed_modality="text_image") == [ + [0.1, 0.2], + [0.3, 0.4], + None, + ] diff --git a/nemo_retriever/tests/test_vllm_embed.py b/nemo_retriever/tests/test_vllm_embed.py index f19aa96ed5..8b93c704a8 100644 --- a/nemo_retriever/tests/test_vllm_embed.py +++ b/nemo_retriever/tests/test_vllm_embed.py @@ -547,40 +547,21 @@ class TestLlamaNemotronEmbed1BV2Embedder: def setup_method(self): self.embedder = _make_text_embedder() - def test_finalize_vectors_all_empty_raises_rows_lost(self): - with pytest.raises(LocalEmbedderRowsLostError): - self.embedder._finalize_vectors([[], []]) + def test_embed_rejects_an_all_empty_result(self): + with patch("nemo_retriever.models.inference.vllm.embed_with_vllm_llm", return_value=[[], []]): + with pytest.raises(LocalEmbedderRowsLostError): + self.embedder.embed(["first", "second"]) - def test_finalize_vectors_rejects_a_missing_row(self): + def test_embed_rejects_a_missing_row(self): drain_errors() - with pytest.raises(LocalEmbedderRowsLostError) as excinfo: - self.embedder._finalize_vectors([[1.0, 0.0], []]) + with patch("nemo_retriever.models.inference.vllm.embed_with_vllm_llm", return_value=[[1.0, 0.0], []]): + with pytest.raises(LocalEmbedderRowsLostError) as excinfo: + self.embedder.embed(["first", "second"]) assert (excinfo.value.lost, excinfo.value.total) == (1, 2) assert excinfo.value.embedder == "LlamaNemotronEmbed1BV2Embedder" assert [error.exc_type for error in drain_errors()] == ["LocalEmbedderRowsLostError"] - def test_report_lost_rows_never_fires_on_a_healthy_batch(self): - from nemo_retriever.models.embed_errors import report_lost_rows - - drain_errors() - for batch in ( - [[0.1] * 2048, [0.2] * 2048], - [[0.0]], - [[0.0, 0.0], [0.0, 0.0]], - [[1.0]] * 256, - ): - assert report_lost_rows(batch, embedder="X") == 0 - assert drain_errors() == [] - - def test_report_lost_rows_does_not_evaluate_truthiness_of_arrays(self): - numpy = pytest.importorskip("numpy") - from nemo_retriever.models.embed_errors import report_lost_rows - - drain_errors() - assert report_lost_rows([numpy.array([1.0, 2.0]), numpy.array([3.0, 4.0])], embedder="X") == 0 - assert drain_errors() == [] - def test_embed_uses_passage_prefix_by_default(self): with patch("nemo_retriever.models.inference.vllm.embed_with_vllm_llm", return_value=[[0.6, 0.8]]) as mock_fn: self.embedder.embed(["hello"])