Repository: flink-extended/dl-on-flink (language: python)

Issue:
Unify Python Environment Configuration of DL on Flink and PyFlink 
<!--

*Thank you very much for contributing to flink-ai-extended. To help the community review your issue or contribution in the best possible way，please take a few minutes to fulfill following items.*

-->

### Describe the feature
Unify Python Environment Configuration of DL on Flink and PyFlink 




Discussion / hints:

Fix patch:
diff --git a/dl-on-flink-framework/src/main/java/org/flinkextended/flink/ml/cluster/node/runner/python/ProcessPythonRunner.java b/dl-on-flink-framework/src/main/java/org/flinkextended/flink/ml/cluster/node/runner/python/ProcessPythonRunner.java
index 0379b8f0..dc6da091 100644
--- a/dl-on-flink-framework/src/main/java/org/flinkextended/flink/ml/cluster/node/runner/python/ProcessPythonRunner.java
+++ b/dl-on-flink-framework/src/main/java/org/flinkextended/flink/ml/cluster/node/runner/python/ProcessPythonRunner.java
@@ -90,17 +90,7 @@ public static int checkPythonEnvironment(String cmd) {
     public void runScript() throws IOException {
         String startupScript = mlContext.getProperties().get(MLConstants.STARTUP_SCRIPT_FILE);
         List<String> args = new ArrayList<>();
-        String pythonVersion =
-                mlContext.getProperties().getOrDefault(MLConstants.PYTHON_VERSION, "");
-        String pythonExec = "python" + pythonVersion;
-        // check if has python2 or python3 environment
-        if (checkPythonEnvironment("which " + pythonExec) != 0) {
-            throw new RuntimeException("Python executable: " + pythonExec + " not found");
-        }
-        String virtualEnv = mlContext.getProperties().getOrDefault(MLConstants.VIRTUAL_ENV_DIR, "");
-        if (!virtualEnv.isEmpty()) {
-            pythonExec = virtualEnv + "/bin/python";
-        }
+        String pythonExec = getPythonExec();
         args.add(pythonExec);
         if (mlContext.startWithStartup()) {
             args.add(startupScript);
@@ -123,6 +113,18 @@ public void runScript() throws IOException {
         runProcess(builder);
     }
 
+    private String getPythonExec() {
+        if (mlContext.getProperties().containsKey(MLConstants.PYTHON_EXEC)) {
+            return mlContext.getProperties().get(MLConstants.PYTHON_EXEC);
+        }
+        String pythonExec = "python";
+        String virtualEnv = mlContext.getProperties().getOrDefault(MLConstants.VIRTUAL_ENV_DIR, "");
+        if (!virtualEnv.isEmpty()) {
+            pythonExec = virtualEnv + "/bin/python";
+        }
+        return pythonExec;
+    }
+
     protected void runProcess(ProcessBuilder builder) throws IOException {
         child = builder.start();
         ProcessLogger processLogger =
@@ -230,6 +232,7 @@ protected void buildProcessBuilder(ProcessBuilder builder) {
             builder.environment()
                     .put(MLConstants.CLASSPATH, Joiner.on(File.pathSeparator).join(stripped));
         }
+        builder.directory(new File(workerDir));
     }
 
     private synchronized void killProcess() {
@@ -278,7 +281,7 @@ protected static String getHadoopClassPath() {
         if (hadoop == null) {
             return null;
         }
-        StringBuffer buffer = new StringBuffer();
+        StringBuilder buffer = new StringBuilder();
         Preconditions.checkState(
                 ShellExec.run(hadoop + " classpath --glob", buffer::append),
                 "Failed to get hadoop class path");
@@ -287,7 +290,7 @@ protected static String getHadoopClassPath() {
 
     protected static String findHadoopBin() {
         String res = null;
-        StringBuffer buffer = new StringBuffer();
+        StringBuilder buffer = new StringBuilder();
         if (ShellExec.run("type -p hadoop", buffer::append, true)) {
             res = buffer.toString();
         } else {
diff --git a/dl-on-flink-framework/src/main/java/org/flinkextended/flink/ml/cluster/rpc/NodeServer.java b/dl-on-flink-framework/src/main/java/org/flinkextended/flink/ml/cluster/rpc/NodeServer.java
index 6b5ae902..f4ac10b7 100644
--- a/dl-on-flink-framework/src/main/java/org/flinkextended/flink/ml/cluster/rpc/NodeServer.java
+++ b/dl-on-flink-framework/src/main/java/org/flinkextended/flink/ml/cluster/rpc/NodeServer.java
@@ -70,7 +70,7 @@ public NodeServer(MLContext mlContext, String jobName) {
         this.mlContext = mlContext;
         this.jobName = jobName;
         idleTimeout =
-                Long.valueOf(
+                Long.parseLong(
                         mlContext
                                 .getProperties()
                                 .getOrDefault(
@@ -226,7 +226,7 @@ public void run() {
 
         // prepare python environment
         prepareRuntimeEnv(mlContext);
-        Future runnerFuture = null;
+        Future<?> runnerFuture = null;
         try {
             // 1. start GRPC server
             this.server =
@@ -292,15 +292,15 @@ public String getDisplayName() {
         return jobName + ":" + mlContext.getIndex();
     }
 
-    private Future startMLRunner() throws Exception {
+    private Future<?> startMLRunner() throws Exception {
         LOG.info("begin start node:" + mlContext.getIdentity());
         runner = MLRunnerFactory.createMLRunner(mlContext, this);
-        Future future = runnerService.submit(runner);
+        Future<?> future = runnerService.submit(runner);
         LOG.info("end start node:" + mlContext.getIdentity());
         return future;
     }
 
-    private void stopMLRunner(Future runnerFuture) {
+    private void stopMLRunner(Future<?> runnerFuture) {
         if (null != runnerService && (!runnerService.isShutdown())) {
             LOG.info("begin stop node:" + mlContext.getIdentity());
             try {
@@ -319,7 +319,7 @@ private void stopMLRunner(Future runnerFuture) {
     }
 
     /** Stop serving requests and shutdown resources. */
-    private synchronized void cleanup(Future runnerFuture) {
+    private synchronized void cleanup(Future<?> runnerFuture) {
         LOG.info("{} run cleanup!", mlContext.getIdentity());
         stopMLRunner(runnerFuture);
         runnerService.shutdownNow();
@@ -335,7 +335,9 @@ private synchronized void cleanup(Future runnerFuture) {
             server.shutdownNow();
             try {
                 if (!server.awaitTermination(2, TimeUnit.MINUTES)) {
-                    LOG.warn("{} timed out waiting for GRPC server to terminate");
+                    LOG.warn(
+                            "{} timed out waiting for GRPC server to terminate",
+                            mlContext.getIdentity());
                 }
             } catch (InterruptedException e) {
                 e.printStackTrace();
diff --git a/dl-on-flink-framework/src/main/java/org/flinkextended/flink/ml/util/MLConstants.java b/dl-on-flink-framework/src/main/java/org/flinkextended/flink/ml/util/MLConstants.java
index 7bcf07c3..d76d73ec 100644
--- a/dl-on-flink-framework/src/main/java/org/flinkextended/flink/ml/util/MLConstants.java
+++ b/dl-on-flink-framework/src/main/java/org/flinkextended/flink/ml/util/MLConstants.java
@@ -109,7 +109,7 @@ public class MLConstants {
     public static final String FAILOVER_RESTART_ALL_STRATEGY = "all";
     public static final String FAILOVER_RESTART_INDIVIDUAL_STRATEGY = "individual";
     public static final String FAILOVER_STRATEGY_DEFAULT = FAILOVER_RESTART_ALL_STRATEGY;
-    public static final String PYTHON_VERSION = "python.version";
+    public static final String PYTHON_EXEC = "python.exec";
 
     public static final String GPU_INFO = "gpu_info";
 }
diff --git a/dl-on-flink-operator/pom.xml b/dl-on-flink-operator/pom.xml
index 9689b6cc..89f0d08b 100644
--- a/dl-on-flink-operator/pom.xml
+++ b/dl-on-flink-operator/pom.xml
@@ -95,6 +95,11 @@
             <artifactId>flink-table-runtime_${scala.major.version}</artifactId>
             <scope>provided</scope>
         </dependency>
+        <dependency>
+            <groupId>org.apache.flink</groupId>
+            <artifactId>flink-python_${scala.major.version}</artifactId>
+            <scope>provided</scope>
+        </dependency>
         <dependency>
             <groupId>org.apache.flink</groupId>
             <artifactId>flink-table-planner_${scala.major.version}</artifactId>
diff --git a/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/client/NodeUtils.java b/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/client/NodeUtils.java
index c6a3d602..50c310ce 100644
--- a/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/client/NodeUtils.java
+++ b/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/client/NodeUtils.java
@@ -21,6 +21,7 @@
 import org.flinkextended.flink.ml.cluster.ClusterConfig;
 import org.flinkextended.flink.ml.cluster.ClusterConfig.Builder;
 import org.flinkextended.flink.ml.operator.ops.NodeOperator;
+import org.flinkextended.flink.ml.operator.ops.inputformat.NodeInputFormat;
 import org.flinkextended.flink.ml.operator.ops.source.NodeSource;
 import org.flinkextended.flink.ml.operator.util.PythonFileUtil;
 import org.flinkextended.flink.ml.operator.util.ReflectionUtils;
@@ -30,12 +31,15 @@
 import org.apache.flink.annotation.Internal;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.api.java.typeutils.RowTypeInfo;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.python.util.PythonConfigUtil;
 import org.apache.flink.streaming.api.datastream.DataStream;
 import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
 import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
 import org.apache.flink.table.api.Schema;
 import org.apache.flink.table.api.StatementSet;
 import org.apache.flink.table.api.Table;
+import org.apache.flink.table.api.TableConfig;
 import org.apache.flink.table.api.TableDescriptor;
 import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
 import org.apache.flink.table.api.bridge.java.internal.StreamTableEnvironmentImpl;
@@ -72,8 +76,10 @@ public static void scheduleNodes(
                 ReflectionUtils.getFieldValue(
                         statementSet, StatementSetImpl.class, "tableEnvironment");
         final StreamExecutionEnvironment env = tEnv.execEnv();
+        final Configuration flinkConfig = mergeConfiguration(env, tEnv.getConfig());
 
-        final DataStream<?> nodeStream = scheduleNodes(env, clusterConfig, DUMMY_TI, nodeType);
+        final DataStream<?> nodeStream =
+                scheduleNodes(env, clusterConfig, DUMMY_TI, nodeType, flinkConfig);
         final Table nodeTable = tEnv.fromDataStream(nodeStream);
         statementSet.addInsert(TableDescriptor.forConnector("blackhole").build(), nodeTable);
     }
@@ -97,13 +103,15 @@ public static Table scheduleNodes(
             String nodeType) {
 
         final StreamExecutionEnvironment env = ((StreamTableEnvironmentImpl) tEnv).execEnv();
+        final Configuration flinkConfig = mergeConfiguration(env, tEnv.getConfig());
 
         final SchemaResolver schemaResolver =
                 ((StreamTableEnvironmentImpl) tEnv).getCatalogManager().getSchemaResolver();
         final RowTypeInfo outTypeInfo =
                 TypeUtil.schemaToRowTypeInfo(outSchema.resolve(schemaResolver));
 
-        final DataStream<?> nodeStream = scheduleNodes(env, clusterConfig, outTypeInfo, nodeType);
+        final DataStream<?> nodeStream =
+                scheduleNodes(env, clusterConfig, outTypeInfo, nodeType, flinkConfig);
         return tEnv.fromDataStream(nodeStream);
     }
 
@@ -125,8 +133,16 @@ public static void scheduleNodes(
                         statementSet, StatementSetImpl.class, "tableEnvironment");
         final StreamExecutionEnvironment env = tEnv.execEnv();
 
+        final Configuration flinkConfig = mergeConfiguration(env, tEnv.getConfig());
+
         final DataStream<?> nodeStream =
-                scheduleNodes(env, tEnv.toDataStream(input), clusterConfig, DUMMY_TI, nodeType);
+                scheduleNodes(
+                        env,
+                        tEnv.toDataStream(input),
+                        clusterConfig,
+                        DUMMY_TI,
+                        nodeType,
+                        flinkConfig);
         final Table nodeTable = tEnv.fromDataStream(nodeStream);
         statementSet.addInsert(TableDescriptor.forConnector("blackhole").build(), nodeTable);
     }
@@ -152,8 +168,16 @@ public static Table scheduleNodes(
         final RowTypeInfo outTypeInfo =
                 TypeUtil.schemaToRowTypeInfo(outSchema.resolve(schemaResolver));
 
+        final Configuration flinkConfig = mergeConfiguration(env, tEnv.getConfig());
+
         final DataStream<?> nodeStream =
-                scheduleNodes(env, tEnv.toDataStream(input), clusterConfig, outTypeInfo, nodeType);
+                scheduleNodes(
+                        env,
+                        tEnv.toDataStream(input),
+                        clusterConfig,
+                        outTypeInfo,
+                        nodeType,
+                        flinkConfig);
         return tEnv.fromDataStream(nodeStream);
     }
 
@@ -185,6 +209,7 @@ public static void scheduleAMNode(StatementSet statementSet, ClusterConfig clust
      * @param clusterConfig The ClusterConfig of the added node.
      * @param outTypeInfo The output type of the node. If null, the output type is Void.
      * @param nodeType The type of node to be added to the cluster.
+     * @param flinkConfig The flink configuration.
      * @return The datastream of the node.
      */
     @Internal
@@ -192,12 +217,15 @@ public static <T> SingleOutputStreamOperator<T> scheduleNodes(
             StreamExecutionEnvironment env,
             ClusterConfig clusterConfig,
             TypeInformation<T> outTypeInfo,
-            String nodeType) {
+            String nodeType,
+            Configuration flinkConfig) {
         final Builder<?> builder = clusterConfig.toBuilder();
         registerFileToFlinkCache(env, clusterConfig.getPythonFilePathsList(), builder);
         ClusterConfig finalClusterConfig = builder.build();
 
-        return env.addSource(NodeSource.createNodeSource(nodeType, finalClusterConfig, outTypeInfo))
+        return env.addSource(
+                        NodeSource.createNodeSource(
+                                nodeType, finalClusterConfig, outTypeInfo, flinkConfig))
                 .setParallelism(finalClusterConfig.getNodeCount(nodeType))
                 .name(nodeType);
     }
@@ -212,6 +240,7 @@ public static <T> SingleOutputStreamOperator<T> scheduleNodes(
      * @param clusterConfig The ClusterConfig of the added node.
      * @param outTypeInfo The output type of the node.
      * @param nodeType The type of node to be added to the cluster.
+     * @param flinkConfig The flink configuration.
      * @return The datastream of the node.
      */
     @Internal
@@ -220,16 +249,33 @@ public static <T> SingleOutputStreamOperator<T> scheduleNodes(
             DataStream<Row> input,
             ClusterConfig clusterConfig,
             TypeInformation<T> outTypeInfo,
-            String nodeType) {
+            String nodeType,
+            Configuration flinkConfig) {
         final Builder<?> builder = clusterConfig.toBuilder();
         registerFileToFlinkCache(env, clusterConfig.getPythonFilePathsList(), builder);
         ClusterConfig finalClusterConfig = builder.build();
 
         return input.transform(
-                        nodeType, outTypeInfo, new NodeOperator<>(nodeType, finalClusterConfig))
+                        nodeType,
+                        outTypeInfo,
+                        new NodeOperator<>(nodeType, finalClusterConfig, flinkConfig))
                 .setParallelism(finalClusterConfig.getNodeCount(nodeType));
     }
 
+    /**
+     * Merge the {@link TableConfig} with the {@link Configuration} in the {@link
+     * StreamExecutionEnvironment} for {@link NodeOperator} and {@link NodeInputFormat}.
+     *
+     * @param env The Flink stream execution environment.
+     * @param tableConfig The table config.
+     * @return The merged configuration.
+     */
+    @Internal
+    public static Configuration mergeConfiguration(
+            StreamExecutionEnvironment env, TableConfig tableConfig) {
+        return PythonConfigUtil.getMergedConfig(env, tableConfig);
+    }
+
     /**
      * Register the file to Flink cache and set the properties to the given cluster config builder.
      */
diff --git a/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/ops/NodeOperator.java b/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/ops/NodeOperator.java
index 36857ac9..d8fc3b7b 100644
--- a/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/ops/NodeOperator.java
+++ b/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/ops/NodeOperator.java
@@ -27,6 +27,7 @@
 import org.flinkextended.flink.ml.operator.util.PythonFileUtil;
 import org.flinkextended.flink.ml.util.MLConstants;
 
+import org.apache.flink.configuration.Configuration;
 import org.apache.flink.iteration.IterationListener;
 import org.apache.flink.streaming.api.operators.AbstractStreamOperator;
 import org.apache.flink.streaming.api.operators.OneInputStreamOperator;
@@ -59,20 +60,31 @@ public class NodeOperator<OUT> extends AbstractStreamOperator<OUT>
 
     private final String nodeType;
     private final ClusterConfig clusterConfig;
+    private final Configuration flinkConfig;
     private MLContext mlContext;
     private DataExchange<Row, OUT> dataExchange;
     private FutureTask<Void> serverFuture;
     private FutureTask<Void> dataExchangeConsumerFuture;
 
     public NodeOperator(String nodeType, ClusterConfig clusterConfig) {
+        this(nodeType, clusterConfig, new Configuration());
+    }
+
+    public NodeOperator(String nodeType, ClusterConfig clusterConfig, Configuration flinkConfig) {
         this.nodeType = nodeType;
         this.clusterConfig = clusterConfig;
+        this.flinkConfig = flinkConfig;
     }
 
     @Override
     public void open() throws Exception {
+        final PythonEnvironmentManager pythonEnvironmentManager =
+                new PythonEnvironmentManager(clusterConfig, flinkConfig);
+        pythonEnvironmentManager.open(getRuntimeContext());
+
         Map<String, String> properties = new HashMap<>(clusterConfig.getProperties());
         properties.put(MLConstants.GPU_INFO, ResourcesUtils.parseGpuInfo(getRuntimeContext()));
+        properties.putAll(pythonEnvironmentManager.getPythonEnvProperties());
 
         mlContext =
                 new MLContext(
diff --git a/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/ops/PythonEnvironmentManager.java b/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/ops/PythonEnvironmentManager.java
new file mode 100644
index 00000000..c541b5a0
--- /dev/null
+++ b/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/ops/PythonEnvironmentManager.java
@@ -0,0 +1,109 @@
+/*
+ * 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.flinkextended.flink.ml.operator.ops;
+
+import org.flinkextended.flink.ml.cluster.ClusterConfig;
+import org.flinkextended.flink.ml.cluster.node.MLContext;
+import org.flinkextended.flink.ml.operator.util.ReflectionUtils;
+import org.flinkextended.flink.ml.util.MLConstants;
+
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.python.PythonConfig;
+import org.apache.flink.python.env.PythonDependencyInfo;
+import org.apache.flink.python.env.beam.ProcessPythonEnvironmentManager;
+import org.apache.flink.streaming.api.operators.StreamingRuntimeContext;
+
+import com.google.common.annotations.VisibleForTesting;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.apache.flink.python.env.beam.ProcessPythonEnvironmentManager.PYTHON_WORKING_DIR;
+
+/**
+ * PythonEnvironmentManager prepare the python execution environment from the pyFlink configuration.
+ * Use the {@link PythonEnvironmentManager#getPythonEnvProperties} to get the properties for the
+ * {@link MLContext}.
+ */
+public class PythonEnvironmentManager {
+
+    private final ClusterConfig clusterConfig;
+    private final PythonConfig pythonConfig;
+    private Map<String, String> pythonEnv;
+
+    public PythonEnvironmentManager(ClusterConfig clusterConfig, Configuration flinkConfig) {
+        this.clusterConfig = clusterConfig;
+        this.pythonConfig = new PythonConfig(flinkConfig);
+    }
+
+    public void open(StreamingRuntimeContext context) throws Exception {
+        final PythonDependencyInfo dependencyInfo =
+                PythonDependencyInfo.create(pythonConfig, context.getDistributedCache());
+
+        ProcessPythonEnvironmentManager pythonEnvironmentManager =
+                new ProcessPythonEnvironmentManager(
+                        dependencyInfo,
+                        context.getTaskManagerRuntimeInfo().getTmpDirectories(),
+                        new HashMap<>(System.getenv()),
+                        context.getJobId());
+        pythonEnvironmentManager.open();
+        pythonEnv =
+                ReflectionUtils.callMethod(
+                        pythonEnvironmentManager,
+                        ProcessPythonEnvironmentManager.class,
+                        "getPythonEnv");
+    }
+
+    @VisibleForTesting
+    String getWorkingDirectory() {
+        final String workDir = pythonEnv.get(PYTHON_WORKING_DIR);
+        if (workDir != null) {
+            return workDir;
+        }
+        return System.getProperty("user.dir");
+    }
+
+    @VisibleForTesting
+    String getPythonPath() {
+        String pythonPath = pythonEnv.get("PYTHONPATH");
+        if (pythonPath == null) {
+            pythonPath = "";
+        }
+        final String originalPath =
+                clusterConfig
+                        .getProperties()
+                        .getOrDefault(
+                                MLConstants.ENV_PROPERTY_PREFIX + MLConstants.PYTHONPATH_ENV, "");
+        return String.join(":", originalPath, pythonPath);
+    }
+
+    @VisibleForTesting
+    String getPythonExec() {
+        return pythonEnv.get("python");
+    }
+
+    public Map<String, String> getPythonEnvProperties() {
+        Map<String, String> properties = new HashMap<>();
+        properties.put(MLConstants.WORK_DIR, getWorkingDirectory());
+        properties.put(
+                MLConstants.ENV_PROPERTY_PREFIX + MLConstants.PYTHONPATH_ENV, getPythonPath());
+        properties.put(MLConstants.PYTHON_EXEC, getPythonExec());
+        return properties;
+    }
+}
diff --git a/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/ops/inputformat/NodeInputFormat.java b/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/ops/inputformat/NodeInputFormat.java
index dce1b2dc..5bfec40e 100644
--- a/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/ops/inputformat/NodeInputFormat.java
+++ b/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/ops/inputformat/NodeInputFormat.java
@@ -22,12 +22,14 @@
 import org.flinkextended.flink.ml.cluster.ExecutionMode;
 import org.flinkextended.flink.ml.cluster.node.MLContext;
 import org.flinkextended.flink.ml.cluster.rpc.NodeServer;
+import org.flinkextended.flink.ml.operator.ops.PythonEnvironmentManager;
 import org.flinkextended.flink.ml.operator.ops.ResourcesUtils;
 import org.flinkextended.flink.ml.operator.util.ColumnInfos;
 import org.flinkextended.flink.ml.util.MLConstants;
 import org.flinkextended.flink.ml.util.MLException;
 
 import org.apache.flink.configuration.Configuration;
+import org.apache.flink.streaming.api.operators.StreamingRuntimeContext;
 import org.apache.flink.util.Preconditions;
 
 import java.io.IOException;
@@ -38,10 +40,17 @@
 public class NodeInputFormat<OUT> extends AbstractNodeInputFormat<OUT> {
 
     private final String nodeType;
+    private final Configuration flinkConfig;
 
     public NodeInputFormat(String nodeType, ClusterConfig clusterConfig) {
+        this(nodeType, clusterConfig, new Configuration());
+    }
+
+    public NodeInputFormat(
+            String nodeType, ClusterConfig clusterConfig, Configuration flinkConfig) {
         super(clusterConfig);
         this.nodeType = nodeType;
+        this.flinkConfig = flinkConfig;
     }
 
     @Override
@@ -66,8 +75,18 @@ public NodeInputSplit[] createInputSplits(int minNumSplits) throws IOException {
 
     @Override
     protected MLContext prepareMLContext(Integer nodeIndex) throws MLException {
+        final PythonEnvironmentManager pythonEnvironmentManager =
+                new PythonEnvironmentManager(clusterConfig, flinkConfig);
+        try {
+            pythonEnvironmentManager.open((StreamingRuntimeContext) getRuntimeContext());
+        } catch (Exception e) {
+            throw new MLException("Fail to open PythonEnvironmentManager", e);
+        }
+
         Map<String, String> properties = new HashMap<>(clusterConfig.getProperties());
         properties.put(MLConstants.GPU_INFO, ResourcesUtils.parseGpuInfo(getRuntimeContext()));
+        properties.putAll(pythonEnvironmentManager.getPythonEnvProperties());
+
         return new MLContext(
                 ExecutionMode.OTHER,
                 nodeType,
diff --git a/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/ops/source/NodeSource.java b/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/ops/source/NodeSource.java
index dc52d2e8..5b83f00f 100644
--- a/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/ops/source/NodeSource.java
+++ b/dl-on-flink-operator/src/main/java/org/flinkextended/flink/ml/operator/ops/source/NodeSource.java
@@ -30,6 +30,7 @@
 import org.apache.flink.api.common.io.InputFormat;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
 import org.apache.flink.api.java.typeutils.ResultTypeQueryable;
+import org.apache.flink.configuration.Configuration;
 import org.apache.flink.streaming.api.functions.source.InputFormatSourceFunction;
 
 /**
@@ -70,11 +71,16 @@ public static <OUT> NodeSource<OUT> createSource(
      * @param nodeType The node type of the node.
      * @param config The ClusterConfig of the node.
      * @param outTI The output type of the node.
+     * @param flinkConfig The flink configuration.
      * @return Flink source that runs the node.
      */
     public static <OUT> NodeSource<OUT> createNodeSource(
-            String nodeType, ClusterConfig config, TypeInformation<OUT> outTI) {
-        final NodeInputFormat<OUT> inputFormat = new NodeInputFormat<>(nodeType, config);
+            String nodeType,
+            ClusterConfig config,
+            TypeInformation<OUT> outTI,
+            Configuration flinkConfig) {
+        final NodeInputFormat<OUT> inputFormat =
+                new NodeInputFormat<>(nodeType, config, flinkConfig);
         return new NodeSource<>(inputFormat, outTI);
     }
 
diff --git a/dl-on-flink-tensorflow-2.x/pom.xml b/dl-on-flink-tensorflow-2.x/pom.xml
index ed8656b9..139a0fce 100644
--- a/dl-on-flink-tensorflow-2.x/pom.xml
+++ b/dl-on-flink-tensorflow-2.x/pom.xml
@@ -63,6 +63,11 @@
             <artifactId>flink-table-runtime_${scala.major.version}</artifactId>
             <scope>provided</scope>
         </dependency>
+        <dependency>
+            <groupId>org.apache.flink</groupId>
+            <artifactId>flink-python_${scala.major.version}</artifactId>
+            <scope>provided</scope>
+        </dependency>
 
         <dependency>
             <groupId>junit</groupId>
diff --git a/dl-on-flink-tensorflow-2.x/python/dl_on_flink_tensorflow/tensorflow_on_flink_mlconf.py b/dl-on-flink-tensorflow-2.x/python/dl_on_flink_tensorflow/tensorflow_on_flink_mlconf.py
index 62dbef3e..a991d582 100644
--- a/dl-on-flink-tensorflow-2.x/python/dl_on_flink_tensorflow/tensorflow_on_flink_mlconf.py
+++ b/dl-on-flink-tensorflow-2.x/python/dl_on_flink_tensorflow/tensorflow_on_flink_mlconf.py
@@ -102,6 +102,6 @@ class MLCONSTANTS(object):
     FAILOVER_RESTART_INDIVIDUAL_STRATEGY = str(ml_constants.FAILOVER_RESTART_INDIVIDUAL_STRATEGY)
     FAILOVER_STRATEGY_DEFAULT = str(ml_constants.FAILOVER_STRATEGY_DEFAULT)
 
-    PYTHON_VERSION = str(ml_constants.PYTHON_VERSION)
+    PYTHON_EXEC = str(ml_constants.PYTHON_EXEC)
 
     GPU_INFO = str(ml_constants.GPU_INFO)
\ No newline at end of file
diff --git a/dl-on-flink-tensorflow-common/pom.xml b/dl-on-flink-tensorflow-common/pom.xml
index b5009d09..6e4a0187 100644
--- a/dl-on-flink-tensorflow-common/pom.xml
+++ b/dl-on-flink-tensorflow-common/pom.xml
@@ -63,6 +63,11 @@
             <artifactId>flink-table-api-java-bridge_${scala.major.version}</artifactId>
             <scope>provided</scope>
         </dependency>
+        <dependency>
+            <groupId>org.apache.flink</groupId>
+            <artifactId>flink-python_${scala.major.version}</artifactId>
+            <scope>provided</scope>
+        </dependency>
         <dependency>
             <groupId>org.apache.flink</groupId>
             <artifactId>flink-ml-iteration_${scala.major.version}</artifactId>
diff --git a/dl-on-flink-tensorflow-common/src/main/java/org/flinkextended/flink/ml/tensorflow/client/TFNodeIterationBody.java b/dl-on-flink-tensorflow-common/src/main/java/org/flinkextended/flink/ml/tensorflow/client/TFNodeIterationBody.java
index d198538d..fb33a2b4 100644
--- a/dl-on-flink-tensorflow-common/src/main/java/org/flinkextended/flink/ml/tensorflow/client/TFNodeIterationBody.java
+++ b/dl-on-flink-tensorflow-common/src/main/java/org/flinkextended/flink/ml/tensorflow/client/TFNodeIterationBody.java
@@ -23,6 +23,7 @@
 import org.apache.flink.annotation.Internal;
 import org.apache.flink.api.common.functions.FlatMapFunction;
 import org.apache.flink.api.common.typeinfo.TypeInformation;
+import org.apache.flink.configuration.Configuration;
 import org.apache.flink.iteration.DataStreamList;
 import org.apache.flink.iteration.IterationBody;
 import org.apache.flink.iteration.IterationBodyResult;
@@ -46,12 +47,17 @@ public class TFNodeIterationBody implements IterationBody {
     private final StreamExecutionEnvironment env;
     private final TFClusterConfig tfClusterConfig;
     private final Integer maxEpoch;
+    private final Configuration flinkConfig;
 
     public TFNodeIterationBody(
-            StreamExecutionEnvironment env, TFClusterConfig tfClusterConfig, Integer maxEpoch) {
+            StreamExecutionEnvironment env,
+            TFClusterConfig tfClusterConfig,
+            Integer maxEpoch,
+            Configuration flinkConfig) {
         this.env = env;
         this.tfClusterConfig = tfClusterConfig;
         this.maxEpoch = maxEpoch;
+        this.flinkConfig = flinkConfig;
     }
 
     @Override
@@ -63,7 +69,8 @@ public IterationBodyResult process(DataStreamList variableStreams, DataStreamLis
                         input,
                         tfClusterConfig,
                         TypeInformation.of(Void.class),
-                        TFClusterConfig.WORKER_NODE_TYPE);
+                        TFClusterConfig.WORKER_NODE_TYPE,
+                        flinkConfig);
         final DataStream<Integer> terminateStream =
                 trainResStream
                         .getSideOutput(new OutputTag<Integer>("termination") {})
diff --git a/dl-on-flink-tensorflow-common/src/main/java/org/flinkextended/flink/ml/tensorflow/client/TFUtils.java b/dl-on-flink-tensorflow-common/src/main/java/org/flinkextended/flink/ml/tensorflow/client/TFUtils.java
index 6403dab2..3fcaa27f 100644
--- a/dl-on-flink-tensorflow-common/src/main/java/org/flinkextended/flink/ml/tensorflow/client/TFUtils.java
+++ b/dl-on-flink-tensorflow-common/src/main/java/org/flinkextended/flink/ml/tensorflow/client/TFUtils.java
@@ -24,6 +24,7 @@
 import org.flinkextended.flink.ml.util.MLConstants;
 
 import org.apache.flink.annotation.Public;
+import org.apache.flink.configuration.Configuration;
 import org.apache.flink.iteration.DataStreamList;
 import org.apache.flink.iteration.IterationConfig;
 import org.apache.flink.iteration.Iterations;
@@ -117,6 +118,7 @@ public static void train(
                         statementSet, StatementSetImpl.class, "tableEnvironment");
         final StreamExecutionEnvironment env = tEnv.execEnv();
         final DataStream<Row> inputDataStream = tEnv.toDataStream(input);
+        final Configuration flinkConfig = NodeUtils.mergeConfiguration(env, tEnv.getConfig());
 
         NodeUtils.scheduleAMNode(statementSet, tfClusterConfig);
 
@@ -130,7 +132,7 @@ public static void train(
                         DataStreamList.of(dummyInitVariable),
                         ReplayableDataStreamList.replay(inputDataStream),
                         iterationConfig,
-                        new TFNodeIterationBody(env, tfClusterConfig, epoch));
+                        new TFNodeIterationBody(env, tfClusterConfig, epoch, flinkConfig));
 
         final DataStream<Integer> trainResDataStream = dataStreamList.get(0);
         statementSet.addInsert(
diff --git a/dl-on-flink-tensorflow/pom.xml b/dl-on-flink-tensorflow/pom.xml
index 8e363ffa..9bd755b9 100644
--- a/dl-on-flink-tensorflow/pom.xml
+++ b/dl-on-flink-tensorflow/pom.xml
@@ -82,6 +82,11 @@
             <artifactId>flink-table-runtime_${scala.major.version}</artifactId>
             <scope>provided</scope>
         </dependency>
+        <dependency>
+            <groupId>org.apache.flink</groupId>
+            <artifactId>flink-python_${scala.major.version}</artifactId>
+            <scope>provided</scope>
+        </dependency>
 
         <dependency>
             <groupId>junit</groupId>
diff --git a/dl-on-flink-tensorflow/python/dl_on_flink_tensorflow/tensorflow_on_flink_mlconf.py b/dl-on-flink-tensorflow/python/dl_on_flink_tensorflow/tensorflow_on_flink_mlconf.py
index 62dbef3e..a991d582 100644
--- a/dl-on-flink-tensorflow/python/dl_on_flink_tensorflow/tensorflow_on_flink_mlconf.py
+++ b/dl-on-flink-tensorflow/python/dl_on_flink_tensorflow/tensorflow_on_flink_mlconf.py
@@ -102,6 +102,6 @@ class MLCONSTANTS(object):
     FAILOVER_RESTART_INDIVIDUAL_STRATEGY = str(ml_constants.FAILOVER_RESTART_INDIVIDUAL_STRATEGY)
     FAILOVER_STRATEGY_DEFAULT = str(ml_constants.FAILOVER_STRATEGY_DEFAULT)
 
-    PYTHON_VERSION = str(ml_constants.PYTHON_VERSION)
+    PYTHON_EXEC = str(ml_constants.PYTHON_EXEC)
 
     GPU_INFO = str(ml_constants.GPU_INFO)
\ No newline at end of file
diff --git a/pom.xml b/pom.xml
index 64f4073b..74ec62c1 100644
--- a/pom.xml
+++ b/pom.xml
@@ -302,6 +302,11 @@
                 <artifactId>flink-core</artifactId>
                 <version>${flink.version}</version>
             </dependency>
+            <dependency>
+                <groupId>org.apache.flink</groupId>
+                <artifactId>flink-python_${scala.major.version}</artifactId>
+                <version>${flink.version}</version>
+            </dependency>
             <dependency>
                 <groupId>org.apache.flink</groupId>
                 <artifactId>flink-ml-iteration_${scala.major.version}</artifactId>


Test patch:
diff --git a/dl-on-flink-operator/src/test/java/org/flinkextended/flink/ml/operator/ops/NodeOperatorTest.java b/dl-on-flink-operator/src/test/java/org/flinkextended/flink/ml/operator/ops/NodeOperatorTest.java
index ce8d56ba..081675cc 100644
--- a/dl-on-flink-operator/src/test/java/org/flinkextended/flink/ml/operator/ops/NodeOperatorTest.java
+++ b/dl-on-flink-operator/src/test/java/org/flinkextended/flink/ml/operator/ops/NodeOperatorTest.java
@@ -24,6 +24,8 @@
 import org.flinkextended.flink.ml.operator.coding.RowCSVCoding;
 import org.flinkextended.flink.ml.util.MLConstants;
 
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.runtime.taskmanager.TaskManagerRuntimeInfo;
 import org.apache.flink.streaming.api.operators.StreamingRuntimeContext;
 import org.apache.flink.streaming.runtime.streamrecord.StreamRecord;
 import org.apache.flink.types.Row;
@@ -38,6 +40,8 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 /** Unit test for {@link NodeOperator}. */
 public class NodeOperatorTest {
@@ -58,6 +62,11 @@ public void setUp() throws Exception {
                         .build();
 
         StreamingRuntimeContext mockRuntimeContext = Mockito.mock(StreamingRuntimeContext.class);
+        final TaskManagerRuntimeInfo taskManagerRuntimeInfo = mock(TaskManagerRuntimeInfo.class);
+        when(taskManagerRuntimeInfo.getTmpDirectories()).thenReturn(new String[] {"/tmp"});
+        when(mockRuntimeContext.getTaskManagerRuntimeInfo()).thenReturn(taskManagerRuntimeInfo);
+        when(mockRuntimeContext.getJobId()).thenReturn(new JobID());
+
         nodeOperator = new TestNodeOperator("worker", config, mockRuntimeContext);
         writtenSb = new StringBuilder();
     }
diff --git a/dl-on-flink-operator/src/test/java/org/flinkextended/flink/ml/operator/ops/PythonEnvironmentManagerTest.java b/dl-on-flink-operator/src/test/java/org/flinkextended/flink/ml/operator/ops/PythonEnvironmentManagerTest.java
new file mode 100644
index 00000000..56d4952c
--- /dev/null
+++ b/dl-on-flink-operator/src/test/java/org/flinkextended/flink/ml/operator/ops/PythonEnvironmentManagerTest.java
@@ -0,0 +1,184 @@
+/*
+ * 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.flinkextended.flink.ml.operator.ops;
+
+import org.flinkextended.flink.ml.cluster.ClusterConfig;
+
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.cache.DistributedCache;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.Path;
+import org.apache.flink.python.util.PythonDependencyUtils;
+import org.apache.flink.runtime.taskmanager.TaskManagerRuntimeInfo;
+import org.apache.flink.streaming.api.operators.StreamingRuntimeContext;
+
+import org.junit.Before;
+import org.junit.Test;
+
+import java.net.URL;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Future;
+
+import static org.apache.flink.python.PythonOptions.PYTHON_EXECUTABLE;
+import static org.apache.flink.python.util.PythonDependencyUtils.PYTHON_ARCHIVES;
+import static org.apache.flink.python.util.PythonDependencyUtils.PYTHON_FILES;
+import static org.apache.flink.python.util.PythonDependencyUtils.PYTHON_REQUIREMENTS_FILE;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+/** Unit test for {@link PythonEnvironmentManager}. */
+public class PythonEnvironmentManagerTest {
+
+    private StreamingRuntimeContext context;
+    private Configuration flinkConfig;
+
+    @Before
+    public void setUp() throws Exception {
+        context = mock(StreamingRuntimeContext.class);
+        final TaskManagerRuntimeInfo taskManagerRuntimeInfo = mock(TaskManagerRuntimeInfo.class);
+        when(taskManagerRuntimeInfo.getTmpDirectories()).thenReturn(new String[] {"/tmp"});
+        when(context.getTaskManagerRuntimeInfo()).thenReturn(taskManagerRuntimeInfo);
+        when(context.getJobId()).thenReturn(new JobID());
+        flinkConfig = new Configuration();
+    }
+
+    @Test
+    public void testArchiveZipFile() throws Exception {
+        final ClusterConfig config =
+                ClusterConfig.newBuilder().setNodeEntry("/tmp/test.py", "main").build();
+        flinkConfig.set(PYTHON_ARCHIVES, Collections.singletonMap("data", "data.zip"));
+        final PythonEnvironmentManager pythonEnvironmentManager =
+                new PythonEnvironmentManager(config, flinkConfig);
+
+        final URL data = Thread.currentThread().getContextClassLoader().getResource("data.zip");
+        assert data != null;
+        final Future<Path> future = CompletableFuture.completedFuture(new Path(data.getPath()));
+        final DistributedCache distributedCache =
+                new DistributedCache(Collections.singletonMap("data", future));
+        when(context.getDistributedCache()).thenReturn(distributedCache);
+
+        pythonEnvironmentManager.open(context);
+        assertTrue(
+                Files.exists(
+                        Paths.get(
+                                pythonEnvironmentManager.getWorkingDirectory(),
+                                "data.zip",
+                                "data.txt")));
+    }
+
+    @Test
+    public void testNoPyFlinkConfig() throws Exception {
+        final ClusterConfig config =
+                ClusterConfig.newBuilder().setNodeEntry("/tmp/test.py", "main").build();
+        final PythonEnvironmentManager pythonEnvironmentManager =
+                new PythonEnvironmentManager(config, flinkConfig);
+
+        pythonEnvironmentManager.open(context);
+        assertEquals(
+                System.getProperty("user.dir"), pythonEnvironmentManager.getWorkingDirectory());
+        assertEquals(":", pythonEnvironmentManager.getPythonPath());
+        assertEquals("python", pythonEnvironmentManager.getPythonExec());
+    }
+
+    @Test
+    public void testPythonFiles() throws Exception {
+        final ClusterConfig config =
+                ClusterConfig.newBuilder().setNodeEntry("/tmp/test.py", "main").build();
+        Map<String, String> pythonFiles = new HashMap<>();
+        pythonFiles.put("row_input.py", "row_input.py");
+        pythonFiles.put("greeter.zip", "greeter.zip");
+        flinkConfig.set(PYTHON_FILES, pythonFiles);
+        final PythonEnvironmentManager pythonEnvironmentManager =
+                new PythonEnvironmentManager(config, flinkConfig);
+
+        final URL greeterPy =
+                Thread.currentThread().getContextClassLoader().getResource("row_input.py");
+        final URL greeterZip =
+                Thread.currentThread().getContextClassLoader().getResource("greeter.zip");
+
+        Map<String, Future<Path>> copyTask = new HashMap<>();
+        copyTask.put(
+                "row_input.py",
+                CompletableFuture.completedFuture(
+                        new Path(Objects.requireNonNull(greeterPy).getPath())));
+        copyTask.put(
+                "greeter.zip",
+                CompletableFuture.completedFuture(
+                        new Path(Objects.requireNonNull(greeterZip).getPath())));
+        final DistributedCache distributedCache = new DistributedCache(copyTask);
+        when(context.getDistributedCache()).thenReturn(distributedCache);
+
+        pythonEnvironmentManager.open(context);
+        assertTrue(checkInPythonPath("row_input.py", pythonEnvironmentManager.getPythonPath()));
+        assertTrue(checkInPythonPath("greeter.py", pythonEnvironmentManager.getPythonPath()));
+    }
+
+    @Test
+    public void testRequirements() throws Exception {
+        final ClusterConfig config =
+                ClusterConfig.newBuilder().setNodeEntry("/tmp/test.py", "main").build();
+        flinkConfig.set(
+                PYTHON_REQUIREMENTS_FILE,
+                Collections.singletonMap(PythonDependencyUtils.FILE, "requirements.txt"));
+        final PythonEnvironmentManager pythonEnvironmentManager =
+                new PythonEnvironmentManager(config, flinkConfig);
+
+        final URL requirements =
+                Thread.currentThread().getContextClassLoader().getResource("requirements.txt");
+        assert requirements != null;
+        final Future<Path> future =
+                CompletableFuture.completedFuture(new Path(requirements.getPath()));
+        final DistributedCache distributedCache =
+                new DistributedCache(Collections.singletonMap("requirements.txt", future));
+        when(context.getDistributedCache()).thenReturn(distributedCache);
+
+        pythonEnvironmentManager.open(context);
+        checkInPythonPath(
+                "dl_on_flink_framework/cluster_config.py",
+                pythonEnvironmentManager.getPythonPath());
+        flinkConfig.set(
+                PYTHON_REQUIREMENTS_FILE,
+                Collections.singletonMap(PythonDependencyUtils.FILE, "requirements.txt"));
+    }
+
+    @Test
+    public void testPythonExec() throws Exception {
+        final ClusterConfig config =
+                ClusterConfig.newBuilder().setNodeEntry("/tmp/test.py", "main").build();
+        flinkConfig.set(PYTHON_EXECUTABLE, "/my_python");
+        final PythonEnvironmentManager pythonEnvironmentManager =
+                new PythonEnvironmentManager(config, flinkConfig);
+        pythonEnvironmentManager.open(context);
+        assertEquals("/my_python", pythonEnvironmentManager.getPythonExec());
+    }
+
+    private boolean checkInPythonPath(String pythonFile, String pythonPath) {
+        return Arrays.stream(pythonPath.split(":"))
+                .anyMatch(path -> Files.exists(Paths.get(path, pythonFile)));
+    }
+}
diff --git a/dl-on-flink-operator/src/test/java/org/flinkextended/flink/ml/operator/ops/inputformat/NodeInputFormatTest.java b/dl-on-flink-operator/src/test/java/org/flinkextended/flink/ml/operator/ops/inputformat/NodeInputFormatTest.java
index d9d5aedf..d31fbd6e 100644
--- a/dl-on-flink-operator/src/test/java/org/flinkextended/flink/ml/operator/ops/inputformat/NodeInputFormatTest.java
+++ b/dl-on-flink-operator/src/test/java/org/flinkextended/flink/ml/operator/ops/inputformat/NodeInputFormatTest.java
@@ -24,8 +24,10 @@
 import org.flinkextended.flink.ml.util.MLConstants;
 import org.flinkextended.flink.ml.util.MLException;
 
-import org.apache.flink.api.common.functions.RuntimeContext;
+import org.apache.flink.api.common.JobID;
 import org.apache.flink.configuration.Configuration;
+import org.apache.flink.runtime.taskmanager.TaskManagerRuntimeInfo;
+import org.apache.flink.streaming.api.operators.StreamingRuntimeContext;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -38,6 +40,7 @@
 import static org.hamcrest.Matchers.isA;
 import static org.junit.Assert.assertEquals;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
 
 /** Unit test for {@link NodeInputFormat}. */
 public class NodeInputFormatTest {
@@ -52,7 +55,12 @@ public void setUp() throws Exception {
                         .setNodeEntry("entry.py", "main")
                         .build();
         nodeInputFormat = new NodeInputFormat<>("worker", clusterConfig);
-        RuntimeContext mockRuntimeContext = mock(RuntimeContext.class);
+        StreamingRuntimeContext mockRuntimeContext = mock(StreamingRuntimeContext.class);
+        final TaskManagerRuntimeInfo taskManagerRuntimeInfo = mock(TaskManagerRuntimeInfo.class);
+        when(taskManagerRuntimeInfo.getTmpDirectories()).thenReturn(new String[] {"/tmp"});
+        when(mockRuntimeContext.getTaskManagerRuntimeInfo()).thenReturn(taskManagerRuntimeInfo);
+        when(mockRuntimeContext.getJobId()).thenReturn(new JobID());
+
         nodeInputFormat.setRuntimeContext(mockRuntimeContext);
     }
 
diff --git a/dl-on-flink-operator/src/test/resources/data.zip b/dl-on-flink-operator/src/test/resources/data.zip
new file mode 100644
index 00000000..79d707d0
Binary files /dev/null and b/dl-on-flink-operator/src/test/resources/data.zip differ
diff --git a/dl-on-flink-operator/src/test/resources/greeter.zip b/dl-on-flink-operator/src/test/resources/greeter.zip
new file mode 100644
index 00000000..1c1aaad4
Binary files /dev/null and b/dl-on-flink-operator/src/test/resources/greeter.zip differ
diff --git a/dl-on-flink-operator/src/test/resources/requirements.txt b/dl-on-flink-operator/src/test/resources/requirements.txt
new file mode 100644
index 00000000..7b7015c0
--- /dev/null
+++ b/dl-on-flink-operator/src/test/resources/requirements.txt
@@ -0,0 +1,1 @@
+dl-on-flink-framework
\ No newline at end of file
diff --git a/dl-on-flink-tensorflow-2.x/python/tests/file.py b/dl-on-flink-tensorflow-2.x/python/tests/file.py
new file mode 100644
index 00000000..ae97dacb
--- /dev/null
+++ b/dl-on-flink-tensorflow-2.x/python/tests/file.py
@@ -0,0 +1,18 @@
+#  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.
+
+def my_func():
+    return True
\ No newline at end of file
diff --git a/dl-on-flink-tensorflow-2.x/python/tests/file.zip b/dl-on-flink-tensorflow-2.x/python/tests/file.zip
new file mode 100644
index 00000000..6905d85d
Binary files /dev/null and b/dl-on-flink-tensorflow-2.x/python/tests/file.zip differ
diff --git a/dl-on-flink-tensorflow-2.x/python/tests/flink_ml_tensorflow/test_pyflink_env.py b/dl-on-flink-tensorflow-2.x/python/tests/flink_ml_tensorflow/test_pyflink_env.py
new file mode 100644
index 00000000..493723cd
--- /dev/null
+++ b/dl-on-flink-tensorflow-2.x/python/tests/flink_ml_tensorflow/test_pyflink_env.py
@@ -0,0 +1,100 @@
+#  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.
+import os.path
+import unittest
+
+from dl_on_flink_framework.context import Context
+from pyflink.datastream import StreamExecutionEnvironment
+from pyflink.table import StreamTableEnvironment
+
+from dl_on_flink_tensorflow.tf_cluster_config import TFClusterConfig
+from dl_on_flink_tensorflow.tf_utils import train
+
+test_folder = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+archive_path = os.path.join(test_folder, "sample.zip")
+python_file_path = os.path.join(test_folder, "file.py")
+zipped_python_file_path = os.path.join(test_folder, "file.zip")
+requirements_file_path = os.path.join(test_folder, "requirements.txt")
+
+
+def verify_zipped_archive(context: Context):
+    path = "sample.zip/sample.csv"
+    print(f"checking if {os.path.abspath(path)} exists")
+    assert os.path.exists(path)
+
+
+def verify_added_python_file(context: Context):
+    from file import my_func
+    assert my_func()
+
+
+def verify_sklearn_installed(context: Context):
+    import sklearn
+    print(sklearn.show_versions())
+
+
+class TestPyFlinkEnv(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        from tests.flink_ml_tensorflow.utils import add_dl_on_flink_jar
+        add_dl_on_flink_jar()
+        super().setUpClass()
+
+    def setUp(self) -> None:
+        super().setUp()
+        from tests.flink_ml_tensorflow.utils import find_jar_path
+        self.env = StreamExecutionEnvironment.get_execution_environment()
+        self.env.add_jars("file://{}".format(find_jar_path()))
+        self.env.set_parallelism(1)
+        self.t_env = StreamTableEnvironment.create(self.env)
+        self.statement_set = self.t_env.create_statement_set()
+
+    def test_add_zipped_archive(self):
+        self.t_env.add_python_archive(archive_path)
+
+        config = TFClusterConfig.new_builder().set_worker_count(1) \
+            .set_node_entry(verify_zipped_archive).build()
+
+        train(self.statement_set, config)
+        self.statement_set.execute().wait()
+
+    def test_add_python_files(self):
+        self.t_env.add_python_file(python_file_path)
+
+        config = TFClusterConfig.new_builder().set_worker_count(1) \
+            .set_node_entry(verify_added_python_file).build()
+
+        train(self.statement_set, config)
+        self.statement_set.execute().wait()
+
+    def test_add_zipped_python_file(self):
+        self.t_env.add_python_file(zipped_python_file_path)
+
+        config = TFClusterConfig.new_builder().set_worker_count(1) \
+            .set_node_entry(verify_added_python_file).build()
+
+        train(self.statement_set, config)
+        self.statement_set.execute().wait()
+
+    def test_requirements(self):
+        self.t_env.set_python_requirements(requirements_file_path)
+
+        config = TFClusterConfig.new_builder().set_worker_count(1) \
+            .set_node_entry(verify_sklearn_installed).build()
+
+        train(self.statement_set, config)
+        self.statement_set.execute().wait()
diff --git a/dl-on-flink-tensorflow-2.x/python/tests/requirements.txt b/dl-on-flink-tensorflow-2.x/python/tests/requirements.txt
new file mode 100644
index 00000000..ff88936c
--- /dev/null
+++ b/dl-on-flink-tensorflow-2.x/python/tests/requirements.txt
@@ -0,0 +1,1 @@
+scikit-learn
\ No newline at end of file
diff --git a/dl-on-flink-tensorflow-2.x/python/tests/sample.zip b/dl-on-flink-tensorflow-2.x/python/tests/sample.zip
new file mode 100644
index 00000000..4b9a4528
Binary files /dev/null and b/dl-on-flink-tensorflow-2.x/python/tests/sample.zip differ
diff --git a/dl-on-flink-tensorflow/python/tests/file.py b/dl-on-flink-tensorflow/python/tests/file.py
new file mode 100644
index 00000000..ae97dacb
--- /dev/null
+++ b/dl-on-flink-tensorflow/python/tests/file.py
@@ -0,0 +1,18 @@
+#  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.
+
+def my_func():
+    return True
\ No newline at end of file
diff --git a/dl-on-flink-tensorflow/python/tests/file.zip b/dl-on-flink-tensorflow/python/tests/file.zip
new file mode 100644
index 00000000..6905d85d
Binary files /dev/null and b/dl-on-flink-tensorflow/python/tests/file.zip differ
diff --git a/dl-on-flink-tensorflow/python/tests/flink_ml_tensorflow/test_flink_ml_api.py b/dl-on-flink-tensorflow/python/tests/flink_ml_tensorflow/test_flink_ml_api.py
index 984e73de..1e317f7d 100644
--- a/dl-on-flink-tensorflow/python/tests/flink_ml_tensorflow/test_flink_ml_api.py
+++ b/dl-on-flink-tensorflow/python/tests/flink_ml_tensorflow/test_flink_ml_api.py
@@ -61,8 +61,7 @@ def setUp(self) -> None:
     def get_tf_config() -> TFConfig:
         work_num = 1
         ps_num = 1
-        prop = {MLCONSTANTS.PYTHON_VERSION: '',
-                MLCONSTANTS.ENCODING_CLASS: 'org.flinkextended.flink.ml.operator.coding.RowCSVCoding',
+        prop = {MLCONSTANTS.ENCODING_CLASS: 'org.flinkextended.flink.ml.operator.coding.RowCSVCoding',
                 MLCONSTANTS.DECODING_CLASS: 'org.flinkextended.flink.ml.operator.coding.RowCSVCoding',
                 'input_types': 'STRING'}
         tf_config = TFConfig(work_num, ps_num, prop, os.path.join(os.path.dirname(__file__), "add_one.py"),
diff --git a/dl-on-flink-tensorflow/python/tests/flink_ml_tensorflow/test_pyflink_env.py b/dl-on-flink-tensorflow/python/tests/flink_ml_tensorflow/test_pyflink_env.py
new file mode 100644
index 00000000..493723cd
--- /dev/null
+++ b/dl-on-flink-tensorflow/python/tests/flink_ml_tensorflow/test_pyflink_env.py
@@ -0,0 +1,100 @@
+#  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.
+import os.path
+import unittest
+
+from dl_on_flink_framework.context import Context
+from pyflink.datastream import StreamExecutionEnvironment
+from pyflink.table import StreamTableEnvironment
+
+from dl_on_flink_tensorflow.tf_cluster_config import TFClusterConfig
+from dl_on_flink_tensorflow.tf_utils import train
+
+test_folder = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+archive_path = os.path.join(test_folder, "sample.zip")
+python_file_path = os.path.join(test_folder, "file.py")
+zipped_python_file_path = os.path.join(test_folder, "file.zip")
+requirements_file_path = os.path.join(test_folder, "requirements.txt")
+
+
+def verify_zipped_archive(context: Context):
+    path = "sample.zip/sample.csv"
+    print(f"checking if {os.path.abspath(path)} exists")
+    assert os.path.exists(path)
+
+
+def verify_added_python_file(context: Context):
+    from file import my_func
+    assert my_func()
+
+
+def verify_sklearn_installed(context: Context):
+    import sklearn
+    print(sklearn.show_versions())
+
+
+class TestPyFlinkEnv(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        from tests.flink_ml_tensorflow.utils import add_dl_on_flink_jar
+        add_dl_on_flink_jar()
+        super().setUpClass()
+
+    def setUp(self) -> None:
+        super().setUp()
+        from tests.flink_ml_tensorflow.utils import find_jar_path
+        self.env = StreamExecutionEnvironment.get_execution_environment()
+        self.env.add_jars("file://{}".format(find_jar_path()))
+        self.env.set_parallelism(1)
+        self.t_env = StreamTableEnvironment.create(self.env)
+        self.statement_set = self.t_env.create_statement_set()
+
+    def test_add_zipped_archive(self):
+        self.t_env.add_python_archive(archive_path)
+
+        config = TFClusterConfig.new_builder().set_worker_count(1) \
+            .set_node_entry(verify_zipped_archive).build()
+
+        train(self.statement_set, config)
+        self.statement_set.execute().wait()
+
+    def test_add_python_files(self):
+        self.t_env.add_python_file(python_file_path)
+
+        config = TFClusterConfig.new_builder().set_worker_count(1) \
+            .set_node_entry(verify_added_python_file).build()
+
+        train(self.statement_set, config)
+        self.statement_set.execute().wait()
+
+    def test_add_zipped_python_file(self):
+        self.t_env.add_python_file(zipped_python_file_path)
+
+        config = TFClusterConfig.new_builder().set_worker_count(1) \
+            .set_node_entry(verify_added_python_file).build()
+
+        train(self.statement_set, config)
+        self.statement_set.execute().wait()
+
+    def test_requirements(self):
+        self.t_env.set_python_requirements(requirements_file_path)
+
+        config = TFClusterConfig.new_builder().set_worker_count(1) \
+            .set_node_entry(verify_sklearn_installed).build()
+
+        train(self.statement_set, config)
+        self.statement_set.execute().wait()
diff --git a/dl-on-flink-tensorflow/python/tests/requirements.txt b/dl-on-flink-tensorflow/python/tests/requirements.txt
new file mode 100644
index 00000000..ff88936c
--- /dev/null
+++ b/dl-on-flink-tensorflow/python/tests/requirements.txt
@@ -0,0 +1,1 @@
+scikit-learn
\ No newline at end of file
diff --git a/dl-on-flink-tensorflow/python/tests/sample.zip b/dl-on-flink-tensorflow/python/tests/sample.zip
new file mode 100644
index 00000000..4b9a4528
Binary files /dev/null and b/dl-on-flink-tensorflow/python/tests/sample.zip differ