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 d37402779..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 @@ -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 { 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..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 @@ -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 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"); } }