mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-23 05:35:30 +00:00
fix(embedding): make transport retry ownership explicit
The store retried a connection failure or a timeout on top of the embedding wrapper and the provider SDK, which already retry both, so a persistent timeout on a batch reached the provider 36 times instead of 12: the OpenAI model retries a batch max_retries + 1 times and the SDK retries each of those attempts. Transport failures stay provider-wide -- the request is never split per item -- but the retry is now owned explicitly instead of implicitly: * a whole-batch request leaves it to the layers below, which already retry it; * an isolated single-item attempt keeps the store's backoff, because that attempt is the last chance to keep the one vector the isolation step is recovering. Rate limiting is unchanged. The transport cases now pin both budgets, and three tests drive a real OpenAIAsEmbedding wrapper and a real openai.AsyncClient over an httpx.MockTransport, so the provider attempt count is measured rather than assumed.
This commit is contained in:
BeiMu-new
parent
6dcc7664a3
commit
13cff385db
2 changed files with 284 additions and 66 deletions
|
|
@ -268,11 +268,17 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
|
|||
self,
|
||||
texts: list[str],
|
||||
*,
|
||||
isolated: bool = False,
|
||||
terminal_failure: list[bool] | None = None,
|
||||
**kwargs,
|
||||
) -> list[list[float] | None] | None:
|
||||
"""Answer a batch, retrying transient failures and isolating input ones.
|
||||
|
||||
``isolated`` marks a single-item attempt made while isolating a rejected batch. Such an
|
||||
attempt keeps the store's transport backoff, because it is the last chance to keep that
|
||||
one vector; every other request leaves transport retries to the embedding wrapper and the
|
||||
provider SDK, which already own them.
|
||||
|
||||
``terminal_failure`` is an optional out-parameter. When the caller passes a list,
|
||||
it is set to ``[True]`` if this call gave up on a failure that no remaining item
|
||||
can explain (a provider-wide error, exhausted quota, or exhausted transient
|
||||
|
|
@ -293,13 +299,22 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
|
|||
if result and len(result) == len(texts):
|
||||
return result
|
||||
except Exception as error:
|
||||
if self._is_transient_transport_failure(error):
|
||||
if self._is_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
|
||||
# item, so this request is never split per item. Which layer retries it is
|
||||
# explicit: an isolated single-item attempt keeps the store's backoff, while
|
||||
# any other request is left to the layers below, which already own that retry
|
||||
# -- the OpenAI model retries a batch max_retries + 1 times and the SDK retries
|
||||
# each of those attempts, so a batch they answer after 12 attempts used to
|
||||
# reach the provider 36 times with this store retrying on top of them.
|
||||
if isolated:
|
||||
if attempt < self.max_retries - 1:
|
||||
await asyncio.sleep(2**attempt)
|
||||
continue
|
||||
self.logger.exception("Embedding request failed: transport error, not splitting")
|
||||
self.is_healthy = False
|
||||
stop()
|
||||
break
|
||||
if self._is_rate_limited(error):
|
||||
if attempt < self.max_retries - 1:
|
||||
delay = 2**attempt
|
||||
|
|
@ -339,8 +354,8 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
|
|||
# below it -- a batch the SDK answered after 12 attempts became 36 once this
|
||||
# store retried on top of it. Retry ownership therefore stays with those
|
||||
# layers: a caller that wants a longer budget configures it there instead of
|
||||
# stacking another silent loop here. Transport failures keep their retry
|
||||
# above because they describe a connection that no lower layer has reported.
|
||||
# stacking another silent loop here. A transport failure is left to them for
|
||||
# the same reason; it only skips the per-item split.
|
||||
self.logger.exception("Embedding request failed: provider-wide error, not splitting")
|
||||
stop()
|
||||
break
|
||||
|
|
@ -378,16 +393,22 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
|
|||
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.
|
||||
Each single-item request goes back through ``_call_with_retry`` with ``isolated=True`` so
|
||||
a rate limit and a transient transport failure still get the store's backoff instead of
|
||||
losing the vector; a provider-wide failure ends the work for the whole batch rather than
|
||||
being asked for every remaining item, and a one-item request never splits again.
|
||||
When the loop stops on a provider-wide failure, ``terminal_failure`` is set for the
|
||||
caller so the vectors recovered before it are not read as a healthy provider.
|
||||
"""
|
||||
results: list[list[float] | None] = []
|
||||
for index, text in enumerate(texts):
|
||||
terminal: list[bool] = []
|
||||
single = await self._call_with_retry([text], terminal_failure=terminal, **kwargs)
|
||||
single = await self._call_with_retry(
|
||||
[text],
|
||||
isolated=True,
|
||||
terminal_failure=terminal,
|
||||
**kwargs,
|
||||
)
|
||||
if single:
|
||||
results.append(single[0])
|
||||
continue
|
||||
|
|
@ -410,12 +431,15 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
|
|||
return results
|
||||
|
||||
@staticmethod
|
||||
def _is_transient_transport_failure(error: Exception) -> bool:
|
||||
"""Return whether an error is a transient transport failure worth retrying.
|
||||
def _is_transport_failure(error: Exception) -> bool:
|
||||
"""Return whether an error describes a transport failure rather than a rejected input.
|
||||
|
||||
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.
|
||||
input can be blamed, so the request is never split per item. Who retries them is explicit
|
||||
in the transport branch of ``_call_with_retry``: an isolated single-item attempt keeps the
|
||||
store's backoff because it is the last chance to keep that vector, while every other
|
||||
request leaves the retry to the layers below, which already own it.
|
||||
"""
|
||||
if isinstance(error, (TimeoutError, ConnectionError, OSError)):
|
||||
return True
|
||||
|
|
|
|||
|
|
@ -6,7 +6,9 @@ import asyncio
|
|||
import json
|
||||
from types import SimpleNamespace
|
||||
|
||||
import httpx
|
||||
import numpy as np
|
||||
import openai
|
||||
|
||||
from reme.components.as_embedding import DashScopeAsEmbedding, OllamaAsEmbedding, OpenAIAsEmbedding
|
||||
from reme.components.embedding_store.base_embedding_store import BaseEmbeddingStore
|
||||
|
|
@ -301,6 +303,108 @@ class TransientThenSuccessAsEmbedding:
|
|||
return [[1.0, 0.0] for _ in texts]
|
||||
|
||||
|
||||
_REAL_OPENAI_ASYNC_CLIENT = openai.AsyncClient
|
||||
_OFFLINE_OPENAI_BASE_URL = "http://127.0.0.1:9/v1"
|
||||
|
||||
|
||||
class CountingOpenAIWrapper(OpenAIAsEmbedding):
|
||||
"""Real OpenAI wrapper that records how often the store asked it for embeddings."""
|
||||
|
||||
def __init__(self, **kwargs) -> None:
|
||||
super().__init__(**kwargs)
|
||||
self.calls = 0
|
||||
|
||||
async def __call__(self, inputs, **kwargs):
|
||||
self.calls += 1
|
||||
return await super().__call__(inputs, **kwargs)
|
||||
|
||||
|
||||
class OfflineOpenAIEndpoint:
|
||||
"""Real openai.AsyncClient wired to an offline httpx transport that counts provider attempts.
|
||||
|
||||
Gives the wrapper-and-SDK layers their real code path, so a regression test can assert how
|
||||
many attempts the provider would see instead of trusting the store's own loop.
|
||||
"""
|
||||
|
||||
def __init__(self, handler, sdk_max_retries: int = 2) -> None:
|
||||
self.requests = 0
|
||||
self.handler = handler
|
||||
self.client = _REAL_OPENAI_ASYNC_CLIENT(
|
||||
api_key="test-key",
|
||||
base_url=_OFFLINE_OPENAI_BASE_URL,
|
||||
http_client=httpx.AsyncClient(transport=httpx.MockTransport(self._handle)),
|
||||
max_retries=sdk_max_retries,
|
||||
)
|
||||
|
||||
def _handle(self, request):
|
||||
self.requests += 1
|
||||
return self.handler(request)
|
||||
|
||||
def install(self, monkeypatch) -> "OfflineOpenAIEndpoint":
|
||||
"""Route every openai.AsyncClient built during this test to the offline client."""
|
||||
monkeypatch.setattr(openai, "AsyncClient", lambda **kwargs: self.client)
|
||||
return self
|
||||
|
||||
|
||||
def offline_openai_endpoint(handler) -> OfflineOpenAIEndpoint:
|
||||
"""Build an offline endpoint whose transport follows ``handler``."""
|
||||
return OfflineOpenAIEndpoint(handler)
|
||||
|
||||
|
||||
def openai_wrapper(name: str) -> CountingOpenAIWrapper:
|
||||
"""Build a real OpenAI embedding wrapper pointed at the offline endpoint."""
|
||||
return CountingOpenAIWrapper(
|
||||
name=name,
|
||||
model="text-embedding-v4",
|
||||
backend="openai",
|
||||
dimensions=2,
|
||||
credential={"base_url": _OFFLINE_OPENAI_BASE_URL, "api_key": "test-key"},
|
||||
retry_delay=0,
|
||||
pass_dimensions=False,
|
||||
)
|
||||
|
||||
|
||||
def embedding_payload(request) -> httpx.Response:
|
||||
"""Answer an embeddings request the way the provider would."""
|
||||
body = json.loads(request.content.decode("utf-8"))
|
||||
inputs = body["input"]
|
||||
if isinstance(inputs, str):
|
||||
inputs = [inputs]
|
||||
return httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"object": "list",
|
||||
"data": [
|
||||
{"object": "embedding", "index": position, "embedding": [0.0, 1.0]} for position in range(len(inputs))
|
||||
],
|
||||
"model": body.get("model", "text-embedding-v4"),
|
||||
"usage": {"prompt_tokens": len(inputs), "total_tokens": len(inputs)},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def context_length_rejection() -> httpx.Response:
|
||||
"""Answer a batch the way an OpenAI-compatible provider reports an over-long input."""
|
||||
return httpx.Response(
|
||||
400,
|
||||
json={
|
||||
"error": {
|
||||
"message": (
|
||||
"This model's maximum context length is 8191 tokens. " "Please reduce the length of the input."
|
||||
),
|
||||
"type": "invalid_request_error",
|
||||
"param": None,
|
||||
"code": "context_length_exceeded",
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def read_timeout(request) -> httpx.Response:
|
||||
"""Fail one attempt the way a read timeout reaches the SDK."""
|
||||
raise httpx.ReadTimeout("simulated read timeout", request=request)
|
||||
|
||||
|
||||
class BadNodeEmbeddingStore(BaseEmbeddingStore):
|
||||
"""Embedding store that returns wrong-dimensional vectors."""
|
||||
|
||||
|
|
@ -1020,38 +1124,37 @@ def test_exhausted_quota_retries_still_do_not_split_the_batch(monkeypatch):
|
|||
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."""
|
||||
def test_per_item_fallback_keeps_a_rate_limited_vector(monkeypatch):
|
||||
"""A 429 on one isolated item still gets the store's backoff and keeps its vector.
|
||||
|
||||
Rate limiting is the failure class the store still owns: no lower layer reports it as a
|
||||
provider-wide error worth its own budget, so the retry lives here. Connection errors and
|
||||
timeouts are left to the wrapper and the SDK instead; see the transport tests above.
|
||||
"""
|
||||
|
||||
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(
|
||||
RateLimitError("Requests are too frequent"),
|
||||
)
|
||||
store = LocalEmbeddingStore(
|
||||
name="t_local_embedding_fallback_rate_limit",
|
||||
max_retries=3,
|
||||
enable_cache=False,
|
||||
)
|
||||
store.as_embedding = provider
|
||||
monkeypatch.setattr(asyncio, "sleep", fake_sleep)
|
||||
|
||||
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"])
|
||||
|
||||
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]
|
||||
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())
|
||||
|
||||
|
|
@ -1099,8 +1202,12 @@ def test_provider_wide_server_errors_are_not_retried_without_splitting(monkeypat
|
|||
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."""
|
||||
def test_sdk_transport_errors_are_provider_wide_and_not_retried_here(monkeypatch):
|
||||
"""SDK connection/timeout errors are provider-wide: never split, and not retried here.
|
||||
|
||||
The wrapper and the SDK already retry a transport failure, so the store answers once and
|
||||
leaves the retry budget to those layers instead of multiplying their attempts.
|
||||
"""
|
||||
|
||||
async def go():
|
||||
sleeps = []
|
||||
|
|
@ -1117,14 +1224,23 @@ def test_sdk_transport_errors_are_retried_without_splitting(monkeypatch):
|
|||
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]
|
||||
# One request only: the wrapper and the SDK already retry a transport failure
|
||||
# (4 wrapper attempts, each retried by the SDK), so another loop here only
|
||||
# multiplies them -- a persistent timeout used to cost 36 attempts instead of 12.
|
||||
assert provider.batch_sizes == [10]
|
||||
assert sleeps == []
|
||||
assert store.is_healthy is False
|
||||
|
||||
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."""
|
||||
def test_isolated_item_keeps_the_store_backoff_for_a_transport_error(monkeypatch):
|
||||
"""An isolated single-item attempt still gets the store's backoff and keeps its vector.
|
||||
|
||||
This fake provider has no wrapper or SDK beneath it, so it pins the store's own budget: the
|
||||
isolated attempt is the last chance to keep that one vector. A whole-batch transport failure
|
||||
is left to the layers below instead; see the earlier transport tests.
|
||||
"""
|
||||
|
||||
async def go():
|
||||
sleeps = []
|
||||
|
|
@ -1132,44 +1248,122 @@ def test_sdk_transport_error_recovers_without_losing_the_batch(monkeypatch):
|
|||
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)
|
||||
for transient in (
|
||||
APIConnectionError("connection reset"),
|
||||
ConnectionError("connection reset"),
|
||||
TimeoutError("request timed out"),
|
||||
):
|
||||
sleeps.clear()
|
||||
|
||||
results = await store._call_with_retry([f"item{i}" for i in range(10)])
|
||||
provider = PerItemLimitThenTransientErrorAsEmbedding(transient)
|
||||
store = LocalEmbeddingStore(
|
||||
name="t_local_embedding_sdk_isolated",
|
||||
max_retries=3,
|
||||
enable_cache=False,
|
||||
)
|
||||
store.as_embedding = provider
|
||||
monkeypatch.setattr(asyncio, "sleep", fake_sleep)
|
||||
|
||||
assert results == [[1.0, 0.0]] * 10
|
||||
assert provider.batch_sizes == [10, 10]
|
||||
assert sleeps == [1.0]
|
||||
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 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."""
|
||||
def test_a_persistent_transport_error_is_not_multiplied_by_the_store(monkeypatch):
|
||||
"""A timeout is asked once: the wrapper and the SDK own the retry budget.
|
||||
|
||||
The provider sees the 4 wrapper attempts x 3 SDK attempts those layers make, while the
|
||||
store's max_retries=3 no longer multiplies them -- a persistent timeout used to cost 36
|
||||
provider attempts instead of 12.
|
||||
"""
|
||||
|
||||
async def go():
|
||||
sleeps = []
|
||||
|
||||
async def fake_sleep(delay):
|
||||
sleeps.append(delay)
|
||||
|
||||
provider = PerItemLimitThenTransientErrorAsEmbedding(APIConnectionError("connection reset"))
|
||||
endpoint = offline_openai_endpoint(read_timeout).install(monkeypatch)
|
||||
provider = openai_wrapper("t_openai_transport_budget")
|
||||
store = LocalEmbeddingStore(
|
||||
name="t_local_embedding_sdk_isolated",
|
||||
name="t_local_embedding_sdk_transport_budget",
|
||||
max_retries=3,
|
||||
enable_cache=False,
|
||||
)
|
||||
store.as_embedding = provider
|
||||
|
||||
assert await store._call_with_retry([f"item{i}" for i in range(3)]) is None
|
||||
assert provider.calls == 1
|
||||
assert endpoint.requests == 12
|
||||
assert store.is_healthy is False
|
||||
|
||||
run(go())
|
||||
|
||||
|
||||
def test_a_transient_transport_error_still_recovers_through_the_sdk(monkeypatch):
|
||||
"""One timeout is answered by the SDK's own retry: the batch survives and never splits."""
|
||||
|
||||
async def go():
|
||||
state = {"timeouts": 1}
|
||||
|
||||
def handler(request):
|
||||
if state["timeouts"] > 0:
|
||||
state["timeouts"] -= 1
|
||||
return read_timeout(request)
|
||||
return embedding_payload(request)
|
||||
|
||||
endpoint = offline_openai_endpoint(handler).install(monkeypatch)
|
||||
provider = openai_wrapper("t_openai_transport_recover")
|
||||
store = LocalEmbeddingStore(
|
||||
name="t_local_embedding_sdk_transport_recover",
|
||||
max_retries=3,
|
||||
enable_cache=False,
|
||||
)
|
||||
store.as_embedding = provider
|
||||
|
||||
results = await store._call_with_retry([f"item{i}" for i in range(3)])
|
||||
|
||||
assert results == [[0.0, 1.0]] * 3
|
||||
assert provider.calls == 1
|
||||
assert endpoint.requests == 2
|
||||
|
||||
run(go())
|
||||
|
||||
|
||||
def test_an_isolated_item_still_recovers_through_the_sdk(monkeypatch):
|
||||
"""After a context-length rejection, a timeout on one remaining item is still recovered.
|
||||
|
||||
The store asks the wrapper once per item; the retry that saves that item belongs to the SDK.
|
||||
"""
|
||||
|
||||
async def go():
|
||||
state = {"timeout_pending": True}
|
||||
|
||||
def handler(request):
|
||||
body = json.loads(request.content.decode("utf-8"))
|
||||
inputs = body["input"]
|
||||
if isinstance(inputs, str):
|
||||
inputs = [inputs]
|
||||
if any(len(item) > 32 for item in inputs):
|
||||
return context_length_rejection()
|
||||
if len(inputs) == 1 and state["timeout_pending"]:
|
||||
state["timeout_pending"] = False
|
||||
return read_timeout(request)
|
||||
return embedding_payload(request)
|
||||
|
||||
endpoint = offline_openai_endpoint(handler).install(monkeypatch)
|
||||
provider = openai_wrapper("t_openai_transport_isolated")
|
||||
store = LocalEmbeddingStore(
|
||||
name="t_local_embedding_sdk_transport_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]
|
||||
assert results == [[0.0, 1.0], None, [0.0, 1.0]]
|
||||
assert provider.calls == 4
|
||||
assert endpoint.requests == 5
|
||||
|
||||
run(go())
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue