fix(embedding): retry vector space changes per request

This commit is contained in:
jinli.yl 2026-08-28 11:31:11 +08:00
parent 0684553223
commit 55905ed595
6 changed files with 98 additions and 68 deletions

View file

@ -173,11 +173,8 @@ curl -s http://127.0.0.1:2333/auto_fin \
When the application uses an MCP service, service-enabled plugin Jobs appear as MCP tools instead. When the application uses an MCP service, service-enabled plugin Jobs appear as MCP tools instead.
To add the plugin to another application config, select it explicitly: Custom application configs must provide the plugin's runtime dependencies, including an `agent_wrapper.default` and
the `search` and `read` Jobs used by Auto Fin.
```bash
reme start config=demo plugins='["auto-fin"]'
```
## Uninstall a plugin ## Uninstall a plugin

View file

@ -167,11 +167,7 @@ curl -s http://127.0.0.1:2333/auto_fin \
当应用使用 MCP service 时,允许对外服务的插件 Job 会显示为 MCP tool。 当应用使用 MCP service 时,允许对外服务的插件 Job 会显示为 MCP tool。
如果需要将插件叠加到其他应用配置,则显式选择该配置: 自定义应用配置需要提供插件的运行依赖,包括 `agent_wrapper.default`,以及 Auto Fin 使用的 `search``read` Jobs。
```bash
reme start config=demo plugins='["auto-fin"]'
```
## 卸载插件 ## 卸载插件

View file

@ -59,11 +59,7 @@ reme start plugins='["auto-fin"]' \
service.backend=http service.backend=http
``` ```
To add Auto Fin to another application instead, select that config explicitly, for example: Custom application configs must provide `agent_wrapper.default` and the `search` and `read` Jobs used by Auto Fin.
```bash
reme start config=demo plugins='["auto-fin"]'
```
## Pipeline ## Pipeline

View file

@ -54,11 +54,7 @@ reme start plugins='["auto-fin"]' \
service.backend=http service.backend=http
``` ```
如果需要将 Auto Fin 叠加到其他应用,则显式选择相应配置,例如: 自定义应用配置需要提供 `agent_wrapper.default`,以及 Auto Fin 使用的 `search``read` Jobs。
```bash
reme start config=demo plugins='["auto-fin"]'
```
## 流程 ## 流程

View file

@ -104,12 +104,25 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
# -- Public API -- # -- Public API --
async def get_embeddings(self, input_text: list[str], **kwargs) -> list[np.ndarray | None]: async def get_embeddings(self, input_text: list[str], **kwargs) -> list[np.ndarray | None]:
await self._sync_cache_space()
texts = [self._truncate(t) for t in input_text] texts = [self._truncate(t) for t in input_text]
results, misses = self._partition_by_cache(texts) for attempt in range(1, _MAX_VECTOR_SPACE_ATTEMPTS + 1):
if misses: await self._sync_cache_space()
await self._fill_misses(misses, results, **kwargs) vector_space_id = self._cache_space
return results results, misses = self._partition_by_cache(texts)
stable = not misses or await self._fill_misses(misses, results, vector_space_id, **kwargs)
if stable and vector_space_id == self.vector_space_id == self._cache_space:
return results
if attempt == _MAX_VECTOR_SPACE_ATTEMPTS:
self.logger.warning(
f"Embedding vector space kept changing while computing a request; "
f"discarding all result(s) after {attempt} attempts",
)
else:
self.logger.info(
f"Embedding vector space changed while computing a request; "
f"discarding all result(s) and retrying ({attempt}/{_MAX_VECTOR_SPACE_ATTEMPTS})",
)
return [None] * len(texts)
# -- Batching -- # -- Batching --
@ -125,31 +138,26 @@ class LocalEmbeddingStore(BaseEmbeddingStore):
misses.append((idx, text, key)) misses.append((idx, text, key))
return results, misses return results, misses
async def _fill_misses(self, misses: list[Miss], results: list[np.ndarray | None], **kwargs) -> None: async def _fill_misses(
self,
misses: list[Miss],
results: list[np.ndarray | None],
vector_space_id: str,
**kwargs,
) -> bool:
"""Fill every miss only while the request remains in one vector space."""
size = self.max_batch_size size = self.max_batch_size
for start in range(0, len(misses), size): for start in range(0, len(misses), size):
if vector_space_id != self.vector_space_id or vector_space_id != self._cache_space:
return False
batch = misses[start : start + size] batch = misses[start : start + size]
for attempt in range(1, _MAX_VECTOR_SPACE_ATTEMPTS + 1): computed = await self._compute_batch(batch, **kwargs)
await self._sync_cache_space() if vector_space_id != self.vector_space_id or vector_space_id != self._cache_space:
vector_space_id = self._cache_space return False
computed = await self._compute_batch(batch, **kwargs) for idx, key, emb in computed:
if vector_space_id != self.vector_space_id or vector_space_id != self._cache_space: results[idx] = emb
if attempt == _MAX_VECTOR_SPACE_ATTEMPTS: self._cache_put(key, emb)
self.logger.warning( return True
f"Embedding vector space kept changing while computing a batch; "
f"discarding {len(computed)} stale result(s) after {attempt} attempts",
)
else:
self.logger.info(
f"Embedding vector space changed while computing a batch; "
f"discarding {len(computed)} stale result(s) and retrying "
f"({attempt}/{_MAX_VECTOR_SPACE_ATTEMPTS})",
)
continue
for idx, key, emb in computed:
results[idx] = emb
self._cache_put(key, emb)
break
async def _compute_batch(self, batch: list[Miss], **kwargs) -> list[tuple[int, str, np.ndarray]]: async def _compute_batch(self, batch: list[Miss], **kwargs) -> list[tuple[int, str, np.ndarray]]:
texts = [text for _, text, _ in batch] texts = [text for _, text, _ in batch]

View file

@ -509,41 +509,79 @@ def test_cache_space_is_rechecked_after_async_load(monkeypatch, tmp_path):
run(go()) run(go())
def test_completed_request_retries_after_vector_space_changes(): def test_whole_request_retries_after_vector_space_changes_between_batches():
"""A request completed by the old provider must not escape into the new vector space.""" """Completed batches must be discarded when a later batch changes vector space."""
async def go(): async def go():
embedding = OpenAIAsEmbedding(name="t_space_write_race", backend="openai", model="v3", dimensions=2) embedding = FakeAsEmbedding()
store = LocalEmbeddingStore(name="t_local_write_race") embedding.vector_space_id = "v3"
store = LocalEmbeddingStore(name="t_local_write_race", max_batch_size=1, enable_cache=False)
store.as_embedding = embedding store.as_embedding = embedding
store._cache_space = embedding.vector_space_id store._cache_space = embedding.vector_space_id
calls = 0 calls = 0
async def compute_after_round_trip(_batch, **_kwargs): async def switch_during_second_batch(batch, **_kwargs):
nonlocal calls nonlocal calls
calls += 1 calls += 1
if calls == 1: if calls == 2:
embedding.model = FakeProviderModel("v4") embedding.vector_space_id = "v4"
return [(0, "key", np.array([3.0, 0.0], dtype=np.float16))] idx, _text, key = batch[0]
return [(0, "key", np.array([4.0, 0.0], dtype=np.float16))] version = 3.0 if calls < 3 else 4.0
return [(idx, key, np.array([version, 0.0], dtype=np.float16))]
store._compute_batch = compute_after_round_trip store._compute_batch = switch_during_second_batch
results = [None] results = await store.get_embeddings(["first", "second"])
await store._fill_misses([(0, "text", "key")], results)
assert calls == 2 assert calls == 4
assert store._cache_space == embedding.vector_space_id assert store._cache_space == embedding.vector_space_id
np.testing.assert_array_equal(results[0], np.array([4.0, 0.0], dtype=np.float16)) for result in results:
np.testing.assert_array_equal(store._cache["key"], np.array([4.0, 0.0], dtype=np.float16)) np.testing.assert_array_equal(result, np.array([4.0, 0.0], dtype=np.float16))
run(go()) run(go())
def test_completed_request_stops_retrying_when_vector_space_keeps_changing(): def test_whole_request_rereads_cache_after_vector_space_changes(monkeypatch, tmp_path):
"""Continuous configuration churn must leave the batch empty instead of blocking forever.""" """A cache hit from the old space must not survive a later provider switch."""
async def go(): async def go():
embedding = OpenAIAsEmbedding(name="t_space_write_churn", backend="openai", model="v3", dimensions=2) monkeypatch.setattr(
LocalEmbeddingStore,
"component_metadata_path",
property(lambda _self: tmp_path),
)
embedding = FakeAsEmbedding()
embedding.vector_space_id = "v3"
store = LocalEmbeddingStore(name="t_local_cache_race", enable_cache=True)
store.as_embedding = embedding
store._cache_space = embedding.vector_space_id
first_key = store._cache_key("first")
store._cache[first_key] = np.array([3.0, 0.0], dtype=np.float16)
calls = 0
async def switch_on_miss(batch, **_kwargs):
nonlocal calls
calls += 1
if calls == 1:
embedding.vector_space_id = "v4"
version = 3.0 if calls == 1 else 4.0
return [(idx, key, np.array([version, 0.0], dtype=np.float16)) for idx, _text, key in batch]
store._compute_batch = switch_on_miss
results = await store.get_embeddings(["first", "second"])
assert calls == 2
for result in results:
np.testing.assert_array_equal(result, np.array([4.0, 0.0], dtype=np.float16))
run(go())
def test_whole_request_stops_retrying_when_vector_space_keeps_changing():
"""Continuous configuration churn must discard the whole request instead of blocking forever."""
async def go():
embedding = FakeAsEmbedding()
embedding.vector_space_id = "v3"
store = LocalEmbeddingStore(name="t_local_write_churn") store = LocalEmbeddingStore(name="t_local_write_churn")
store.as_embedding = embedding store.as_embedding = embedding
store._cache_space = embedding.vector_space_id store._cache_space = embedding.vector_space_id
@ -552,12 +590,11 @@ def test_completed_request_stops_retrying_when_vector_space_keeps_changing():
async def change_space_every_time(_batch, **_kwargs): async def change_space_every_time(_batch, **_kwargs):
nonlocal calls nonlocal calls
calls += 1 calls += 1
embedding.model = FakeProviderModel(f"v{calls + 3}") embedding.vector_space_id = f"v{calls + 3}"
return [(0, "key", np.array([float(calls), 0.0], dtype=np.float16))] return [(0, "key", np.array([float(calls), 0.0], dtype=np.float16))]
store._compute_batch = change_space_every_time store._compute_batch = change_space_every_time
results = [None] results = await store.get_embeddings(["text"])
await store._fill_misses([(0, "text", "key")], results)
assert calls == 3 assert calls == 3
assert results == [None] assert results == [None]