fix(embedding): keep one bad item from dropping its whole batch

A provider-side request error rejects an entire batch, so a single over-limit
item used to cost the vectors of every item beside it. QwenPaw #8040 shows two
poison chunks turning into exactly twenty lost vectors, while the backfill loop
still reported full success.

Retry a rejected multi-item request one item at a time so only the item that
actually fails is left without a vector. Single-item requests keep their
previous behaviour, and provider-level failures (429 / quota) keep using the
existing retry paths instead of being split.

Also make the backfill loop report what it actually filled, and name the failing
chunk ids in the final reindex error. `processed` counts attempts, so it claimed
success after a partial loss.
This commit is contained in:
BeiMu-new 2026-09-30 17:06:04 +08:00
parent 638c81e88e
commit ab96ab680a
4 changed files with 176 additions and 2 deletions

View file

@ -215,10 +215,33 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
continue
self.logger.exception("Embedding request failed")
self.is_healthy = False
if len(texts) > 1:
return await self._embed_individually(texts, **kwargs)
return None
self.is_healthy = False
return None
async def _embed_individually(self, texts: list[str], **kwargs) -> list[list[float] | None]:
"""Retry a rejected batch one item at a time so one bad item cannot drop the rest.
Provider-side request errors, such as a per-item context-length limit, reject the
whole batch and would otherwise cost every vector in it. Isolating each item bounds
the damage to the item that actually failed and keeps the remaining vectors usable.
"""
results: list[list[float] | None] = []
for text in texts:
try:
result = await self.as_embedding([text], **kwargs)
except Exception as error:
self.logger.error(
f"Embedding request failed for a single item (chars={len(text)}): "
f"{type(error).__name__}: {error}",
)
results.append(None)
continue
results.append(result[0] if result and len(result) == 1 else None)
return results
@staticmethod
def _is_rate_limited(error: Exception) -> bool:
"""Recognize an OpenAI-compatible 429 response without importing a provider SDK."""

View file

@ -43,6 +43,7 @@ _SEARCH_FILTER_RESERVED_KEYS = frozenset(
"strict_date_filter",
},
)
_FAILED_CHUNK_PREVIEW = 5
@R.register("local")
@ -440,7 +441,12 @@ class LocalFileStore(BaseFileStore):
if chunk.text and not self._embedding_dim_matches(chunk.embedding)
]
if missing:
raise RuntimeError(f"embedding reindex incomplete: {len(missing)} chunks failed")
preview = ", ".join(missing[:_FAILED_CHUNK_PREVIEW])
if len(missing) > _FAILED_CHUNK_PREVIEW:
preview += f" (+{len(missing) - _FAILED_CHUNK_PREVIEW} more)"
raise RuntimeError(
f"embedding reindex incomplete: {len(missing)} chunks failed: {preview}",
)
await self._finalize_embedding_reindex()
await self._dump_owned_state()
if self.embedding_store is not None:
@ -544,6 +550,7 @@ class LocalFileStore(BaseFileStore):
return
processed = 0
filled = 0
batch_count = 0
embedding_started_at = time.monotonic()
next_percent = _PROGRESS_LOG_PERCENT_STEP
@ -552,10 +559,11 @@ class LocalFileStore(BaseFileStore):
await self.embedding_store.get_node_embeddings(batch)
self._drop_stale_embeddings(batch, "backfill")
processed += len(batch)
filled += sum(1 for chunk in batch if chunk.embedding is not None)
batch_count += 1
next_percent = self._log_progress("embedding backfill", processed, total, next_percent)
self.logger.info(
f"{self.name}: embedding batches complete: processed={processed}/{total}, "
f"{self.name}: embedding batches complete: processed={processed}/{total}, filled={filled}, "
f"batches={batch_count}, elapsed={time.monotonic() - embedding_started_at:.3f}s",
)
except asyncio.CancelledError:

View file

@ -118,6 +118,16 @@ class RecoveringEmbeddingStore(CountingFakeEmbeddingStore):
return await super().get_node_embeddings(nodes, **kwargs)
class PartialEmbeddingStore(FakeEmbeddingStore):
"""Fake provider that leaves chunks whose text contains "beta" without a vector."""
async def get_node_embeddings(self, nodes: list[FileChunk], **_kwargs) -> list[FileChunk]:
for chunk_node in nodes:
if "beta" not in chunk_node.text:
chunk_node.embedding = self._embed(chunk_node.text)
return nodes
class HealthCountingEmbeddingStore(FakeEmbeddingStore):
"""Fake provider that records eager health checks."""
@ -838,6 +848,41 @@ def test_verified_resume_without_chunks_supersedes_inflight_health_check():
run(go())
@pytest.mark.parametrize("store_factory", [_new_local_store, _new_faiss_store, _new_zvec_store])
def test_embedding_reindex_error_names_the_chunks_without_vectors(store_factory):
"""The final reindex error must name the chunks that still have no vector."""
async def go():
with tempfile.TemporaryDirectory() as tmp, temp_chdir(tmp):
store = store_factory("t_reindex_failure_ids")
await store.start()
await set_chunks_with_graph(
store,
{
"good": chunk("good", "a.md", "alpha text"),
"bad": chunk("bad", "b.md", "beta text"),
},
)
store.embedding_store = PartialEmbeddingStore()
_ensure_zvec_collection(store)
if isinstance(store, FaissLocalFileStore):
store._rebuild_index()
elif isinstance(store, ZvecLocalFileStore):
store._rebuild_collection()
await store.require_embedding_rebuild()
with pytest.raises(RuntimeError) as failure:
await store.reindex("embedding")
message = str(failure.value)
assert "1 chunks failed" in message
assert "bad" in message
assert "good" not in message
await store.close()
run(go())
@pytest.mark.parametrize("store_factory", [_new_local_store, _new_faiss_store, _new_zvec_store])
def test_embedding_reindex_is_explicit_and_keeps_vectors_disabled_until_success(
store_factory,

View file

@ -126,6 +126,38 @@ class QuotaThenSuccessAsEmbedding:
return [[1.0, 0.0] for _ in texts]
class ContextLengthError(Exception):
"""OpenAI-compatible 400 raised when one item exceeds the provider's per-item window."""
status_code = 400
body = {"error": {"message": "the input length exceeds the context length"}}
class PerItemLimitAsEmbedding:
"""Reject a whole request when any single text is over the provider's per-item limit.
Mirrors an OpenAI-compatible gateway: one over-limit item makes the provider answer
HTTP 400 for the entire batch, so a client that gives up on the batch loses every
vector in it.
"""
dimensions = 2
vector_space_id = "fakespace000"
max_chars = 32
def __init__(self):
self.batch_sizes: list[int] = []
def initialize_model(self):
"""Mirror the real component's idempotent initialization hook."""
async def __call__(self, texts: list[str], **_kwargs):
self.batch_sizes.append(len(texts))
if any(len(text) > self.max_chars for text in texts):
raise ContextLengthError("the input length exceeds the context length")
return [[1.0, 0.0] for _ in texts]
class BadNodeEmbeddingStore(BaseEmbeddingStore):
"""Embedding store that returns wrong-dimensional vectors."""
@ -730,3 +762,69 @@ def test_start_ignores_cache_file_without_vector_space_tag(monkeypatch, tmp_path
assert not store._cache
run(go())
def test_batch_failure_falls_back_to_per_item_embedding():
"""One over-limit item must cost one vector, not its whole batch."""
async def go():
provider = PerItemLimitAsEmbedding()
store = LocalEmbeddingStore(name="t_local_embedding_poison_batch", enable_cache=False)
store.as_embedding = provider
results = await store._call_with_retry(["ok", "x" * 64, "also ok"])
assert results == [[1.0, 0.0], None, [1.0, 0.0]]
assert provider.batch_sizes == [3, 1, 1, 1]
run(go())
def test_poison_chunk_keeps_its_batch_mates():
"""A rejected chunk must not clear the vectors of the chunks beside it."""
async def go():
store = LocalEmbeddingStore(name="t_local_embedding_poison_nodes", enable_cache=False)
store.as_embedding = PerItemLimitAsEmbedding()
nodes = [EmbNode(text="first"), EmbNode(text="y" * 64), EmbNode(text="third")]
await store.get_node_embeddings(nodes)
assert nodes[0].embedding is not None
assert nodes[1].embedding is None
assert nodes[2].embedding is not None
run(go())
def test_single_item_failure_still_reports_a_failed_batch():
"""A one-item request has nothing to isolate, so the existing contract stands."""
async def go():
store = LocalEmbeddingStore(name="t_local_embedding_single_poison", enable_cache=False)
store.as_embedding = PerItemLimitAsEmbedding()
assert await store._call_with_retry(["x" * 64]) is None
run(go())
def test_rate_limited_batch_is_not_split_into_per_item_requests(monkeypatch):
"""A 429 is a provider-level failure; splitting it would multiply the load."""
async def go():
sleeps = []
async def fake_sleep(delay):
sleeps.append(delay)
store = LocalEmbeddingStore(name="t_local_embedding_rate_limit_batch", max_retries=2)
embedding = AlwaysRateLimitedAsEmbedding()
store.as_embedding = embedding
monkeypatch.setattr(asyncio, "sleep", fake_sleep)
assert await store._call_with_retry(["a", "b"]) is None
assert embedding.calls == 2
assert sleeps == [1.0]
run(go())