This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new ea5e5ea8 [plan][python] Remove split async vector-store query
workaround (#1148)
ea5e5ea8 is described below
commit ea5e5ea89c1e76f8ea18ad766f9ee14d67a3ec44
Author: Wenjin Xie <[email protected]>
AuthorDate: Mon Sep 21 22:44:45 2026 +0800
[plan][python] Remove split async vector-store query workaround (#1148)
Generated-by: Codex CLI 0.153.4 (GPT-6)
Co-authored-by: Codex <[email protected]>
---
.../api/vectorstores/python/PythonVectorStore.java | 47 ---------
.../test/VectorStoreCrossLanguageTest.java | 24 +++--
.../plan/actions/ContextRetrievalAction.java | 114 ++++-----------------
.../flink_agents/api/vector_stores/vector_store.py | 10 --
.../vector_stores/chroma/chroma_vector_store.py | 18 ----
5 files changed, 38 insertions(+), 175 deletions(-)
diff --git
a/api/src/main/java/org/apache/flink/agents/api/vectorstores/python/PythonVectorStore.java
b/api/src/main/java/org/apache/flink/agents/api/vectorstores/python/PythonVectorStore.java
index 4a393adc..7edf17c4 100644
---
a/api/src/main/java/org/apache/flink/agents/api/vectorstores/python/PythonVectorStore.java
+++
b/api/src/main/java/org/apache/flink/agents/api/vectorstores/python/PythonVectorStore.java
@@ -168,53 +168,6 @@ public class PythonVectorStore extends BaseVectorStore
implements PythonResource
}
}
- /** Embed query text via the configured model (no numpy, so it stays on
the async pool). */
- public float[] embedQuery(String text) {
- return getEmbeddingModel().embed(text);
- }
-
- /**
- * Convert the raw embedding to the Python store's native vector form.
This runs the numpy
- * conversion on the mailbox thread: numpy releases/re-acquires the GIL
during the copy, and
- * pemja keeps a single PyThreadState, so doing it on an async worker
thread can stall the
- * interpreter (observed as a hang in CI; benign with spare cores
locally). The returned Python
- * object is forwarded back into {@link #queryNormalized}. See
- * https://github.com/apache/flink-agents/issues/844.
- */
- public Object normalizeEmbedding(float[] embedding) {
- List<Float> embeddingList = new ArrayList<>(embedding.length);
- for (float v : embedding) {
- embeddingList.add(v);
- }
- Map<String, Object> kwargs = new HashMap<>();
- kwargs.put("embeddings", embeddingList);
- return adapter.callMethod(vectorStore, "_normalize_embeddings",
kwargs);
- }
-
- /** Query with a pre-normalized embedding; numpy-free, so it stays on the
async pool. */
- @SuppressWarnings("unchecked")
- public List<Document> queryNormalized(
- Object normalizedEmbedding,
- int limit,
- @Nullable String collection,
- @Nullable Map<String, Object> filters,
- Map<String, Object> args) {
- Map<String, Object> kwargs = new HashMap<>(args);
- kwargs.put("embedding", normalizedEmbedding);
- kwargs.put("limit", limit);
- if (collection != null) {
- kwargs.put("collection_name", collection);
- }
- if (filters != null) {
- kwargs.put("filters", filters);
- }
- try (PythonObjectScope scope = new PythonObjectScope()) {
- Object pythonDocuments =
- scope.own(adapter.callMethod(vectorStore,
"_query_embedding", kwargs));
- return adapter.fromPythonDocuments((List<PyObject>)
pythonDocuments);
- }
- }
-
@Override
@SuppressWarnings("unchecked")
public List<String> addEmbedding(
diff --git
a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/VectorStoreCrossLanguageTest.java
b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/VectorStoreCrossLanguageTest.java
index 554462d9..eb5a8824 100644
---
a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/VectorStoreCrossLanguageTest.java
+++
b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/VectorStoreCrossLanguageTest.java
@@ -19,6 +19,7 @@
package org.apache.flink.agents.resource.test;
import org.apache.flink.agents.api.AgentsExecutionEnvironment;
+import org.apache.flink.agents.api.agents.AgentExecutionOptions;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
@@ -26,8 +27,9 @@ import
org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.util.CloseableIterator;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Assumptions;
+import org.junit.jupiter.api.Timeout;
import org.junit.jupiter.params.ParameterizedTest;
-import org.junit.jupiter.params.provider.ValueSource;
+import org.junit.jupiter.params.provider.CsvSource;
import java.io.IOException;
import java.util.Map;
@@ -44,22 +46,26 @@ public class VectorStoreCrossLanguageTest {
}
@ParameterizedTest
- @ValueSource(strings = {"JAVA", "PYTHON"})
- public void testVectorStoreIntegration(String embeddingType) throws
Exception {
+ @CsvSource({"JAVA, false", "JAVA, true", "PYTHON, false", "PYTHON, true"})
+ @Timeout(120)
+ public void testVectorStoreIntegration(String embeddingType, boolean
ragAsync)
+ throws Exception {
System.setProperty("EMBEDDING_TYPE", embeddingType);
Assumptions.assumeTrue(ollamaReady, "Ollama Server information is not
provided");
StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
env.setParallelism(1);
- final DataStreamSource<String> inputStream = env.fromData("What is
Apache Flink");
+ final DataStreamSource<String> inputStream =
+ env.fromData("What is Apache Flink", "What is ChromaDB");
final AgentsExecutionEnvironment agentEnv =
AgentsExecutionEnvironment.getExecutionEnvironment(env);
+ agentEnv.getConfig().set(AgentExecutionOptions.RAG_ASYNC, ragAsync);
+ agentEnv.getConfig().set(AgentExecutionOptions.NUM_ASYNC_THREADS, 2);
final DataStream<Object> outputStream =
- agentEnv.fromDataStream(
- inputStream, (KeySelector<String, String>)
value -> "orderKey")
+ agentEnv.fromDataStream(inputStream, (KeySelector<String,
String>) value -> value)
.apply(new VectorStoreCrossLanguageAgent())
.toDataStream();
@@ -67,7 +73,11 @@ public class VectorStoreCrossLanguageTest {
agentEnv.execute();
- checkResult(results);
+ try (results) {
+ checkResult(results);
+ checkResult(results);
+ Assertions.assertFalse(results.hasNext(), "Unexpected extra
retrieval response");
+ }
}
@SuppressWarnings("unchecked")
diff --git
a/plan/src/main/java/org/apache/flink/agents/plan/actions/ContextRetrievalAction.java
b/plan/src/main/java/org/apache/flink/agents/plan/actions/ContextRetrievalAction.java
index 02610d08..900c149f 100644
---
a/plan/src/main/java/org/apache/flink/agents/plan/actions/ContextRetrievalAction.java
+++
b/plan/src/main/java/org/apache/flink/agents/plan/actions/ContextRetrievalAction.java
@@ -25,9 +25,7 @@ import org.apache.flink.agents.api.context.RunnerContext;
import org.apache.flink.agents.api.event.ContextRetrievalRequestEvent;
import org.apache.flink.agents.api.event.ContextRetrievalResponseEvent;
import org.apache.flink.agents.api.resource.ResourceType;
-import org.apache.flink.agents.api.resource.python.PythonObjectScope;
import org.apache.flink.agents.api.vectorstores.BaseVectorStore;
-import org.apache.flink.agents.api.vectorstores.Document;
import org.apache.flink.agents.api.vectorstores.VectorStoreQuery;
import org.apache.flink.agents.api.vectorstores.VectorStoreQueryResult;
import org.apache.flink.agents.api.vectorstores.python.PythonVectorStore;
@@ -73,36 +71,27 @@ public class ContextRetrievalAction {
contextRetrievalRequestEvent.getQuery(),
contextRetrievalRequestEvent.getMaxResults());
- final VectorStoreQueryResult result;
- if (ragAsync && vectorStore instanceof PythonVectorStore) {
- // A Python store's query path runs numpy, which can stall on
the async pool
- // (pemja keeps a single PyThreadState; numpy
releasing/re-acquiring the GIL on a
- // worker thread hangs intermittently — seen in CI, fine
locally with spare cores).
- // Keep only that numpy step on the mailbox thread; embed and
query stay async.
- result = queryPythonAsync((PythonVectorStore) vectorStore,
vectorStoreQuery, ctx);
- } else {
- DurableCallable<VectorStoreQueryResult> callable =
- new DurableCallable<VectorStoreQueryResult>() {
- @Override
- public String getId() {
- return "rag-async";
- }
-
- @Override
- public Class<VectorStoreQueryResult>
getResultClass() {
- return VectorStoreQueryResult.class;
- }
-
- @Override
- public VectorStoreQueryResult call() throws
Exception {
- return vectorStore.query(vectorStoreQuery);
- }
- };
- result =
- ragAsync
- ? ctx.durableExecuteAsync(callable).await()
- : ctx.durableExecute(callable);
- }
+ DurableCallable<VectorStoreQueryResult> callable =
+ new DurableCallable<VectorStoreQueryResult>() {
+ @Override
+ public String getId() {
+ return "rag-async";
+ }
+
+ @Override
+ public Class<VectorStoreQueryResult> getResultClass() {
+ return VectorStoreQueryResult.class;
+ }
+
+ @Override
+ public VectorStoreQueryResult call() throws Exception {
+ return vectorStore.query(vectorStoreQuery);
+ }
+ };
+ final VectorStoreQueryResult result =
+ ragAsync
+ ? ctx.durableExecuteAsync(callable).await()
+ : ctx.durableExecute(callable);
ctx.sendEvent(
new ContextRetrievalResponseEvent(
@@ -111,65 +100,4 @@ public class ContextRetrievalAction {
result.getDocuments()));
}
}
-
- /**
- * Run a Python vector-store RAG query while keeping numpy off the async
pool: embed async,
- * normalize the embedding synchronously on the mailbox thread (numpy on a
worker thread can
- * stall under pemja's single PyThreadState), then query async with the
pre-normalized vector.
- * See https://github.com/apache/flink-agents/issues/844.
- */
- private static VectorStoreQueryResult queryPythonAsync(
- PythonVectorStore store, VectorStoreQuery query, RunnerContext
ctx) throws Exception {
- final float[] embedding =
- ctx.durableExecuteAsync(
- new DurableCallable<float[]>() {
- @Override
- public String getId() {
- return "rag-embed";
- }
-
- @Override
- public Class<float[]> getResultClass() {
- return float[].class;
- }
-
- @Override
- public float[] call() {
- return
store.embedQuery(query.getQueryText());
- }
- })
- .await();
-
- try (PythonObjectScope scope = new PythonObjectScope()) {
- final Object normalized =
scope.own(store.normalizeEmbedding(embedding));
-
- final List<Document> documents =
- ctx.durableExecuteAsync(
- new DurableCallable<List<Document>>() {
- @Override
- public String getId() {
- return "rag-query";
- }
-
- @SuppressWarnings("unchecked")
- @Override
- public Class<List<Document>>
getResultClass() {
- return (Class<List<Document>>)
(Class<?>) List.class;
- }
-
- @Override
- public List<Document> call() {
- return store.queryNormalized(
- normalized,
- query.getLimit(),
- query.getCollection(),
- query.getFilters(),
- store.getStoreKwargs());
- }
- })
- .await();
-
- return new VectorStoreQueryResult(documents);
- }
- }
}
diff --git a/python/flink_agents/api/vector_stores/vector_store.py
b/python/flink_agents/api/vector_stores/vector_store.py
index 2cf04c49..d16d00c7 100644
--- a/python/flink_agents/api/vector_stores/vector_store.py
+++ b/python/flink_agents/api/vector_stores/vector_store.py
@@ -397,16 +397,6 @@ class BaseVectorStore(Resource, ABC):
**kwargs: Vector store specific parameters.
"""
- @staticmethod
- def _normalize_embeddings(embeddings: list[float]) -> Any:
- """Pre-process a query embedding before the search call.
-
- Hook for backends whose query path performs CPU/numpy work that must
run
- on the mailbox thread rather than an async cross-language worker (see
the
- ChromaDB override). Default is identity.
- """
- return embeddings
-
@abstractmethod
def _query_embedding(
self,
diff --git
a/python/flink_agents/integrations/vector_stores/chroma/chroma_vector_store.py
b/python/flink_agents/integrations/vector_stores/chroma/chroma_vector_store.py
index cab7ec3d..e65ba922 100644
---
a/python/flink_agents/integrations/vector_stores/chroma/chroma_vector_store.py
+++
b/python/flink_agents/integrations/vector_stores/chroma/chroma_vector_store.py
@@ -22,7 +22,6 @@ from typing import Any, Dict, Generator, List
import chromadb
from chromadb import ClientAPI as ChromaClient
from chromadb import CloudClient
-from chromadb.api.types import normalize_embeddings
from chromadb.config import Settings
from pydantic import Field
from typing_extensions import override
@@ -317,20 +316,6 @@ class ChromaVectorStore(CollectionManageableVectorStore):
metadatas=[doc.metadata for doc in documents],
)
- @staticmethod
- @override
- def _normalize_embeddings(embeddings: List[float]) -> Any:
- """Convert the raw embedding to chroma's numpy form on the caller's
thread.
-
- chroma's query normally runs this np.array conversion internally. numpy
- releases and re-acquires the GIL during the copy, which can stall on an
- async pemja worker thread (pemja keeps a single PyThreadState) — seen
as a
- hang in CI, benign locally with spare cores. Running it here lets the
- mailbox thread do the numpy step, so the async query sees a ready
ndarray.
- See https://github.com/apache/flink-agents/issues/844.
- """
- return normalize_embeddings([embeddings])[0]
-
@override
def _query_embedding(
self,
@@ -340,9 +325,6 @@ class ChromaVectorStore(CollectionManageableVectorStore):
filters: Dict[str, Any] | None = None,
**kwargs: Any,
) -> List[Document]:
- # ``embedding`` may be a pre-normalized ndarray (async path) or a raw
list
- # (sync path); chroma takes the ndarray branch for the former,
avoiding any
- # numpy work on this thread.
collection = self._resolve_collection(collection_name, kwargs)
results = collection.query(
query_embeddings=[embedding],
