fix(embed): fail the run when embed rows arrive with no embedding - #2551
Open
rhdong wants to merge 4 commits into
Open
fix(embed): fail the run when embed rows arrive with no embedding#2551rhdong wants to merge 4 commits into
rhdong wants to merge 4 commits into
Conversation
Contributor
Greptile SummaryThe PR makes missing embeddings fatal instead of allowing an apparently successful but incomplete index.
|
| Filename | Overview |
|---|---|
| nemo_retriever/src/nemo_retriever/common/vdb/lancedb_collections.py | Completes the prior fix by rejecting empty list or tuple embeddings before any collection catalog or table mutation. |
| nemo_retriever/src/nemo_retriever/common/vdb/lancedb.py | Rejects empty dense embeddings before fixed-table creation and reports detailed row counters. |
| nemo_retriever/src/nemo_retriever/models/inference/main_text_embed.py | Adds fatal checks for empty or short embedding results instead of padding or discarding failed rows. |
| nemo_retriever/src/nemo_retriever/models/inference/runtime.py | Propagates recognized embedding-engine failures rather than converting them into empty embeddings. |
| nemo_retriever/src/nemo_retriever/models/embed_errors.py | Defines the embedding-specific failures used to abort incomplete ingestion. |
| nemo_retriever/tests/test_lancedb_collections.py | Covers empty list and tuple embeddings on the collection path and verifies no document is persisted. |
| nemo_retriever/tests/test_lancedb_write_policy.py | Exercises fixed-table empty-vector rejection and existing malformed-vector policies. |
| nemo_retriever/tests/test_embed_engine_failure_propagation.py | Covers propagation of embedding-engine startup and runtime failures. |
| nemo_retriever/tests/test_multimodal_embed.py | Covers multimodal response-cardinality failures. |
| docs/docs/extraction/troubleshoot.md | Documents fatal local embedding failures, writer behavior, and the remaining distinction between empty and None embeddings. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart LR
A[Embedding runtime] --> B{Vectors complete?}
B -- No --> C[Raise embedding failure]
B -- Yes --> D[Canonical records]
D --> E{Write mode}
E -- Fixed table --> F[Fixed-table empty-vector guard]
E -- Collection --> G[Collection empty-vector guard]
F --> H[LanceDB index]
G --> I[Scoped collection table]
Reviews (10): Last reviewed commit: "test(embed): exercise loss guards throug..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
rhdong
force-pushed
the
fix/incomplete-index-silent-success
branch
3 times, most recently
from
August 19, 2026 00:28
3a2721d to
c28212d
Compare
Collaborator
|
Overall, this direction makes sense and addresses the failure mode we saw. I checked out the PR and ran the focused tests locally: 146 passed and one failed. One scope question: is this intentionally scoped to local vLLM failures? Some missing embeddings can still arrive as Non-blocking: most of the diff size is tests rather than docs—the two new test files are a little over 1,000 lines and cover several unchanged behaviors/known gaps. Could we consolidate them around the core regression cases? |
rhdong
force-pushed
the
fix/incomplete-index-silent-success
branch
2 times, most recently
from
August 19, 2026 15:46
ea9aef0 to
9f7a404
Compare
Author
|
Hi @jioffe502, thanks again for the review. I’ve pushed the follow-up changes through 872aa35:
The five focused suites now pass with 179 passed and 3 skipped on both LanceDB versions. Could you please take another look when you have a chance? Thanks! |
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 <hrong@nvidia.com>
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 <hrong@nvidia.com>
Replace redundant private-helper checks with public runtime, embedder, and writer coverage. Clarify the dense-only and None-conversion behavior in the shipped documentation.
rhdong
force-pushed
the
fix/incomplete-index-silent-success
branch
from
August 20, 2026 20:00
872aa35 to
5998966
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
Rows whose embedding failed were silently excluded from the index, and the run still reported success.
exceptcaught the refusal and turned the whole batch into{"embedding": []}.success: truewithexit_code 0.embed_workers: 3: the index held 41,181 of 79,234 rows, and bo767 nDCG@10 was 0.2598 against an expected ~0.75. Every log line looked healthy.Root cause
The engine is refused permission to start. This is a capacity check, not an allocation failure.
embedding_endpointat a separate embedding service starts no in-process engine and makes the failure disappear.gpu_memory_utilization x total. That threshold scales with card size; the memory the run already holds when engines start does not.ValueError: Free memory ... is less than desired GPU memory utilization. When free memory collapses to ~0.16 GiB,cudaMallocfails before that check and vLLM reportsCUDA error: out of memoryinstead. Both end inRuntimeError: Engine core initialization failed, which is the string the classifier matches, so both are covered.Fix
Five guards, each covering a failure shape the others cannot see.
models/inference/runtime.pyre-raises instead of returning empty embeddings[]for every row, raising nothingmodels/inference/main_text_embed.pytreats an empty batch result as fatalmain_text_embed._multimodal_callable_runnercompares vectors received against rows submitted with an image, instead of padding the shortfall withNone_finalize_vectorsstops discarding its own count of failed rowscommon/vdb/lancedb.pyrefuses to build an index from rows with an empty embeddingThe writer guard carries the correctness guarantee for the shape it owns: it keys on the row, not the cause, so an empty embedding is refused whatever produced it - any backend, any GPU, endpoint path included. It is not a guarantee over every way a row can lose its vector. A missing embedding delivered as
Nonestays a counted, non-fatal drop, becauseoperators/embed/text_embed.py:107writesNoneon purpose for a blank-text row it chose not to embed; makingNonefatal would fail ingests that work today. The runtime classifier buys no correctness at all and only shortens time-to-failure, because the vector-store write is a single terminal global-batch stage and no earlier write exists to fail at. Do not delete either one as redundant.on_bad_vectorsis untouched. It governs malformed vectors, while an empty embedding means the embed stage produced nothing: a separate category, counted asempty_embedding.Verification
exit_code 10in ~107 s, no index directory, 0 rows reached the writerexit_code 10after ~1 h 48 m,lancedb/left empty: no table, no.lancefilesThe full suite passes on Linux with the same failure set as
main, and allpre-commithooks pass.Residual risk
OutOfMemoryErroris deliberately not fatal, because on the HuggingFace local backend an OOM is per-batch and recoverable. That backend was never exercised on hardware, so the protection rests on code reasoning and unit tests.BaseExceptionhandlers convert the new error to{"embedding": None}. Neither is on the shipped route, but one exports an identically named function, so future rewiring could lose the protection. Tests pin the gap.Noneis not covered. FourBaseExceptionhandlers inoperators/embed/text_embed.pyandcommon/modality/pipeline/embedding.pyrewrite a failed batch to{"embedding": None}, which the writers deliberately ignore. Neither route has a shipped importer, and tests pin that; closing it properly means separating "not embedded on purpose" from "embedding lost" at the source, which is out of scope here.Known gaps, not addressed here
Out of scope, and the reason one missing check became a silent 48% loss.
drain_errors()has no production callers. The failure path callsreport_error(), so the exception is stored, but nothing innemo_retriever/srcreads it back.has_embeddingis written but never read as a condition. It is set in six places, including correctly toFalsehere, and no code branches on it.Metric gates count quantity, not quality. The failing run passed
files==767,pages==54730,query_count==991, and nothing asked whether those pages produced usable vectors.The writer guard has not fired on hardware at this commit. In every end-to-end run the
runtime classifier aborted first, and the one run that did reach the writer used an
earlier revision of
lancedb.py. That path is covered by unit tests only.Checklist