org.apache.zeppelin.submarine.hadoop.YarnClient.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.zeppelin.submarine.hadoop.YarnClient.java

Source

/*
 * Licensed 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.zeppelin.submarine.hadoop;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.security.Principal;
import java.security.PrivilegedAction;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
import java.util.stream.Collectors;

import com.google.gson.Gson;
import com.google.gson.JsonArray;
import com.google.gson.JsonElement;
import com.google.gson.JsonIOException;
import com.google.gson.JsonObject;
import com.google.gson.JsonParser;
import com.google.gson.JsonPrimitive;
import com.google.gson.JsonSyntaxException;
import com.google.gson.reflect.TypeToken;
import org.apache.commons.lang.StringUtils;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.security.UserGroupInformation;

import org.apache.http.HttpResponse;
import org.apache.http.auth.AuthSchemeProvider;
import org.apache.http.auth.AuthScope;
import org.apache.http.auth.Credentials;
import org.apache.http.client.HttpClient;
import org.apache.http.client.config.AuthSchemes;
import org.apache.http.client.config.CookieSpecs;
import org.apache.http.client.config.RequestConfig;
import org.apache.http.client.methods.HttpDelete;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.client.methods.HttpUriRequest;
import org.apache.http.config.Lookup;
import org.apache.http.config.RegistryBuilder;
import org.apache.http.impl.auth.SPNegoSchemeFactory;
import org.apache.http.impl.client.BasicCredentialsProvider;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClientBuilder;
import org.apache.zeppelin.conf.ZeppelinConfiguration;
import org.apache.zeppelin.submarine.commons.SubmarineConstants;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.security.auth.Subject;
import javax.security.auth.kerberos.KerberosPrincipal;
import javax.security.auth.login.AppConfigurationEntry;
import javax.security.auth.login.LoginContext;

public class YarnClient {
    private Logger LOGGER = LoggerFactory.getLogger(YarnClient.class);

    private Configuration hadoopConf;
    private String yarnWebHttpAddr;
    private String principal = "";
    private String keytab = "";

    public static final String YARN_REST_APPATTEMPTS = "appAttempts";
    public static final String YARN_REST_CONTAINER = "container";
    public static final String YARN_REST_APPATTEMPT = "appAttempt";
    public static final String YARN_REST_APPATTEMPTID = "appAttemptId";

    public static final String YARN_REST_EXPOSEDPORTS = "EXPOSEDPORTS";
    public static final String CONTAINER_IP = "CONTAINER_IP";
    public static final String CONTAINER_PORT = "CONTAINER_PORT";
    public static final String HOST_IP = "HOST_IP";
    public static final String HOST_PORT = "HOST_PORT";

    String SERVICE_PATH = "/services/{service_name}";

    private boolean hadoopSecurityEnabled = true; // simple or kerberos

    public YarnClient(Properties properties) {
        this.hadoopConf = new Configuration();

        String hadoopAuthType = properties.getProperty(SubmarineConstants.ZEPPELIN_SUBMARINE_AUTH_TYPE, "kerberos");
        if (StringUtils.equals(hadoopAuthType, "simple")) {
            hadoopSecurityEnabled = false;
        }

        yarnWebHttpAddr = properties.getProperty(SubmarineConstants.YARN_WEB_HTTP_ADDRESS, "");
        boolean isSecurityEnabled = UserGroupInformation.isSecurityEnabled();
        if (isSecurityEnabled || hadoopSecurityEnabled) {
            String krb5conf = properties.getProperty(SubmarineConstants.SUBMARINE_HADOOP_KRB5_CONF, "");
            if (StringUtils.isEmpty(krb5conf)) {
                krb5conf = "/etc/krb5.conf";
                System.setProperty("java.security.krb5.conf", krb5conf);
            }

            String keytab = properties.getProperty(SubmarineConstants.SUBMARINE_HADOOP_KEYTAB, "");
            String principal = properties.getProperty(SubmarineConstants.SUBMARINE_HADOOP_PRINCIPAL, "");

            ZeppelinConfiguration zConf = ZeppelinConfiguration.create();
            if (StringUtils.isEmpty(keytab)) {
                keytab = zConf.getString(ZeppelinConfiguration.ConfVars.ZEPPELIN_SERVER_KERBEROS_KEYTAB);
            }
            if (StringUtils.isEmpty(principal)) {
                principal = zConf.getString(ZeppelinConfiguration.ConfVars.ZEPPELIN_SERVER_KERBEROS_PRINCIPAL);
            }
            if (StringUtils.isBlank(keytab) || StringUtils.isBlank(principal)) {
                throw new RuntimeException(
                        "keytab and principal can not be empty, keytab: " + keytab + ", principal: " + principal);
            }

            this.principal = principal;
            this.keytab = keytab;
            if (LOGGER.isDebugEnabled()) {
                System.setProperty("sun.security.spnego.debug", "true");
                System.setProperty("sun.security.krb5.debug", "true");
            }
        }
    }

    // http://yarn-web-http-address/app/v1/services/{service_name}
    public void deleteService(String serviceName) {
        String appUrl = this.yarnWebHttpAddr + "/app/v1/services/" + serviceName + "?_="
                + System.currentTimeMillis();

        InputStream inputStream = null;
        try {
            HttpResponse response = callRestUrl(appUrl, principal, HTTP.DELETE);
            inputStream = response.getEntity().getContent();
            String result = new BufferedReader(new InputStreamReader(inputStream)).lines()
                    .collect(Collectors.joining(System.lineSeparator()));
            if (response.getStatusLine().getStatusCode() != 200 /*success*/) {
                LOGGER.warn("Status code " + response.getStatusLine().getStatusCode());
                LOGGER.warn("message is :" + Arrays.deepToString(response.getAllHeaders()));
                LOGGER.warn("result\n" + result);
            }
        } catch (Exception exp) {
            exp.printStackTrace();
        } finally {
            try {
                if (null != inputStream) {
                    inputStream.close();
                }
            } catch (Exception e) {
                LOGGER.error(e.getMessage(), e);
            }
        }
    }

    // http://yarn-web-http-address/app/v1/services/{appIdOrName}
    // test/resources/app-v1-services-app_name.json
    public Map<String, Object> getAppServices(String appIdOrName) {
        Map<String, Object> mapStatus = new HashMap<>();
        String appUrl = this.yarnWebHttpAddr + "/app/v1/services/" + appIdOrName + "?_="
                + System.currentTimeMillis();

        InputStream inputStream = null;
        try {
            HttpResponse response = callRestUrl(appUrl, principal, HTTP.GET);
            inputStream = response.getEntity().getContent();
            String result = new BufferedReader(new InputStreamReader(inputStream)).lines()
                    .collect(Collectors.joining(System.lineSeparator()));
            if (response.getStatusLine().getStatusCode() != 200 /*success*/
                    && response.getStatusLine().getStatusCode() != 404 /*Not found*/) {
                LOGGER.warn("Status code " + response.getStatusLine().getStatusCode());
                LOGGER.warn("message is :" + Arrays.deepToString(response.getAllHeaders()));
                LOGGER.warn("result\n" + result);
            }

            // parse app status json
            mapStatus = parseAppServices(result);
        } catch (Exception exp) {
            exp.printStackTrace();
        } finally {
            try {
                if (null != inputStream) {
                    inputStream.close();
                }
            } catch (Exception e) {
                LOGGER.error(e.getMessage(), e);
            }
        }

        return mapStatus;
    }

    // http://yarn-web-http-address/ws/v1/cluster/apps/{appId}
    // test/resources/ws-v1-cluster-apps-application_id-failed.json
    // test/resources/ws-v1-cluster-apps-application_id-finished.json
    // test/resources/ws-v1-cluster-apps-application_id-running.json
    public Map<String, Object> getClusterApps(String appId) {
        Map<String, Object> appAttempts = new HashMap<>();
        String appUrl = this.yarnWebHttpAddr + "/ws/v1/cluster/apps/" + appId + "?_=" + System.currentTimeMillis();

        InputStream inputStream = null;
        try {
            HttpResponse response = callRestUrl(appUrl, principal, HTTP.GET);
            inputStream = response.getEntity().getContent();
            String result = new BufferedReader(new InputStreamReader(inputStream)).lines()
                    .collect(Collectors.joining(System.lineSeparator()));
            if (response.getStatusLine().getStatusCode() != 200 /*success*/) {
                LOGGER.warn("Status code " + response.getStatusLine().getStatusCode());
                LOGGER.warn("message is :" + Arrays.deepToString(response.getAllHeaders()));
                LOGGER.warn("result\n" + result);
            }
            // parse app status json
            appAttempts = parseClusterApps(result);

            return appAttempts;
        } catch (Exception exp) {
            exp.printStackTrace();
        } finally {
            try {
                if (null != inputStream) {
                    inputStream.close();
                }
            } catch (Exception e) {
                LOGGER.error(e.getMessage(), e);
            }
        }

        return appAttempts;
    }

    public Map<String, Object> parseClusterApps(String jsonContent) {
        Map<String, Object> appAttempts = new HashMap<>();

        try {
            JsonParser jsonParser = new JsonParser();
            JsonObject jsonObject = (JsonObject) jsonParser.parse(jsonContent);

            JsonObject jsonAppAttempts = jsonObject.get("app").getAsJsonObject();
            if (null == jsonAppAttempts) {
                return appAttempts;
            }
            for (Map.Entry<String, JsonElement> entry : jsonAppAttempts.entrySet()) {
                String key = entry.getKey();
                if (null != entry.getValue() && entry.getValue() instanceof JsonPrimitive) {
                    Object value = entry.getValue().getAsString();
                    appAttempts.put(key, value);
                }
            }
        } catch (JsonIOException e) {
            LOGGER.error(e.getMessage(), e);
        } catch (JsonSyntaxException e) {
            LOGGER.error(e.getMessage(), e);
        }

        return appAttempts;
    }

    // http://yarn-web-http-address/ws/v1/cluster/apps/{appId}/appattempts
    // test/resources/ws-v1-cluster-apps-application_id-appattempts.json
    public List<Map<String, Object>> getAppAttempts(String appId) {
        List<Map<String, Object>> appAttempts = new ArrayList<>();
        String appUrl = this.yarnWebHttpAddr + "/ws/v1/cluster/apps/" + appId + "/appattempts?_="
                + System.currentTimeMillis();

        InputStream inputStream = null;
        try {
            HttpResponse response = callRestUrl(appUrl, principal, HTTP.GET);
            inputStream = response.getEntity().getContent();
            String result = new BufferedReader(new InputStreamReader(inputStream)).lines()
                    .collect(Collectors.joining(System.lineSeparator()));
            if (response.getStatusLine().getStatusCode() != 200 /*success*/) {
                LOGGER.warn("Status code " + response.getStatusLine().getStatusCode());
                LOGGER.warn("message is :" + Arrays.deepToString(response.getAllHeaders()));
                LOGGER.warn("result\n" + result);
            }

            // parse app status json
            appAttempts = parseAppAttempts(result);
        } catch (Exception exp) {
            exp.printStackTrace();
        } finally {
            try {
                if (null != inputStream) {
                    inputStream.close();
                }
            } catch (Exception e) {
                LOGGER.error(e.getMessage(), e);
            }
        }

        return appAttempts;
    }

    // http://yarn-web-http-address/ws/v1/cluster/apps/{appId}/appattempts/{appAttemptId}/containers
    // test/resources/ws-v1-cluster-apps-application_id-appattempts-appattempt_id-containers.json
    public List<Map<String, Object>> getAppAttemptsContainers(String appId, String appAttemptId) {
        List<Map<String, Object>> appAttemptsContainers = new ArrayList<>();
        String appUrl = this.yarnWebHttpAddr + "/ws/v1/cluster/apps/" + appId + "/appattempts/" + appAttemptId
                + "/containers?_=" + System.currentTimeMillis();

        InputStream inputStream = null;
        try {
            HttpResponse response = callRestUrl(appUrl, principal, HTTP.GET);
            inputStream = response.getEntity().getContent();
            String result = new BufferedReader(new InputStreamReader(inputStream)).lines()
                    .collect(Collectors.joining(System.lineSeparator()));
            if (response.getStatusLine().getStatusCode() != 200 /*success*/) {
                LOGGER.warn("Status code " + response.getStatusLine().getStatusCode());
                LOGGER.warn("message is :" + Arrays.deepToString(response.getAllHeaders()));
                LOGGER.warn("result\n" + result);
            }

            // parse app status json
            appAttemptsContainers = parseAppAttemptsContainers(result);
        } catch (Exception exp) {
            exp.printStackTrace();
        } finally {
            try {
                if (null != inputStream) {
                    inputStream.close();
                }
            } catch (Exception e) {
                LOGGER.error(e.getMessage(), e);
            }
        }

        return appAttemptsContainers;
    }

    public List<Map<String, Object>> getAppAttemptsContainersExportPorts(String appId) {
        List<Map<String, Object>> listExportPorts = new ArrayList<>();

        // appId -> appAttemptId
        List<Map<String, Object>> listAppAttempts = getAppAttempts(appId);
        for (Map<String, Object> mapAppAttempts : listAppAttempts) {
            if (mapAppAttempts.containsKey(YARN_REST_APPATTEMPTID)) {
                String appAttemptId = (String) mapAppAttempts.get(YARN_REST_APPATTEMPTID);
                List<Map<String, Object>> exportPorts = getAppAttemptsContainers(appId, appAttemptId);
                if (exportPorts.size() > 0) {
                    listExportPorts.addAll(exportPorts);
                }
            }
        }

        return listExportPorts;
    }

    // Kerberos authentication for simulated curling
    private static HttpClient buildSpengoHttpClient() {
        HttpClientBuilder builder = HttpClientBuilder.create();
        Lookup<AuthSchemeProvider> authSchemeRegistry = RegistryBuilder.<AuthSchemeProvider>create()
                .register(AuthSchemes.SPNEGO, new SPNegoSchemeFactory(true)).build();
        builder.setDefaultAuthSchemeRegistry(authSchemeRegistry);
        BasicCredentialsProvider credentialsProvider = new BasicCredentialsProvider();
        credentialsProvider.setCredentials(new AuthScope(null, -1, null), new Credentials() {
            @Override
            public Principal getUserPrincipal() {
                return null;
            }

            @Override
            public String getPassword() {
                return null;
            }
        });
        builder.setDefaultCredentialsProvider(credentialsProvider);

        // Avoid output WARN: Cookie rejected
        RequestConfig globalConfig = RequestConfig.custom().setCookieSpec(CookieSpecs.IGNORE_COOKIES).build();
        builder.setDefaultRequestConfig(globalConfig);

        CloseableHttpClient httpClient = builder.build();

        return httpClient;
    }

    public HttpResponse callRestUrl(final String url, final String userId, HTTP operation) {
        if (LOGGER.isDebugEnabled()) {
            LOGGER.debug(String.format("Calling YarnClient %s %s %s", this.principal, this.keytab, url));
        }
        javax.security.auth.login.Configuration config = new javax.security.auth.login.Configuration() {
            @SuppressWarnings("serial")
            @Override
            public AppConfigurationEntry[] getAppConfigurationEntry(String name) {
                return new AppConfigurationEntry[] { new AppConfigurationEntry(
                        "com.sun.security.auth.module.Krb5LoginModule",
                        AppConfigurationEntry.LoginModuleControlFlag.REQUIRED, new HashMap<String, Object>() {
                            {
                                put("useTicketCache", "false");
                                put("useKeyTab", "true");
                                put("keyTab", keytab);
                                // Krb5 in GSS API needs to be refreshed so it does not throw the error
                                // Specified version of key is not available
                                put("refreshKrb5Config", "true");
                                put("principal", principal);
                                put("storeKey", "true");
                                put("doNotPrompt", "true");
                                put("isInitiator", "true");
                                if (LOGGER.isDebugEnabled()) {
                                    put("debug", "true");
                                }
                            }
                        }) };
            }
        };

        Set<Principal> principals = new HashSet<Principal>(1);
        principals.add(new KerberosPrincipal(userId));
        Subject sub = new Subject(false, principals, new HashSet<Object>(), new HashSet<Object>());
        try {
            // Authentication module: Krb5Login
            LoginContext loginContext = new LoginContext("Krb5Login", sub, null, config);
            loginContext.login();
            Subject serviceSubject = loginContext.getSubject();
            return Subject.doAs(serviceSubject, new PrivilegedAction<HttpResponse>() {
                HttpResponse httpResponse = null;

                @Override
                public HttpResponse run() {
                    try {
                        HttpUriRequest request = null;
                        switch (operation) {
                        case DELETE:
                            request = new HttpDelete(url);
                            break;
                        case POST:
                            request = new HttpPost(url);
                            break;
                        default:
                            request = new HttpGet(url);
                            break;
                        }

                        HttpClient spengoClient = buildSpengoHttpClient();
                        httpResponse = spengoClient.execute(request);
                        return httpResponse;
                    } catch (IOException e) {
                        LOGGER.error(e.getMessage(), e);
                    }
                    return httpResponse;
                }
            });
        } catch (Exception e) {
            LOGGER.error(e.getMessage(), e);
        }
        return null;
    }

    private Map<String, Object> parseAppServices(String appJson) {
        Map<String, Object> mapStatus = new HashMap<>();

        try {
            JsonParser jsonParser = new JsonParser();
            JsonObject jsonObject = (JsonObject) jsonParser.parse(appJson);

            JsonElement elementAppId = jsonObject.get("id");
            JsonElement elementAppState = jsonObject.get("state");
            JsonElement elementAppName = jsonObject.get("name");

            String appId = (elementAppId == null) ? "" : elementAppId.getAsString();
            String appState = (elementAppState == null) ? "" : elementAppState.getAsString();
            String appName = (elementAppName == null) ? "" : elementAppName.getAsString();

            if (!StringUtils.isEmpty(appId)) {
                mapStatus.put(SubmarineConstants.YARN_APPLICATION_ID, appId);
            }
            if (!StringUtils.isEmpty(appName)) {
                mapStatus.put(SubmarineConstants.YARN_APPLICATION_NAME, appName);
            }
            if (!StringUtils.isEmpty(appState)) {
                mapStatus.put(SubmarineConstants.YARN_APPLICATION_STATUS, appState);
            }
        } catch (JsonIOException e) {
            LOGGER.error(e.getMessage(), e);
        } catch (JsonSyntaxException e) {
            LOGGER.error(e.getMessage(), e);
        }

        return mapStatus;
    }

    // appJson format : submarine/src/test/resources/appAttempts.json
    public List<Map<String, Object>> parseAppAttempts(String jsonContent) {
        List<Map<String, Object>> appAttempts = new ArrayList<>();

        try {
            JsonParser jsonParser = new JsonParser();
            JsonObject jsonObject = (JsonObject) jsonParser.parse(jsonContent);

            JsonObject jsonAppAttempts = jsonObject.get(YARN_REST_APPATTEMPTS).getAsJsonObject();
            if (null == jsonAppAttempts) {
                return appAttempts;
            }
            JsonArray jsonAppAttempt = jsonAppAttempts.get(YARN_REST_APPATTEMPT).getAsJsonArray();
            if (null == jsonAppAttempt) {
                return appAttempts;
            }
            for (int i = 0; i < jsonAppAttempt.size(); i++) {
                Map<String, Object> mapAppAttempt = new HashMap<>();

                JsonObject jsonParagraph = jsonAppAttempt.get(i).getAsJsonObject();

                JsonElement jsonElement = jsonParagraph.get("id");
                String id = (jsonElement == null) ? "" : jsonElement.getAsString();
                mapAppAttempt.put("id", id);

                jsonElement = jsonParagraph.get(YARN_REST_APPATTEMPTID);
                String appAttemptId = (jsonElement == null) ? "" : jsonElement.getAsString();
                mapAppAttempt.put(YARN_REST_APPATTEMPTID, appAttemptId);

                appAttempts.add(mapAppAttempt);
            }
        } catch (JsonIOException e) {
            LOGGER.error(e.getMessage(), e);
        } catch (JsonSyntaxException e) {
            LOGGER.error(e.getMessage(), e);
        }

        return appAttempts;
    }

    // appJson format : submarine/src/test/resources/appAttempts.json
    public List<Map<String, Object>> parseAppAttemptsContainers(String jsonContent) {
        List<Map<String, Object>> appContainers = new ArrayList<>();

        try {
            JsonParser jsonParser = new JsonParser();
            JsonObject jsonObject = (JsonObject) jsonParser.parse(jsonContent);

            JsonArray jsonContainers = jsonObject.get(YARN_REST_CONTAINER).getAsJsonArray();
            for (int i = 0; i < jsonContainers.size(); i++) {
                String hostIp = "";

                JsonObject jsonContainer = jsonContainers.get(i).getAsJsonObject();

                JsonElement jsonElement = jsonContainer.get("nodeId");
                String nodeId = (jsonElement == null) ? "" : jsonElement.getAsString();
                String[] nodeIdParts = nodeId.split(":");
                if (nodeIdParts.length == 2) {
                    hostIp = nodeIdParts[0];
                }

                jsonElement = jsonContainer.get("exposedPorts");
                String exposedPorts = (jsonElement == null) ? "" : jsonElement.getAsString();

                Gson gson = new Gson();
                Map<String, List<Map<String, String>>> listExposedPorts = gson.fromJson(exposedPorts,
                        new TypeToken<Map<String, List<Map<String, String>>>>() {
                        }.getType());
                if (null == listExposedPorts) {
                    continue;
                }
                for (Map.Entry<String, List<Map<String, String>>> entry : listExposedPorts.entrySet()) {
                    String containerPort = entry.getKey();
                    String[] containerPortParts = containerPort.split("/");
                    if (containerPortParts.length == 2) {
                        List<Map<String, String>> hostIps = entry.getValue();
                        for (Map<String, String> hostAttrib : hostIps) {
                            Map<String, Object> containerExposedPort = new HashMap<>();
                            String hostPort = hostAttrib.get("HostPort");
                            containerExposedPort.put(HOST_IP, hostIp);
                            containerExposedPort.put(HOST_PORT, hostPort);
                            containerExposedPort.put(CONTAINER_PORT, containerPortParts[0]);
                            appContainers.add(containerExposedPort);
                        }
                    }
                }
            }
        } catch (JsonIOException e) {
            LOGGER.error(e.getMessage(), e);
        } catch (JsonSyntaxException e) {
            LOGGER.error(e.getMessage(), e);
        }

        return appContainers;
    }

    public List<Map<String, Object>> getAppExportPorts(String name) {
        // Query the IP and port of the submarine interpreter process through the yarn client
        Map<String, Object> mapAppStatus = getAppServices(name);
        if (mapAppStatus.containsKey(SubmarineConstants.YARN_APPLICATION_ID)
                && mapAppStatus.containsKey(SubmarineConstants.YARN_APPLICATION_NAME)
                && mapAppStatus.containsKey(SubmarineConstants.YARN_APPLICATION_STATUS)) {
            String appId = mapAppStatus.get(SubmarineConstants.YARN_APPLICATION_ID).toString();
            String appStatus = mapAppStatus.get(SubmarineConstants.YARN_APPLICATION_STATUS).toString();

            // if (StringUtils.equals(appStatus, SubmarineJob.YarnApplicationState.RUNNING.toString())) {
            List<Map<String, Object>> mapAppAttempts = getAppAttemptsContainersExportPorts(appId);
            return mapAppAttempts;
            //}
        }

        return new ArrayList<Map<String, Object>>() {
        };
    }

    public enum HTTP {
        GET, POST, DELETE;
    }
}