From fb14b7e048886b2ccc27e7220c4257965b035c88 Mon Sep 17 00:00:00 2001 From: joeyutong Date: Tue, 25 Aug 2026 20:49:58 +0800 Subject: [PATCH 1/2] [runtime][python] Release per-action Pemja objects Close action-scoped PyObject handles after interpreter ownership is established and remove completed awaitables from interpreter globals. Generated-by: Codex (GPT-5) Co-Authored-By: Claude Code AI-Model: gpt-5 AI-Contributed/Feature: 37/37 AI-Contributed/UT: 155/155 --- .../python/utils/PythonActionExecutor.java | 37 +++-- .../utils/PythonActionExecutorTest.java | 155 +++++++++++++++++- 2 files changed, 175 insertions(+), 17 deletions(-) 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 d37402779..a4c949b73 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 @@ -119,9 +119,9 @@ public void open() throws Exception { /** * Execute the Python function, which may return a Python coroutine (awaitable) that needs to be * processed in the future. Due to an issue in Pemja regarding incorrect object reference - * counting, this may lead to garbage collection of the object. To prevent this, we use the set - * and get methods to manually increment the object's reference count, then return the name of - * the Python awaitable variable. + * counting, this may lead to garbage collection of the object. To prevent this, we store the + * awaitable in the interpreter globals, then return the name of that variable. The temporary + * Java wrapper can be closed after the interpreter takes ownership of its own reference. * * @return The name of the Python awaitable variable. It may be null if the Python function does * not return a coroutine. @@ -131,10 +131,10 @@ public String executePythonFunction(PythonFunction function, Event event) throws 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); + try (PyObject pythonEventObject = + (PyObject) interpreter.invoke(CONVERT_JSON_TO_PYTHON_EVENT, eventJson); + PyObject calledResult = + (PyObject) function.call(pythonEventObject, pythonRunnerContext)) { if (calledResult == null) { return null; } else { @@ -189,16 +189,21 @@ public Object getOutputFromOutputEvent(String eventJson) { * interpreter's context * @return true if the awaitable has completed; false otherwise */ - public boolean callPythonAwaitable(String pythonAwaitableRef) { + public boolean callPythonAwaitable(String pythonAwaitableRef) throws Exception { // 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]; + try (PyObject pythonAwaitable = (PyObject) 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); + boolean finished = (boolean) ((Object[]) invokeResult)[0]; + if (finished) { + interpreter.exec("del " + pythonAwaitableRef); + } + return finished; + } } public void close() throws Exception { 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 baf584fa4..ee8bc0b36 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 @@ -17,18 +17,34 @@ */ package org.apache.flink.agents.runtime.python.utils; +import org.apache.flink.agents.api.InputEvent; +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; +import org.mockito.ArgumentCaptor; +import org.mockito.InOrder; import pemja.core.PythonInterpreter; +import pemja.core.object.PyObject; import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.ArgumentMatchers.same; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.inOrder; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; class PythonActionExecutorTest { + private static final String CONVERT_JSON_TO_PYTHON_EVENT = + "python_java_utils.convert_json_to_python_event"; + private static final String CALL_PYTHON_AWAITABLE = "function.call_python_awaitable"; + @Test void resolvesPickledPythonKeyTextFromPyFlinkKeyRow() throws Exception { PythonInterpreter interpreter = mock(PythonInterpreter.class); @@ -87,8 +103,145 @@ void propagatesDecodeFailure() throws Exception { .hasMessage("bad pickle"); } + @Test + void closesPythonEventAfterSynchronousAction() throws Exception { + PythonInterpreter interpreter = mock(PythonInterpreter.class); + PythonRunnerContextImpl runnerContext = mock(PythonRunnerContextImpl.class); + PythonActionExecutor executor = newExecutor(interpreter, runnerContext); + PythonFunction function = mock(PythonFunction.class); + PyObject pythonEvent = mock(PyObject.class); + when(interpreter.invoke(same(CONVERT_JSON_TO_PYTHON_EVENT), anyString())) + .thenReturn(pythonEvent); + when(function.call(same(pythonEvent), isNull())).thenReturn(null); + + assertThat(executor.executePythonFunction(function, new InputEvent(1L))).isNull(); + + verify(function).setInterpreter(interpreter); + verify(pythonEvent).close(); + } + + @Test + void closesTemporaryWrappersAfterStoringAwaitable() throws Exception { + PythonInterpreter interpreter = mock(PythonInterpreter.class); + PythonRunnerContextImpl runnerContext = mock(PythonRunnerContextImpl.class); + PythonActionExecutor executor = newExecutor(interpreter, runnerContext); + PythonFunction function = mock(PythonFunction.class); + PyObject pythonEvent = mock(PyObject.class); + PyObject pythonAwaitable = mock(PyObject.class); + when(interpreter.invoke(same(CONVERT_JSON_TO_PYTHON_EVENT), anyString())) + .thenReturn(pythonEvent); + when(function.call(same(pythonEvent), isNull())).thenReturn(pythonAwaitable); + + String pythonAwaitableRef = executor.executePythonFunction(function, new InputEvent(1L)); + + ArgumentCaptor refCaptor = ArgumentCaptor.forClass(String.class); + verify(interpreter).set(refCaptor.capture(), same(pythonAwaitable)); + assertThat(pythonAwaitableRef) + .isEqualTo(refCaptor.getValue()) + .startsWith("python_awaitable_"); + InOrder closeOrder = inOrder(interpreter, pythonAwaitable, pythonEvent); + closeOrder.verify(interpreter).set(pythonAwaitableRef, pythonAwaitable); + closeOrder.verify(pythonAwaitable).close(); + closeOrder.verify(pythonEvent).close(); + } + + @Test + void closesPythonEventWhenActionFails() throws Exception { + PythonInterpreter interpreter = mock(PythonInterpreter.class); + PythonRunnerContextImpl runnerContext = mock(PythonRunnerContextImpl.class); + PythonActionExecutor executor = newExecutor(interpreter, runnerContext); + PythonFunction function = mock(PythonFunction.class); + PyObject pythonEvent = mock(PyObject.class); + RuntimeException failure = new RuntimeException("action failed"); + when(interpreter.invoke(same(CONVERT_JSON_TO_PYTHON_EVENT), anyString())) + .thenReturn(pythonEvent); + when(function.call(same(pythonEvent), isNull())).thenThrow(failure); + + assertThatThrownBy(() -> executor.executePythonFunction(function, new InputEvent(1L))) + .isInstanceOf(PythonActionExecutor.PythonActionExecutionException.class) + .hasCause(failure); + verify(pythonEvent).close(); + verify(runnerContext).drainEvents(null); + } + + @Test + void closesAwaitableAndEventWhenStoringAwaitableFails() throws Exception { + PythonInterpreter interpreter = mock(PythonInterpreter.class); + PythonRunnerContextImpl runnerContext = mock(PythonRunnerContextImpl.class); + PythonActionExecutor executor = newExecutor(interpreter, runnerContext); + PythonFunction function = mock(PythonFunction.class); + PyObject pythonEvent = mock(PyObject.class); + PyObject pythonAwaitable = mock(PyObject.class); + RuntimeException failure = new RuntimeException("set failed"); + when(interpreter.invoke(same(CONVERT_JSON_TO_PYTHON_EVENT), anyString())) + .thenReturn(pythonEvent); + when(function.call(same(pythonEvent), isNull())).thenReturn(pythonAwaitable); + doThrow(failure).when(interpreter).set(anyString(), same(pythonAwaitable)); + + assertThatThrownBy(() -> executor.executePythonFunction(function, new InputEvent(1L))) + .isInstanceOf(PythonActionExecutor.PythonActionExecutionException.class) + .hasCause(failure); + verify(pythonAwaitable).close(); + verify(pythonEvent).close(); + } + + @Test + void closesRetrievedAwaitableWhileItIsPending() throws Exception { + PythonInterpreter interpreter = mock(PythonInterpreter.class); + PythonActionExecutor executor = newExecutor(interpreter); + PyObject pythonAwaitable = mock(PyObject.class); + String pythonAwaitableRef = "python_awaitable_1"; + when(interpreter.get(pythonAwaitableRef)).thenReturn(pythonAwaitable); + when(interpreter.invoke(CALL_PYTHON_AWAITABLE, pythonAwaitable)) + .thenReturn(new Object[] {false, null}); + + assertThat(executor.callPythonAwaitable(pythonAwaitableRef)).isFalse(); + + verify(pythonAwaitable).close(); + verify(interpreter, never()).exec(anyString()); + } + + @Test + void deletesCompletedAwaitableAndClosesRetrievedWrapper() throws Exception { + PythonInterpreter interpreter = mock(PythonInterpreter.class); + PythonActionExecutor executor = newExecutor(interpreter); + PyObject pythonAwaitable = mock(PyObject.class); + String pythonAwaitableRef = "python_awaitable_1"; + when(interpreter.get(pythonAwaitableRef)).thenReturn(pythonAwaitable); + when(interpreter.invoke(CALL_PYTHON_AWAITABLE, pythonAwaitable)) + .thenReturn(new Object[] {true, null}); + + assertThat(executor.callPythonAwaitable(pythonAwaitableRef)).isTrue(); + + InOrder closeOrder = inOrder(interpreter, pythonAwaitable); + closeOrder.verify(interpreter).invoke(CALL_PYTHON_AWAITABLE, pythonAwaitable); + closeOrder.verify(interpreter).exec("del " + pythonAwaitableRef); + closeOrder.verify(pythonAwaitable).close(); + } + + @Test + void closesRetrievedAwaitableWhenPollingFails() throws Exception { + PythonInterpreter interpreter = mock(PythonInterpreter.class); + PythonActionExecutor executor = newExecutor(interpreter); + PyObject pythonAwaitable = mock(PyObject.class); + String pythonAwaitableRef = "python_awaitable_1"; + RuntimeException failure = new RuntimeException("poll failed"); + when(interpreter.get(pythonAwaitableRef)).thenReturn(pythonAwaitable); + when(interpreter.invoke(CALL_PYTHON_AWAITABLE, pythonAwaitable)).thenThrow(failure); + + assertThatThrownBy(() -> executor.callPythonAwaitable(pythonAwaitableRef)) + .isSameAs(failure); + verify(pythonAwaitable).close(); + verify(interpreter, never()).exec(anyString()); + } + private static PythonActionExecutor newExecutor(PythonInterpreter interpreter) throws Exception { - return new PythonActionExecutor(interpreter, null, null, null, "test-job"); + return newExecutor(interpreter, null); + } + + private static PythonActionExecutor newExecutor( + PythonInterpreter interpreter, PythonRunnerContextImpl runnerContext) throws Exception { + return new PythonActionExecutor(interpreter, null, null, runnerContext, "test-job"); } } From 2114dc046a50992baa78c514577ae87e8a9c2699 Mon Sep 17 00:00:00 2001 From: joeyutong Date: Wed, 26 Aug 2026 16:54:57 +0800 Subject: [PATCH 2/2] [runtime][python] Avoid materializing ignored action results Return only the awaitable completion state across Pemja because Action yielded and returned values are not consumed by the runtime. Co-Authored-By: Claude Code AI-Model: gpt-5 AI-Contributed/Feature: 29/29 AI-Contributed/UT: 17/17 --- python/flink_agents/plan/function.py | 23 +++++++------------ .../flink_agents/plan/tests/test_function.py | 11 +++++++++ .../python/utils/PythonActionExecutor.java | 6 ++--- .../utils/PythonActionExecutorTest.java | 6 ++--- 4 files changed, 24 insertions(+), 22 deletions(-) diff --git a/python/flink_agents/plan/function.py b/python/flink_agents/plan/function.py index 7988084ae..073780ccd 100644 --- a/python/flink_agents/plan/function.py +++ b/python/flink_agents/plan/function.py @@ -414,28 +414,21 @@ def get_python_function_cache_keys() -> List[Tuple[str, str]]: ) -def call_python_awaitable(awaitable: Any) -> Tuple[bool, Any]: - """Invokes the next step of a Python coroutine or generator and returns whether - it is done, along with the yielded or returned value. +def call_python_awaitable(awaitable: Any) -> bool: + """Invoke the next step of a Python coroutine or generator. Args: awaitable: A Python coroutine or generator object that can be driven by the send() method. Returns: - Tuple[bool, Any]: - - The first element is a boolean flag indicating whether the awaitable - has finished: - * False: The awaitable has more values to yield. - * True: The awaitable has completed. - - The second element is either: - * The value yielded by the awaitable (when not exhausted), or - * The return value of the awaitable (when it has finished). + True if the awaitable has completed, otherwise False. Yielded and returned + values are discarded because Actions communicate through emitted Events. """ try: - result = awaitable.send(None) - except StopIteration as e: - return True, e.value if hasattr(e, "value") else None + awaitable.send(None) + except StopIteration: + return True except RuntimeError as e: err_msg = str(e) if ( @@ -448,4 +441,4 @@ def call_python_awaitable(awaitable: Any) -> Tuple[bool, Any]: logger.exception("Error in awaitable execution") raise else: - return False, result + return False diff --git a/python/flink_agents/plan/tests/test_function.py b/python/flink_agents/plan/tests/test_function.py index 5c9ec2f79..232d54b7f 100644 --- a/python/flink_agents/plan/tests/test_function.py +++ b/python/flink_agents/plan/tests/test_function.py @@ -29,6 +29,7 @@ JavaFunction, PythonFunction, _is_function_cacheable, + call_python_awaitable, call_python_function, clear_python_function_cache, get_python_function_cache_keys, @@ -226,6 +227,16 @@ def test_call_python_function_basic() -> None: assert result == 10 +def test_call_python_awaitable_returns_only_completion_state() -> None: + def awaitable_with_values() -> Generator[Any, None, Any]: + yield object() + return object() + + awaitable = awaitable_with_values() + assert call_python_awaitable(awaitable) is False + assert call_python_awaitable(awaitable) is True + + def test_call_python_function_caching() -> None: """Test that call_python_function reuses cached instances.""" # Clear cache before testing 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 a4c949b73..623a960e0 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 @@ -190,15 +190,15 @@ public Object getOutputFromOutputEvent(String eventJson) { * @return true if the awaitable has completed; false otherwise */ public boolean callPythonAwaitable(String pythonAwaitableRef) throws Exception { - // Calling awaitable.send(None) in Python returns a tuple of (finished, output). + // Python discards yielded/returned values because Actions communicate through Events. try (PyObject pythonAwaitable = (PyObject) 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); - boolean finished = (boolean) ((Object[]) invokeResult)[0]; + checkState(invokeResult instanceof Boolean); + boolean finished = (boolean) invokeResult; if (finished) { interpreter.exec("del " + pythonAwaitableRef); } 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 ee8bc0b36..db5dc849d 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 @@ -192,8 +192,7 @@ void closesRetrievedAwaitableWhileItIsPending() throws Exception { PyObject pythonAwaitable = mock(PyObject.class); String pythonAwaitableRef = "python_awaitable_1"; when(interpreter.get(pythonAwaitableRef)).thenReturn(pythonAwaitable); - when(interpreter.invoke(CALL_PYTHON_AWAITABLE, pythonAwaitable)) - .thenReturn(new Object[] {false, null}); + when(interpreter.invoke(CALL_PYTHON_AWAITABLE, pythonAwaitable)).thenReturn(false); assertThat(executor.callPythonAwaitable(pythonAwaitableRef)).isFalse(); @@ -208,8 +207,7 @@ void deletesCompletedAwaitableAndClosesRetrievedWrapper() throws Exception { PyObject pythonAwaitable = mock(PyObject.class); String pythonAwaitableRef = "python_awaitable_1"; when(interpreter.get(pythonAwaitableRef)).thenReturn(pythonAwaitable); - when(interpreter.invoke(CALL_PYTHON_AWAITABLE, pythonAwaitable)) - .thenReturn(new Object[] {true, null}); + when(interpreter.invoke(CALL_PYTHON_AWAITABLE, pythonAwaitable)).thenReturn(true); assertThat(executor.callPythonAwaitable(pythonAwaitableRef)).isTrue();