diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index 4fad8361..75b8dbdf 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -14,6 +14,27 @@ from ..as_embedding import BaseAsEmbedding Miss = tuple[int, str, str] # (result_index, text, cache_key) _MAX_VECTOR_SPACE_ATTEMPTS = 3 +_PROVIDER_WIDE_STATUS_CODES = frozenset({401, 403, 404}) + +# Transport failures raised by the OpenAI SDK and httpx. They do not inherit from the built-in +# TimeoutError / ConnectionError / OSError, so the built-in check alone lets them slip into the +# generic handler and fan out into per-item requests. Matching on class names (across the MRO, +# so subclasses such as APITimeoutError are covered) keeps this module free of SDK imports. +_TRANSPORT_FAILURE_TYPES = frozenset( + { + "APIConnectionError", + "APITimeoutError", + "ConnectError", + "ConnectTimeout", + "PoolTimeout", + "ReadError", + "ReadTimeout", + "RemoteProtocolError", + "WriteError", + "WriteTimeout", + }, +) + @R.register("local") class LocalEmbeddingStore(BaseEmbeddingStore): @@ -190,10 +211,14 @@ class LocalEmbeddingStore(BaseEmbeddingStore): result = await self.as_embedding(texts, **kwargs) if result and len(result) == len(texts): return result - except (TimeoutError, ConnectionError, OSError): - if attempt < self.max_retries - 1: - await asyncio.sleep(2**attempt) except Exception as error: + if self._is_transient_transport_failure(error): + # A connection failure or a timeout describes the endpoint, not any single + # item: retry the very same request with the usual backoff instead of + # fanning out into per-item requests that would fail the same way. + if attempt < self.max_retries - 1: + await asyncio.sleep(2**attempt) + continue if self._is_rate_limited(error): if attempt < self.max_retries - 1: delay = 2**attempt @@ -213,12 +238,93 @@ class LocalEmbeddingStore(BaseEmbeddingStore): ) await asyncio.sleep(self.quota_retry_delay) continue + if self._is_insufficient_quota(error): + # Quota is a service-level failure: no single item is to blame, so splitting the + # batch cannot isolate anything and would only multiply the rejected requests. + # This also covers the case where the quota retries above are exhausted. + self.logger.exception("Embedding request failed: quota exhausted") + self.is_healthy = False + return None + if self._is_provider_wide_failure(error): + # Credentials, permissions and server outages describe the endpoint, not any + # single item: isolation cannot fix them, so retry the transient ones with the + # usual backoff and give up on the rest instead of multiplying the requests. + if self._is_transient_provider_failure(error) and attempt < self.max_retries - 1: + delay = 2**attempt + self.logger.warning(f"Embedding provider error; retrying in {delay:.1f}s") + await asyncio.sleep(delay) + continue + self.logger.exception("Embedding request failed: provider-wide error, not splitting") + break 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. + Each single-item request goes back through ``_call_with_retry`` so a transient + failure (429, connection error, timeout) still gets the existing backoff and retry + policy instead of losing the vector; a one-item request never splits again. + """ + results: list[list[float] | None] = [] + for text in texts: + single = await self._call_with_retry([text], **kwargs) + if single and len(single) == 1: + results.append(single[0]) + continue + self.logger.error(f"Embedding request failed for a single item (chars={len(text)})") + results.append(None) + return results + + @staticmethod + def _is_transient_transport_failure(error: Exception) -> bool: + """Return whether an error is a transient transport failure worth retrying. + + Covers the built-in socket errors as well as the provider/httpx exceptions that describe + a connection failure or a timeout. Transport failures are provider-wide: no individual + input can be blamed, so the same request is retried and never split per item. + """ + if isinstance(error, (TimeoutError, ConnectionError, OSError)): + return True + return any(cls.__name__ in _TRANSPORT_FAILURE_TYPES for cls in type(error).__mro__) + + @staticmethod + def _is_provider_wide_failure(error: Exception) -> bool: + """Return whether an error describes the provider rather than a single input. + + Authentication, permission, missing model or route and server errors belong to the + endpoint, the credentials or the deployment: no individual input can be blamed, so + splitting the batch cannot isolate anything and would only multiply the rejected + requests. Rate limits and quota exhaustion have their own branches above and never + reach this classifier. A context-length ``400`` stays input-dependent on purpose, + because that is the one status an over-limit item can explain by itself. + """ + status = getattr(error, "status_code", None) + if isinstance(status, int) and (status in _PROVIDER_WIDE_STATUS_CODES or 500 <= status < 600): + return True + return type(error).__name__ in ( + "AuthenticationError", + "PermissionDeniedError", + "NotFoundError", + "InternalServerError", + ) + + @staticmethod + def _is_transient_provider_failure(error: Exception) -> bool: + """Return whether a provider-wide failure is worth retrying with the usual backoff.""" + status = getattr(error, "status_code", None) + if isinstance(status, int): + return 500 <= status < 600 + return type(error).__name__ == "InternalServerError" + @staticmethod def _is_rate_limited(error: Exception) -> bool: """Recognize an OpenAI-compatible 429 response without importing a provider SDK.""" diff --git a/reme/components/file_store/local_file_store.py b/reme/components/file_store/local_file_store.py index 243917f8..13e95378 100644 --- a/reme/components/file_store/local_file_store.py +++ b/reme/components/file_store/local_file_store.py @@ -29,6 +29,7 @@ _EMBEDDING_F16_DTYPE = np.dtype(" _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: @@ -530,6 +536,7 @@ class LocalFileStore(BaseFileStore): return processed = 0 + filled = 0 batch_count = 0 embedding_started_at = time.monotonic() next_percent = _PROGRESS_LOG_PERCENT_STEP @@ -538,10 +545,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: diff --git a/tests/unit/test_file_store_consistency.py b/tests/unit/test_file_store_consistency.py index b7750017..9e6af475 100644 --- a/tests/unit/test_file_store_consistency.py +++ b/tests/unit/test_file_store_consistency.py @@ -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, diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 51b7bf63..8d3053e5 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -126,6 +126,151 @@ 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 AlwaysInsufficientQuotaAsEmbedding: + """Always fail with quota exhaustion, recording every requested batch size.""" + + dimensions = 2 + + def __init__(self): + self.batch_sizes: list[int] = [] + + async def __call__(self, texts: list[str], **_kwargs): + self.batch_sizes.append(len(texts)) + raise InsufficientQuotaError("quota exhausted") + + +class PerItemLimitThenTransientErrorAsEmbedding: + """Reject an over-limit batch, then fail the first single-item request once. + + Mirrors the review case: the batch is rejected because one item exceeds the per-item + window, and a sibling item then hits a transient provider failure while retried alone. + """ + + dimensions = 2 + vector_space_id = "fakespace000" + max_chars = 32 + + def __init__(self, transient: Exception): + self.transient = transient + self.batch_sizes: list[int] = [] + self.transient_failures = 0 + + 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") + if len(texts) == 1 and self.transient_failures == 0: + self.transient_failures += 1 + raise self.transient + return [[1.0, 0.0] for _ in texts] + + +class AuthenticationError(Exception): + """OpenAI-compatible 401 error used without importing the provider SDK.""" + + status_code = 401 + + +class PermissionDeniedError(Exception): + """OpenAI-compatible 403 error used without importing the provider SDK.""" + + status_code = 403 + + +class InternalServerError(Exception): + """OpenAI-compatible 500 error used without importing the provider SDK.""" + + status_code = 500 + + +class NotFoundError(Exception): + """OpenAI-compatible 404 error, e.g. a missing model or endpoint.""" + + status_code = 404 + + +class APIError(Exception): + """Base of the OpenAI SDK error hierarchy, reproduced without importing the SDK.""" + + +class APIConnectionError(APIError): + """OpenAI SDK connection failure: it does not inherit from the built-in ConnectionError.""" + + +class APITimeoutError(APIConnectionError): + """OpenAI SDK timeout, a subclass of APIConnectionError.""" + + +class AlwaysProviderWideFailureAsEmbedding: + """Always fail with the same provider-wide error, recording every requested batch size.""" + + dimensions = 2 + + def __init__(self, error: Exception): + self.error = error + self.batch_sizes: list[int] = [] + + async def __call__(self, texts: list[str], **_kwargs): + self.batch_sizes.append(len(texts)) + raise self.error + + +class TransientThenSuccessAsEmbedding: + """Fail the first `failures` calls with one error, then return valid embeddings.""" + + dimensions = 2 + vector_space_id = "fakespace000" + + def __init__(self, error: Exception, failures: int = 1): + self.error = error + self.failures = failures + 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 len(self.batch_sizes) <= self.failures: + raise self.error + return [[1.0, 0.0] for _ in texts] + + class BadNodeEmbeddingStore(BaseEmbeddingStore): """Embedding store that returns wrong-dimensional vectors.""" @@ -730,3 +875,307 @@ 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()) + + +def test_insufficient_quota_batch_is_not_split_into_per_item_requests(monkeypatch): + """Quota exhaustion is a service-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_quota_batch", max_retries=2) + embedding = AlwaysInsufficientQuotaAsEmbedding() + store.as_embedding = embedding + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + assert await store._call_with_retry(["a", "b", "c"]) is None + assert embedding.batch_sizes == [3] + assert store.is_healthy is False + assert not sleeps + + run(go()) + + +def test_exhausted_quota_retries_still_do_not_split_the_batch(monkeypatch): + """Once the opt-in quota retries run out, the store gives up instead of splitting.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + store = LocalEmbeddingStore( + name="t_local_embedding_quota_batch_exhausted", + max_retries=3, + quota_retry_delay=5.0, + ) + embedding = AlwaysInsufficientQuotaAsEmbedding() + store.as_embedding = embedding + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + assert await store._call_with_retry(["a", "b"]) is None + assert embedding.batch_sizes == [2, 2, 2] + assert sleeps == [5.0, 5.0] + + run(go()) + + +def test_per_item_fallback_keeps_a_transiently_failed_vector(monkeypatch): + """A 429 / connection error / timeout on one isolated item must not lose its vector.""" + + async def go(): + transients = ( + RateLimitError("Requests are too frequent"), + ConnectionError("connection reset"), + TimeoutError("request timed out"), + ) + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + for transient in transients: + sleeps.clear() + + provider = PerItemLimitThenTransientErrorAsEmbedding(transient) + store = LocalEmbeddingStore( + name="t_local_embedding_fallback_transient", + max_retries=3, + enable_cache=False, + ) + store.as_embedding = provider + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + 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, 1] + assert provider.transient_failures == 1 + assert sleeps == [1.0] + + run(go()) + + +def test_provider_wide_credentials_errors_are_not_split_into_per_item_requests(): + """401/403 belong to the credentials: splitting ten items only multiplies the failures.""" + + async def go(): + for error in ( + AuthenticationError("invalid api key"), + PermissionDeniedError("no access to model"), + ): + provider = AlwaysProviderWideFailureAsEmbedding(error) + store = LocalEmbeddingStore(name="t_local_embedding_provider_wide", max_retries=3) + store.as_embedding = provider + + assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None + assert provider.batch_sizes == [10] + assert store.is_healthy is False + + run(go()) + + +def test_provider_wide_server_errors_are_retried_without_splitting(monkeypatch): + """A 500 is transient, so it is retried with the usual backoff and never split per item.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + provider = AlwaysProviderWideFailureAsEmbedding(InternalServerError("upstream unavailable")) + store = LocalEmbeddingStore(name="t_local_embedding_provider_server", max_retries=3) + store.as_embedding = provider + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None + assert provider.batch_sizes == [10, 10, 10] + assert sleeps == [1.0, 2.0] + assert store.is_healthy is False + + run(go()) + + +def test_sdk_transport_errors_are_retried_without_splitting(monkeypatch): + """SDK connection/timeout errors are provider-wide: retry the batch, never split it.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + for error in (APIConnectionError("connection reset"), APITimeoutError("timed out")): + sleeps.clear() + provider = AlwaysProviderWideFailureAsEmbedding(error) + store = LocalEmbeddingStore(name="t_local_embedding_sdk_transport", max_retries=3) + store.as_embedding = provider + + assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None + assert provider.batch_sizes == [10, 10, 10] + assert sleeps == [1.0, 2.0] + + run(go()) + + +def test_sdk_transport_error_recovers_without_losing_the_batch(monkeypatch): + """One SDK transport failure retries the original batch and keeps every vector.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + provider = TransientThenSuccessAsEmbedding(APITimeoutError("timed out")) + store = LocalEmbeddingStore(name="t_local_embedding_sdk_recover", max_retries=3) + store.as_embedding = provider + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + results = await store._call_with_retry([f"item{i}" for i in range(10)]) + + assert results == [[1.0, 0.0]] * 10 + assert provider.batch_sizes == [10, 10] + assert sleeps == [1.0] + + run(go()) + + +def test_isolated_item_retries_an_sdk_transport_error(monkeypatch): + """After a context-length rejection, one isolated item still retries an SDK transport error.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + provider = PerItemLimitThenTransientErrorAsEmbedding(APIConnectionError("connection reset")) + store = LocalEmbeddingStore( + name="t_local_embedding_sdk_isolated", + max_retries=3, + enable_cache=False, + ) + store.as_embedding = provider + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + 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, 1] + assert provider.transient_failures == 1 + assert sleeps == [1.0] + + run(go()) + + +def test_provider_wide_statuses_exclude_input_dependent_bad_request(): + """404 is provider-wide, but an over-limit 400 must stay input-dependent.""" + + async def go(): + store = LocalEmbeddingStore(name="t_local_embedding_classifier") + for error in ( + AuthenticationError("bad key"), + PermissionDeniedError("no access"), + NotFoundError("model_not_found"), + InternalServerError("boom"), + ): + assert store._is_provider_wide_failure(error) is True + over_limit = ContextLengthError("the input length exceeds the context length") + assert store._is_provider_wide_failure(over_limit) is False + + run(go()) + + +def test_provider_wide_not_found_errors_are_not_split(monkeypatch): + """A missing model or endpoint is provider-wide: one request, then failure, never a split.""" + + async def go(): + sleeps = [] + + async def fake_sleep(delay): + sleeps.append(delay) + + provider = AlwaysProviderWideFailureAsEmbedding(NotFoundError("model_not_found")) + store = LocalEmbeddingStore(name="t_local_embedding_not_found", max_retries=3) + store.as_embedding = provider + monkeypatch.setattr(asyncio, "sleep", fake_sleep) + + assert await store._call_with_retry([f"item{i}" for i in range(10)]) is None + assert provider.batch_sizes == [10] + assert not sleeps + + run(go())