[runtime][python] Release per-action Pemja objects

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

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

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) {
-        // 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];
+    public boolean callPythonAwaitable(String pythonAwaitableRef) throws Exception {
+        // 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 instanceof Boolean);
+            boolean finished = (boolean) invokeResult;
+            if (finished) {
+                interpreter.exec("del " + pythonAwaitableRef);
+            }
+            return finished;
+        }
     }
 
     public void close() throws Exception {

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,143 @@ 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<String> 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(false);
+
+        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(true);
+
+        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");
     }
 }