Skip to content

[runtime][python] Release Pemja objects when closing action executor #944

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
joeyutong wants to merge 1 commit into
base: main
Choose a base branch
from
+175 −10
Open

[runtime][python] Release Pemja objects when closing action executor #944

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 26 additions & 10 deletions runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
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.util.ExceptionUtils;
import pemja.core.PythonInterpreter;
import pemja.core.object.PyObject;

Expand Down Expand Up @@ -188,17 +189,32 @@ public boolean callPythonAwaitable(String pythonAwaitableRef) {
}

public void close() throws Exception {
if (interpreter != null) {
if (pythonAsyncThreadPool != null) {
interpreter.invoke(CLOSE_ASYNC_THREAD_POOL, pythonAsyncThreadPool);
}
PyObject asyncThreadPool = pythonAsyncThreadPool;

@weiqingy weiqingy Aug 1, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: The copy-then-null reads as defensive style, but it looks load-bearing. PyObject.close() in pemja 0.5.7 is an unguarded decRef(tState, pyobject) with no null check and no double-close flag, so clearing the fields first is the only thing stopping a repeated close() from decrementing a second time on an already-released handle. Someone later tidying this into closePythonObject(CLOSE_ASYNC_THREAD_POOL, pythonAsyncThreadPool) would drop that quietly, and the only assertion that would notice is the three-line tail of releasesBothPythonObjectsWhenLogicalCleanupFails.

Would a short comment here save the next reader that trip? There is precedent right next door at ActionExecutionOperator.java:471 (// Must close before pythonInterpreter since cached resources may hold Python references.).

Something like this, if it helps:

// Clear the fields before releasing: PyObject.close() is an unguarded native decRef,
// so a repeated close() must not reach the same handle twice.

PyObject runnerContext = pythonRunnerContext;
pythonAsyncThreadPool = null;
pythonRunnerContext = null;

Exception exception = null;
try {
closePythonObject(CLOSE_ASYNC_THREAD_POOL, asyncThreadPool);
} catch (Exception e) {
exception = ExceptionUtils.firstOrSuppressed(e, exception);
}
try {
closePythonObject(CLOSE_FLINK_RUNNER_CONTEXT, runnerContext);
} catch (Exception e) {
exception = ExceptionUtils.firstOrSuppressed(e, exception);
}

if (exception != null) {
throw exception;

@weiqingy weiqingy Aug 1, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Combining both failures and rethrowing is the right call. What I keep looking at is what happens to this exception one frame up:

// PythonBridgeManager.close(), lines 292-302
if (pythonActionExecutor != null) { pythonActionExecutor.close(); }
if (pythonInterpreter != null) { pythonInterpreter.close(); }
if (pythonEnvironmentManager != null) { pythonEnvironmentManager.close(); }

That is a plain sequence, so on exactly the failure path this PR is built for, pythonInterpreter.close() never runs, and that is the release that tears down the interpreter owning every handle still outstanding. ActionExecutionOperator.close() (lines 469-489) has the same shape across its five closes.

To be clear, this is pre-existing. The old close() threw on a failed interpreter.invoke too, so nothing has regressed here. But given the stated goal is "releases both handles even if one cleanup operation fails", how do you see that goal holding one frame up? Carrying the same firstOrSuppressed pattern into PythonBridgeManager.close() would make the guarantee end-to-end, though I may be missing a reason the interpreter is fine to leak on that path.

}
}

if (pythonRunnerContext != null) {
try {
interpreter.invoke(CLOSE_FLINK_RUNNER_CONTEXT, pythonRunnerContext);
} finally {
pythonRunnerContext = null;
}
private void closePythonObject(String closeFunction, PyObject pythonObject) throws Exception {
if (pythonObject != null) {
try (pythonObject) {

@weiqingy weiqingy Aug 1, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is the shape the whole fix turns on: logical cleanup inside, native release on the way out. Two other Pemja handles in the same lifecycle still have the pre-PR shape.

Mem0LongTermMemory.close() (Mem0LongTermMemory.java:137-140) calls adapter.callMethod(pyMem0, "close", Map.of()) and never pyMem0.close(), which is the old PythonActionExecutor.close() exactly. It is live rather than dead code: RunnerContextImpl.java:339-345 calls ltm.close(), and the handle comes from PythonBridgeManager.java:237.

PythonResourceAdapterImpl.pythonResourceContext (PythonResourceAdapterImpl.java:86,99) is built as interpreter.invoke(GET_RESOURCE_CONTEXT, this), so it is a Java object handed into Python, the same JNI-global-ref pattern #942 describes. That class has no close() at all, and PythonBridgeManager.close() never touches the adapter.

I read the PR and #942 as deliberately scoped to the action executor, so I am not suggesting you widen this one. Is a follow-up issue the plan for the sibling handles, or is there something that already releases those two that I have missed?

interpreter.invoke(closeFunction, pythonObject);
}
}
}
Expand Down
149 changes: 149 additions & 0 deletions .../src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,149 @@
/*
* 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 com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.flink.agents.api.agents.AgentExecutionOptions;
import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl;
import org.junit.jupiter.api.Test;
import org.mockito.InOrder;
import pemja.core.PythonInterpreter;
import pemja.core.object.PyObject;

import java.util.HashMap;

import static org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.Mockito.clearInvocations;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.inOrder;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.when;

/** Tests for {@link PythonActionExecutor}. */
class PythonActionExecutorTest {

private static final String CREATE_ASYNC_THREAD_POOL =
"flink_runner_context.create_async_thread_pool";
private static final String CLOSE_ASYNC_THREAD_POOL =
"flink_runner_context.close_async_thread_pool";
private static final String CREATE_FLINK_RUNNER_CONTEXT =
"flink_runner_context.create_flink_runner_context";
private static final String CLOSE_FLINK_RUNNER_CONTEXT =
"flink_runner_context.close_flink_runner_context";

@Test
void releasesPythonObjectsAfterLogicalCleanup() throws Exception {
TestFixture fixture = createOpenedExecutor();
clearInvocations(fixture.interpreter, fixture.asyncThreadPool, fixture.runnerContextObject);

fixture.executor.close();

InOrder closeOrder =
inOrder(fixture.interpreter, fixture.asyncThreadPool, fixture.runnerContextObject);
closeOrder
.verify(fixture.interpreter)
.invoke(CLOSE_ASYNC_THREAD_POOL, fixture.asyncThreadPool);
closeOrder.verify(fixture.asyncThreadPool).close();
closeOrder
.verify(fixture.interpreter)
.invoke(CLOSE_FLINK_RUNNER_CONTEXT, fixture.runnerContextObject);
closeOrder.verify(fixture.runnerContextObject).close();
}

@Test
void releasesBothPythonObjectsWhenLogicalCleanupFails() throws Exception {
TestFixture fixture = createOpenedExecutor();
RuntimeException asyncFailure = new RuntimeException("async cleanup failed");
RuntimeException contextFailure = new RuntimeException("context cleanup failed");
doThrow(asyncFailure)
.when(fixture.interpreter)
.invoke(CLOSE_ASYNC_THREAD_POOL, fixture.asyncThreadPool);
doThrow(contextFailure)
.when(fixture.interpreter)
.invoke(CLOSE_FLINK_RUNNER_CONTEXT, fixture.runnerContextObject);

assertThatThrownBy(fixture.executor::close)
.isSameAs(asyncFailure)
.hasSuppressedException(contextFailure);
InOrder closeOrder =
inOrder(fixture.interpreter, fixture.asyncThreadPool, fixture.runnerContextObject);
closeOrder
.verify(fixture.interpreter)
.invoke(CLOSE_ASYNC_THREAD_POOL, fixture.asyncThreadPool);
closeOrder.verify(fixture.asyncThreadPool).close();
closeOrder
.verify(fixture.interpreter)
.invoke(CLOSE_FLINK_RUNNER_CONTEXT, fixture.runnerContextObject);
closeOrder.verify(fixture.runnerContextObject).close();

clearInvocations(fixture.interpreter, fixture.asyncThreadPool, fixture.runnerContextObject);
fixture.executor.close();
verifyNoInteractions(
fixture.interpreter, fixture.asyncThreadPool, fixture.runnerContextObject);
}

private static TestFixture createOpenedExecutor() throws Exception {
PythonInterpreter interpreter = mock(PythonInterpreter.class);
PythonRunnerContextImpl runnerContext = mock(PythonRunnerContextImpl.class);
JavaResourceAdapter resourceAdapter = mock(JavaResourceAdapter.class);
PyObject asyncThreadPool = mock(PyObject.class);
PyObject runnerContextObject = mock(PyObject.class);
AgentPlan plan = new AgentPlan(new HashMap<>(), new HashMap<>());
String planJson = new ObjectMapper().writeValueAsString(plan);
String jobIdentifier = "job-1";

when(interpreter.invoke(
CREATE_ASYNC_THREAD_POOL,
plan.getConfig().get(AgentExecutionOptions.NUM_ASYNC_THREADS)))
.thenReturn(asyncThreadPool);
when(interpreter.invoke(
CREATE_FLINK_RUNNER_CONTEXT,
runnerContext,
planJson,
asyncThreadPool,
resourceAdapter,
jobIdentifier))
.thenReturn(runnerContextObject);

PythonActionExecutor executor =
new PythonActionExecutor(
interpreter, plan, resourceAdapter, runnerContext, jobIdentifier);
executor.open();
return new TestFixture(interpreter, asyncThreadPool, runnerContextObject, executor);
}

private static final class TestFixture {
private final PythonInterpreter interpreter;
private final PyObject asyncThreadPool;
private final PyObject runnerContextObject;
private final PythonActionExecutor executor;

private TestFixture(
PythonInterpreter interpreter,
PyObject asyncThreadPool,
PyObject runnerContextObject,
PythonActionExecutor executor) {
this.interpreter = interpreter;
this.asyncThreadPool = asyncThreadPool;
this.runnerContextObject = runnerContextObject;
this.executor = executor;
}
}
}
Loading