mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-23 05:35:30 +00:00
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:
parent
638c81e88e
commit
ab96ab680a
4 changed files with 176 additions and 2 deletions
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue