[FLINK-31286][python] Make sure Python processes are cleaned up when …

…TaskManager crashes This closes apache#22065.
zchovan · Mar 2, 2023 · ab4e85e · ab4e85e
1 parent 95dd542
commit ab4e85e
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/flink-python/src/main/java/org/apache/flink/python/PythonFunctionRunner.java b/flink-python/src/main/java/org/apache/flink/python/PythonFunctionRunner.java
@@ -24,7 +24,7 @@
 
 /** The base interface of runner which is responsible for the execution of Python functions. */
 @Internal
-public interface PythonFunctionRunner {
+public interface PythonFunctionRunner extends AutoCloseable {
 
     /**
      * Prepares the Python function runner, such as preparing the Python execution environment, etc.

diff --git a/...ain/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java b/...ain/java/org/apache/flink/streaming/api/runners/python/beam/BeamPythonFunctionRunner.java
@@ -38,6 +38,7 @@
 import org.apache.flink.streaming.api.operators.python.process.timer.TimerRegistration;
 import org.apache.flink.streaming.api.runners.python.beam.state.BeamStateRequestHandler;
 import org.apache.flink.util.Preconditions;
+import org.apache.flink.util.ShutdownHookUtil;
 import org.apache.flink.util.function.LongFunctionWithException;
 
 import org.apache.beam.model.fnexecution.v1.BeamFnApi;
@@ -186,6 +187,8 @@ public abstract class BeamPythonFunctionRunner implements PythonFunctionRunner {
     /** The shared resource among Python operators of the same slot. */
     private transient OpaqueMemoryResource<PythonSharedResources> sharedResources;
 
+    private transient Thread shutdownHook;
+
     public BeamPythonFunctionRunner(
             String taskName,
             ProcessPythonEnvironmentManager environmentManager,
@@ -287,6 +290,10 @@ public void open(ReadableConfig config) throws Exception {
                             jobBundleFactory, createPythonExecutionEnvironment(config, -1));
         }
         progressHandler = getProgressHandler(flinkMetricContainer);
+
+        shutdownHook =
+                ShutdownHookUtil.addShutdownHook(
+                        this, BeamPythonFunctionRunner.class.getSimpleName(), LOG);
     }
 
     @Override
@@ -311,6 +318,12 @@ public void close() throws Exception {
         } finally {
             sharedResources = null;
         }
+
+        if (shutdownHook != null) {
+            ShutdownHookUtil.removeShutdownHook(
+                    shutdownHook, BeamPythonFunctionRunner.class.getSimpleName(), LOG);
+            shutdownHook = null;
+        }
     }
 
     @Override