From 930fddc4d25f3977fb05a80c27faada3800bf1ce Mon Sep 17 00:00:00 2001 From: BeiMu-new <2690322907@qq.com> Date: Fri, 9 Oct 2026 13:24:22 +0800 Subject: [PATCH] fix(embedding): retry SDK transport failures instead of splitting the batch Third review follow-up on #579. `_call_with_retry` only caught the built-in `TimeoutError`, `ConnectionError` and `OSError`. `openai.APIConnectionError` and `openai.APITimeoutError` do not inherit from any of them (they go APIError -> OpenAIError -> Exception), so they reached the generic handler: `_is_provider_wide_failure()` did not recognize them either, and a multi-item request fanned out into per-item requests that failed for exactly the same reason. The same gap meant a single-item request got no ReMe-level retry at all after a transport failure. Transport failures are now recognized by a dedicated classifier, `_is_transient_transport_failure()`, covering the built-in socket errors as well as the SDK/httpx exception types (matched across the MRO, so subclasses such as `APITimeoutError` are included) without importing the provider SDK. They retry the original request with the existing backoff and never split; the retry budget is unchanged. Regression coverage: - test_sdk_transport_errors_are_retried_without_splitting APIConnectionError / APITimeoutError -> batch sizes [10, 10, 10], sleeps [1.0, 2.0], then failure with no single-item requests. - test_sdk_transport_error_recovers_without_losing_the_batch one APITimeoutError then success -> [[1.0, 0.0]] * 10, batch sizes [10, 10]. - test_isolated_item_retries_an_sdk_transport_error context-length rejection, then APIConnectionError on one isolated item -> that item's vector is retained. The existing context-length, 401/403, 5xx, quota and rate-limit tests are retained. Local verification with the repo's own hook arguments: pylint -> Your code has been rated at 10.00/10, exit 0 black -> 2 files would be left unchanged flake8 -> clean Branch-order discriminator check (classifiers extracted from the production file and executed against each error type): case old new 400 context-length SPLIT SPLIT 401 / 403 SPLIT STOP 500 SPLIT RETRY SDK APIConnectionError SPLIT RETRY SDK APITimeoutError SPLIT RETRY builtin ConnectionError RETRY RETRY builtin TimeoutError RETRY RETRY 429 RETRY RETRY insufficient_quota STOP STOP --- .../embedding_store/local_embedding_store.py | 41 ++++++- tests/unit/test_local_embedding_store.py | 108 ++++++++++++++++++ 2 files changed, 146 insertions(+), 3 deletions(-) diff --git a/reme/components/embedding_store/local_embedding_store.py b/reme/components/embedding_store/local_embedding_store.py index d5e26e23..c0f1e257 100644 --- a/reme/components/embedding_store/local_embedding_store.py +++ b/reme/components/embedding_store/local_embedding_store.py @@ -14,6 +14,25 @@ from ..as_embedding import BaseAsEmbedding Miss = tuple[int, str, str] # (result_index, text, cache_key) _MAX_VECTOR_SPACE_ATTEMPTS = 3 +# 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 +209,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 @@ -259,6 +282,18 @@ class LocalEmbeddingStore(BaseEmbeddingStore): 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. diff --git a/tests/unit/test_local_embedding_store.py b/tests/unit/test_local_embedding_store.py index 926b872c..9e01215a 100644 --- a/tests/unit/test_local_embedding_store.py +++ b/tests/unit/test_local_embedding_store.py @@ -218,6 +218,18 @@ class InternalServerError(Exception): status_code = 500 +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.""" @@ -232,6 +244,27 @@ class AlwaysProviderWideFailureAsEmbedding: 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.""" @@ -1026,3 +1059,78 @@ def test_provider_wide_server_errors_are_retried_without_splitting(monkeypatch): 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())