From ee12ec21cc41cbbe3c97df49a620b78d666c55b5 Mon Sep 17 00:00:00 2001 From: WenjinXie Date: Thu, 3 Sep 2026 12:42:33 +0800 Subject: [PATCH 1/3] [runtime][java] Make Python interpreters thread-confined Create an interpreter lazily for each caller thread and keep Python invocation and result conversion on the same interpreter. Preserve reentrant callbacks and add deterministic concurrent cross-language ChatModel E2E coverage for Java and Python agents. Generated-by: Codex 0.144.5 (GPT-5) Co-authored-by: Codex --- ...ConcurrentChatModelCrossLanguageAgent.java | 108 +++++++++ .../ConcurrentChatModelCrossLanguageTest.java | 69 ++++++ .../resource/test/Mem0LongTermMemoryTest.java | 2 - ...current_chat_model_cross_language_agent.py | 109 +++++++++ ...ncurrent_chat_model_cross_language_test.py | 99 ++++++++ .../runtime/operator/PythonBridgeManager.java | 47 ++-- .../python/utils/JavaResourceAdapter.java | 13 +- .../python/utils/PythonActionExecutor.java | 87 +++---- .../utils/PythonInterpreterManager.java | 171 ++++++++++++++ .../utils/PythonResourceAdapterImpl.java | 40 ++-- .../operator/PythonBridgeManagerTest.java | 32 +-- .../python/utils/JavaResourceAdapterTest.java | 53 +++++ .../utils/PythonActionExecutorTest.java | 78 ++++++- .../utils/PythonInterpreterManagerTest.java | 218 ++++++++++++++++++ .../utils/PythonResourceAdapterImplTest.java | 12 +- 15 files changed, 1032 insertions(+), 106 deletions(-) create mode 100644 e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/ConcurrentChatModelCrossLanguageAgent.java create mode 100644 e2e-test/flink-agents-end-to-end-tests-resource-cross-language/src/test/java/org/apache/flink/agents/resource/test/ConcurrentChatModelCrossLanguageTest.java create mode 100644 python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/concurrent_chat_model_cross_language_agent.py create mode 100644 python/flink_agents/e2e_tests/e2e_tests_resource_cross_language/concurrent_chat_model_cross_language_test.py create mode 100644 runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManager.java create mode 100644 runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManagerTest.java 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 000000000..f5558c2e0 --- /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 messages, List tools, Map 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 000000000..4f909c8ec --- /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 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 outputStream = + agentsEnv + .fromDataStream(inputStream, (KeySelector) value -> value) + .apply(new ConcurrentChatModelCrossLanguageAgent()) + .toDataStream(); + + CloseableIterator results = outputStream.collectAsync(); + agentsEnv.execute(); + + List 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 0aa31fd5f..c6dce1550 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.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 @@ * MILVUS_URI} env var * */ -@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 000000000..b90b96e0f --- /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 000000000..609bdd1cd --- /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/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 ba44a889a..b36f6703f 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 @@ -32,6 +32,7 @@ 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 +58,7 @@ * *
    *
  • The {@link PythonEnvironmentManager} that prepares dependencies and the Pemja runtime. - *
  • The {@link PythonInterpreter} obtained from that environment. + *
  • The thread-confined Python interpreters obtained from that environment. *
  • The {@link PythonActionExecutor} (when the plan contains Python actions or Mem0). *
  • The {@link PythonRunnerContextImpl} consumed by Python actions. *
  • The Java/Python resource adapters that bridge resource lookups across languages. @@ -70,7 +71,7 @@ * 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}. * *

    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 @@ -81,7 +82,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; @@ -99,12 +101,14 @@ class PythonBridgeManager implements AutoCloseable { *

    Scans the agent plan for any {@link PythonFunction} action or {@link * PythonResourceProvider}. If neither is present, 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 - * PythonRunnerContextImpl}, wires the Java/Python resource adapters, and conditionally - * initializes the Python action executor and the Python resource adapter (each only when the - * corresponding component is present in the plan). The generation guard runs immediately after - * interpreter construction and before any user module import. + * PythonEnvironmentManager}, opens an owner {@link PythonInterpreter}, refreshes the shared + * import state for the current dependency generation, and creates a {@link + * PythonInterpreterManager} that lazily binds a separate interpreter to every calling thread. + * It then constructs the shared {@link PythonRunnerContextImpl}, wires the Java/Python resource + * adapters, and conditionally initializes the Python action executor and the Python resource + * adapter (each only when the corresponding component is present in the plan). The 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. @@ -158,12 +162,13 @@ void open( 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); @@ -173,6 +178,12 @@ void open( 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); pythonRunnerContext = new PythonRunnerContextImpl( metricGroup, @@ -184,7 +195,7 @@ void open( javaResourceAdapter = new JavaResourceAdapter( resourceCache.getResourceContext(), - pythonInterpreter, + pythonInterpreterManager, userCodeClassLoader); if (containPythonResource || mem0Configured) { initPythonResourceAdapter(agentPlan, resourceCache); @@ -242,7 +253,8 @@ private boolean isMem0Configured(AgentPlan agentPlan) { */ 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 +282,7 @@ private void initPythonActionExecutor(AgentPlan agentPlan, String jobIdentifier) throws Exception { pythonActionExecutor = new PythonActionExecutor( - pythonInterpreter, + pythonInterpreterManager, agentPlan, javaResourceAdapter, pythonRunnerContext, @@ -282,7 +294,9 @@ private void initPythonResourceAdapter(AgentPlan agentPlan, ResourceCache resour throws Exception { pythonResourceAdapter = new PythonResourceAdapterImpl( - resourceCache.getResourceContext(), pythonInterpreter, javaResourceAdapter); + resourceCache.getResourceContext(), + pythonInterpreterManager, + javaResourceAdapter); pythonResourceAdapter.open(); PythonMCPResourceDiscovery.discoverPythonMCPResources( agentPlan.getResourceProviders(), pythonResourceAdapter, resourceCache); @@ -333,7 +347,8 @@ public void close() throws Exception { 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 2585145b8..7ef6155b2 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.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; @@ -42,7 +41,7 @@ public class JavaResourceAdapter { private final ResourceContext resourceContext; - private final transient PythonInterpreter interpreter; + private final transient PythonInterpreterManager interpreterManager; /** * Class loader used to resolve Java tool methods declared by name. Captured at construction @@ -54,10 +53,10 @@ public class JavaResourceAdapter { public JavaResourceAdapter( ResourceContext resourceContext, - PythonInterpreter interpreter, + PythonInterpreterManager interpreterManager, ClassLoader userCodeClassLoader) { this.resourceContext = resourceContext; - this.interpreter = interpreter; + this.interpreterManager = interpreterManager; this.userCodeClassLoader = userCodeClassLoader; } @@ -98,12 +97,12 @@ public List getSkillDirs(List skillNames) throws Exception { public ChatMessage fromPythonChatMessage(Object pythonChatMessage) { // 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."); + if (interpreterManager == null) { + throw new IllegalStateException("Python interpreter manager is not set."); } String roleValue = (String) - interpreter.invoke( + interpreterManager.invoke( "python_java_utils.update_java_chat_message", pythonChatMessage, chatMessage); 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 f39958e07..19f35c0c5 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 @@ -26,7 +26,6 @@ import org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl; 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; @@ -37,11 +36,6 @@ /** 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"; @@ -57,6 +51,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); @@ -71,7 +66,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; @@ -80,13 +75,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; @@ -98,17 +93,15 @@ public PyObject getPythonRunnerContext() { } 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), @@ -129,23 +122,31 @@ public void open() throws Exception { */ public String executePythonFunction(PythonFunction function, Event event) throws Exception { runnerContext.checkNoPendingEvents(); - function.setInterpreter(interpreter); - String eventJson = new ObjectMapper().writeValueAsString(event); - Object pythonEventObject = interpreter.invoke(CONVERT_JSON_TO_PYTHON_EVENT, eventJson); try { - Object calledResult = 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; - } + return interpreterManager.withInterpreter( + interpreter -> { + Object pythonEventObject = + interpreter.invoke(CONVERT_JSON_TO_PYTHON_EVENT, eventJson); + Object calledResult = + 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, + // and reference retention 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); @@ -155,7 +156,7 @@ public String executePythonFunction(PythonFunction function, Event event) throws 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); } @@ -170,14 +171,14 @@ public String resolveKeyText(Object flinkKey, boolean pythonKeyIsPickled) { 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); } /** @@ -192,14 +193,20 @@ public Object getOutputFromOutputEvent(String eventJson) { */ public boolean callPythonAwaitable(String pythonAwaitableRef) { // Calling awaitable.send(None) in Python returns a tuple of (finished, output). - Object pythonAwaitable = interpreter.get(pythonAwaitableRef); - checkState( - pythonAwaitable != null, - "Python awaitable '%s' not found in interpreter. ", - pythonAwaitableRef); - Object invokeResult = interpreter.invoke(CALL_PYTHON_AWAITABLE, pythonAwaitable); - checkState(invokeResult.getClass().isArray() && ((Object[]) invokeResult).length == 2); - return (boolean) ((Object[]) invokeResult)[0]; + return interpreterManager.withInterpreter( + interpreter -> { + Object pythonAwaitable = interpreter.get(pythonAwaitableRef); + checkState( + pythonAwaitable != null, + "Python awaitable '%s' not found in interpreter. ", + pythonAwaitableRef); + Object invokeResult = + interpreter.invoke(CALL_PYTHON_AWAITABLE, pythonAwaitable); + checkState( + invokeResult.getClass().isArray() + && ((Object[]) invokeResult).length == 2); + return (boolean) ((Object[]) invokeResult)[0]; + }); } @Override @@ -209,7 +216,7 @@ public void close() throws Exception { // 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; } @@ -240,7 +247,7 @@ public void close() throws Exception { 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 000000000..29c39d666 --- /dev/null +++ b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManager.java @@ -0,0 +1,171 @@ +/* + * 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.util.ExceptionUtils; +import pemja.core.PythonInterpreter; + +import java.util.List; +import java.util.Objects; +import java.util.concurrent.CopyOnWriteArrayList; +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. + * + *

    A {@link PythonInterpreter} is never shared by concurrently executing threads. The first + * bridge call on a thread creates an interpreter and binds it to that thread for the lifetime of + * this manager. Nested Python-to-Java-to-Python calls therefore reuse the same interpreter and run + * inline, while calls made by different async workers use different interpreters. + * + *

    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. 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 { + + 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 interpreterFactory; + private final Consumer interpreterInitializer; + private final List interpreters = new CopyOnWriteArrayList<>(); + private final ThreadLocal threadInterpreter = new ThreadLocal<>(); + private final ReentrantReadWriteLock lifecycleLock = new ReentrantReadWriteLock(); + + private volatile boolean closed; + + public PythonInterpreterManager( + PythonInterpreter ownerInterpreter, Supplier interpreterFactory) { + this(ownerInterpreter, interpreterFactory, PythonInterpreterManager::initializeInterpreter); + } + + PythonInterpreterManager( + PythonInterpreter ownerInterpreter, + Supplier interpreterFactory, + Consumer interpreterInitializer) { + this.interpreterFactory = Objects.requireNonNull(interpreterFactory); + this.interpreterInitializer = Objects.requireNonNull(interpreterInitializer); + + PythonInterpreter initializedOwner = initialize(Objects.requireNonNull(ownerInterpreter)); + interpreters.add(initializedOwner); + threadInterpreter.set(initializedOwner); + } + + /** Executes an operation using the interpreter bound to the calling thread. */ + public T withInterpreter(Function operation) { + lifecycleLock.readLock().lock(); + try { + if (closed) { + throw new IllegalStateException("Python interpreter manager is already closed."); + } + return Objects.requireNonNull(operation).apply(currentInterpreter()); + } 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 currentInterpreter() { + PythonInterpreter interpreter = threadInterpreter.get(); + if (interpreter != null) { + return interpreter; + } + + PythonInterpreter created = initialize(interpreterFactory.get()); + interpreters.add(created); + threadInterpreter.set(created); + return created; + } + + 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); + } + + @Override + public void close() throws Exception { + lifecycleLock.writeLock().lock(); + try { + if (closed) { + return; + } + closed = true; + threadInterpreter.remove(); + + Throwable firstFailure = null; + for (int i = interpreters.size() - 1; i >= 0; i--) { + try { + interpreters.get(i).close(); + } catch (Throwable closeFailure) { + firstFailure = ExceptionUtils.firstOrSuppressed(closeFailure, firstFailure); + } + } + interpreters.clear(); + + if (firstFailure != null) { + ExceptionUtils.rethrowException(firstFailure); + } + } finally { + lifecycleLock.writeLock().unlock(); + } + } +} 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 31c897134..58e6a057d 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 @@ -30,7 +30,6 @@ 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; @@ -40,8 +39,6 @@ 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"; @@ -81,22 +78,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 @@ -138,18 +134,18 @@ private Object toPythonResource(String resourceType, Resource resource) { 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 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 @@ -158,7 +154,9 @@ public ChatMessage fromPythonChatMessage(Object pythonChatMessage) { 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; } @@ -167,7 +165,7 @@ public ChatMessage fromPythonChatMessage(Object pythonChatMessage) { public Object toPythonDocuments(List documents) { List 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; } @@ -188,7 +186,7 @@ public List fromPythonDocuments(List pythonDocuments) { @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 @@ -201,26 +199,26 @@ public VectorStoreQueryResult fromPythonVectorStoreQueryResult( @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 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 @@ -234,7 +232,7 @@ public Map getPythonToolMetadata( @SuppressWarnings("unchecked") Map result = (Map) - interpreter.invoke( + interpreterManager.invoke( GET_PYTHON_TOOL_METADATA, module, qualName, injectedArgs); if (result == null) { throw new IllegalStateException( @@ -245,6 +243,6 @@ public Map getPythonToolMetadata( @Override public Object invokePythonTool(String module, String qualName, Map kwargs) { - return interpreter.invoke(INVOKE_PYTHON_TOOL, module, qualName, kwargs); + return interpreterManager.invoke(INVOKE_PYTHON_TOOL, module, qualName, kwargs); } } 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 fa25afffb..04c4535fc 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.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 @@ void closeAttemptsAllResourcesAndSuppressesLaterFailures() throws Exception { 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 @@ void closeAttemptsAllResourcesAndSuppressesLaterFailures() throws Exception { 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 @@ void closeAttemptsAllResourcesAndSuppressesLaterFailures() throws Exception { 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 @@ void openIsNoOpWhenPlanHasNeitherPythonActionsNorResources() throws Exception { 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 @@ void closeReleasesInterpreterAndEnvironmentWhenActionExecutorFails() throws Exce // 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 @@ void closeReleasesInterpreterAndEnvironmentWhenActionExecutorFails() throws Exce 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 @@ void closeReportsFirstFailureWithLaterOnesSuppressed() throws Exception { .close(); setField(bridge, "pythonActionExecutor", actionExecutor); - setField(bridge, "pythonInterpreter", interpreter); + setField(bridge, "pythonInterpreterManager", interpreterManager); setField(bridge, "pythonEnvironmentManager", environmentManager); assertThatThrownBy(bridge::close) @@ -177,7 +177,7 @@ void closeReportsFirstFailureWithLaterOnesSuppressed() throws Exception { .extracting(Throwable::getMessage) .containsExactly("environment manager close failed")); - verify(interpreter).close(); + verify(interpreterManager).close(); } /** @@ -189,13 +189,13 @@ void closeReportsFirstFailureWithLaterOnesSuppressed() throws Exception { 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 @@ void closeReleasesInterpreterAndEnvironmentWhenActionExecutorThrowsError() throw .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 f1d4a82ef..5b46af4df 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,16 +22,69 @@ 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; +import pemja.core.PythonInterpreter; import java.util.List; import java.util.Map; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.any; +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; class JavaResourceAdapterTest { + @Test + void convertsPythonChatMessageWithCallingThreadsInterpreter() throws Exception { + PythonInterpreter owner = mock(PythonInterpreter.class); + PythonInterpreter worker = mock(PythonInterpreter.class); + PythonInterpreterManager manager = + new PythonInterpreterManager(owner, () -> worker, ignored -> {}); + JavaResourceAdapter adapter = + new JavaResourceAdapter( + null, manager, Thread.currentThread().getContextClassLoader()); + Object pythonChatMessage = new Object(); + ExecutorService executor = Executors.newSingleThreadExecutor(); + + when(worker.invoke( + eq("python_java_utils.update_java_chat_message"), + eq(pythonChatMessage), + any(ChatMessage.class))) + .thenAnswer( + invocation -> { + invocation.getArgument(2).setContent("hello"); + return "user"; + }); + + try { + ChatMessage converted = + executor.submit(() -> adapter.fromPythonChatMessage(pythonChatMessage)) + .get(5, TimeUnit.SECONDS); + + assertThat(converted.getRole()).isEqualTo(MessageRole.USER); + assertThat(converted.getContent()).isEqualTo("hello"); + verify(worker) + .invoke( + eq("python_java_utils.update_java_chat_message"), + eq(pythonChatMessage), + any(ChatMessage.class)); + verifyNoInteractions(owner); + } finally { + executor.shutdownNow(); + manager.close(); + } + } + @Test void getJavaToolMetadataHidesInjectedArgsAndReturnsAnnotatedDeclaration() throws Exception { JavaResourceAdapter adapter = 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 5bdac5dcc..33335b46e 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,8 +18,10 @@ 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.agents.AgentExecutionOptions; import org.apache.flink.agents.plan.AgentPlan; +import org.apache.flink.agents.plan.PythonFunction; import org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl; import org.apache.flink.types.Row; import org.junit.jupiter.api.Test; @@ -29,6 +31,10 @@ import java.lang.reflect.Field; import java.util.HashMap; +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; @@ -52,6 +58,60 @@ class PythonActionExecutorTest { private static final String CLOSE_FLINK_RUNNER_CONTEXT = "flink_runner_context.close_flink_runner_context"; + @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); + Object pythonEvent = new Object(); + Object awaitable = new Object(); + 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 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(awaitable); + when(worker.get(org.mockito.ArgumentMatchers.anyString())).thenReturn(awaitable); + when(worker.invoke("function.call_python_awaitable", awaitable)) + .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(), awaitable); + verifyNoInteractions(owner); + } finally { + executorService.shutdownNow(); + manager.close(); + } + } + @Test void resolvesPickledPythonKeyTextFromPyFlinkKeyRow() throws Exception { PythonInterpreter interpreter = mock(PythonInterpreter.class); @@ -245,7 +305,8 @@ void releasesBothPythonObjectsWhenLogicalCleanupFails() throws Exception { private static PythonActionExecutor newExecutor(PythonInterpreter interpreter) throws Exception { - return new PythonActionExecutor(interpreter, null, null, null, "test-job"); + return new PythonActionExecutor( + newInterpreterManager(interpreter), null, null, null, "test-job"); } private static TestFixture createOpenedExecutor() throws Exception { @@ -273,11 +334,24 @@ private static TestFixture createOpenedExecutor() throws Exception { 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 000000000..b2747b531 --- /dev/null +++ b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManagerTest.java @@ -0,0 +1,218 @@ +/* + * 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.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.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 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 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 first = + executor.submit( + () -> + manager.withInterpreter( + interpreter -> { + ready.countDown(); + await(release); + PythonInterpreter nested = + manager.withInterpreter(value -> value); + assertThat(nested).isSameAs(interpreter); + return interpreter; + })); + Future 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 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 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 f372821a5..c40b2a1de 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 @@ -45,18 +45,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(); } } @@ -87,7 +96,6 @@ void testOpen() { pythonResourceAdapter.open(); - verify(mockInterpreter).exec(PythonResourceAdapterImpl.PYTHON_IMPORTS); verify(mockInterpreter) .invoke(PythonResourceAdapterImpl.GET_RESOURCE_CONTEXT, pythonResourceAdapter); } From f456e11f8d0187bbc16a6a3eacd5457bf7b91ca2 Mon Sep 17 00:00:00 2001 From: WenjinXie Date: Thu, 3 Sep 2026 16:53:14 +0800 Subject: [PATCH 2/3] [runtime][java] Route Python callbacks through interpreter workers Keep managed Java async workers on thread-local interpreters while routing Python-created and other unmanaged callers through bounded, affinity-preserving callback lanes. Close every interpreter on its owning worker and extend cross-language concurrency coverage for reverse callbacks. Generated-by: Codex 0.144.5 (GPT-5) Co-authored-by: Codex --- ...current_chat_model_cross_language_agent.py | 16 +- .../async/AsyncExecutorThreadFactory.java | 55 +++- .../async/ContinuationActionExecutor.java | 3 + .../operator/ActionExecutionOperator.java | 3 +- .../operator/ActionTaskContextManager.java | 7 +- .../runtime/operator/PythonBridgeManager.java | 25 +- .../utils/PythonInterpreterManager.java | 288 ++++++++++++++++-- .../async/ContinuationActionExecutor.java | 23 +- .../async/AsyncExecutorThreadFactoryTest.java | 25 ++ .../utils/PythonInterpreterManagerTest.java | 196 ++++++++++++ 10 files changed, 604 insertions(+), 37 deletions(-) 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 index b90b96e0f..f7c9edfb9 100644 --- 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 @@ -16,6 +16,7 @@ # limitations under the License. ################################################################################# import threading +from concurrent.futures import ThreadPoolExecutor from typing import Any, List, Sequence from pydantic import PrivateAttr @@ -57,9 +58,22 @@ def chat( message = "Timed out waiting for concurrent cross-language chat request." raise RuntimeError(message) from error + # Run a Python -> Java -> Python conversion from a Python-created thread. This + # mirrors Mem0's callback path and verifies that the Java bridge routes the + # reverse Python invocation to a bounded callback worker instead of creating a + # second Pemja thread state on this CPython thread. + if self.resource_context is None: + message = "The Python chat connection has no resource context." + raise RuntimeError(message) + java_adapter = self.resource_context._j_resource_adapter + with ThreadPoolExecutor(max_workers=1) as executor: + java_message = executor.submit( + java_adapter.fromPythonChatMessage, messages[-1] + ).result(timeout=30) + return ChatMessage( role=MessageRole.ASSISTANT, - content=f"python-connection:{messages[-1].content}", + content=f"python-connection:{java_message.getContent()}", ) 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 788a9c3e4..14edc1e23 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 @@ public final class AsyncExecutorThreadFactory implements ThreadFactory { private static final String NAME_PREFIX = "flink-agents-java-async-"; + private static final ThreadLocal ASYNC_EXECUTOR_THREAD = new ThreadLocal<>(); private final ThreadFactory delegate = Executors.defaultThreadFactory(); + private final Runnable threadCleanup; + private final Queue 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 f6134da8a..6a3692548 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 f82d8047a..0014a1d18 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 @@ -220,7 +220,8 @@ public void open() throws Exception { // 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 64e9e9ed9..06162a9ad 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 @@ -84,11 +84,16 @@ class ActionTaskContextManager implements AutoCloseable { private final ContinuationActionExecutor continuationActionExecutor; ActionTaskContextManager(int numAsyncThreads) { + this(numAsyncThreads, () -> {}); + } + + ActionTaskContextManager(int numAsyncThreads, Runnable asyncThreadCleanup) { this.actionTaskMemoryContexts = new HashMap<>(); this.continuationContexts = new HashMap<>(); this.pythonAwaitableRefs = new HashMap<>(); this.activeReportedExecutionsByActionExecutionId = 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 b36f6703f..9c6b6f6e0 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; @@ -103,12 +104,12 @@ class PythonBridgeManager implements AutoCloseable { * #isInitialized()} stays {@code false}. Otherwise it builds the {@link * PythonEnvironmentManager}, opens an owner {@link PythonInterpreter}, refreshes the shared * import state for the current dependency generation, and creates a {@link - * PythonInterpreterManager} that lazily binds a separate interpreter to every calling thread. - * It then constructs the shared {@link PythonRunnerContextImpl}, wires the Java/Python resource - * adapters, and conditionally initializes the Python action executor and the Python resource - * adapter (each only when the corresponding component is present in the plan). The generation - * guard runs immediately after owner-interpreter construction and before any user module - * import. + * 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 action executor and the Python resource adapter (each only when the + * corresponding component is present in the plan). The 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. @@ -183,7 +184,10 @@ void open( // construction fails, its constructor releases the owner itself. initializingPythonInterpreter = null; pythonInterpreterManager = - new PythonInterpreterManager(ownerInterpreter, env::getInterpreter); + new PythonInterpreterManager( + ownerInterpreter, + env::getInterpreter, + agentPlan.getConfig().get(AgentExecutionOptions.NUM_ASYNC_THREADS)); pythonRunnerContext = new PythonRunnerContextImpl( metricGroup, @@ -332,6 +336,13 @@ boolean isInitialized() { 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 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 index 29c39d666..c007126bf 100644 --- 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 @@ -17,12 +17,23 @@ */ 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.List; +import java.util.ArrayList; +import java.util.Map; import java.util.Objects; -import java.util.concurrent.CopyOnWriteArrayList; +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; @@ -31,18 +42,25 @@ /** * Owns the Pemja interpreters used by one operator subtask. * - *

    A {@link PythonInterpreter} is never shared by concurrently executing threads. The first - * bridge call on a thread creates an interpreter and binds it to that thread for the lifetime of - * this manager. Nested Python-to-Java-to-Python calls therefore reuse the same interpreter and run - * inline, while calls made by different async workers use different interpreters. + *

    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. * *

    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. Each interpreter has its own globals, however, so the - * bridge modules must be imported for every newly created interpreter. + * 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" @@ -50,37 +68,95 @@ public final class PythonInterpreterManager implements AutoCloseable { private final Supplier interpreterFactory; private final Consumer interpreterInitializer; - private final List interpreters = new CopyOnWriteArrayList<>(); + private final Thread ownerThread; + private final Map interpreters = new ConcurrentHashMap<>(); private final ThreadLocal threadInterpreter = new ThreadLocal<>(); + private final ThreadLocal callbackLane; + private final ExecutorService[] callbackExecutors; + private final Queue callbackThreads = new ConcurrentLinkedQueue<>(); + private final Queue workerCloseFailures = new ConcurrentLinkedQueue<>(); private final ReentrantReadWriteLock lifecycleLock = new ReentrantReadWriteLock(); private volatile boolean closed; public PythonInterpreterManager( PythonInterpreter ownerInterpreter, Supplier interpreterFactory) { - this(ownerInterpreter, interpreterFactory, PythonInterpreterManager::initializeInterpreter); + this( + ownerInterpreter, + interpreterFactory, + PythonInterpreterManager::initializeInterpreter, + DEFAULT_CALLBACK_WORKERS); + } + + public PythonInterpreterManager( + PythonInterpreter ownerInterpreter, + Supplier interpreterFactory, + int callbackWorkerCount) { + this( + ownerInterpreter, + interpreterFactory, + PythonInterpreterManager::initializeInterpreter, + callbackWorkerCount); } PythonInterpreterManager( PythonInterpreter ownerInterpreter, Supplier interpreterFactory, Consumer interpreterInitializer) { + this( + ownerInterpreter, + interpreterFactory, + interpreterInitializer, + DEFAULT_CALLBACK_WORKERS); + } + + PythonInterpreterManager( + PythonInterpreter ownerInterpreter, + Supplier interpreterFactory, + Consumer 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.add(initializedOwner); + 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 using the interpreter bound to the calling thread. */ + /** + * Executes an operation on a thread-confined interpreter. + * + *

    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 withInterpreter(Function operation) { lifecycleLock.readLock().lock(); try { if (closed) { throw new IllegalStateException("Python interpreter manager is already closed."); } - return Objects.requireNonNull(operation).apply(currentInterpreter()); + Function 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(); } @@ -110,18 +186,47 @@ public void set(String name, Object value) { }); } - private PythonInterpreter currentInterpreter() { + private PythonInterpreter createAndBindInterpreter() { PythonInterpreter interpreter = threadInterpreter.get(); if (interpreter != null) { return interpreter; } PythonInterpreter created = initialize(interpreterFactory.get()); - interpreters.add(created); + interpreters.put(Thread.currentThread(), created); threadInterpreter.set(created); return created; } + private T executeOnCallbackWorker(Function operation) { + int lane = callbackLane.get(); + Future result = + callbackExecutors[lane].submit(() -> operation.apply(createAndBindInterpreter())); + return awaitResult(result); + } + + private static T awaitResult(Future 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 { @@ -141,31 +246,164 @@ 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 void close() throws Exception { + 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; - threadInterpreter.remove(); + } finally { + lifecycleLock.writeLock().unlock(); + } - Throwable firstFailure = null; - for (int i = interpreters.size() - 1; i >= 0; i--) { + 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 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 { - interpreters.get(i).close(); + entry.getValue().close(); } catch (Throwable closeFailure) { - firstFailure = ExceptionUtils.firstOrSuppressed(closeFailure, firstFailure); + workerCloseFailures.add(closeFailure); } } - interpreters.clear(); + } - if (firstFailure != null) { - ExceptionUtils.rethrowException(firstFailure); + 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 { - lifecycleLock.writeLock().unlock(); + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + } + + private static void awaitThreadExit(Iterable 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/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java b/runtime/src/main/java21/org/apache/flink/agents/runtime/async/ContinuationActionExecutor.java index 6818441b5..e34649ed8 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 @@ private long getDeadlineNanos(Duration timeout) { 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 125d18122..a19b04770 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 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 @@ /** 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/python/utils/PythonInterpreterManagerTest.java b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonInterpreterManagerTest.java index b2747b531..5df8d9f25 100644 --- 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 @@ -17,6 +17,7 @@ */ 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; @@ -27,6 +28,7 @@ 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; @@ -40,6 +42,180 @@ /** Defect-oriented concurrency tests for {@link PythonInterpreterManager}. */ class PythonInterpreterManagerTest { + @Test + void routesUnmanagedThreadCallsAwayFromTheCallerThread() throws Exception { + PythonInterpreter owner = mock(PythonInterpreter.class); + AtomicReference callerThread = new AtomicReference<>(); + AtomicReference 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 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 creationThread = new AtomicReference<>(); + AtomicReference invocationThread = new AtomicReference<>(); + AtomicReference 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 creationThread = new AtomicReference<>(); + AtomicReference 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); @@ -184,6 +360,26 @@ void reentrantCallUsesSameInterpreter() throws Exception { } } + @Test + void reentrantCallbackCallUsesSameInterpreter() throws Exception { + PythonInterpreter owner = mock(PythonInterpreter.class); + PythonInterpreter callback = mock(PythonInterpreter.class); + AtomicReference 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); From 2db48c356bcdf08d35617ab3bde4d35451c8f812 Mon Sep 17 00:00:00 2001 From: WenjinXie Date: Thu, 3 Sep 2026 21:27:43 +0800 Subject: [PATCH 3/3] [runtime][java][python] Avoid nested callback for chat messages Extract Python ChatMessage fields before crossing into Java so message conversion no longer re-enters Python through the interpreter manager. Keep the concurrent E2E focused on supported overlapping calls and document the nested callback thread limitation. Generated-by: Codex 0.144.5 (GPT-5) Co-authored-by: Codex --- .../docs/development/workflow_agent.md | 14 +++++ ...current_chat_model_cross_language_agent.py | 16 +---- .../runtime/java/java_chat_model.py | 29 +++++---- .../runtime/tests/test_java_chat_model.py | 63 +++++++++++++++++++ .../runtime/operator/PythonBridgeManager.java | 4 +- .../python/utils/JavaResourceAdapter.java | 36 +++++------ .../python/utils/JavaResourceAdapterTest.java | 57 ++++------------- 7 files changed, 122 insertions(+), 97 deletions(-) create mode 100644 python/flink_agents/runtime/tests/test_java_chat_model.py diff --git a/docs/content/docs/development/workflow_agent.md b/docs/content/docs/development/workflow_agent.md index b8166cd0e..90cde7dbc 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/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 index f7c9edfb9..b90b96e0f 100644 --- 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 @@ -16,7 +16,6 @@ # limitations under the License. ################################################################################# import threading -from concurrent.futures import ThreadPoolExecutor from typing import Any, List, Sequence from pydantic import PrivateAttr @@ -58,22 +57,9 @@ def chat( message = "Timed out waiting for concurrent cross-language chat request." raise RuntimeError(message) from error - # Run a Python -> Java -> Python conversion from a Python-created thread. This - # mirrors Mem0's callback path and verifies that the Java bridge routes the - # reverse Python invocation to a bounded callback worker instead of creating a - # second Pemja thread state on this CPython thread. - if self.resource_context is None: - message = "The Python chat connection has no resource context." - raise RuntimeError(message) - java_adapter = self.resource_context._j_resource_adapter - with ThreadPoolExecutor(max_workers=1) as executor: - java_message = executor.submit( - java_adapter.fromPythonChatMessage, messages[-1] - ).result(timeout=30) - return ChatMessage( role=MessageRole.ASSISTANT, - content=f"python-connection:{java_message.getContent()}", + content=f"python-connection:{messages[-1].content}", ) diff --git a/python/flink_agents/runtime/java/java_chat_model.py b/python/flink_agents/runtime/java/java_chat_model.py index 15df31578..43aca50eb 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.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 @@ def chat( """ 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 @@ def chat( ] 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 @@ def chat( """ # 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 000000000..87d4ae91d --- /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/operator/PythonBridgeManager.java b/runtime/src/main/java/org/apache/flink/agents/runtime/operator/PythonBridgeManager.java index 9c6b6f6e0..35c0381c6 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 @@ -198,9 +198,7 @@ void open( javaResourceAdapter = new JavaResourceAdapter( - resourceCache.getResourceContext(), - pythonInterpreterManager, - userCodeClassLoader); + resourceCache.getResourceContext(), userCodeClassLoader); if (containPythonResource || mem0Configured) { initPythonResourceAdapter(agentPlan, resourceCache); } 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 7ef6155b2..f3f8dd8c8 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 @@ -41,8 +41,6 @@ public class JavaResourceAdapter { private final ResourceContext resourceContext; - private final transient PythonInterpreterManager interpreterManager; - /** * Class loader used to resolve Java tool methods declared by name. Captured at construction * (the operator passes its {@code RuntimeContext.getUserCodeClassLoader()}) because pemja @@ -51,12 +49,8 @@ public class JavaResourceAdapter { */ private final transient ClassLoader userCodeClassLoader; - public JavaResourceAdapter( - ResourceContext resourceContext, - PythonInterpreterManager interpreterManager, - ClassLoader userCodeClassLoader) { + public JavaResourceAdapter(ResourceContext resourceContext, ClassLoader userCodeClassLoader) { this.resourceContext = resourceContext; - this.interpreterManager = interpreterManager; this.userCodeClassLoader = userCodeClassLoader; } @@ -91,23 +85,23 @@ public List getSkillDirs(List skillNames) throws Exception { * 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 + *

    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> toolCalls, + Map extraArgs) { // TODO: Delete this method after the pemja findClass method is fixed. - ChatMessage chatMessage = new ChatMessage(); - if (interpreterManager == null) { - throw new IllegalStateException("Python interpreter manager is not set."); - } - String roleValue = - (String) - interpreterManager.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/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 5b46af4df..49cce424f 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 @@ -26,69 +26,34 @@ import org.apache.flink.agents.api.chat.messages.MessageRole; import org.apache.flink.agents.api.tools.ToolParameterSource; import org.junit.jupiter.api.Test; -import pemja.core.PythonInterpreter; import java.util.List; import java.util.Map; -import java.util.concurrent.ExecutorService; -import java.util.concurrent.Executors; -import java.util.concurrent.TimeUnit; import static org.assertj.core.api.Assertions.assertThat; -import static org.mockito.ArgumentMatchers.any; -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; class JavaResourceAdapterTest { @Test - void convertsPythonChatMessageWithCallingThreadsInterpreter() throws Exception { - PythonInterpreter owner = mock(PythonInterpreter.class); - PythonInterpreter worker = mock(PythonInterpreter.class); - PythonInterpreterManager manager = - new PythonInterpreterManager(owner, () -> worker, ignored -> {}); + void buildsJavaChatMessageFromExtractedPythonFields() { JavaResourceAdapter adapter = - new JavaResourceAdapter( - null, manager, Thread.currentThread().getContextClassLoader()); - Object pythonChatMessage = new Object(); - ExecutorService executor = Executors.newSingleThreadExecutor(); + new JavaResourceAdapter(null, Thread.currentThread().getContextClassLoader()); + List> toolCalls = List.of(Map.of("id", "call-1", "type", "function")); + Map extraArgs = Map.of("reasoning", "brief"); - when(worker.invoke( - eq("python_java_utils.update_java_chat_message"), - eq(pythonChatMessage), - any(ChatMessage.class))) - .thenAnswer( - invocation -> { - invocation.getArgument(2).setContent("hello"); - return "user"; - }); + ChatMessage converted = + adapter.fromPythonChatMessage("user", "hello", toolCalls, extraArgs); - try { - ChatMessage converted = - executor.submit(() -> adapter.fromPythonChatMessage(pythonChatMessage)) - .get(5, TimeUnit.SECONDS); - - assertThat(converted.getRole()).isEqualTo(MessageRole.USER); - assertThat(converted.getContent()).isEqualTo("hello"); - verify(worker) - .invoke( - eq("python_java_utils.update_java_chat_message"), - eq(pythonChatMessage), - any(ChatMessage.class)); - verifyNoInteractions(owner); - } finally { - executor.shutdownNow(); - manager.close(); - } + 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 metadata = adapter.getJavaToolMetadata(