mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-13 04:01:35 +00:00
fix(embedding): stop quota failures and transient item errors from multiplying
Review follow-up on the two P2 findings in #579: * quota exhaustion no longer falls into the per-item split. The store now gives up explicitly when the provider reports insufficient_quota, including after the opt-in quota retries are spent, because no single item is to blame and splitting only multiplied the rejected requests (3 items -> 4 requests). * the per-item fallback goes back through `_call_with_retry`, so a 429, connection error or timeout on an isolated item keeps the existing backoff and retry policy instead of silently dropping that item's vector. A one-item request still cannot split again. Regression tests cover both findings: a quota failure stays a single request with the default policy and with exhausted quota retries, and a transient failure on one item of a split batch still ends with every healthy sibling embedded. python -m pytest tests/unit/test_local_embedding_store.py tests/unit/test_file_store_consistency.py -q -> 83 passed, 41 skipped
This commit is contained in:
parent
49c0e7a215
commit
d5dc9b224d
2 changed files with 138 additions and 9 deletions
|
|
@ -213,6 +213,13 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
|
||||||
)
|
)
|
||||||
await asyncio.sleep(self.quota_retry_delay)
|
await asyncio.sleep(self.quota_retry_delay)
|
||||||
continue
|
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
|
||||||
self.logger.exception("Embedding request failed")
|
self.logger.exception("Embedding request failed")
|
||||||
self.is_healthy = False
|
self.is_healthy = False
|
||||||
if len(texts) > 1:
|
if len(texts) > 1:
|
||||||
|
|
@ -227,19 +234,18 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
|
||||||
Provider-side request errors, such as a per-item context-length limit, reject the
|
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
|
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.
|
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] = []
|
results: list[list[float] | None] = []
|
||||||
for text in texts:
|
for text in texts:
|
||||||
try:
|
single = await self._call_with_retry([text], **kwargs)
|
||||||
result = await self.as_embedding([text], **kwargs)
|
if single and len(single) == 1:
|
||||||
except Exception as error:
|
results.append(single[0])
|
||||||
self.logger.error(
|
|
||||||
f"Embedding request failed for a single item (chars={len(text)}): "
|
|
||||||
f"{type(error).__name__}: {error}",
|
|
||||||
)
|
|
||||||
results.append(None)
|
|
||||||
continue
|
continue
|
||||||
results.append(result[0] if result and len(result) == 1 else None)
|
self.logger.error(f"Embedding request failed for a single item (chars={len(text)})")
|
||||||
|
results.append(None)
|
||||||
return results
|
return results
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
|
|
|
||||||
|
|
@ -158,6 +158,48 @@ class PerItemLimitAsEmbedding:
|
||||||
return [[1.0, 0.0] for _ in texts]
|
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 BadNodeEmbeddingStore(BaseEmbeddingStore):
|
class BadNodeEmbeddingStore(BaseEmbeddingStore):
|
||||||
"""Embedding store that returns wrong-dimensional vectors."""
|
"""Embedding store that returns wrong-dimensional vectors."""
|
||||||
|
|
||||||
|
|
@ -828,3 +870,84 @@ def test_rate_limited_batch_is_not_split_into_per_item_requests(monkeypatch):
|
||||||
assert sleeps == [1.0]
|
assert sleeps == [1.0]
|
||||||
|
|
||||||
run(go())
|
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"),
|
||||||
|
)
|
||||||
|
for transient in transients:
|
||||||
|
sleeps = []
|
||||||
|
|
||||||
|
async def fake_sleep(delay):
|
||||||
|
sleeps.append(delay)
|
||||||
|
|
||||||
|
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())
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue