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 4140c638 [runtime][java] Make embedded Python interpreters
thread-confined (#1092)
4140c638 is described below
commit 4140c6385dfad2bbf20fdc6d391ff47df0984e82
Author: Wenjin Xie <[email protected]>
AuthorDate: Tue Sep 15 16:20:11 2026 +0800
[runtime][java] Make embedded Python interpreters thread-confined (#1092)
Generated-by: Codex 0.144.5 (GPT-5)
Co-authored-by: Codex <[email protected]>
---
.../PythonCollectionManageableVectorStore.java | 4 +-
.../PythonCollectionManageableVectorStoreTest.java | 5 +-
docs/content/docs/development/workflow_agent.md | 14 +
.../ConcurrentChatModelCrossLanguageAgent.java | 108 ++++++
.../test/ConcurrentChatModelCrossLanguageTest.java | 69 ++++
.../resource/test/Mem0LongTermMemoryTest.java | 2 -
.../concurrent_chat_model_cross_language_agent.py | 109 ++++++
.../concurrent_chat_model_cross_language_test.py | 99 +++++
.../flink_agents/runtime/java/java_chat_model.py | 29 +-
.../runtime/tests/test_java_chat_model.py | 63 ++++
.../runtime/async/AsyncExecutorThreadFactory.java | 55 ++-
.../runtime/async/ContinuationActionExecutor.java | 3 +
.../runtime/operator/ActionExecutionOperator.java | 3 +-
.../runtime/operator/ActionTaskContextManager.java | 7 +-
.../runtime/operator/PythonBridgeManager.java | 56 ++-
.../runtime/python/utils/JavaResourceAdapter.java | 37 +-
.../runtime/python/utils/PythonActionExecutor.java | 145 +++++---
.../python/utils/PythonInterpreterManager.java | 409 ++++++++++++++++++++
.../python/utils/PythonResourceAdapterImpl.java | 40 +-
.../runtime/async/ContinuationActionExecutor.java | 23 +-
.../async/AsyncExecutorThreadFactoryTest.java | 25 ++
.../runtime/operator/PythonBridgeManagerTest.java | 32 +-
.../python/utils/JavaResourceAdapterTest.java | 24 +-
.../python/utils/PythonActionExecutorTest.java | 112 +++++-
.../python/utils/PythonInterpreterManagerTest.java | 414 +++++++++++++++++++++
.../utils/PythonResourceAdapterImplTest.java | 12 +-
26 files changed, 1725 insertions(+), 174 deletions(-)
diff --git
a/api/src/main/java/org/apache/flink/agents/api/vectorstores/python/PythonCollectionManageableVectorStore.java
b/api/src/main/java/org/apache/flink/agents/api/vectorstores/python/PythonCollectionManageableVectorStore.java
index 26cce423..6fb1e315 100644
---
a/api/src/main/java/org/apache/flink/agents/api/vectorstores/python/PythonCollectionManageableVectorStore.java
+++
b/api/src/main/java/org/apache/flink/agents/api/vectorstores/python/PythonCollectionManageableVectorStore.java
@@ -70,6 +70,8 @@ public class PythonCollectionManageableVectorStore extends
PythonVectorStore
@Override
public void deleteCollection(String name) throws Exception {
- this.vectorStore.invokeMethod("delete_collection", name);
+ Map<String, Object> kwargs = new HashMap<>();
+ kwargs.put("name", name);
+ this.adapter.callMethod(vectorStore, "delete_collection", kwargs);
}
}
diff --git
a/api/src/test/java/org/apache/flink/agents/api/vectorstores/python/PythonCollectionManageableVectorStoreTest.java
b/api/src/test/java/org/apache/flink/agents/api/vectorstores/python/PythonCollectionManageableVectorStoreTest.java
index 566db3fe..e407bcc4 100644
---
a/api/src/test/java/org/apache/flink/agents/api/vectorstores/python/PythonCollectionManageableVectorStoreTest.java
+++
b/api/src/test/java/org/apache/flink/agents/api/vectorstores/python/PythonCollectionManageableVectorStoreTest.java
@@ -41,6 +41,7 @@ import static org.mockito.ArgumentMatchers.argThat;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.when;
public class PythonCollectionManageableVectorStoreTest {
@@ -157,7 +158,9 @@ public class PythonCollectionManageableVectorStoreTest {
vectorStore.deleteCollection(collectionName);
- verify(mockVectorStore).invokeMethod("delete_collection",
collectionName);
+ verify(mockAdapter)
+ .callMethod(mockVectorStore, "delete_collection",
Map.of("name", collectionName));
+ verifyNoInteractions(mockVectorStore);
}
@Test
diff --git a/docs/content/docs/development/workflow_agent.md
b/docs/content/docs/development/workflow_agent.md
index b8166cd0..90cde7db 100644
--- a/docs/content/docs/development/workflow_agent.md
+++ b/docs/content/docs/development/workflow_agent.md
@@ -657,6 +657,20 @@ To use async execution on JDK 21+, user should append jvm
option `--add-exports=
{{< /tab >}}
{{< /tabs >}}
+{{< hint warning >}}
+Async actions in either language may call cross-language resources. Within one
async execution,
+keep the cross-language call chain to one reverse callback. Only the following
call shapes are
+supported:
+
+- Python action → Java resource → Python callback
+- Java action → Python resource → Java callback
+
+For example, a Python action may use a Java `ChatModelSetup` backed by a Python
+`ChatModelConnection`. Multiple or deeper cross-language callbacks, such as
+Python→Java→Python→Java or Java→Python→Java→Python, are not recommended and
may deadlock
+when an outer call occupies an async or interpreter worker while waiting for
another worker.
+{{< /hint >}}
+
### Cross-language Actions
An action declared in one language can dispatch its body to the other language
by setting a `target` on the decorator/annotation. The decorated function or
annotated method then acts as a stub — it should raise so direct calls outside
the framework fail loud.
diff --git
a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/ConcurrentChatModelCrossLanguageAgent.java
b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/ConcurrentChatModelCrossLanguageAgent.java
new file mode 100644
index 00000000..f5558c2e
--- /dev/null
+++
b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/ConcurrentChatModelCrossLanguageAgent.java
@@ -0,0 +1,108 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.agents.resource.test;
+
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.EventType;
+import org.apache.flink.agents.api.InputEvent;
+import org.apache.flink.agents.api.OutputEvent;
+import org.apache.flink.agents.api.agents.Agent;
+import org.apache.flink.agents.api.annotation.Action;
+import org.apache.flink.agents.api.annotation.ChatModelConnection;
+import org.apache.flink.agents.api.annotation.ChatModelSetup;
+import org.apache.flink.agents.api.chat.messages.ChatMessage;
+import org.apache.flink.agents.api.chat.messages.MessageRole;
+import org.apache.flink.agents.api.chat.model.BaseChatModelConnection;
+import org.apache.flink.agents.api.context.RunnerContext;
+import org.apache.flink.agents.api.event.ChatRequestEvent;
+import org.apache.flink.agents.api.event.ChatResponseEvent;
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import org.apache.flink.agents.api.resource.ResourceName;
+import org.apache.flink.agents.api.tools.Tool;
+
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+/** Java agent exercising concurrent Java-to-Python-to-Java chat calls. */
+public class ConcurrentChatModelCrossLanguageAgent extends Agent {
+
+ /** Java connection that only returns after two chat requests overlap. */
+ public static class OverlappingJavaChatModelConnection extends
BaseChatModelConnection {
+ private final CountDownLatch concurrentCalls = new CountDownLatch(2);
+
+ public OverlappingJavaChatModelConnection(
+ ResourceDescriptor descriptor, ResourceContext
resourceContext) {
+ super(descriptor, resourceContext);
+ }
+
+ @Override
+ public ChatMessage chat(
+ List<ChatMessage> messages, List<Tool> tools, Map<String,
Object> modelParams) {
+ concurrentCalls.countDown();
+ try {
+ if (!concurrentCalls.await(30, TimeUnit.SECONDS)) {
+ throw new IllegalStateException(
+ "Timed out waiting for concurrent cross-language
chat request.");
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IllegalStateException(
+ "Interrupted while waiting for concurrent
cross-language chat request.", e);
+ }
+
+ ChatMessage request = messages.get(messages.size() - 1);
+ return new ChatMessage(
+ MessageRole.ASSISTANT, "java-connection:" +
request.getContent());
+ }
+ }
+
+ @ChatModelConnection
+ public static ResourceDescriptor overlappingJavaConnection() {
+ return ResourceDescriptor.Builder.newBuilder(
+ OverlappingJavaChatModelConnection.class.getName())
+ .build();
+ }
+
+ @ChatModelSetup
+ public static ResourceDescriptor pythonChatModel() {
+ return
ResourceDescriptor.Builder.newBuilder(ResourceName.ChatModel.PYTHON_WRAPPER_SETUP)
+ .addInitialArgument("pythonClazz",
ResourceName.ChatModel.Python.OLLAMA_SETUP)
+ .addInitialArgument("connection", "overlappingJavaConnection")
+ .addInitialArgument("model", "mock-model")
+ .addInitialArgument("extract_reasoning", false)
+ .build();
+ }
+
+ @Action(EventType.InputEvent)
+ public static void requestChat(Event event, RunnerContext ctx) {
+ String input = String.valueOf(InputEvent.fromEvent(event).getInput());
+ ctx.sendEvent(
+ new ChatRequestEvent(
+ "pythonChatModel", List.of(new
ChatMessage(MessageRole.USER, input))));
+ }
+
+ @Action(EventType.ChatResponseEvent)
+ public static void emitResponse(Event event, RunnerContext ctx) {
+ ChatResponseEvent response = ChatResponseEvent.fromEvent(event);
+ ctx.sendEvent(new OutputEvent(response.getResponse().getContent()));
+ }
+}
diff --git
a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/ConcurrentChatModelCrossLanguageTest.java
b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/ConcurrentChatModelCrossLanguageTest.java
new file mode 100644
index 00000000..4f909c8e
--- /dev/null
+++
b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/ConcurrentChatModelCrossLanguageTest.java
@@ -0,0 +1,69 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+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.environment.StreamExecutionEnvironment;
+import org.apache.flink.util.CloseableIterator;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** E2E coverage for concurrent calls through a Python setup and Java
connection. */
+public class ConcurrentChatModelCrossLanguageTest {
+
+ @Test
+ public void testConcurrentPythonSetupWithJavaConnection() throws Exception
{
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(1);
+
+ DataStream<String> inputStream = env.fromData("first-request",
"second-request");
+ AgentsExecutionEnvironment agentsEnv =
+ AgentsExecutionEnvironment.getExecutionEnvironment(env);
+ agentsEnv.getConfig().set(AgentExecutionOptions.NUM_ASYNC_THREADS, 2);
+ agentsEnv.getConfig().set(AgentExecutionOptions.CHAT_ASYNC, true);
+
+ DataStream<Object> outputStream =
+ agentsEnv
+ .fromDataStream(inputStream, (KeySelector<String,
String>) value -> value)
+ .apply(new ConcurrentChatModelCrossLanguageAgent())
+ .toDataStream();
+
+ CloseableIterator<Object> results = outputStream.collectAsync();
+ agentsEnv.execute();
+
+ List<String> responses = new ArrayList<>();
+ while (results.hasNext()) {
+ responses.add(String.valueOf(results.next()));
+ }
+
+ assertThat(responses).hasSize(2);
+ assertThat(new HashSet<>(responses))
+ .isEqualTo(
+ Set.of("java-connection:first-request",
"java-connection:second-request"));
+ }
+}
diff --git
a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/Mem0LongTermMemoryTest.java
b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/Mem0LongTermMemoryTest.java
index 0aa31fd5..c6dce155 100644
---
a/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/Mem0LongTermMemoryTest.java
+++
b/e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/Mem0LongTermMemoryTest.java
@@ -27,7 +27,6 @@ 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.Disabled;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.ValueSource;
@@ -61,7 +60,6 @@ import static
org.apache.flink.agents.resource.test.Mem0LongTermMemoryAgent.OLLA
* MILVUS_URI} env var
* </ul>
*/
-@Disabled("Disabled until #1087 is resolved")
public class Mem0LongTermMemoryTest {
private final boolean embeddingReady;
diff --git
a/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/concurrent_chat_model_cross_language_agent.py
b/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/concurrent_chat_model_cross_language_agent.py
new file mode 100644
index 00000000..b90b96e0
--- /dev/null
+++
b/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/concurrent_chat_model_cross_language_agent.py
@@ -0,0 +1,109 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#################################################################################
+import threading
+from typing import Any, List, Sequence
+
+from pydantic import PrivateAttr
+from typing_extensions import override
+
+from flink_agents.api.agents.agent import Agent
+from flink_agents.api.agents.types import OutputSchema
+from flink_agents.api.chat_message import ChatMessage, MessageRole
+from flink_agents.api.chat_models.chat_model import BaseChatModelConnection
+from flink_agents.api.decorators import action, chat_model_connection,
chat_model_setup
+from flink_agents.api.events.chat_event import ChatRequestEvent,
ChatResponseEvent
+from flink_agents.api.events.event import Event, InputEvent, OutputEvent
+from flink_agents.api.events.event_type import EventType
+from flink_agents.api.resource import ResourceDescriptor, ResourceName
+from flink_agents.api.runner_context import RunnerContext
+from flink_agents.api.tools.tool import Tool
+
+
+class OverlappingPythonChatModelConnection(BaseChatModelConnection):
+ """Python connection that only returns after two chat requests overlap."""
+
+ _concurrent_calls: threading.Barrier = PrivateAttr(
+ default_factory=lambda: threading.Barrier(2)
+ )
+
+ @override
+ def chat(
+ self,
+ messages: Sequence[ChatMessage],
+ tools: List[Tool] | None = None,
+ output_schema: OutputSchema | None = None,
+ **kwargs: Any,
+ ) -> ChatMessage:
+ """Echo the request after observing another in-flight call."""
+ self._reject_unsupported_output_schema(output_schema)
+ try:
+ self._concurrent_calls.wait(timeout=30)
+ except threading.BrokenBarrierError as error:
+ message = "Timed out waiting for concurrent cross-language chat
request."
+ raise RuntimeError(message) from error
+
+ return ChatMessage(
+ role=MessageRole.ASSISTANT,
+ content=f"python-connection:{messages[-1].content}",
+ )
+
+
+class ConcurrentChatModelCrossLanguageAgent(Agent):
+ """Python agent exercising concurrent Python-to-Java-to-Python chat
calls."""
+
+ @chat_model_connection
+ @staticmethod
+ def overlapping_python_connection() -> ResourceDescriptor:
+ """Declare the Python connection used behind the Java setup."""
+ return ResourceDescriptor(
+ clazz=(
+ f"{OverlappingPythonChatModelConnection.__module__}."
+ f"{OverlappingPythonChatModelConnection.__name__}"
+ )
+ )
+
+ @chat_model_setup
+ @staticmethod
+ def java_chat_model() -> ResourceDescriptor:
+ """Declare a Java setup backed by the Python connection."""
+ return ResourceDescriptor(
+ clazz=ResourceName.ChatModel.JAVA_WRAPPER_SETUP,
+ java_clazz=ResourceName.ChatModel.Java.OLLAMA_SETUP,
+ connection="overlapping_python_connection",
+ model="mock-model",
+ extract_reasoning=False,
+ )
+
+ @action(EventType.InputEvent)
+ @staticmethod
+ def request_chat(event: Event, ctx: RunnerContext) -> None:
+ """Send one chat request per input key."""
+ input_value = str(InputEvent.from_event(event).input)
+ ctx.send_event(
+ ChatRequestEvent(
+ model="java_chat_model",
+ messages=[ChatMessage(role=MessageRole.USER,
content=input_value)],
+ )
+ )
+
+ @action(EventType.ChatResponseEvent)
+ @staticmethod
+ def emit_response(event: Event, ctx: RunnerContext) -> None:
+ """Emit the cross-language chat response."""
+ response = ChatResponseEvent.from_event(event).response
+ ctx.send_event(OutputEvent(output=response.content))
diff --git
a/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/concurrent_chat_model_cross_language_test.py
b/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/concurrent_chat_model_cross_language_test.py
new file mode 100644
index 00000000..609bdd1c
--- /dev/null
+++
b/python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/concurrent_chat_model_cross_language_test.py
@@ -0,0 +1,99 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#################################################################################
+import os
+import sys
+import sysconfig
+from pathlib import Path
+
+from pyflink.common import Encoder, WatermarkStrategy
+from pyflink.common.typeinfo import Types
+from pyflink.datastream import RuntimeExecutionMode, StreamExecutionEnvironment
+from pyflink.datastream.connectors.file_system import (
+ FileSource,
+ StreamFormat,
+ StreamingFileSink,
+)
+
+from flink_agents.api.core_options import AgentExecutionOptions
+from flink_agents.api.execution_environment import AgentsExecutionEnvironment
+from
flink_agents.e2e_tests.e2e_tests_resource_cross_language.concurrent_chat_model_cross_language_agent
import (
+ ConcurrentChatModelCrossLanguageAgent,
+)
+
+current_dir = Path(__file__).parent
+
+os.environ["PYTHONPATH"] = sysconfig.get_paths()["purelib"]
+
+
+def test_concurrent_java_setup_with_python_connection(tmp_path: Path) -> None:
+ """Run two overlapping chat requests through Java setup and Python
connection."""
+ env = StreamExecutionEnvironment.get_execution_environment()
+ env.set_runtime_mode(RuntimeExecutionMode.STREAMING)
+ env.set_parallelism(1)
+ env.set_python_executable(sys.executable)
+
+ input_datastream = env.from_source(
+ source=FileSource.for_record_stream_format(
+ StreamFormat.text_line_format(),
+ f"file:///{current_dir}/../resources/java_chat_module_input",
+ ).build(),
+ watermark_strategy=WatermarkStrategy.no_watermarks(),
+ source_name="concurrent_chat_inputs",
+ ).map(lambda value: str(value))
+
+ agents_env = AgentsExecutionEnvironment.get_execution_environment(env=env)
+ agents_env.get_config().set(AgentExecutionOptions.NUM_ASYNC_THREADS, 2)
+ agents_env.get_config().set(AgentExecutionOptions.CHAT_ASYNC, True)
+ output_datastream = (
+ agents_env.from_datastream(
+ input=input_datastream,
+ key_selector=lambda value: value,
+ )
+ .apply(ConcurrentChatModelCrossLanguageAgent())
+ .to_datastream()
+ )
+
+ result_dir = tmp_path / "results"
+ result_dir.mkdir(parents=True, exist_ok=True)
+ output_datastream.map(
+ lambda value: str(value).replace("\n", "").replace("\r", ""),
+ Types.STRING(),
+ ).add_sink(
+ StreamingFileSink.for_row_format(
+ base_path=str(result_dir.absolute()),
+ encoder=Encoder.simple_string_encoder(),
+ ).build()
+ )
+
+ agents_env.execute()
+
+ responses: list[str] = []
+ for file in result_dir.iterdir():
+ if file.is_dir():
+ for child in file.iterdir():
+ with child.open() as result_file:
+ responses.extend(line.strip() for line in result_file if
line.strip())
+ elif file.is_file():
+ with file.open() as result_file:
+ responses.extend(line.strip() for line in result_file if
line.strip())
+
+ assert len(responses) == 2
+ assert set(responses) == {
+ "python-connection:calculate the sum of 1 and 2.",
+ "python-connection:Tell me a joke about cats.",
+ }
diff --git a/python/flink_agents/runtime/java/java_chat_model.py
b/python/flink_agents/runtime/java/java_chat_model.py
index 15df3157..43aca50e 100644
--- a/python/flink_agents/runtime/java/java_chat_model.py
+++ b/python/flink_agents/runtime/java/java_chat_model.py
@@ -30,6 +30,21 @@ from flink_agents.api.tools.tool import Tool
from flink_agents.runtime.java.java_resource_wrapper import (
set_java_resource_metric_group,
)
+from flink_agents.runtime.python_java_utils import (
+ from_java_chat_message,
+ normalize_tool_call_id,
+)
+
+
+def _to_java_chat_message(j_resource_adapter: Any, message: ChatMessage) ->
Any:
+ """Build a Java message from fields extracted on the Python calling
thread."""
+ tool_calls = [normalize_tool_call_id(call) for call in message.tool_calls]
+ return j_resource_adapter.fromPythonChatMessage(
+ message.role.value,
+ message.content,
+ tool_calls,
+ message.extra_args,
+ )
class JavaChatModelConnectionImpl(JavaChatModelConnection):
@@ -82,7 +97,7 @@ class JavaChatModelConnectionImpl(JavaChatModelConnection):
"""
self._reject_unsupported_output_schema(output_schema)
java_messages = [
- self._j_resource_adapter.fromPythonChatMessage(message)
+ _to_java_chat_message(self._j_resource_adapter, message)
for message in messages
]
java_tools = [
@@ -91,11 +106,6 @@ class JavaChatModelConnectionImpl(JavaChatModelConnection):
]
j_response_message = self._j_resource.chat(java_messages, java_tools,
kwargs)
- # Convert Java response back to Python format
- from flink_agents.runtime.python_java_utils import (
- from_java_chat_message,
- )
-
return from_java_chat_message(j_response_message)
@override
@@ -180,16 +190,11 @@ class JavaChatModelSetupImpl(JavaChatModelSetup):
"""
# Convert Python messages to Java format
java_messages = [
- self._j_resource_adapter.fromPythonChatMessage(message)
+ _to_java_chat_message(self._j_resource_adapter, message)
for message in messages
]
j_response_message = self._j_resource.chat(
java_messages, prompt_args or {}, kwargs
)
- # Convert Java response back to Python format
- from flink_agents.runtime.python_java_utils import (
- from_java_chat_message,
- )
-
return from_java_chat_message(j_response_message)
diff --git a/python/flink_agents/runtime/tests/test_java_chat_model.py
b/python/flink_agents/runtime/tests/test_java_chat_model.py
new file mode 100644
index 00000000..87d4ae91
--- /dev/null
+++ b/python/flink_agents/runtime/tests/test_java_chat_model.py
@@ -0,0 +1,63 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#################################################################################
+from typing import Any
+
+from flink_agents.api.chat_message import ChatMessage, MessageRole
+from flink_agents.runtime.java.java_chat_model import _to_java_chat_message
+
+
+class _JavaResourceAdapter:
+ def __init__(self) -> None:
+ self.arguments: tuple[Any, ...] | None = None
+ self.result = object()
+
+ def fromPythonChatMessage(self, *arguments: Any) -> Any:
+ self.arguments = arguments
+ return self.result
+
+
+def test_to_java_chat_message_extracts_java_safe_fields() -> None:
+ adapter = _JavaResourceAdapter()
+ message = ChatMessage(
+ role=MessageRole.ASSISTANT,
+ content="hello",
+ tool_calls=[
+ {
+ "id": 7,
+ "type": "function",
+ "function": {"name": "lookup", "arguments": "{}"},
+ }
+ ],
+ extra_args={"reasoning": "brief"},
+ )
+
+ result = _to_java_chat_message(adapter, message)
+
+ assert result is adapter.result
+ assert adapter.arguments == (
+ "assistant",
+ "hello",
+ [
+ {
+ "id": "7",
+ "type": "function",
+ "function": {"name": "lookup", "arguments": "{}"},
+ }
+ ],
+ {"reasoning": "brief"},
+ )
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/async/AsyncExecutorThreadFactory.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/async/AsyncExecutorThreadFactory.java
index 788a9c3e..14edc1e2 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/async/AsyncExecutorThreadFactory.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/async/AsyncExecutorThreadFactory.java
@@ -17,6 +17,9 @@
*/
package org.apache.flink.agents.runtime.async;
+import java.util.Objects;
+import java.util.Queue;
+import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadFactory;
@@ -35,13 +38,63 @@ import java.util.concurrent.ThreadFactory;
public final class AsyncExecutorThreadFactory implements ThreadFactory {
private static final String NAME_PREFIX = "flink-agents-java-async-";
+ private static final ThreadLocal<Boolean> ASYNC_EXECUTOR_THREAD = new
ThreadLocal<>();
private final ThreadFactory delegate = Executors.defaultThreadFactory();
+ private final Runnable threadCleanup;
+ private final Queue<Thread> createdThreads = new ConcurrentLinkedQueue<>();
+
+ public AsyncExecutorThreadFactory() {
+ this(() -> {});
+ }
+
+ public AsyncExecutorThreadFactory(Runnable threadCleanup) {
+ this.threadCleanup = Objects.requireNonNull(threadCleanup);
+ }
+
+ /** Returns whether the calling thread is owned by the Flink Agents Java
async executor. */
+ public static boolean isAsyncExecutorThread() {
+ return Boolean.TRUE.equals(ASYNC_EXECUTOR_THREAD.get());
+ }
@Override
public Thread newThread(Runnable runnable) {
- Thread thread = delegate.newThread(runnable);
+ Thread thread =
+ delegate.newThread(
+ () -> {
+ ASYNC_EXECUTOR_THREAD.set(true);
+ try {
+ runnable.run();
+ } finally {
+ try {
+ threadCleanup.run();
+ } finally {
+ ASYNC_EXECUTOR_THREAD.remove();
+ }
+ }
+ });
thread.setName(NAME_PREFIX + thread.getName());
+ createdThreads.add(thread);
return thread;
}
+
+ /** Waits until every worker created by this factory has run its exit
cleanup. */
+ public void awaitThreadExit() {
+ boolean interrupted = false;
+ try {
+ for (Thread thread : createdThreads) {
+ while (thread.isAlive()) {
+ try {
+ thread.join();
+ } catch (InterruptedException e) {
+ interrupted = true;
+ }
+ }
+ }
+ } finally {
+ if (interrupted) {
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
}
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java
index f6134da8..6a369254 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java
@@ -37,6 +37,9 @@ public class ContinuationActionExecutor {
/** Creates a new ContinuationActionExecutor. */
public ContinuationActionExecutor(int numAsyncThreads) {}
+ /** JDK 11 fallback has no worker threads, so the cleanup hook is never
needed. */
+ public ContinuationActionExecutor(int numAsyncThreads, Runnable
threadCleanup) {}
+
/**
* Executes the action. In JDK 11, this simply runs the action
synchronously.
*
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
index 1db7421b..94c050a9 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionExecutionOperator.java
@@ -245,7 +245,8 @@ public class ActionExecutionOperator<IN, OUT> extends
AbstractStreamOperator<OUT
// init context manager for runner context creation and memory contexts
contextManager =
new ActionTaskContextManager(
-
agentPlan.getConfig().get(AgentExecutionOptions.NUM_ASYNC_THREADS));
+
agentPlan.getConfig().get(AgentExecutionOptions.NUM_ASYNC_THREADS),
+ pythonBridge::releaseCurrentThreadInterpreter);
mailboxProcessor = getMailboxProcessor();
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java
index fb107a87..30211b11 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/ActionTaskContextManager.java
@@ -77,8 +77,13 @@ class ActionTaskContextManager implements AutoCloseable {
private ContinuationActionExecutor continuationActionExecutor;
ActionTaskContextManager(int numAsyncThreads) {
+ this(numAsyncThreads, () -> {});
+ }
+
+ ActionTaskContextManager(int numAsyncThreads, Runnable asyncThreadCleanup)
{
this.actionTaskContexts = new HashMap<>();
- this.continuationActionExecutor = new
ContinuationActionExecutor(numAsyncThreads);
+ this.continuationActionExecutor =
+ new ContinuationActionExecutor(numAsyncThreads,
asyncThreadCleanup);
}
/**
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
index ba5b77c7..e8b4d646 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java
@@ -17,6 +17,7 @@
*/
package org.apache.flink.agents.runtime.operator;
+import org.apache.flink.agents.api.agents.AgentExecutionOptions;
import org.apache.flink.agents.api.memory.LongTermMemoryOptions;
import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.plan.JavaFunction;
@@ -32,6 +33,7 @@ import
org.apache.flink.agents.runtime.metrics.FlinkAgentsMetricGroupImpl;
import org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl;
import org.apache.flink.agents.runtime.python.utils.JavaResourceAdapter;
import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor;
+import org.apache.flink.agents.runtime.python.utils.PythonInterpreterManager;
import org.apache.flink.agents.runtime.python.utils.PythonResourceAdapterImpl;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.JobID;
@@ -57,7 +59,7 @@ import static
org.apache.flink.agents.plan.actions.Utils.supportAsync;
*
* <ul>
* <li>The {@link PythonEnvironmentManager} that prepares dependencies and
the Pemja runtime.
- * <li>The {@link PythonInterpreter} obtained from that environment.
+ * <li>The thread-confined Python interpreters obtained from that
environment.
* <li>The {@link PythonActionExecutor} (when the plan contains Python
actions, Python-owned
* resources, or Mem0).
* <li>The {@link PythonRunnerContextImpl} consumed by Python actions.
@@ -71,7 +73,7 @@ import static
org.apache.flink.agents.plan.actions.Utils.supportAsync;
* in that case all accessors return {@code null} and {@link #isInitialized()}
returns {@code
* false}. {@link #close()} closes the owned resources in the reverse order of
creation: {@code
* longTermMemory} → {@code pythonActionExecutor} → {@code
pythonResourceAdapter} → {@code
- * pythonInterpreter} → {@code pythonEnvironmentManager}.
+ * pythonInterpreterManager} → {@code pythonEnvironmentManager}.
*
* <p>Design constraint: package-private; no manager-to-manager held
references. Other managers
* receive what they need (e.g. the Python runner context, the action
executor) via method
@@ -82,7 +84,8 @@ class PythonBridgeManager implements AutoCloseable {
private static final Logger LOG =
LoggerFactory.getLogger(PythonBridgeManager.class);
private PythonEnvironmentManager pythonEnvironmentManager;
- private PythonInterpreter pythonInterpreter;
+ private PythonInterpreter initializingPythonInterpreter;
+ private PythonInterpreterManager pythonInterpreterManager;
private PythonActionExecutor pythonActionExecutor;
private PythonRunnerContextImpl pythonRunnerContext;
private PythonResourceAdapterImpl pythonResourceAdapter;
@@ -100,14 +103,16 @@ class PythonBridgeManager implements AutoCloseable {
* <p>Scans the agent plan for any {@link PythonFunction} action or
Python-owned resource
* provider. If neither is present and Mem0 is not configured, this method
is a no-op and {@link
* #isInitialized()} stays {@code false}. Otherwise it builds the {@link
- * PythonEnvironmentManager}, opens an embedded {@link PythonInterpreter},
refreshes the shared
- * import state for the current dependency generation, constructs the
shared {@link
+ * PythonEnvironmentManager}, opens an owner {@link PythonInterpreter},
refreshes the shared
+ * import state for the current dependency generation, and creates a {@link
+ * PythonInterpreterManager} that binds interpreters to managed Java
workers and routes other
+ * callers through bounded callback workers. It then constructs the shared
{@link
* PythonRunnerContextImpl}, wires the Java/Python resource adapters, and
conditionally
* initializes the Python resource adapter (when Python-owned resources or
Mem0 are present) and
* the Python action executor (when Python actions, Python-owned
resources, or Mem0 are present,
* since the executor is also the bridge that materializes Python-owned
resources). The
- * generation guard runs immediately after interpreter construction and
before any user module
- * import.
+ * generation guard runs immediately after owner-interpreter construction
and before any user
+ * module import.
*
* @param agentPlan the agent plan describing actions and resources.
* @param resourceCache the resource cache visible to both languages.
@@ -157,12 +162,13 @@ class PythonBridgeManager implements AutoCloseable {
dependencyInfo, tmpDirs, new
HashMap<>(System.getenv()), jobId);
pythonEnvironmentManager.open();
EmbeddedPythonEnvironment env =
pythonEnvironmentManager.createEnvironment();
- pythonInterpreter = env.getInterpreter();
+ PythonInterpreter ownerInterpreter = env.getInterpreter();
+ initializingPythonInterpreter = ownerInterpreter;
String dependencyGeneration =
pythonEnvironmentManager.getBaseDirectory();
String pythonPath = env.getEnv().get("PYTHONPATH");
boolean dependencyGenerationChanged =
PythonDependencyGenerationManager.ensurePythonDependencyGeneration(
- pythonInterpreter,
+ ownerInterpreter,
jobId,
dependencyGeneration,
pythonPath == null ? "" : pythonPath);
@@ -172,6 +178,15 @@ class PythonBridgeManager implements AutoCloseable {
dependencyGeneration,
jobId);
}
+ // Transfer ownership only after dependency generation is active.
If generation setup
+ // fails, close() can still release initializingPythonInterpreter;
if manager
+ // construction fails, its constructor releases the owner itself.
+ initializingPythonInterpreter = null;
+ pythonInterpreterManager =
+ new PythonInterpreterManager(
+ ownerInterpreter,
+ env::getInterpreter,
+
agentPlan.getConfig().get(AgentExecutionOptions.NUM_ASYNC_THREADS));
pythonRunnerContext =
new PythonRunnerContextImpl(
metricGroup,
@@ -182,9 +197,7 @@ class PythonBridgeManager implements AutoCloseable {
javaResourceAdapter =
new JavaResourceAdapter(
- resourceCache.getResourceContext(),
- pythonInterpreter,
- userCodeClassLoader);
+ resourceCache.getResourceContext(),
userCodeClassLoader);
if (containPythonResource || mem0Configured) {
initPythonResourceAdapter(agentPlan, resourceCache);
}
@@ -242,7 +255,8 @@ class PythonBridgeManager implements AutoCloseable {
*/
private void wireLongTermMemory(AgentPlan agentPlan, Runnable
mailboxThreadChecker) {
PyObject pyCtx = pythonActionExecutor.getPythonRunnerContext();
- Object pyLtm =
pythonInterpreter.invoke("python_java_utils.get_long_term_memory", pyCtx);
+ Object pyLtm =
+
pythonInterpreterManager.invoke("python_java_utils.get_long_term_memory",
pyCtx);
if (pyLtm == null) {
throw new IllegalStateException(
String.format(
@@ -270,7 +284,7 @@ class PythonBridgeManager implements AutoCloseable {
throws Exception {
pythonActionExecutor =
new PythonActionExecutor(
- pythonInterpreter,
+ pythonInterpreterManager,
agentPlan,
javaResourceAdapter,
pythonRunnerContext,
@@ -282,7 +296,9 @@ class PythonBridgeManager implements AutoCloseable {
throws Exception {
pythonResourceAdapter =
new PythonResourceAdapterImpl(
- resourceCache.getResourceContext(), pythonInterpreter,
javaResourceAdapter);
+ resourceCache.getResourceContext(),
+ pythonInterpreterManager,
+ javaResourceAdapter);
pythonResourceAdapter.open();
PythonMCPResourceDiscovery.discoverPythonMCPResources(
agentPlan.getResourceProviders(), pythonResourceAdapter,
resourceCache);
@@ -318,6 +334,13 @@ class PythonBridgeManager implements AutoCloseable {
return initialized;
}
+ /** Releases a Python interpreter bound to the current managed Java async
worker. */
+ void releaseCurrentThreadInterpreter() {
+ if (pythonInterpreterManager != null) {
+ pythonInterpreterManager.releaseCurrentThreadInterpreter();
+ }
+ }
+
@Override
public void close() throws Exception {
// Close every component even when an earlier one fails, so a failing
action executor
@@ -333,7 +356,8 @@ class PythonBridgeManager implements AutoCloseable {
longTermMemory,
pythonActionExecutor,
pythonResourceAdapter,
- pythonInterpreter,
+ pythonInterpreterManager,
+ initializingPythonInterpreter,
pythonEnvironmentManager
}) {
if (closeable == null) {
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/JavaResourceAdapter.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/JavaResourceAdapter.java
index 4a03a827..2affcf85 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/JavaResourceAdapter.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/JavaResourceAdapter.java
@@ -30,7 +30,6 @@ import org.apache.flink.agents.api.tools.ToolResponse;
import org.apache.flink.agents.api.vectorstores.Document;
import org.apache.flink.agents.plan.tools.FunctionTool;
import org.apache.flink.agents.plan.tools.ToolMetadataFactory;
-import pemja.core.PythonInterpreter;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
@@ -44,8 +43,6 @@ public class JavaResourceAdapter {
private final ResourceContext resourceContext;
- private final transient PythonInterpreter interpreter;
-
/**
* Class loader used to resolve Java tool methods declared by name.
Captured at construction
* (the operator passes its {@code
RuntimeContext.getUserCodeClassLoader()}) because pemja
@@ -54,12 +51,8 @@ public class JavaResourceAdapter {
*/
private final transient ClassLoader userCodeClassLoader;
- public JavaResourceAdapter(
- ResourceContext resourceContext,
- PythonInterpreter interpreter,
- ClassLoader userCodeClassLoader) {
+ public JavaResourceAdapter(ResourceContext resourceContext, ClassLoader
userCodeClassLoader) {
this.resourceContext = resourceContext;
- this.interpreter = interpreter;
this.userCodeClassLoader = userCodeClassLoader;
}
@@ -94,23 +87,23 @@ public class JavaResourceAdapter {
* Convert a Python chat message to a Java chat message. This method is
intended for use by the
* Python interpreter.
*
- * @param pythonChatMessage the Python chat message
+ * <p>The Python caller extracts the message fields before crossing into
Java. Keeping this
+ * method Java-only avoids an unnecessary Python→Java→Python callback
while constructing the
+ * Java value.
+ *
+ * @param roleValue the Python message role value
+ * @param content the message content
+ * @param toolCalls the normalized tool calls
+ * @param extraArgs additional message arguments
* @return the Java chat message
*/
- public ChatMessage fromPythonChatMessage(Object pythonChatMessage) {
+ public ChatMessage fromPythonChatMessage(
+ String roleValue,
+ String content,
+ List<Map<String, Object>> toolCalls,
+ Map<String, Object> extraArgs) {
// TODO: Delete this method after the pemja findClass method is fixed.
- ChatMessage chatMessage = new ChatMessage();
- if (interpreter == null) {
- throw new IllegalStateException("Python interpreter is not set.");
- }
- String roleValue =
- (String)
- interpreter.invoke(
- "python_java_utils.update_java_chat_message",
- pythonChatMessage,
- chatMessage);
- chatMessage.setRole(MessageRole.fromValue(roleValue));
- return chatMessage;
+ return new ChatMessage(MessageRole.fromValue(roleValue), content,
toolCalls, extraArgs);
}
/**
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
index 4012bd11..3f97b961 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
@@ -31,7 +31,6 @@ import
org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl;
import org.apache.flink.agents.runtime.python.resource.PythonRuntimeResource;
import org.apache.flink.types.Row;
import org.apache.flink.util.ExceptionUtils;
-import pemja.core.PythonInterpreter;
import pemja.core.object.PyObject;
import java.io.IOException;
@@ -45,11 +44,6 @@ import static org.apache.flink.util.Preconditions.checkState;
/** Execute the corresponding Python action in the agent. */
public class PythonActionExecutor implements AutoCloseable {
- private static final String PYTHON_IMPORTS =
- "from flink_agents.plan import function\n"
- + "from flink_agents.runtime import flink_runner_context\n"
- + "from flink_agents.runtime import python_java_utils";
-
// =========== RUNNER CONTEXT ===========
private static final String CREATE_FLINK_RUNNER_CONTEXT =
"flink_runner_context.create_flink_runner_context";
@@ -87,6 +81,7 @@ public class PythonActionExecutor implements AutoCloseable {
// =========== PYTHON AWAITABLE ===========
private static final String CALL_PYTHON_AWAITABLE =
"function.call_python_awaitable";
+ private static final String CALL_PYTHON_FUNCTION =
"function.call_python_function";
private static final String PYTHON_AWAITABLE_VAR_NAME_PREFIX =
"python_awaitable_";
private static final AtomicLong PYTHON_AWAITABLE_VAR_ID = new
AtomicLong(0);
@@ -101,7 +96,7 @@ public class PythonActionExecutor implements AutoCloseable {
private static final String GET_OUTPUT_FROM_OUTPUT_EVENT =
"python_java_utils.get_output_from_output_event";
- private final PythonInterpreter interpreter;
+ private final PythonInterpreterManager interpreterManager;
private final AgentPlan agentPlan;
private final PythonRunnerContextImpl runnerContext;
private final JavaResourceAdapter javaResourceAdapter;
@@ -110,13 +105,13 @@ public class PythonActionExecutor implements
AutoCloseable {
private PyObject pythonRunnerContext;
public PythonActionExecutor(
- PythonInterpreter interpreter,
+ PythonInterpreterManager interpreterManager,
AgentPlan agentPlan,
JavaResourceAdapter javaResourceAdapter,
PythonRunnerContextImpl runnerContext,
String jobIdentifier)
throws JsonProcessingException {
- this.interpreter = interpreter;
+ this.interpreterManager = interpreterManager;
this.agentPlan = agentPlan;
this.runnerContext = runnerContext;
this.javaResourceAdapter = javaResourceAdapter;
@@ -136,7 +131,7 @@ public class PythonActionExecutor implements AutoCloseable {
@SuppressWarnings("unchecked")
public Map<String, Resource> eagerMaterialize(ResourceType type) {
Object pythonResources =
- interpreter.invoke(EAGER_MATERIALIZE, pythonRunnerContext,
type.getValue());
+ interpreterManager.invoke(EAGER_MATERIALIZE,
pythonRunnerContext, type.getValue());
if (pythonResources == null) {
return Collections.emptyMap();
}
@@ -158,68 +153,66 @@ public class PythonActionExecutor implements
AutoCloseable {
*/
public boolean addTaskLifecycleListener(PyObject pythonListener) {
Object registered =
- interpreter.invoke(
+ interpreterManager.invoke(
ADD_TASK_LIFECYCLE_LISTENER, pythonRunnerContext,
pythonListener);
return Boolean.TRUE.equals(registered);
}
/** Forwards {@code onRecordStart} to the Python runtime lifecycle
listeners. */
public void notifyRecordStart(Object key) {
- interpreter.invoke(NOTIFY_RECORD_START, pythonRunnerContext, key);
+ interpreterManager.invoke(NOTIFY_RECORD_START, pythonRunnerContext,
key);
}
/** Forwards {@code onActionPrepared} to the Python runtime lifecycle
listeners. */
public void notifyActionPrepared(ActionTask task) {
- interpreter.invoke(NOTIFY_ACTION_PREPARED, pythonRunnerContext, task);
+ interpreterManager.invoke(NOTIFY_ACTION_PREPARED, pythonRunnerContext,
task);
}
/** Forwards {@code onActionStarted} to the Python runtime lifecycle
listeners. */
public void notifyActionStarted(ActionTask task) {
- interpreter.invoke(NOTIFY_ACTION_STARTED, pythonRunnerContext, task);
+ interpreterManager.invoke(NOTIFY_ACTION_STARTED, pythonRunnerContext,
task);
}
/** Forwards {@code onActionTransferred} to the Python runtime lifecycle
listeners. */
public void notifyActionTransferred(ActionTask fromTask, ActionTask
toTask) {
- interpreter.invoke(NOTIFY_ACTION_TRANSFERRED, pythonRunnerContext,
fromTask, toTask);
+ interpreterManager.invoke(NOTIFY_ACTION_TRANSFERRED,
pythonRunnerContext, fromTask, toTask);
}
/** Forwards {@code onActionFinishing} to the Python runtime lifecycle
listeners. */
public void notifyActionFinishing(ActionTask task) {
- interpreter.invoke(NOTIFY_ACTION_FINISHING, pythonRunnerContext, task);
+ interpreterManager.invoke(NOTIFY_ACTION_FINISHING,
pythonRunnerContext, task);
}
/** Forwards {@code onActionFinished} to the Python runtime lifecycle
listeners. */
public void notifyActionFinished(ActionTask task) {
- interpreter.invoke(NOTIFY_ACTION_FINISHED, pythonRunnerContext, task);
+ interpreterManager.invoke(NOTIFY_ACTION_FINISHED, pythonRunnerContext,
task);
}
/** Forwards {@code onActionReused} to the Python runtime lifecycle
listeners. */
public void notifyActionReused(ActionTask task) {
- interpreter.invoke(NOTIFY_ACTION_REUSED, pythonRunnerContext, task);
+ interpreterManager.invoke(NOTIFY_ACTION_REUSED, pythonRunnerContext,
task);
}
/** Forwards {@code onActionFailed} to the Python runtime lifecycle
listeners. */
public void notifyActionFailed(ActionTask task, Throwable error) {
- interpreter.invoke(NOTIFY_ACTION_FAILED, pythonRunnerContext, task,
error);
+ interpreterManager.invoke(NOTIFY_ACTION_FAILED, pythonRunnerContext,
task, error);
}
/** Forwards {@code onRecordFinished} to the Python runtime lifecycle
listeners. */
public void notifyRecordFinished(Object key) {
- interpreter.invoke(NOTIFY_RECORD_FINISHED, pythonRunnerContext, key);
+ interpreterManager.invoke(NOTIFY_RECORD_FINISHED, pythonRunnerContext,
key);
}
public void open() throws Exception {
- interpreter.exec(PYTHON_IMPORTS);
-
pythonAsyncThreadPool =
(PyObject)
- interpreter.invoke(
+ interpreterManager.invoke(
CREATE_ASYNC_THREAD_POOL,
agentPlan.getConfig().get(AgentExecutionOptions.NUM_ASYNC_THREADS));
pythonRunnerContext =
(PyObject)
- interpreter.invoke(
+ interpreterManager.invoke(
CREATE_FLINK_RUNNER_CONTEXT,
runnerContext,
new
ObjectMapper().writeValueAsString(agentPlan),
@@ -240,23 +233,42 @@ public class PythonActionExecutor implements
AutoCloseable {
*/
public String executePythonFunction(PythonFunction function, Event event)
throws Exception {
runnerContext.checkNoPendingEvents();
- function.setInterpreter(interpreter);
-
String eventJson = new ObjectMapper().writeValueAsString(event);
- try (PyObject pythonEventObject =
- (PyObject)
interpreter.invoke(CONVERT_JSON_TO_PYTHON_EVENT, eventJson);
- PyObject calledResult =
- (PyObject) function.call(pythonEventObject,
pythonRunnerContext)) {
- if (calledResult == null) {
- return null;
- } else {
- // must be a coroutine (awaitable)
- String pythonAwaitableRef =
- PYTHON_AWAITABLE_VAR_NAME_PREFIX
- + PYTHON_AWAITABLE_VAR_ID.incrementAndGet();
- interpreter.set(pythonAwaitableRef, calledResult);
- return pythonAwaitableRef;
- }
+ try {
+ return interpreterManager.withInterpreter(
+ interpreter -> {
+ try (PythonObjectScope scope = new
PythonObjectScope()) {
+ PyObject pythonEventObject =
+ scope.own(
+ (PyObject)
+ interpreter.invoke(
+
CONVERT_JSON_TO_PYTHON_EVENT,
+ eventJson));
+ PyObject calledResult =
+ scope.own(
+ (PyObject)
+ interpreter.invoke(
+
CALL_PYTHON_FUNCTION,
+
function.getModule(),
+
function.getQualName(),
+ new Object[] {
+
pythonEventObject,
+
pythonRunnerContext
+ }));
+ if (calledResult == null) {
+ return null;
+ }
+
+ // The result must be a coroutine (awaitable).
Keep conversion,
+ // invocation, reference retention, and temporary
reference cleanup on
+ // the same thread-confined interpreter.
+ String pythonAwaitableRef =
+ PYTHON_AWAITABLE_VAR_NAME_PREFIX
+ +
PYTHON_AWAITABLE_VAR_ID.incrementAndGet();
+ interpreter.set(pythonAwaitableRef, calledResult);
+ return pythonAwaitableRef;
+ }
+ });
} catch (Exception e) {
runnerContext.drainEvents(null);
throw new PythonActionExecutionException("Failed to execute Python
action", e);
@@ -266,7 +278,7 @@ public class PythonActionExecutor implements AutoCloseable {
public Event wrapToInputEvent(Object eventData) throws IOException {
checkState(eventData instanceof byte[]);
// wrap_to_input_event returns a JSON string
- Object result = interpreter.invoke(WRAP_TO_INPUT_EVENT, eventData);
+ Object result = interpreterManager.invoke(WRAP_TO_INPUT_EVENT,
eventData);
checkState(result instanceof String);
return Event.fromJson((String) result);
}
@@ -281,14 +293,14 @@ public class PythonActionExecutor implements
AutoCloseable {
String keySerialization =
pythonKeyIsPickled ? PICKLED_KEY_SERIALIZATION :
EXPLICIT_KEY_SERIALIZATION;
return (String)
- interpreter.invoke(
+ interpreterManager.invoke(
CONVERT_TO_PYTHON_KEY_TEXT, (byte[]) logicalKey,
keySerialization);
}
return String.valueOf(logicalKey);
}
public Object getOutputFromOutputEvent(String eventJson) {
- return interpreter.invoke(GET_OUTPUT_FROM_OUTPUT_EVENT, eventJson);
+ return interpreterManager.invoke(GET_OUTPUT_FROM_OUTPUT_EVENT,
eventJson);
}
/**
@@ -302,23 +314,32 @@ public class PythonActionExecutor implements
AutoCloseable {
* @return true if the awaitable has completed; false otherwise
*/
public boolean callPythonAwaitable(String pythonAwaitableRef) {
- try (PythonObjectScope scope = new PythonObjectScope()) {
- PyObject pythonAwaitable = scope.own((PyObject)
interpreter.get(pythonAwaitableRef));
- checkState(
- pythonAwaitable != null,
- "Python awaitable '%s' not found in interpreter.",
- pythonAwaitableRef);
- // Actions communicate through Events, so this caller consumes
only the completion flag.
- Object invokeResult =
- scope.own(interpreter.invoke(CALL_PYTHON_AWAITABLE,
pythonAwaitable));
- checkState(invokeResult instanceof Object[] && ((Object[])
invokeResult).length == 2);
- Object[] result = (Object[]) invokeResult;
- boolean finished = (boolean) result[0];
- if (finished) {
- interpreter.exec("del " + pythonAwaitableRef);
- }
- return finished;
- }
+ // Calling awaitable.send(None) in Python returns a tuple of
(finished, output).
+ return interpreterManager.withInterpreter(
+ interpreter -> {
+ try (PythonObjectScope scope = new PythonObjectScope()) {
+ PyObject pythonAwaitable =
+ scope.own((PyObject)
interpreter.get(pythonAwaitableRef));
+ checkState(
+ pythonAwaitable != null,
+ "Python awaitable '%s' not found in
interpreter.",
+ pythonAwaitableRef);
+ // Actions communicate through Events, so this caller
consumes only the
+ // completion flag.
+ Object invokeResult =
+ scope.own(
+
interpreter.invoke(CALL_PYTHON_AWAITABLE, pythonAwaitable));
+ checkState(
+ invokeResult instanceof Object[]
+ && ((Object[]) invokeResult).length ==
2);
+ Object[] result = (Object[]) invokeResult;
+ boolean finished = (boolean) result[0];
+ if (finished) {
+ interpreter.exec("del " + pythonAwaitableRef);
+ }
+ return finished;
+ }
+ });
}
@Override
@@ -328,7 +349,7 @@ public class PythonActionExecutor implements AutoCloseable {
// resource cache unreleased, and PythonBridgeManager closes the
interpreter right behind
// us, so there is no later chance to run it. The first failure is
rethrown with the later
// one suppressed, matching the ladders in the managers above.
- if (interpreter == null) {
+ if (interpreterManager == null) {
return;
}
@@ -359,7 +380,7 @@ public class PythonActionExecutor implements AutoCloseable {
private void closePythonObject(String closeFunction, PyObject
pythonObject) throws Exception {
if (pythonObject != null) {
try (pythonObject) {
- interpreter.invoke(closeFunction, pythonObject);
+ interpreterManager.invoke(closeFunction, pythonObject);
}
}
}
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManager.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManager.java
new file mode 100644
index 00000000..c007126b
--- /dev/null
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManager.java
@@ -0,0 +1,409 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.flink.agents.runtime.python.utils;
+
+import org.apache.flink.agents.runtime.async.AsyncExecutorThreadFactory;
+import org.apache.flink.util.ExceptionUtils;
+import pemja.core.PythonInterpreter;
+
+import java.util.ArrayList;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Queue;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+/**
+ * Owns the Pemja interpreters used by one operator subtask.
+ *
+ * <p>A {@link PythonInterpreter} is never shared by concurrently executing
threads. The owner
+ * mailbox thread and Flink Agents' managed Java async workers bind an
interpreter to themselves.
+ * Calls from every other thread are routed to a bounded set of Java callback
workers, each of which
+ * owns its interpreter. This distinction is required for threads created by
CPython (for example
+ * Mem0 or Python's async executor): creating a Pemja interpreter there would
attach a second {@code
+ * PyThreadState} to the same native thread.
+ *
+ * <p>Pemja's {@code MULTI_THREAD} interpreters share the same CPython main
interpreter, so Python
+ * resource objects can still be initialized once and passed as opaque handles
to calls made by any
+ * of these thread-confined interpreters. Calls from one external thread
always use the same
+ * callback worker, keeping multi-step conversions involving opaque Python
handles on one
+ * interpreter. Each interpreter has its own globals, however, so the bridge
modules must be
+ * imported for every newly created interpreter.
+ */
+public final class PythonInterpreterManager implements AutoCloseable {
+
+ private static final int DEFAULT_CALLBACK_WORKERS = 2;
+ private static final AtomicInteger CALLBACK_POOL_ID = new AtomicInteger();
+
+ static final String PYTHON_IMPORTS =
+ "from flink_agents.plan import function\n"
+ + "from flink_agents.runtime import flink_runner_context\n"
+ + "from flink_agents.runtime import python_java_utils";
+
+ private final Supplier<PythonInterpreter> interpreterFactory;
+ private final Consumer<PythonInterpreter> interpreterInitializer;
+ private final Thread ownerThread;
+ private final Map<Thread, PythonInterpreter> interpreters = new
ConcurrentHashMap<>();
+ private final ThreadLocal<PythonInterpreter> threadInterpreter = new
ThreadLocal<>();
+ private final ThreadLocal<Integer> callbackLane;
+ private final ExecutorService[] callbackExecutors;
+ private final Queue<Thread> callbackThreads = new
ConcurrentLinkedQueue<>();
+ private final Queue<Throwable> workerCloseFailures = new
ConcurrentLinkedQueue<>();
+ private final ReentrantReadWriteLock lifecycleLock = new
ReentrantReadWriteLock();
+
+ private volatile boolean closed;
+
+ public PythonInterpreterManager(
+ PythonInterpreter ownerInterpreter, Supplier<PythonInterpreter>
interpreterFactory) {
+ this(
+ ownerInterpreter,
+ interpreterFactory,
+ PythonInterpreterManager::initializeInterpreter,
+ DEFAULT_CALLBACK_WORKERS);
+ }
+
+ public PythonInterpreterManager(
+ PythonInterpreter ownerInterpreter,
+ Supplier<PythonInterpreter> interpreterFactory,
+ int callbackWorkerCount) {
+ this(
+ ownerInterpreter,
+ interpreterFactory,
+ PythonInterpreterManager::initializeInterpreter,
+ callbackWorkerCount);
+ }
+
+ PythonInterpreterManager(
+ PythonInterpreter ownerInterpreter,
+ Supplier<PythonInterpreter> interpreterFactory,
+ Consumer<PythonInterpreter> interpreterInitializer) {
+ this(
+ ownerInterpreter,
+ interpreterFactory,
+ interpreterInitializer,
+ DEFAULT_CALLBACK_WORKERS);
+ }
+
+ PythonInterpreterManager(
+ PythonInterpreter ownerInterpreter,
+ Supplier<PythonInterpreter> interpreterFactory,
+ Consumer<PythonInterpreter> interpreterInitializer,
+ int callbackWorkerCount) {
+ if (callbackWorkerCount <= 0) {
+ throw new IllegalArgumentException("Callback worker count must be
greater than zero.");
+ }
+ this.interpreterFactory = Objects.requireNonNull(interpreterFactory);
+ this.interpreterInitializer =
Objects.requireNonNull(interpreterInitializer);
+ this.ownerThread = Thread.currentThread();
+
+ PythonInterpreter initializedOwner =
initialize(Objects.requireNonNull(ownerInterpreter));
+ interpreters.put(ownerThread, initializedOwner);
+ threadInterpreter.set(initializedOwner);
+ callbackExecutors = createCallbackExecutors(callbackWorkerCount);
+ AtomicInteger nextCallbackLane = new AtomicInteger();
+ callbackLane =
+ ThreadLocal.withInitial(
+ () ->
+ Math.floorMod(
+ nextCallbackLane.getAndIncrement(),
+ callbackExecutors.length));
+ }
+
+ /**
+ * Executes an operation on a thread-confined interpreter.
+ *
+ * <p>The owner and managed Java async threads execute inline. An
unmanaged caller is commonly a
+ * CPython-created thread, so it is assigned to a stable callback lane
instead of receiving a
+ * second Python thread state itself.
+ */
+ public <T> T withInterpreter(Function<PythonInterpreter, T> operation) {
+ lifecycleLock.readLock().lock();
+ try {
+ if (closed) {
+ throw new IllegalStateException("Python interpreter manager is
already closed.");
+ }
+ Function<PythonInterpreter, T> checkedOperation =
Objects.requireNonNull(operation);
+ PythonInterpreter current = threadInterpreter.get();
+ if (current != null) {
+ return checkedOperation.apply(current);
+ }
+ if (AsyncExecutorThreadFactory.isAsyncExecutorThread()) {
+ return checkedOperation.apply(createAndBindInterpreter());
+ }
+ return executeOnCallbackWorker(checkedOperation);
+ } finally {
+ lifecycleLock.readLock().unlock();
+ }
+ }
+
+ public void exec(String code) {
+ withInterpreter(
+ interpreter -> {
+ interpreter.exec(code);
+ return null;
+ });
+ }
+
+ public Object invoke(String name, Object... args) {
+ return withInterpreter(interpreter -> interpreter.invoke(name, args));
+ }
+
+ public Object get(String name) {
+ return withInterpreter(interpreter -> interpreter.get(name));
+ }
+
+ public void set(String name, Object value) {
+ withInterpreter(
+ interpreter -> {
+ interpreter.set(name, value);
+ return null;
+ });
+ }
+
+ private PythonInterpreter createAndBindInterpreter() {
+ PythonInterpreter interpreter = threadInterpreter.get();
+ if (interpreter != null) {
+ return interpreter;
+ }
+
+ PythonInterpreter created = initialize(interpreterFactory.get());
+ interpreters.put(Thread.currentThread(), created);
+ threadInterpreter.set(created);
+ return created;
+ }
+
+ private <T> T executeOnCallbackWorker(Function<PythonInterpreter, T>
operation) {
+ int lane = callbackLane.get();
+ Future<T> result =
+ callbackExecutors[lane].submit(() ->
operation.apply(createAndBindInterpreter()));
+ return awaitResult(result);
+ }
+
+ private static <T> T awaitResult(Future<T> result) {
+ boolean interrupted = false;
+ try {
+ while (true) {
+ try {
+ return result.get();
+ } catch (InterruptedException e) {
+ // A bridge invocation cannot be abandoned while close()
may reclaim its native
+ // thread state. Finish the accepted operation, then
restore interruption.
+ interrupted = true;
+ } catch (ExecutionException e) {
+ ExceptionUtils.rethrow(e.getCause());
+ throw new AssertionError("Unreachable after rethrowing
callback failure");
+ }
+ }
+ } finally {
+ if (interrupted) {
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ private PythonInterpreter initialize(PythonInterpreter interpreter) {
+ Objects.requireNonNull(interpreter);
+ try {
+ interpreterInitializer.accept(interpreter);
+ return interpreter;
+ } catch (Throwable initializationFailure) {
+ try {
+ interpreter.close();
+ } catch (Throwable closeFailure) {
+ initializationFailure.addSuppressed(closeFailure);
+ }
+ throw initializationFailure;
+ }
+ }
+
+ private static void initializeInterpreter(PythonInterpreter interpreter) {
+ interpreter.exec(PYTHON_IMPORTS);
+ }
+
+ /** Releases the interpreter owned by the current managed Java worker. */
+ public void releaseCurrentThreadInterpreter() {
+ if (!AsyncExecutorThreadFactory.isAsyncExecutorThread()) {
+ return;
+ }
+ lifecycleLock.writeLock().lock();
+ try {
+ closeCurrentInterpreterAndRecordFailure();
+ } finally {
+ lifecycleLock.writeLock().unlock();
+ }
+ }
+
+ private void closeCurrentInterpreterAndRecordFailure() {
+ Thread thread = Thread.currentThread();
+ PythonInterpreter interpreter = threadInterpreter.get();
+ threadInterpreter.remove();
+ PythonInterpreter registered = interpreters.remove(thread);
+ if (registered == null) {
+ return;
+ }
+ if (registered != interpreter) {
+ workerCloseFailures.add(
+ new IllegalStateException(
+ "The interpreter registered for worker "
+ + thread.getName()
+ + " did not match its thread-local
interpreter."));
+ }
+ try {
+ registered.close();
+ } catch (Throwable closeFailure) {
+ workerCloseFailures.add(closeFailure);
+ }
+ }
+
+ private ExecutorService[] createCallbackExecutors(int callbackWorkerCount)
{
+ int poolId = CALLBACK_POOL_ID.incrementAndGet();
+ ThreadFactory delegate = Executors.defaultThreadFactory();
+ ExecutorService[] executors = new ExecutorService[callbackWorkerCount];
+ for (int i = 0; i < callbackWorkerCount; i++) {
+ int workerId = i + 1;
+ executors[i] =
+ Executors.newSingleThreadExecutor(
+ runnable -> {
+ Thread thread =
+ delegate.newThread(
+ () -> {
+ try {
+ runnable.run();
+ } finally {
+
closeCurrentInterpreterAndRecordFailure();
+ }
+ });
+ thread.setName(
+ "flink-agents-python-callback-pool-"
+ + poolId
+ + "-thread-"
+ + workerId);
+ callbackThreads.add(thread);
+ return thread;
+ });
+ }
+ return executors;
+ }
+
+ @Override
+ public synchronized void close() throws Exception {
+ if (Thread.currentThread() != ownerThread) {
+ throw new IllegalStateException(
+ "Python interpreter manager must be closed by its owner
thread.");
+ }
+
+ lifecycleLock.writeLock().lock();
+ try {
+ if (closed) {
+ return;
+ }
+ closed = true;
+ } finally {
+ lifecycleLock.writeLock().unlock();
+ }
+
+ for (ExecutorService callbackExecutor : callbackExecutors) {
+ callbackExecutor.shutdown();
+ }
+ for (ExecutorService callbackExecutor : callbackExecutors) {
+ awaitTermination(callbackExecutor);
+ }
+ awaitThreadExit(callbackThreads);
+
+ closeCurrentInterpreterAndRecordFailure();
+ callbackLane.remove();
+
+ // Managed worker executors are closed before this manager. A dead
thread cannot retain a
+ // live GILState TLS slot, so this is only a defensive fallback for a
worker whose exit hook
+ // was unable to run. Never delete a live thread's Python state from
the owner thread.
+ for (Map.Entry<Thread, PythonInterpreter> entry :
+ new ArrayList<>(interpreters.entrySet())) {
+ Thread thread = entry.getKey();
+ if (thread.isAlive()) {
+ workerCloseFailures.add(
+ new IllegalStateException(
+ "Python interpreter worker is still alive
during manager close: "
+ + thread.getName()));
+ continue;
+ }
+ if (interpreters.remove(thread) == entry.getValue()) {
+ try {
+ entry.getValue().close();
+ } catch (Throwable closeFailure) {
+ workerCloseFailures.add(closeFailure);
+ }
+ }
+ }
+
+ Throwable firstFailure = null;
+ Throwable closeFailure;
+ while ((closeFailure = workerCloseFailures.poll()) != null) {
+ firstFailure = ExceptionUtils.firstOrSuppressed(closeFailure,
firstFailure);
+ }
+ if (firstFailure != null) {
+ ExceptionUtils.rethrowException(firstFailure);
+ }
+ }
+
+ private static void awaitTermination(ExecutorService executor) {
+ boolean interrupted = false;
+ try {
+ while (!executor.isTerminated()) {
+ try {
+ executor.awaitTermination(1, TimeUnit.DAYS);
+ } catch (InterruptedException e) {
+ interrupted = true;
+ }
+ }
+ } finally {
+ if (interrupted) {
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+
+ private static void awaitThreadExit(Iterable<Thread> threads) {
+ boolean interrupted = false;
+ try {
+ for (Thread thread : threads) {
+ while (thread.isAlive()) {
+ try {
+ thread.join();
+ } catch (InterruptedException e) {
+ interrupted = true;
+ }
+ }
+ }
+ } finally {
+ if (interrupted) {
+ Thread.currentThread().interrupt();
+ }
+ }
+ }
+}
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImpl.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImpl.java
index 73f206ab..ae840b3c 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImpl.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImpl.java
@@ -31,7 +31,6 @@ import org.apache.flink.agents.api.tools.Tool;
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 pemja.core.PythonInterpreter;
import pemja.core.object.PyObject;
import java.util.ArrayList;
@@ -41,8 +40,6 @@ import java.util.Map;
public class PythonResourceAdapterImpl implements PythonResourceAdapter,
AutoCloseable {
- static final String PYTHON_IMPORTS = "from flink_agents.runtime import
python_java_utils";
-
static final String JAVA_RESOURCE = "j_resource";
static final String JAVA_RESOURCE_ADAPTER = "j_resource_adapter";
@@ -82,22 +79,21 @@ public class PythonResourceAdapterImpl implements
PythonResourceAdapter, AutoClo
static final String INVOKE_PYTHON_TOOL = PYTHON_MODULE_PREFIX +
"invoke_python_tool";
private final ResourceContext resourceContext;
- private final PythonInterpreter interpreter;
+ private final PythonInterpreterManager interpreterManager;
private final JavaResourceAdapter javaResourceAdapter;
private PyObject pythonResourceContext;
public PythonResourceAdapterImpl(
ResourceContext resourceContext,
- PythonInterpreter interpreter,
+ PythonInterpreterManager interpreterManager,
JavaResourceAdapter javaResourceAdapter) {
this.resourceContext = resourceContext;
- this.interpreter = interpreter;
+ this.interpreterManager = interpreterManager;
this.javaResourceAdapter = javaResourceAdapter;
}
public void open() {
- interpreter.exec(PYTHON_IMPORTS);
- pythonResourceContext = (PyObject)
interpreter.invoke(GET_RESOURCE_CONTEXT, this);
+ pythonResourceContext = (PyObject)
interpreterManager.invoke(GET_RESOURCE_CONTEXT, this);
}
@Override
@@ -139,18 +135,18 @@ public class PythonResourceAdapterImpl implements
PythonResourceAdapter, AutoClo
kwargs.put(JAVA_RESOURCE, resource);
kwargs.put(JAVA_RESOURCE_ADAPTER, javaResourceAdapter);
kwargs.put(RESOURCE_CONTEXT_KEY, pythonResourceContext);
- return interpreter.invoke(FROM_JAVA_RESOURCE, resourceType, kwargs);
+ return interpreterManager.invoke(FROM_JAVA_RESOURCE, resourceType,
kwargs);
}
@Override
public PyObject initPythonResource(String module, String clazz,
Map<String, Object> kwargs) {
kwargs.put(RESOURCE_CONTEXT_KEY, pythonResourceContext);
- return (PyObject) interpreter.invoke(CREATE_RESOURCE, module, clazz,
kwargs);
+ return (PyObject) interpreterManager.invoke(CREATE_RESOURCE, module,
clazz, kwargs);
}
@Override
public Object toPythonChatMessage(ChatMessage message) {
- return interpreter.invoke(FROM_JAVA_CHAT_MESSAGE, message);
+ return interpreterManager.invoke(FROM_JAVA_CHAT_MESSAGE, message);
}
@Override
@@ -159,7 +155,9 @@ public class PythonResourceAdapterImpl implements
PythonResourceAdapter, AutoClo
ChatMessage chatMessage = new ChatMessage();
String roleValue =
- (String) interpreter.invoke(TO_JAVA_CHAT_MESSAGE,
pythonChatMessage, chatMessage);
+ (String)
+ interpreterManager.invoke(
+ TO_JAVA_CHAT_MESSAGE, pythonChatMessage,
chatMessage);
chatMessage.setRole(MessageRole.fromValue(roleValue));
return chatMessage;
}
@@ -168,7 +166,7 @@ public class PythonResourceAdapterImpl implements
PythonResourceAdapter, AutoClo
public Object toPythonDocuments(List<Document> documents) {
List<Object> pythonDocuments = new ArrayList<>();
for (Document document : documents) {
- pythonDocuments.add(interpreter.invoke(FROM_JAVA_DOCUMENT,
document));
+ pythonDocuments.add(interpreterManager.invoke(FROM_JAVA_DOCUMENT,
document));
}
return pythonDocuments;
}
@@ -189,7 +187,7 @@ public class PythonResourceAdapterImpl implements
PythonResourceAdapter, AutoClo
@Override
public Object toPythonVectorStoreQuery(VectorStoreQuery query) {
- return interpreter.invoke(FROM_JAVA_VECTOR_STORE_QUERY, query);
+ return interpreterManager.invoke(FROM_JAVA_VECTOR_STORE_QUERY, query);
}
@Override
@@ -206,26 +204,26 @@ public class PythonResourceAdapterImpl implements
PythonResourceAdapter, AutoClo
@Override
public Object convertToPythonTool(Tool tool) {
- return interpreter.invoke(FROM_JAVA_TOOL, tool);
+ return interpreterManager.invoke(FROM_JAVA_TOOL, tool);
}
private Object convertToPythonPrompt(Prompt prompt) {
- return interpreter.invoke(FROM_JAVA_PROMPT, prompt);
+ return interpreterManager.invoke(FROM_JAVA_PROMPT, prompt);
}
@Override
public Object callMethod(Object obj, String methodName, Map<String,
Object> kwargs) {
- return interpreter.invoke(CALL_METHOD, obj, methodName, kwargs);
+ return interpreterManager.invoke(CALL_METHOD, obj, methodName, kwargs);
}
@Override
public void setMetricGroup(Object pythonResource, FlinkAgentsMetricGroup
metricGroup) {
- interpreter.invoke(SET_METRIC_GROUP, pythonResource, metricGroup);
+ interpreterManager.invoke(SET_METRIC_GROUP, pythonResource,
metricGroup);
}
@Override
public Object invoke(String name, Object... args) {
- return interpreter.invoke(name, args);
+ return interpreterManager.invoke(name, args);
}
@Override
@@ -239,7 +237,7 @@ public class PythonResourceAdapterImpl implements
PythonResourceAdapter, AutoClo
@SuppressWarnings("unchecked")
Map<String, String> result =
(Map<String, String>)
- interpreter.invoke(
+ interpreterManager.invoke(
GET_PYTHON_TOOL_METADATA, module, qualName,
injectedArgs);
if (result == null) {
throw new IllegalStateException(
@@ -250,6 +248,6 @@ public class PythonResourceAdapterImpl implements
PythonResourceAdapter, AutoClo
@Override
public Object invokePythonTool(String module, String qualName, Map<String,
Object> kwargs) {
- return interpreter.invoke(INVOKE_PYTHON_TOOL, module, qualName,
kwargs);
+ return interpreterManager.invoke(INVOKE_PYTHON_TOOL, module, qualName,
kwargs);
}
}
diff --git
a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java
b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java
index 6818441b..e34649ed 100644
---
a/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java
+++
b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java
@@ -49,11 +49,17 @@ public class ContinuationActionExecutor {
private static final ContinuationScope SCOPE = new
ContinuationScope("FlinkAgentsAction");
private final ExecutorService asyncExecutor;
+ private final AsyncExecutorThreadFactory asyncThreadFactory;
public ContinuationActionExecutor(int numAsyncThreads) {
+ this(numAsyncThreads, () -> {});
+ }
+
+ public ContinuationActionExecutor(int numAsyncThreads, Runnable
threadCleanup) {
LOG.info("Initialize fixed thread pool for async task with {}
threads", numAsyncThreads);
+ this.asyncThreadFactory = new
AsyncExecutorThreadFactory(threadCleanup);
this.asyncExecutor =
- Executors.newFixedThreadPool(numAsyncThreads, new
AsyncExecutorThreadFactory());
+ Executors.newFixedThreadPool(numAsyncThreads,
asyncThreadFactory);
}
/**
@@ -339,6 +345,21 @@ public class ContinuationActionExecutor {
public void close() {
asyncExecutor.shutdownNow();
+ boolean interrupted = false;
+ try {
+ while (!asyncExecutor.isTerminated()) {
+ try {
+ asyncExecutor.awaitTermination(1,
java.util.concurrent.TimeUnit.DAYS);
+ } catch (InterruptedException e) {
+ interrupted = true;
+ }
+ }
+ } finally {
+ if (interrupted) {
+ Thread.currentThread().interrupt();
+ }
+ }
+ asyncThreadFactory.awaitThreadExit();
}
/**
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/async/AsyncExecutorThreadFactoryTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/async/AsyncExecutorThreadFactoryTest.java
index 125d1812..a19b0477 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/async/AsyncExecutorThreadFactoryTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/async/AsyncExecutorThreadFactoryTest.java
@@ -23,6 +23,7 @@ import org.junit.jupiter.api.Test;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import static org.assertj.core.api.Assertions.assertThat;
@@ -30,6 +31,30 @@ import static org.assertj.core.api.Assertions.assertThat;
/** Tests for {@link AsyncExecutorThreadFactory} thread naming. */
class AsyncExecutorThreadFactoryTest {
+ @Test
+ void marksManagedWorkersAndRunsCleanupBeforeRemovingTheMarker() throws
Exception {
+ AtomicBoolean workerMarked = new AtomicBoolean();
+ AtomicBoolean cleanupMarked = new AtomicBoolean();
+ AsyncExecutorThreadFactory threadFactory =
+ new AsyncExecutorThreadFactory(
+ () ->
+ cleanupMarked.set(
+
AsyncExecutorThreadFactory.isAsyncExecutorThread()));
+ ExecutorService executor =
Executors.newSingleThreadExecutor(threadFactory);
+ try {
+ workerMarked.set(
+
executor.submit(AsyncExecutorThreadFactory::isAsyncExecutorThread).get());
+ } finally {
+ executor.shutdownNow();
+ assertThat(executor.awaitTermination(5,
TimeUnit.SECONDS)).isTrue();
+ threadFactory.awaitThreadExit();
+ }
+
+ assertThat(workerMarked.get()).isTrue();
+ assertThat(cleanupMarked.get()).isTrue();
+
assertThat(AsyncExecutorThreadFactory.isAsyncExecutorThread()).isFalse();
+ }
+
@Test
@DisplayName("Threads carry the descriptive flink-agents-java-async name")
void testThreadNamesCarryDescriptivePrefix() throws Exception {
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/PythonBridgeManagerTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/PythonBridgeManagerTest.java
index fa25afff..04c4535f 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/operator/PythonBridgeManagerTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/operator/PythonBridgeManagerTest.java
@@ -23,12 +23,12 @@ import org.apache.flink.agents.plan.actions.Action;
import org.apache.flink.agents.runtime.env.PythonEnvironmentManager;
import org.apache.flink.agents.runtime.memory.Mem0LongTermMemory;
import org.apache.flink.agents.runtime.python.utils.PythonActionExecutor;
+import org.apache.flink.agents.runtime.python.utils.PythonInterpreterManager;
import org.apache.flink.agents.runtime.python.utils.PythonResourceAdapterImpl;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.JobID;
import org.junit.jupiter.api.Test;
import org.mockito.InOrder;
-import pemja.core.PythonInterpreter;
import java.lang.reflect.Field;
import java.util.List;
@@ -50,7 +50,7 @@ class PythonBridgeManagerTest {
Mem0LongTermMemory longTermMemory = mock(Mem0LongTermMemory.class);
PythonActionExecutor actionExecutor = mock(PythonActionExecutor.class);
PythonResourceAdapterImpl resourceAdapter =
mock(PythonResourceAdapterImpl.class);
- PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ PythonInterpreterManager interpreterManager =
mock(PythonInterpreterManager.class);
PythonEnvironmentManager environmentManager =
mock(PythonEnvironmentManager.class);
RuntimeException actionExecutorFailure =
new RuntimeException("action executor close failed");
@@ -61,12 +61,12 @@ class PythonBridgeManagerTest {
RuntimeException resourceAdapterFailure =
new RuntimeException("resource adapter close failed");
doThrow(resourceAdapterFailure).when(resourceAdapter).close();
- doThrow(interpreterFailure).when(interpreter).close();
+ doThrow(interpreterFailure).when(interpreterManager).close();
doThrow(environmentFailure).when(environmentManager).close();
setField(bridge, "longTermMemory", longTermMemory);
setField(bridge, "pythonActionExecutor", actionExecutor);
setField(bridge, "pythonResourceAdapter", resourceAdapter);
- setField(bridge, "pythonInterpreter", interpreter);
+ setField(bridge, "pythonInterpreterManager", interpreterManager);
setField(bridge, "pythonEnvironmentManager", environmentManager);
assertThatThrownBy(bridge::close)
@@ -79,12 +79,12 @@ class PythonBridgeManagerTest {
longTermMemory,
actionExecutor,
resourceAdapter,
- interpreter,
+ interpreterManager,
environmentManager);
closeOrder.verify(longTermMemory).close();
closeOrder.verify(actionExecutor).close();
closeOrder.verify(resourceAdapter).close();
- closeOrder.verify(interpreter).close();
+ closeOrder.verify(interpreterManager).close();
closeOrder.verify(environmentManager).close();
}
@@ -128,14 +128,14 @@ class PythonBridgeManagerTest {
void closeReleasesInterpreterAndEnvironmentWhenActionExecutorFails()
throws Exception {
PythonBridgeManager bridge = new PythonBridgeManager();
PythonActionExecutor actionExecutor = mock(PythonActionExecutor.class);
- PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ PythonInterpreterManager interpreterManager =
mock(PythonInterpreterManager.class);
PythonEnvironmentManager environmentManager =
mock(PythonEnvironmentManager.class);
doThrow(new IllegalStateException("action executor close failed"))
.when(actionExecutor)
.close();
setField(bridge, "pythonActionExecutor", actionExecutor);
- setField(bridge, "pythonInterpreter", interpreter);
+ setField(bridge, "pythonInterpreterManager", interpreterManager);
setField(bridge, "pythonEnvironmentManager", environmentManager);
assertThatThrownBy(bridge::close)
@@ -144,9 +144,9 @@ class PythonBridgeManagerTest {
// Contract 3: a lone failure arrives with nothing attached to
it.
.satisfies(thrown ->
assertThat(thrown.getSuppressed()).isEmpty());
- InOrder inOrder = inOrder(actionExecutor, interpreter,
environmentManager);
+ InOrder inOrder = inOrder(actionExecutor, interpreterManager,
environmentManager);
inOrder.verify(actionExecutor).close();
- inOrder.verify(interpreter).close();
+ inOrder.verify(interpreterManager).close();
inOrder.verify(environmentManager).close();
}
@@ -155,7 +155,7 @@ class PythonBridgeManagerTest {
void closeReportsFirstFailureWithLaterOnesSuppressed() throws Exception {
PythonBridgeManager bridge = new PythonBridgeManager();
PythonActionExecutor actionExecutor = mock(PythonActionExecutor.class);
- PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ PythonInterpreterManager interpreterManager =
mock(PythonInterpreterManager.class);
PythonEnvironmentManager environmentManager =
mock(PythonEnvironmentManager.class);
doThrow(new IllegalStateException("action executor close failed"))
.when(actionExecutor)
@@ -165,7 +165,7 @@ class PythonBridgeManagerTest {
.close();
setField(bridge, "pythonActionExecutor", actionExecutor);
- setField(bridge, "pythonInterpreter", interpreter);
+ setField(bridge, "pythonInterpreterManager", interpreterManager);
setField(bridge, "pythonEnvironmentManager", environmentManager);
assertThatThrownBy(bridge::close)
@@ -177,7 +177,7 @@ class PythonBridgeManagerTest {
.extracting(Throwable::getMessage)
.containsExactly("environment manager
close failed"));
- verify(interpreter).close();
+ verify(interpreterManager).close();
}
/**
@@ -189,13 +189,13 @@ class PythonBridgeManagerTest {
void closeReleasesInterpreterAndEnvironmentWhenActionExecutorThrowsError()
throws Exception {
PythonBridgeManager bridge = new PythonBridgeManager();
PythonActionExecutor actionExecutor = mock(PythonActionExecutor.class);
- PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ PythonInterpreterManager interpreterManager =
mock(PythonInterpreterManager.class);
PythonEnvironmentManager environmentManager =
mock(PythonEnvironmentManager.class);
OutOfMemoryError failure = new OutOfMemoryError("action executor close
failed");
doThrow(failure).when(actionExecutor).close();
setField(bridge, "pythonActionExecutor", actionExecutor);
- setField(bridge, "pythonInterpreter", interpreter);
+ setField(bridge, "pythonInterpreterManager", interpreterManager);
setField(bridge, "pythonEnvironmentManager", environmentManager);
// The Error reaches the caller unchanged rather than wrapped in an
Exception, and with
@@ -204,7 +204,7 @@ class PythonBridgeManagerTest {
.isSameAs(failure)
.satisfies(thrown ->
assertThat(thrown.getSuppressed()).isEmpty());
- verify(interpreter).close();
+ verify(interpreterManager).close();
verify(environmentManager).close();
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/JavaResourceAdapterTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/JavaResourceAdapterTest.java
index dd8ddd39..a7cb735a 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/JavaResourceAdapterTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/JavaResourceAdapterTest.java
@@ -22,6 +22,8 @@ import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.flink.agents.api.annotation.Tool;
import org.apache.flink.agents.api.annotation.ToolParam;
+import org.apache.flink.agents.api.chat.messages.ChatMessage;
+import org.apache.flink.agents.api.chat.messages.MessageRole;
import org.apache.flink.agents.api.tools.ToolParameterSource;
import org.junit.jupiter.api.Test;
@@ -32,10 +34,26 @@ import static org.assertj.core.api.Assertions.assertThat;
public class JavaResourceAdapterTest {
+ @Test
+ void buildsJavaChatMessageFromExtractedPythonFields() {
+ JavaResourceAdapter adapter =
+ new JavaResourceAdapter(null,
Thread.currentThread().getContextClassLoader());
+ List<Map<String, Object>> toolCalls = List.of(Map.of("id", "call-1",
"type", "function"));
+ Map<String, Object> extraArgs = Map.of("reasoning", "brief");
+
+ ChatMessage converted =
+ adapter.fromPythonChatMessage("user", "hello", toolCalls,
extraArgs);
+
+ assertThat(converted.getRole()).isEqualTo(MessageRole.USER);
+ assertThat(converted.getContent()).isEqualTo("hello");
+ assertThat(converted.getToolCalls()).isEqualTo(toolCalls);
+ assertThat(converted.getExtraArgs()).isEqualTo(extraArgs);
+ }
+
@Test
void getJavaToolMetadataHidesInjectedArgsAndReturnsAnnotatedDeclaration()
throws Exception {
JavaResourceAdapter adapter =
- new JavaResourceAdapter(null, null,
Thread.currentThread().getContextClassLoader());
+ new JavaResourceAdapter(null,
Thread.currentThread().getContextClassLoader());
Map<String, String> metadata =
adapter.getJavaToolMetadata(
@@ -61,7 +79,7 @@ public class JavaResourceAdapterTest {
@Test
void invokeJavaToolPreservesSuccessResponseForPythonCaller() throws
Exception {
JavaResourceAdapter adapter =
- new JavaResourceAdapter(null, null,
Thread.currentThread().getContextClassLoader());
+ new JavaResourceAdapter(null,
Thread.currentThread().getContextClassLoader());
Map<String, Object> result =
adapter.invokeJavaTool(
@@ -87,7 +105,7 @@ public class JavaResourceAdapterTest {
@Test
void invokeJavaToolPreservesErrorResponseForPythonCaller() throws
Exception {
JavaResourceAdapter adapter =
- new JavaResourceAdapter(null, null,
Thread.currentThread().getContextClassLoader());
+ new JavaResourceAdapter(null,
Thread.currentThread().getContextClassLoader());
Map<String, Object> result =
adapter.invokeJavaTool(
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java
index b2ac5060..63575170 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java
@@ -18,6 +18,7 @@
package org.apache.flink.agents.runtime.python.utils;
import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.flink.agents.api.Event;
import org.apache.flink.agents.api.InputEvent;
import org.apache.flink.agents.api.agents.AgentExecutionOptions;
import org.apache.flink.agents.plan.AgentPlan;
@@ -34,11 +35,14 @@ import java.lang.reflect.Field;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.ArgumentMatchers.anyString;
-import static org.mockito.ArgumentMatchers.isNull;
import static org.mockito.ArgumentMatchers.same;
import static org.mockito.Mockito.clearInvocations;
import static org.mockito.Mockito.doThrow;
@@ -63,6 +67,65 @@ class PythonActionExecutorTest {
private static final String CONVERT_JSON_TO_PYTHON_EVENT =
"python_java_utils.convert_json_to_python_event";
private static final String CALL_PYTHON_AWAITABLE =
"function.call_python_awaitable";
+ private static final String CALL_PYTHON_FUNCTION =
"function.call_python_function";
+
+ @Test
+ void
keepsActionConversionInvocationAndAwaitableOnCallingThreadsInterpreter() throws
Exception {
+ PythonInterpreter owner = mock(PythonInterpreter.class);
+ PythonInterpreter worker = mock(PythonInterpreter.class);
+ PythonInterpreterManager manager =
+ new PythonInterpreterManager(owner, () -> worker, ignored ->
{});
+ PythonRunnerContextImpl runnerContext =
mock(PythonRunnerContextImpl.class);
+ PyObject pythonRunnerContext = mock(PyObject.class);
+ PyObject pythonEvent = mock(PyObject.class);
+ PyObject createdAwaitable = mock(PyObject.class);
+ PyObject retrievedAwaitable = mock(PyObject.class);
+ PythonFunction function = new PythonFunction("test_module",
"test_action");
+ Event event = new Event("test_event");
+ String eventJson = new ObjectMapper().writeValueAsString(event);
+ ExecutorService executorService = Executors.newSingleThreadExecutor();
+ AtomicReference<String> awaitableRef = new AtomicReference<>();
+
+ when(worker.invoke("python_java_utils.convert_json_to_python_event",
eventJson))
+ .thenReturn(pythonEvent);
+ when(worker.invoke(
+ "function.call_python_function",
+ "test_module",
+ "test_action",
+ new Object[] {pythonEvent, pythonRunnerContext}))
+ .thenReturn(createdAwaitable);
+
when(worker.get(org.mockito.ArgumentMatchers.anyString())).thenReturn(retrievedAwaitable);
+ when(worker.invoke("function.call_python_awaitable",
retrievedAwaitable))
+ .thenReturn(new Object[] {false, null});
+
+ PythonActionExecutor actionExecutor =
+ new PythonActionExecutor(manager, null, null, runnerContext,
"test-job");
+ setField(actionExecutor, "pythonRunnerContext", pythonRunnerContext);
+ try {
+ boolean finished =
+ executorService
+ .submit(
+ () -> {
+ String ref =
+
actionExecutor.executePythonFunction(
+ function, event);
+ awaitableRef.set(ref);
+ return
actionExecutor.callPythonAwaitable(ref);
+ })
+ .get(5, TimeUnit.SECONDS);
+
+ assertThat(finished).isFalse();
+ assertThat(awaitableRef.get()).startsWith("python_awaitable_");
+ verify(worker).set(awaitableRef.get(), createdAwaitable);
+ verify(pythonEvent).close();
+ verify(createdAwaitable).close();
+ verify(retrievedAwaitable).close();
+ verifyNoInteractions(owner);
+ } finally {
+ executorService.shutdownNow();
+ manager.close();
+ }
+ }
@Test
void resolvesPickledPythonKeyTextFromPyFlinkKeyRow() throws Exception {
@@ -260,15 +323,19 @@ class PythonActionExecutorTest {
PythonInterpreter interpreter = mock(PythonInterpreter.class);
PythonRunnerContextImpl runnerContext =
mock(PythonRunnerContextImpl.class);
PythonActionExecutor executor = newExecutor(interpreter,
runnerContext);
- PythonFunction function = mock(PythonFunction.class);
+ PythonFunction function = new PythonFunction("test_module",
"test_action");
PyObject pythonEvent = mock(PyObject.class);
when(interpreter.invoke(same(CONVERT_JSON_TO_PYTHON_EVENT),
anyString()))
.thenReturn(pythonEvent);
- when(function.call(same(pythonEvent), isNull())).thenReturn(null);
+ when(interpreter.invoke(
+ CALL_PYTHON_FUNCTION,
+ "test_module",
+ "test_action",
+ new Object[] {pythonEvent, null}))
+ .thenReturn(null);
assertThat(executor.executePythonFunction(function, new
InputEvent(1L))).isNull();
- verify(function).setInterpreter(interpreter);
verify(pythonEvent).close();
}
@@ -277,12 +344,17 @@ class PythonActionExecutorTest {
PythonInterpreter interpreter = mock(PythonInterpreter.class);
PythonRunnerContextImpl runnerContext =
mock(PythonRunnerContextImpl.class);
PythonActionExecutor executor = newExecutor(interpreter,
runnerContext);
- PythonFunction function = mock(PythonFunction.class);
+ PythonFunction function = new PythonFunction("test_module",
"test_action");
PyObject pythonEvent = mock(PyObject.class);
PyObject pythonAwaitable = mock(PyObject.class);
when(interpreter.invoke(same(CONVERT_JSON_TO_PYTHON_EVENT),
anyString()))
.thenReturn(pythonEvent);
- when(function.call(same(pythonEvent),
isNull())).thenReturn(pythonAwaitable);
+ when(interpreter.invoke(
+ CALL_PYTHON_FUNCTION,
+ "test_module",
+ "test_action",
+ new Object[] {pythonEvent, null}))
+ .thenReturn(pythonAwaitable);
String pythonAwaitableRef = executor.executePythonFunction(function,
new InputEvent(1L));
@@ -291,10 +363,12 @@ class PythonActionExecutorTest {
assertThat(pythonAwaitableRef)
.isEqualTo(refCaptor.getValue())
.startsWith("python_awaitable_");
- InOrder closeOrder = inOrder(interpreter, pythonAwaitable,
pythonEvent);
- closeOrder.verify(interpreter).set(pythonAwaitableRef,
pythonAwaitable);
- closeOrder.verify(pythonAwaitable).close();
- closeOrder.verify(pythonEvent).close();
+ InOrder awaitableCloseOrder = inOrder(interpreter, pythonAwaitable);
+ awaitableCloseOrder.verify(interpreter).set(pythonAwaitableRef,
pythonAwaitable);
+ awaitableCloseOrder.verify(pythonAwaitable).close();
+ InOrder eventCloseOrder = inOrder(interpreter, pythonEvent);
+ eventCloseOrder.verify(interpreter).set(pythonAwaitableRef,
pythonAwaitable);
+ eventCloseOrder.verify(pythonEvent).close();
}
@Test
@@ -343,7 +417,8 @@ class PythonActionExecutorTest {
private static PythonActionExecutor newExecutor(
PythonInterpreter interpreter, PythonRunnerContextImpl
runnerContext) throws Exception {
- return new PythonActionExecutor(interpreter, null, null,
runnerContext, "test-job");
+ return new PythonActionExecutor(
+ newInterpreterManager(interpreter), null, null, runnerContext,
"test-job");
}
private static TestFixture createOpenedExecutor() throws Exception {
@@ -371,11 +446,24 @@ class PythonActionExecutorTest {
PythonActionExecutor executor =
new PythonActionExecutor(
- interpreter, plan, resourceAdapter, runnerContext,
jobIdentifier);
+ newInterpreterManager(interpreter),
+ plan,
+ resourceAdapter,
+ runnerContext,
+ jobIdentifier);
executor.open();
return new TestFixture(interpreter, asyncThreadPool,
runnerContextObject, executor);
}
+ private static PythonInterpreterManager
newInterpreterManager(PythonInterpreter interpreter) {
+ return new PythonInterpreterManager(
+ interpreter,
+ () -> {
+ throw new AssertionError("unexpected worker interpreter");
+ },
+ ignored -> {});
+ }
+
private static final class TestFixture {
private final PythonInterpreter interpreter;
private final PyObject asyncThreadPool;
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManagerTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManagerTest.java
new file mode 100644
index 00000000..5df8d9f2
--- /dev/null
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManagerTest.java
@@ -0,0 +1,414 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.flink.agents.runtime.python.utils;
+
+import org.apache.flink.agents.runtime.async.AsyncExecutorThreadFactory;
+import org.junit.jupiter.api.Test;
+import pemja.core.PythonInterpreter;
+
+import java.util.Queue;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+
+/** Defect-oriented concurrency tests for {@link PythonInterpreterManager}. */
+class PythonInterpreterManagerTest {
+
+ @Test
+ void routesUnmanagedThreadCallsAwayFromTheCallerThread() throws Exception {
+ PythonInterpreter owner = mock(PythonInterpreter.class);
+ AtomicReference<Thread> callerThread = new AtomicReference<>();
+ AtomicReference<Thread> interpreterCreationThread = new
AtomicReference<>();
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+
+ try (PythonInterpreterManager manager =
+ new PythonInterpreterManager(
+ owner,
+ () -> {
+
interpreterCreationThread.set(Thread.currentThread());
+ return mock(PythonInterpreter.class);
+ },
+ ignored -> {})) {
+ executor.submit(
+ () -> {
+ callerThread.set(Thread.currentThread());
+ manager.invoke("callback");
+ })
+ .get(5, TimeUnit.SECONDS);
+
+ assertThat(interpreterCreationThread.get())
+ .as("an unmanaged caller must not create a Pemja thread
state on itself")
+ .isNotSameAs(callerThread.get());
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ void boundsInterpretersCreatedForTransientUnmanagedThreads() throws
Exception {
+ PythonInterpreter owner = mock(PythonInterpreter.class);
+ AtomicInteger created = new AtomicInteger();
+ AtomicReference<Throwable> failure = new AtomicReference<>();
+
+ try (PythonInterpreterManager manager =
+ new PythonInterpreterManager(
+ owner,
+ () -> {
+ created.incrementAndGet();
+ return mock(PythonInterpreter.class);
+ },
+ ignored -> {})) {
+ for (int i = 0; i < 20; i++) {
+ Thread caller =
+ new Thread(
+ () -> {
+ try {
+ manager.invoke("callback");
+ } catch (Throwable t) {
+ failure.set(t);
+ }
+ });
+ caller.start();
+ caller.join(5000);
+ assertThat(caller.isAlive()).isFalse();
+ }
+
+ assertThat(failure.get()).isNull();
+ assertThat(created.get())
+ .as("transient Python-originated callers must reuse
bounded callback workers")
+ .isBetween(1, 2);
+ }
+ }
+
+ @Test
+ void managedJavaAsyncWorkerCreatesUsesAndClosesInterpreterOnItself()
throws Exception {
+ PythonInterpreter owner = mock(PythonInterpreter.class);
+ PythonInterpreter worker = mock(PythonInterpreter.class);
+ AtomicReference<Thread> creationThread = new AtomicReference<>();
+ AtomicReference<Thread> invocationThread = new AtomicReference<>();
+ AtomicReference<Thread> closeThread = new AtomicReference<>();
+
+ doAnswer(
+ invocation -> {
+ invocationThread.set(Thread.currentThread());
+ return null;
+ })
+ .when(worker)
+ .invoke("worker-call");
+ doAnswer(
+ invocation -> {
+ closeThread.set(Thread.currentThread());
+ return null;
+ })
+ .when(worker)
+ .close();
+
+ try (PythonInterpreterManager manager =
+ new PythonInterpreterManager(
+ owner,
+ () -> {
+ creationThread.set(Thread.currentThread());
+ return worker;
+ },
+ ignored -> {})) {
+ AsyncExecutorThreadFactory threadFactory =
+ new
AsyncExecutorThreadFactory(manager::releaseCurrentThreadInterpreter);
+ ExecutorService executor =
Executors.newSingleThreadExecutor(threadFactory);
+ try {
+ executor.submit(() -> manager.invoke("worker-call")).get(5,
TimeUnit.SECONDS);
+ } finally {
+ executor.shutdownNow();
+ assertThat(executor.awaitTermination(5,
TimeUnit.SECONDS)).isTrue();
+ threadFactory.awaitThreadExit();
+ }
+
+ assertThat(invocationThread.get()).isSameAs(creationThread.get());
+ verify(worker).close();
+ assertThat(closeThread.get()).isSameAs(creationThread.get());
+ }
+ }
+
+ @Test
+ void callbackInterpreterIsClosedByItsOwningWorker() throws Exception {
+ PythonInterpreter owner = mock(PythonInterpreter.class);
+ PythonInterpreter callback = mock(PythonInterpreter.class);
+ AtomicReference<Thread> creationThread = new AtomicReference<>();
+ AtomicReference<Thread> closeThread = new AtomicReference<>();
+ ExecutorService caller = Executors.newSingleThreadExecutor();
+
+ doAnswer(
+ invocation -> {
+ closeThread.set(Thread.currentThread());
+ return null;
+ })
+ .when(callback)
+ .close();
+
+ PythonInterpreterManager manager =
+ new PythonInterpreterManager(
+ owner,
+ () -> {
+ creationThread.set(Thread.currentThread());
+ return callback;
+ },
+ ignored -> {});
+ try {
+ caller.submit(() -> manager.invoke("callback")).get(5,
TimeUnit.SECONDS);
+ manager.close();
+
+ verify(callback).close();
+ assertThat(closeThread.get()).isSameAs(creationThread.get());
+ } finally {
+ manager.close();
+ caller.shutdownNow();
+ }
+ }
+
+ @Test
+ void rejectsCloseFromANonOwnerThread() throws Exception {
+ PythonInterpreter owner = mock(PythonInterpreter.class);
+ ExecutorService caller = Executors.newSingleThreadExecutor();
+
+ try (PythonInterpreterManager manager =
+ new PythonInterpreterManager(owner, () ->
mock(PythonInterpreter.class))) {
+ Future<?> close =
+ caller.submit(
+ () -> {
+ manager.close();
+ return null;
+ });
+
+ assertThatThrownBy(() -> close.get(5, TimeUnit.SECONDS))
+ .hasCauseInstanceOf(IllegalStateException.class)
+ .hasRootCauseMessage(
+ "Python interpreter manager must be closed by its
owner thread.");
+ assertThatCode(() ->
manager.invoke("still-open")).doesNotThrowAnyException();
+ } finally {
+ caller.shutdownNow();
+ }
+ }
+
+ @Test
+ void reusesOwnerInterpreterOnCreatingThread() throws Exception {
+ PythonInterpreter owner = mock(PythonInterpreter.class);
+ PythonInterpreter unused = mock(PythonInterpreter.class);
+
+ try (PythonInterpreterManager manager =
+ new PythonInterpreterManager(owner, () -> unused, ignored ->
{})) {
+ manager.invoke("first");
+ manager.invoke("second");
+
+ verify(owner).invoke("first");
+ verify(owner).invoke("second");
+ verify(unused, never()).invoke("first");
+ verify(unused, never()).invoke("second");
+ }
+ }
+
+ @Test
+ void initializesEveryThreadConfinedInterpreter() throws Exception {
+ PythonInterpreter owner = mock(PythonInterpreter.class);
+ PythonInterpreter worker = mock(PythonInterpreter.class);
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+
+ try (PythonInterpreterManager manager = new
PythonInterpreterManager(owner, () -> worker)) {
+ executor.submit(() -> manager.invoke("worker-call")).get(5,
TimeUnit.SECONDS);
+
+ verify(owner).exec(PythonInterpreterManager.PYTHON_IMPORTS);
+ verify(worker).exec(PythonInterpreterManager.PYTHON_IMPORTS);
+ verify(worker).invoke("worker-call");
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ void bindsDifferentInterpretersToDifferentThreads() throws Exception {
+ PythonInterpreter owner = mock(PythonInterpreter.class);
+ Queue<PythonInterpreter> created = new ConcurrentLinkedQueue<>();
+ ExecutorService executor = Executors.newFixedThreadPool(2);
+
+ try (PythonInterpreterManager manager =
+ new PythonInterpreterManager(
+ owner,
+ () -> {
+ PythonInterpreter interpreter =
mock(PythonInterpreter.class);
+ created.add(interpreter);
+ return interpreter;
+ },
+ ignored -> {})) {
+ CountDownLatch ready = new CountDownLatch(2);
+ CountDownLatch release = new CountDownLatch(1);
+ Future<PythonInterpreter> first =
+ executor.submit(
+ () ->
+ manager.withInterpreter(
+ interpreter -> {
+ ready.countDown();
+ await(release);
+ PythonInterpreter nested =
+
manager.withInterpreter(value -> value);
+
assertThat(nested).isSameAs(interpreter);
+ return interpreter;
+ }));
+ Future<PythonInterpreter> second =
+ executor.submit(
+ () ->
+ manager.withInterpreter(
+ interpreter -> {
+ ready.countDown();
+ await(release);
+ PythonInterpreter nested =
+
manager.withInterpreter(value -> value);
+
assertThat(nested).isSameAs(interpreter);
+ return interpreter;
+ }));
+
+ assertThat(ready.await(5, TimeUnit.SECONDS)).isTrue();
+ release.countDown();
+
+ assertThat(first.get(5,
TimeUnit.SECONDS)).isNotSameAs(second.get(5, TimeUnit.SECONDS));
+ assertThat(created).hasSize(2);
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ void doesNotSerializeCallsMadeByDifferentThreads() throws Exception {
+ PythonInterpreter owner = mock(PythonInterpreter.class);
+ CountDownLatch entered = new CountDownLatch(2);
+ CountDownLatch release = new CountDownLatch(1);
+ ExecutorService executor = Executors.newFixedThreadPool(2);
+
+ try (PythonInterpreterManager manager =
+ new PythonInterpreterManager(
+ owner,
+ () -> {
+ PythonInterpreter interpreter =
mock(PythonInterpreter.class);
+ doAnswer(
+ invocation -> {
+ entered.countDown();
+ await(release);
+ return null;
+ })
+ .when(interpreter)
+ .invoke("blocking");
+ return interpreter;
+ },
+ ignored -> {})) {
+ Future<?> first = executor.submit(() ->
manager.invoke("blocking"));
+ Future<?> second = executor.submit(() ->
manager.invoke("blocking"));
+
+ assertThat(entered.await(5, TimeUnit.SECONDS))
+ .as("both interpreter calls should overlap")
+ .isTrue();
+ release.countDown();
+ first.get(5, TimeUnit.SECONDS);
+ second.get(5, TimeUnit.SECONDS);
+ } finally {
+ release.countDown();
+ executor.shutdownNow();
+ }
+ }
+
+ @Test
+ void reentrantCallUsesSameInterpreter() throws Exception {
+ PythonInterpreter owner = mock(PythonInterpreter.class);
+ AtomicReference<PythonInterpreterManager> managerRef = new
AtomicReference<>();
+ doAnswer(invocation ->
managerRef.get().invoke("inner")).when(owner).invoke("outer");
+
+ try (PythonInterpreterManager manager =
+ new PythonInterpreterManager(
+ owner,
+ () -> {
+ throw new AssertionError("reentrant call created
another interpreter");
+ },
+ ignored -> {})) {
+ managerRef.set(manager);
+
+ assertThatCode(() ->
manager.invoke("outer")).doesNotThrowAnyException();
+ verify(owner).invoke("inner");
+ }
+ }
+
+ @Test
+ void reentrantCallbackCallUsesSameInterpreter() throws Exception {
+ PythonInterpreter owner = mock(PythonInterpreter.class);
+ PythonInterpreter callback = mock(PythonInterpreter.class);
+ AtomicReference<PythonInterpreterManager> managerRef = new
AtomicReference<>();
+ ExecutorService caller = Executors.newSingleThreadExecutor();
+ doAnswer(invocation ->
managerRef.get().invoke("inner")).when(callback).invoke("outer");
+
+ try (PythonInterpreterManager manager =
+ new PythonInterpreterManager(owner, () -> callback, ignored ->
{})) {
+ managerRef.set(manager);
+
+ caller.submit(() -> manager.invoke("outer")).get(5,
TimeUnit.SECONDS);
+ verify(callback).invoke("inner");
+ verify(owner, never()).invoke("inner");
+ } finally {
+ caller.shutdownNow();
+ }
+ }
+
+ @Test
+ void closesEveryInterpreterAndRejectsLaterCalls() throws Exception {
+ PythonInterpreter owner = mock(PythonInterpreter.class);
+ PythonInterpreter worker = mock(PythonInterpreter.class);
+ ExecutorService executor = Executors.newSingleThreadExecutor();
+ PythonInterpreterManager manager =
+ new PythonInterpreterManager(owner, () -> worker, ignored ->
{});
+
+ try {
+ executor.submit(() -> manager.invoke("worker-call")).get(5,
TimeUnit.SECONDS);
+ manager.close();
+
+ verify(owner).close();
+ verify(worker).close();
+ assertThatThrownBy(() -> manager.invoke("after-close"))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("already closed");
+ } finally {
+ manager.close();
+ executor.shutdownNow();
+ }
+ }
+
+ private static void await(CountDownLatch latch) {
+ try {
+ assertThat(latch.await(5, TimeUnit.SECONDS)).isTrue();
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new AssertionError(e);
+ }
+ }
+}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImplTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImplTest.java
index b2225fdb..41bf28ed 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImplTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImplTest.java
@@ -47,18 +47,27 @@ public class PythonResourceAdapterImplTest {
@Mock private ResourceContext resourceContext;
private PythonResourceAdapterImpl pythonResourceAdapter;
+ private PythonInterpreterManager interpreterManager;
private AutoCloseable mocks;
@BeforeEach
void setUp() throws Exception {
mocks = MockitoAnnotations.openMocks(this);
+ interpreterManager =
+ new PythonInterpreterManager(
+ mockInterpreter,
+ () -> {
+ throw new AssertionError("unexpected worker
interpreter");
+ },
+ ignored -> {});
pythonResourceAdapter =
- new PythonResourceAdapterImpl(resourceContext,
mockInterpreter, null);
+ new PythonResourceAdapterImpl(resourceContext,
interpreterManager, null);
}
@AfterEach
void tearDown() throws Exception {
if (mocks != null) {
+ interpreterManager.close();
mocks.close();
}
}
@@ -89,7 +98,6 @@ public class PythonResourceAdapterImplTest {
pythonResourceAdapter.open();
- verify(mockInterpreter).exec(PythonResourceAdapterImpl.PYTHON_IMPORTS);
verify(mockInterpreter)
.invoke(PythonResourceAdapterImpl.GET_RESOURCE_CONTEXT,
pythonResourceAdapter);
}
