Search in sources :

Example 1 with PythonProcessShutdownHook

use of org.apache.flink.client.python.PythonEnvUtils.PythonProcessShutdownHook in project flink by apache.

the class PythonFunctionFactory method createPythonFunctionFactory.

static PythonFunctionFactory createPythonFunctionFactory(ReadableConfig config) throws ExecutionException, InterruptedException, IOException {
    Map<String, Object> entryPoint;
    PythonProcessShutdownHook shutdownHook = null;
    if (getGatewayServer() == null) {
        GatewayServer gatewayServer = null;
        Process pythonProcess = null;
        String tmpDir = null;
        try {
            gatewayServer = startGatewayServer();
            List<String> commands = new ArrayList<>();
            commands.add("-m");
            commands.add("pyflink.pyflink_callback_server");
            tmpDir = System.getProperty("java.io.tmpdir") + File.separator + "pyflink" + File.separator + UUID.randomUUID();
            pythonProcess = launchPy4jPythonClient(gatewayServer, config, commands, null, tmpDir, false);
            entryPoint = (Map<String, Object>) gatewayServer.getGateway().getEntryPoint();
            int i = 0;
            while (!entryPoint.containsKey("PythonFunctionFactory")) {
                if (!pythonProcess.isAlive()) {
                    throw new RuntimeException("Python callback server start failed!");
                }
                try {
                    Thread.sleep(CHECK_INTERVAL);
                } catch (InterruptedException e) {
                    throw new RuntimeException("Interrupted while waiting for the python process to start.", e);
                }
                i++;
                if (i > TIMEOUT_MILLIS / CHECK_INTERVAL) {
                    throw new RuntimeException("Python callback server start failed!");
                }
            }
        } catch (Throwable e) {
            try {
                if (gatewayServer != null) {
                    gatewayServer.shutdown();
                }
            } catch (Throwable e2) {
            // ignore, do not swallow the origin exception.
            }
            try {
                if (pythonProcess != null) {
                    shutdownPythonProcess(pythonProcess, TIMEOUT_MILLIS);
                }
            } catch (Throwable e3) {
            // ignore, do not swallow the origin exception.
            }
            if (tmpDir != null) {
                FileUtils.deleteDirectoryQuietly(new File(tmpDir));
            }
            throw e;
        }
        shutdownHook = new PythonProcessShutdownHook(pythonProcess, gatewayServer, tmpDir);
        Runtime.getRuntime().addShutdownHook(shutdownHook);
    } else {
        entryPoint = (Map<String, Object>) getGatewayServer().getGateway().getEntryPoint();
    }
    return new PythonFunctionFactoryImpl((PythonFunctionFactory) entryPoint.get("PythonFunctionFactory"), shutdownHook);
}
Also used : ArrayList(java.util.ArrayList) PythonEnvUtils.shutdownPythonProcess(org.apache.flink.client.python.PythonEnvUtils.shutdownPythonProcess) PythonProcessShutdownHook(org.apache.flink.client.python.PythonEnvUtils.PythonProcessShutdownHook) GatewayServer(py4j.GatewayServer) PythonEnvUtils.startGatewayServer(org.apache.flink.client.python.PythonEnvUtils.startGatewayServer) PythonEnvUtils.getGatewayServer(org.apache.flink.client.python.PythonEnvUtils.getGatewayServer) File(java.io.File)

Aggregations

File (java.io.File)1 ArrayList (java.util.ArrayList)1 PythonProcessShutdownHook (org.apache.flink.client.python.PythonEnvUtils.PythonProcessShutdownHook)1 PythonEnvUtils.getGatewayServer (org.apache.flink.client.python.PythonEnvUtils.getGatewayServer)1 PythonEnvUtils.shutdownPythonProcess (org.apache.flink.client.python.PythonEnvUtils.shutdownPythonProcess)1 PythonEnvUtils.startGatewayServer (org.apache.flink.client.python.PythonEnvUtils.startGatewayServer)1 GatewayServer (py4j.GatewayServer)1